Update generate.py

generate.py

@@ -1,14 +1,14 @@
 import torch
 import torch.nn.functional as F
-from transformers import GPT2Tokenizer
 from evo_decoder import EvoDecoder
-from search_utils import web_search
+from transformers import GPT2Tokenizer
 
-# 🔧 Load model and tokenizer
+# ✅ Device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# ✅ Load tokenizer and model
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 tokenizer.pad_token = tokenizer.eos_token
-
 model = EvoDecoder(
     vocab_size=tokenizer.vocab_size,
     d_model=256,
@@ -17,18 +17,24 @@ model = EvoDecoder(
     dim_feedforward=512
 ).to(device)
 
+# ✅ Load trained weights
 model.load_state_dict(torch.load("evo_decoder.pt", map_location=device))
 model.eval()
 
+# ✅ Response Generator
 @torch.no_grad()
-def generate_response(question, context="", use_rag=False, temperature=1.0):
-    if not context and use_rag:
-        context = web_search(question)
+def generate_response(prompt, max_length=128, temperature=1.0, external_context=""):
+    model.eval()
+
+    # ✅ Force prompt into SQuAD-style format Evo was trained on
+    if external_context:
+        full_prompt = f"Context: {external_context}\nQuestion: {prompt}\nAnswer:"
+    else:
+        full_prompt = f"Question: {prompt}\nAnswer:"
 
-    prompt = f"Context: {context}\nQuestion: {question}\nAnswer:"
-    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
+    input_ids = tokenizer.encode(full_prompt, return_tensors="pt").to(device)
 
-    for _ in range(128):
+    for _ in range(max_length):
         logits = model(input_ids)
         logits = logits[:, -1, :] / temperature
         probs = F.softmax(logits, dim=-1)
@@ -38,5 +44,5 @@ def generate_response(question, context="", use_rag=False, temperature=1.0):
         if next_token.item() == tokenizer.eos_token_id:
             break
 
-    output = tokenizer.decode(input_ids[0], skip_special_tokens=True)
-    return output[len(prompt):].strip()
+    output = tokenizer.decode(input_ids.squeeze(), skip_special_tokens=True)
+    return output[len(full_prompt):].strip()