Update llm.py

llm.py

@@ -1,10 +1,48 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline
 
-tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
-model = AutoModelForCausalLM.from_pretrained("distilgpt2")
+# Load conversational model (small version for demo)
+qa_pipeline = pipeline(
+    "text-generation",
+    model="microsoft/DialoGPT-medium",
+    max_new_tokens=200
+)
 
-def generate_answer(context, question):
-    prompt = f"Context:\n{context}\n\nQuestion: {question}\nAnswer:"
-    inputs = tokenizer.encode(prompt, return_tensors='pt', truncation=True, max_length=1024)
-    outputs = model.generate(inputs, max_new_tokens=60, do_sample=True, pad_token_id=tokenizer.eos_token_id)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+def generate_answer(question, context, search_results):
+    # Build context-aware prompt
+    background = "\n".join([
+        f"Source {i+1}: {res['title']}\nContent: {res['content'][:1000]}"
+        for i, res in enumerate(search_results)
+    ])
+    
+    prompt = f"""
+[Conversation History]
+{context}
+
+[Research Materials]
+{background}
+
+[User Question]
+{question}
+
+[Assistant Response Guidelines]
+1. Answer conversationally like a helpful assistant
+2. Acknowledge if previous context exists
+3. Cite sources using numbers like [1], [2]
+4. Be concise but comprehensive
+5. If unsure, say so and suggest alternative approaches
+
+Assistant:
+"""
+    # Generate response
+    response = qa_pipeline(
+        prompt,
+        do_sample=True,
+        temperature=0.7,
+        top_k=50,
+        pad_token_id=50256  # DialoGPT pad token
+    )[0]['generated_text'].split("Assistant:")[-1].strip()
+    
+    return {
+        "response": response,
+        "sources": search_results
+    }