Update generate.py

generate.py

@@ -1,29 +1,31 @@
 import torch
+import torch.nn.functional as F
+from evo_model import EvoDecoder
 from transformers import GPT2Tokenizer
-from evo_model import EvoDecoderModel
 
 # Load tokenizer
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-tokenizer.pad_token = tokenizer.eos_token
 
-# Model configuration (must match evo_model.py and trained weights)
-vocab_size = tokenizer.vocab_size
-model = EvoDecoderModel(vocab_size=vocab_size, d_model=256, nhead=4, num_layers=3)
-model.load_state_dict(torch.load("evo_decoder.pt", map_location=torch.device("cpu")))
+# Load trained model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = EvoDecoder(vocab_size=tokenizer.vocab_size, d_model=512, nhead=8, num_layers=6).to(device)
+model.load_state_dict(torch.load("evo_decoder.pt", map_location=device))
 model.eval()
 
-def generate_response(prompt, max_length=100):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    generated = input_ids.clone()
+@torch.no_grad()
+def generate_response(prompt, max_length=50, temperature=1.0):
+    model.eval()
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
 
-    with torch.no_grad():
-        for _ in range(max_length):
-            output = model(generated, memory=None)
-            next_token_logits = output[:, -1, :]
-            next_token = torch.argmax(next_token_logits, dim=-1).unsqueeze(0)
-            generated = torch.cat((generated, next_token), dim=1)
+    for _ in range(max_length):
+        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
+        if next_token.item() == tokenizer.eos_token_id:
+            break
 
-    return tokenizer.decode(generated[0], skip_special_tokens=True)
+    output = tokenizer.decode(input_ids.squeeze(), skip_special_tokens=True)
+    return output[len(prompt):].strip()