File size: 1,734 Bytes
2b45a2a 1c98e83 2b45a2a 1c98e83 2b45a2a 1c98e83 2b45a2a 1c98e83 2b45a2a 1c98e83 6648cc8 1c98e83 2b45a2a 1c98e83 2b45a2a 1c98e83 2b45a2a 1c98e83 2b45a2a 1c98e83 |
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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
# generate.py — Generates response using EvoDecoderModel with GPT2 tokenizer and top-k/p sampling
import torch
from transformers import GPT2Tokenizer
from evo_model import EvoDecoderModel
# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 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()
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, :].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_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)
input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)
# 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)
return output.split("Assistant:")[-1].strip()
|