generate.py

# 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()