Update generate.py

generate.py

@@ -1,14 +1,14 @@
 import torch
 import torch.nn.functional as F
-from evo_decoder import EvoDecoder
 from transformers import GPT2Tokenizer
+from evo_decoder import EvoDecoder
+from search_utils import web_search  # Optional RAG fallback
+
+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
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 model = EvoDecoder(
     vocab_size=tokenizer.vocab_size,
     d_model=256,
@@ -21,17 +21,12 @@ model.load_state_dict(torch.load("evo_decoder.pt", map_location=device))
 model.eval()
 
 @torch.no_grad()
-def generate_response(prompt, max_length=50, temperature=1.0, external_context=""):
-    """
-    Generate text using EvoDecoder with optional external RAG context.
-    """
-    model.eval()
-    
-    # Combine external context with prompt
-    full_prompt = (external_context.strip() + "\n\n" + prompt.strip()) if external_context else prompt.strip()
-
-    # Truncate if input too long
-    input_ids = tokenizer.encode(full_prompt, return_tensors="pt", truncation=True, max_length=256).to(device)
+def generate_response(question, context="", use_rag=False, temperature=1.0, max_length=100):
+    if use_rag and not context:
+        context = web_search(question)
+
+    prompt = f"Context: {context}\nQuestion: {question}\nAnswer:"
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
 
     for _ in range(max_length):
         logits = model(input_ids)
@@ -39,14 +34,8 @@ def generate_response(prompt, max_length=50, temperature=1.0, external_context="
         probs = F.softmax(logits, dim=-1)
         next_token = torch.multinomial(probs, num_samples=1)
         input_ids = torch.cat((input_ids, next_token), dim=1)
-
-        # Break on EOS
         if next_token.item() == tokenizer.eos_token_id:
             break
 
-        # Prevent overflow
-        if input_ids.shape[1] >= 256:
-            break
-
-    output = tokenizer.decode(input_ids.squeeze(), skip_special_tokens=True)
-    return output[len(full_prompt):].strip()
+    output = tokenizer.decode(input_ids[0], skip_special_tokens=True)
+    return output[len(prompt):].strip()