File size: 3,713 Bytes
e97f4e2 |
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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import numpy as np
import torch
import torch.nn.functional as F
import tiktoken
from dataclasses import dataclass
from model import GPT
class GPT2Inference:
""" To generate text sequences using a trained GPT2 model """
def __init__(self, model, token_encoder, device):
self.model = model
self.token_encoder = token_encoder
self.device = device
self.device_type = 'cuda' if device.startswith('cuda') else 'cpu'
def generate_sequences(self, prompt, num_seq=5, max_tokens=50):
self.model.eval()
tokens = self.token_encoder.encode(prompt)
tokens = torch.tensor(tokens, dtype=torch.long) # (n,) n : current sequence length
tokens = tokens.unsqueeze(0).repeat(num_seq, 1) # (1,n) --> (num_seq, n)
gen_tokens = tokens.to(self.device)
# create a different rng generator so as not to impact the global rng state used for training
sample_rng = torch.Generator(device=self.device).manual_seed(42)
# generate new tokens one token at a time until the sequence length becomes 'max_tokens'
while gen_tokens.shape[-1] <= max_tokens:
with torch.no_grad():
with torch.autocast(device_type=self.device_type, dtype=torch.bfloat16):
logits, loss = self.model(gen_tokens) # (num_seq, n, vocab_size)
logits = logits[:, -1, :] # (num_seq, vocab_size)
probs = F.softmax(logits, dim=-1) # (num_seq, vocab_size)
# take top-k 50 probs
topk_probs, topk_indices = torch.topk(probs, 50, dim=-1) # (num_seq, 50), (num_seq, 50)
# sample a token from top-50 probabilities
ix = torch.multinomial(topk_probs, num_samples=1, generator=sample_rng) # (num_seq, 1)
next_tok = torch.gather(topk_indices, -1, ix) # (num_seq, 1)
gen_tokens = torch.cat([gen_tokens, next_tok], dim=1)
# decode generated tokens and print generated text
for i in range(num_seq):
tokens = gen_tokens[i, :max_tokens].tolist()
gen_text = self.token_encoder.decode(tokens)
print(f"> sample {i}: {gen_text}")
def parse_args():
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--prompt', type=str, default="Hello, I am a language model,")
parser.add_argument('--num_seq', type=int, default=5)
parser.add_argument('--max_tokens', type=int, default=50)
args = parser.parse_args()
return args
@dataclass
class GPTConfig:
context_length: int = 1024 # max context / sequence length
vocab_size: int = 50257 # number of tokens: 50000 BPE merges + 256 bytes tokens + 1 <endoftext> token
num_layers: int = 12
embd_size: int = 768 # embedding dim
num_heads: int = 12
def inference(args=None):
if args is None:
args = parse_args()
device = 'cpu'
if torch.cuda.is_available():
device = 'cuda'
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
device = 'mps' # for apple macbook GPUs
print(f'using device: {device}')
model_path = './logs/model_95364.pt'
checkpoint = torch.load(model_path, weights_only=False)
print(f"loaded model from: {model_path}")
# print(checkpoint['model'].keys())
model = GPT(config=checkpoint['config'])
model.load_state_dict(checkpoint['model'])
model = model.to(device)
token_encoder = tiktoken.get_encoding('gpt2')
generator = GPT2Inference(model, token_encoder, device)
generator.generate_sequences(args.prompt, args.num_seq, args.max_tokens)
if __name__ == '__main__':
inference()
|