Update llm.py

llm.py

@@ -1,20 +1,12 @@
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+# llm.py
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
-model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
+# load model once
+tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+model = AutoModelForCausalLM.from_pretrained("distilgpt2")
 
 def generate_answer(context, question):
-    prompt = f"""
-You are a helpful AI assistant.
-Context:
-{context}
-Question: {question}
-Answer as a helpful paragraph:"""
-
-    inputs = tokenizer(prompt, return_tensors='pt', truncation=True, max_length=512)
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=100,
-        do_sample=False
-    )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+    prompt = f"Context:\n{context}\n\nQuestion: {question}\nAnswer:"
+    inputs = tokenizer.encode(prompt, return_tensors='pt', max_length=1024, truncation=True)
+    outputs = model.generate(inputs, max_new_tokens=50, do_sample=True)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()