Update generate.py

generate.py

@@ -1,24 +1,31 @@
 # generate.py
 import torch
-from transformers import BertTokenizer
+from transformers import AutoTokenizer
 from evo_model import EvoDecoderModel
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 vocab_size = tokenizer.vocab_size
 
-model = EvoDecoderModel(vocab_size=vocab_size)
+model = EvoDecoderModel(vocab_size=vocab_size).to(device)
 model.load_state_dict(torch.load("evo_decoder_model.pt", map_location=device))
-model.to(device)
 model.eval()
 
-def generate_response(prompt, max_length=128, use_web=False):
+def generate_response(prompt, use_web=False):
+    # Tokenize
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=128)
+    input_ids = inputs["input_ids"].to(device)
+
+    # Predict
     with torch.no_grad():
-        input_ids = tokenizer(prompt, return_tensors="pt").input_ids[:, :128].to(device)
         logits = model(input_ids)
-        next_token_logits = logits[:, -1, :]
-        next_token_id = torch.argmax(next_token_logits, dim=-1)
 
-        full_ids = torch.cat([input_ids, next_token_id.unsqueeze(0)], dim=1)
-        return tokenizer.decode(full_ids[0], skip_special_tokens=True)
+    # Take last token's logits and get predicted token
+    next_token_logits = logits[0, -1]  # shape: (vocab_size,)
+    predicted_token_id = torch.argmax(next_token_logits).item()
+
+    # Decode to word
+    predicted_token = tokenizer.decode([predicted_token_id])
+    return predicted_token