File size: 1,088 Bytes
1c98e83 ccae0a9 7dc4300 1c98e83 ccae0a9 1c98e83 ccae0a9 5259900 ccae0a9 1c98e83 ccae0a9 d7a4aba ccae0a9 7dc4300 ccae0a9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
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)
|