Update generate.py

generate.py

@@ -1,28 +1,29 @@
 import torch
-import torch.nn.functional as F
+from transformers import GPT2Tokenizer
 from evo_model import EvoDecoderModel
 
-def load_model(vocab_size, model_path="evo_decoder.pt", device="cpu"):
-    model = EvoDecoderModel(vocab_size)
-    model.load_state_dict(torch.load(model_path, map_location=device))
-    model.to(device)
-    model.eval()
-    return model
+# Load tokenizer
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+tokenizer.pad_token = tokenizer.eos_token
 
-def generate_text(model, tokenizer, prompt, max_new_tokens=50, temperature=1.0, device="cpu"):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
+# 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")))
+model.eval()
 
-    for _ in range(max_new_tokens):
-        with torch.no_grad():
-            logits = model(input_ids)
-        next_token_logits = logits[:, -1, :] / temperature
-        probs = F.softmax(next_token_logits, dim=-1)
-        next_token_id = torch.multinomial(probs, num_samples=1)
+def generate_response(prompt, max_length=100):
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    generated = input_ids.clone()
 
-        input_ids = torch.cat([input_ids, next_token_id], dim=-1)
+    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)
 
-        if next_token_id.item() == tokenizer.eos_token_id:
-            break
+            if next_token.item() == tokenizer.eos_token_id:
+                break
 
-    output_text = tokenizer.decode(input_ids[0], skip_special_tokens=True)
-    return output_text
+    return tokenizer.decode(generated[0], skip_special_tokens=True)