Update llm.py

llm.py

@@ -1,19 +1,23 @@
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
-model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
+tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
 
 def generate_answer(context, question):
-    prompt = f"""Context:
+    prompt = f"""
+You are a helpful AI assistant.
+
+Context:
 {context}
 
-Based on the above context, answer the question:
-{question}
-"""
+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=80,
+        max_new_tokens=100,
         do_sample=False
     )
     return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()