import torch
from transformers import GPT2Tokenizer
from evo_model import EvoDecoderModel

# Load tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

# Model configuration (must match evo_model.py and trained weights)
vocab_size = tokenizer.vocab_size
model = EvoDecoderModel(vocab_size=vocab_size, d_model=256, nhead=4, num_layers=3)
model.load_state_dict(torch.load("evo_decoder.pt", map_location=torch.device("cpu")))
model.eval()

def generate_response(prompt, max_length=100):
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    generated = input_ids.clone()

    with torch.no_grad():
        for _ in range(max_length):
            output = model(generated, memory=None)
            next_token_logits = output[:, -1, :]
            next_token = torch.argmax(next_token_logits, dim=-1).unsqueeze(0)
            generated = torch.cat((generated, next_token), dim=1)

            if next_token.item() == tokenizer.eos_token_id:
                break

    return tokenizer.decode(generated[0], skip_special_tokens=True)