Update llm.py

llm.py

@@ -1,23 +1,20 @@
-# llm.py
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
-model = AutoModelForCausalLM.from_pretrained("distilgpt2")
-
-# Fix: add pad_token_id if missing
-if tokenizer.pad_token_id is None:
-    tokenizer.pad_token_id = tokenizer.eos_token_id
+tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
+model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
 
 def generate_answer(context, question):
-    prompt = f"Context:\n{context}\n\nQuestion: {question}\nAnswer:"
-    # Limit to last N chars if prompt is too long
-    prompt = prompt[-1000:]  
+    prompt = f"""Context:
+{context}
+
+Based on the above context, answer the question:
+Question: {question}
+Answer:"""
 
-    inputs = tokenizer(prompt, return_tensors='pt', truncation=True, max_length=1024)
+    inputs = tokenizer(prompt, return_tensors='pt', truncation=True, max_length=512)
     outputs = model.generate(
-        inputs["input_ids"],
-        max_new_tokens=50,
-        do_sample=True,
-        pad_token_id=tokenizer.eos_token_id  # fix warning
+        **inputs,
+        max_new_tokens=80,
+        do_sample=False  # deterministic
     )
     return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()