Update generate.py

generate.py

@@ -0,0 +1,37 @@
+# generate.py — Loads EvoDecoder model and generates responses
+import torch
+from transformers import AutoTokenizer
+from evo_model import EvoDecoderModel
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+vocab_size = tokenizer.vocab_size
+
+model = EvoDecoderModel(vocab_size=vocab_size).to(device)
+model.load_state_dict(torch.load("evo_decoder_model.pt", map_location=device))
+model.eval()
+
+def generate_response(prompt, max_length=100, top_k=40):
+    input_text = f"User: {prompt}\nAssistant:"
+    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
+
+    for _ in range(max_length):
+        with torch.no_grad():
+            logits = model(input_ids)
+            next_token_logits = logits[:, -1, :]
+
+            # Top-k sampling
+            top_k_probs, top_k_indices = torch.topk(next_token_logits, top_k)
+            probs = torch.softmax(top_k_probs, dim=-1)
+            next_token = top_k_indices[torch.multinomial(probs, 1)]
+
+        input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)
+
+        # Optional EOS stop condition
+        if next_token.item() == tokenizer.eos_token_id:
+            break
+
+    output = tokenizer.decode(input_ids[0], skip_special_tokens=True)
+    return output.split("Assistant:")[-1].strip()