File size: 1,813 Bytes
17f3846 |
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 46 47 48 49 50 51 52 |
import torch
from torch.nn import functional as F
from gpt_class import GPTConfig, GPT
# Assuming tiktoken is correctly imported and functions as expected
import tiktoken
# Setup device
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load model
state_dict = torch.load('model_51999.pt', map_location=device)
config = state_dict['config']
model = GPT(config)
model.load_state_dict(state_dict['model'])
model.to(device)
model.eval()
# Set seed for reproducibility
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)
# Get tokenizer
tokenizer = tiktoken.get_encoding("gpt2")
def Generate(model, tokenizer, example, num_return_sequences, max_length):
model.eval()
tokens = tokenizer.encode(example)
tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).repeat(num_return_sequences, 1)
tokens = tokens.to(device)
sample_rng = torch.Generator(device=device)
xgen = tokens
while xgen.size(1) < max_length:
with torch.no_grad():
with torch.autocast(device_type=device):
logits, _ = model(xgen) # Assumes model returns logits and optional loss
logits = logits[:, -1, :] # Get last token logits
probs = F.softmax(logits, dim=-1)
topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)
ix = torch.multinomial(topk_probs, 1, generator=sample_rng)
xcol = torch.gather(topk_indices, -1, ix)
xgen = torch.cat((xgen, xcol), dim=1)
# Generate output for each sequence
for i in range(num_return_sequences):
tokens = xgen[i, :max_length].tolist()
decoded = tokenizer.decode(tokens)
print(f"Sample {i+1}: {decoded}")
# Generate text
Generate(model, tokenizer, example="It is raining outside and", num_return_sequences=4, max_length=64)
|