Update llm.py

llm.py

@@ -1,23 +1,29 @@
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import pipeline
 
-tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
-model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
+qa_pipeline = pipeline(
+    "text2text-generation",
+    model="google/flan-t5-large",
+    device=0 if torch.cuda.is_available() else -1
+)
 
 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
+    You are a helpful research assistant. Use the context to answer the question.
+    
+    Context:
+    {context}
+    
+    Question: {question}
+    
+    Answer in a comprehensive paragraph with key details:
+    """
+    
+    result = qa_pipeline(
+        prompt,
+        max_length=512,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+    
+    return result[0]['generated_text'].strip()