Update generate.py

generate.py

@@ -2,10 +2,10 @@ import torch
 import torch.nn.functional as F
 from transformers import GPT2Tokenizer
 from evo_decoder import EvoDecoder
-from search_utils import web_search  # Optional RAG fallback
+from search_utils import web_search
 
+# 🔧 Load model and tokenizer
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 tokenizer.pad_token = tokenizer.eos_token
 
@@ -21,19 +21,20 @@ model.load_state_dict(torch.load("evo_decoder.pt", map_location=device))
 model.eval()
 
 @torch.no_grad()
-def generate_response(question, context="", use_rag=False, temperature=1.0, max_length=100):
-    if use_rag and not context:
+def generate_response(question, context="", use_rag=False, temperature=1.0):
+    if not context and use_rag:
         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):
+    for _ in range(128):
         logits = model(input_ids)
         logits = logits[:, -1, :] / temperature
         probs = F.softmax(logits, dim=-1)
         next_token = torch.multinomial(probs, num_samples=1)
         input_ids = torch.cat((input_ids, next_token), dim=1)
+
         if next_token.item() == tokenizer.eos_token_id:
             break