Update generate.py

generate.py

@@ -1,14 +1,17 @@
-# generate.py — Generates responses from EvoDecoderModel with Top-k sampling
+# generate.py — Generates response using EvoDecoderModel with GPT2 tokenizer and top-k/p sampling
 import torch
-from transformers import AutoTokenizer
+from transformers import GPT2Tokenizer
 from evo_model import EvoDecoderModel
 
+# Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+# Load GPT2 tokenizer (better for decoding tasks)
+tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
+tokenizer.pad_token = tokenizer.eos_token  # Safe fallback
 vocab_size = tokenizer.vocab_size
 
+# Load trained EvoDecoder model
 model = EvoDecoderModel(vocab_size=vocab_size).to(device)
 model.load_state_dict(torch.load("evo_decoder_model.pt", map_location=device))
 model.eval()
@@ -20,19 +23,21 @@ def generate_response(prompt, max_length=100, top_k=40):
     for _ in range(max_length):
         with torch.no_grad():
             logits = model(input_ids)
-            next_token_logits = logits[:, -1, :]
+            next_token_logits = logits[:, -1, :].squeeze(0)
+
+            # Apply repetition penalty
+            for token_id in set(input_ids.view(-1).tolist()):
+                next_token_logits[token_id] *= 0.8
 
             # Top-k sampling
-            top_k_probs, top_k_indices = torch.topk(next_token_logits, top_k)
-            probs = torch.softmax(top_k_probs.squeeze(0), dim=-1)
-            sampled_index = torch.multinomial(probs, 1).item()
-            next_token = top_k_indices[0, sampled_index]
+            top_k_logits, top_k_indices = torch.topk(next_token_logits, k=top_k)
+            probs = torch.softmax(top_k_logits, dim=-1)
+            next_token = top_k_indices[torch.multinomial(probs, num_samples=1)].unsqueeze(0)
 
-        # Reshape next_token to match input_ids shape
-        next_token = next_token.unsqueeze(0).unsqueeze(0)  # Shape: (1, 1)
-        input_ids = torch.cat([input_ids, next_token], dim=1)
+        input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)
 
-        if next_token.item() == tokenizer.eos_token_id:
+        # Stop on EOS
+        if tokenizer.eos_token_id and next_token.item() == tokenizer.eos_token_id:
             break
 
     output = tokenizer.decode(input_ids[0], skip_special_tokens=True)