Update llm.py

llm.py

@@ -1,74 +1,41 @@
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-import time
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Load CPU-optimized model
+tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+model = AutoModelForCausalLM.from_pretrained("distilgpt2")
+
+def generate_answer(context, question, max_new_tokens=100):
+    """Generate answer with CPU optimizations"""
+    # Create concise prompt
+    prompt = f"""Based on the context, answer the question conversationally.
 
-# Option 1: Ultra-fast CPU model (best for response speed)
-def generate_answer_fast(context, question):
-    """
-    Uses DistilGPT-2 (smallest version) with optimizations
-    ~2-5 seconds on CPU for 60 tokens
-    """
-    tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
-    model = AutoModelForCausalLM.from_pretrained("distilgpt2")
-    
-    prompt = f"""You are a helpful AI assistant. Using the context, answer the question conversationally.
-    
 Context:
-{context[:1500]}
+{context[:1000]}
 
 Question: {question}
 Answer:"""
     
-    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
+    # Tokenize with truncation
+    inputs = tokenizer(
+        prompt, 
+        return_tensors="pt",
+        max_length=512,
+        truncation=True
+    )
     
-    # Fast generation parameters for CPU
+    # Generate with CPU-optimized settings
     outputs = model.generate(
         inputs.input_ids,
-        max_new_tokens=80,
-        num_beams=1,            # Disable beam search (faster)
-        do_sample=False,         # Disable sampling (faster)
+        max_new_tokens=max_new_tokens,
+        num_beams=1,                # Faster than beam search
+        do_sample=True,             # More natural responses
+        temperature=0.7,            # Balance creativity/focus
+        top_k=40,                   # Focus on likely tokens
+        top_p=0.9,                  # Nucleus sampling
         pad_token_id=tokenizer.eos_token_id,
-        temperature=0.7,
-        top_k=20,
         early_stopping=True
     )
     
     # Extract only the new text
     full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return full_text.split("Answer:")[-1].strip()
-
-# Option 2: Better quality CPU model (balance between speed & quality)
-def generate_answer_quality(context, question):
-    """
-    Uses DialoGPT-small - conversational but still CPU-friendly
-    ~5-10 seconds on CPU for 80 tokens
-    """
-    qa_pipeline = pipeline(
-        "text-generation",
-        model="microsoft/DialoGPT-small",
-        tokenizer="microsoft/DialoGPT-small"
-    )
-    
-    prompt = f"""Context: {context[:1200]}
-Question: {question}
-Assistant:"""
-    
-    response = qa_pipeline(
-        prompt,
-        max_new_tokens=80,
-        num_beams=1,
-        temperature=0.8,
-        top_k=30,
-        do_sample=True,
-        pad_token_id=50256,  # DialoGPT pad token
-        no_repeat_ngram_size=2
-    )
-    
-    return response[0]['generated_text'].split("Assistant:")[-1].strip()
-
-# Choose one based on priority
-def generate_answer(context, question):
-    start_time = time.time()
-    result = generate_answer_fast(context, question)  # For fastest response
-    # result = generate_answer_quality(context, question)  # For better conversation
-    print(f"Generation took: {time.time() - start_time:.2f}s")
-    return result
+    return full_text.split("Answer:")[-1].strip()