Update generate.py

generate.py

@@ -3,18 +3,22 @@ import torch.nn.functional as F
 from evo_model import EvoDecoder
 from transformers import GPT2Tokenizer
 
-# Load tokenizer
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 
-# 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 = EvoDecoder(
+    vocab_size=tokenizer.vocab_size,
+    d_model=256,
+    nhead=4,
+    num_layers=3,
+    dim_feedforward=1024
+).to(device)
+
 model.load_state_dict(torch.load("evo_decoder.pt", map_location=device))
 model.eval()
 
 @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)
 
     for _ in range(max_length):