Update llm.py

llm.py

@@ -1,40 +1,20 @@
-from transformers import pipeline
-import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-qa_pipeline = pipeline(
-    "text2text-generation",
-    model="google/flan-t5-base",
-    device=device
-)
+tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
 
 def generate_answer(context, question):
-    # Handle empty context
-    if context == "No relevant context found.":
-        prompt = f"""
-        You are a helpful AI assistant. Answer the question based on your general knowledge.
-        
-        Question: {question}
-        
-        Answer as a helpful paragraph:
-        """
-    else:
-        prompt = f"""
-        You are a helpful AI assistant. Use the context to answer the question.
-        
-        Context:
-        {context}
-        
-        Question: {question}
-        
-        Answer as a comprehensive paragraph with key details:
-        """
-    
-    result = qa_pipeline(
-        prompt,
-        max_length=400,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9
+    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 result[0]['generated_text'].strip()
+    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()