Update generate.py

generate.py

@@ -1,43 +1,40 @@
-# generate.py — Generates response using EvoDecoderModel with GPT2 tokenizer and top-k/p sampling
+# generate.py — Generates EvoDecoder responses with optional live web context
 import torch
-from transformers import GPT2Tokenizer
+from transformers import AutoTokenizer
 from evo_model import EvoDecoderModel
+from search_utils import web_search
 
-# Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load GPT2 tokenizer (better for decoding tasks)
-tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
-tokenizer.pad_token = tokenizer.eos_token  # Safe fallback
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 vocab_size = tokenizer.vocab_size
 
-# Load trained EvoDecoder model
 model = EvoDecoderModel(vocab_size=vocab_size).to(device)
 model.load_state_dict(torch.load("evo_decoder_model.pt", map_location=device))
 model.eval()
 
-def generate_response(prompt, max_length=100, top_k=40):
-    input_text = f"User: {prompt}\nAssistant:"
-    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
+def generate_response(prompt, use_web=False, max_length=100, top_k=40):
+    if use_web:
+        context = web_search(prompt)
+        prompt = f"Relevant Info: {context}\nUser: {prompt}\nAssistant:"
+    else:
+        prompt = f"User: {prompt}\nAssistant:"
+
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
 
     for _ in range(max_length):
         with torch.no_grad():
             logits = model(input_ids)
-            next_token_logits = logits[:, -1, :].squeeze(0)
-
-            # Apply repetition penalty
-            for token_id in set(input_ids.view(-1).tolist()):
-                next_token_logits[token_id] *= 0.8
-
-            # Top-k sampling
-            top_k_logits, top_k_indices = torch.topk(next_token_logits, k=top_k)
-            probs = torch.softmax(top_k_logits, dim=-1)
-            next_token = top_k_indices[torch.multinomial(probs, num_samples=1)].unsqueeze(0)
+            next_token_logits = logits[:, -1, :]
+            top_k_probs, top_k_indices = torch.topk(next_token_logits, top_k)
+            probs = torch.softmax(top_k_probs, dim=-1)
+            next_token = top_k_indices[0, torch.multinomial(probs, 1)]
 
-        input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)
+        next_token = next_token.unsqueeze(0).unsqueeze(0)  # (1,1)
+        input_ids = torch.cat([input_ids, next_token], dim=1)
 
-        # Stop on EOS
-        if tokenizer.eos_token_id and next_token.item() == tokenizer.eos_token_id:
+        if next_token.item() == tokenizer.eos_token_id:
             break
 
     output = tokenizer.decode(input_ids[0], skip_special_tokens=True)