Update generate.py

generate.py

@@ -1,47 +1,30 @@
-# generate.py — EvoDecoder response generation with optional DuckDuckGo RAG
+# generate.py
 import torch
-from transformers import AutoTokenizer
+from transformers import BertTokenizer
 from evo_model import EvoDecoderModel
-from search_utils import web_search
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-vocab_size = tokenizer.vocab_size
+# Load tokenizer
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
 
-model = EvoDecoderModel(vocab_size=vocab_size).to(device)
+# Initialize model architecture
+vocab_size = tokenizer.vocab_size
+model = EvoDecoderModel(vocab_size=vocab_size)
 model.load_state_dict(torch.load("evo_decoder_model.pt", map_location=device))
+model.to(device)
 model.eval()
 
-def generate_response(prompt, use_web=False, max_length=100, top_k=40):
-    # Augment with web context if enabled
-    context = ""
-    if use_web:
-        web_context = web_search(prompt)
-        context += f"Relevant Info: {web_context}\n"
-
-    input_text = context + f"User: {prompt}\nAssistant:"
-    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
-
-    # Truncate to avoid exceeding model's positional encoding limit
-    if input_ids.size(1) > 512:
-        input_ids = input_ids[:, -512:]
-
-    for _ in range(max_length):
-        with torch.no_grad():
-            logits = model(input_ids)
-            next_token_logits = logits[:, -1, :]
-
-            # Top-k sampling
-            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).item()].unsqueeze(0).unsqueeze(0)
+def generate_response(prompt, max_length=128, use_web=False):
+    with torch.no_grad():
+        input_ids = tokenizer(prompt, return_tensors="pt", padding=False, truncation=True).input_ids.to(device)
+        input_ids = input_ids[:, :128]  # ✅ clip to trained length
 
-        input_ids = torch.cat([input_ids, next_token], dim=1)
+        logits = model(input_ids)
+        next_token_logits = logits[:, -1, :]  # take last token's logits
+        predicted_id = torch.argmax(next_token_logits, dim=-1)
 
-        if next_token.item() == tokenizer.eos_token_id:
-            break
+        output_ids = torch.cat([input_ids, predicted_id.unsqueeze(0)], dim=1)
+        decoded = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
-    output = tokenizer.decode(input_ids[0], skip_special_tokens=True)
-    return output.split("Assistant:")[-1].strip()
+        return decoded