mirror of
https://github.com/osmarks/nanogpt-experiments.git
synced 2024-11-10 20:09:58 +00:00
allow sample.py to init from a pretrained gpt2 checkpoints as well, in similar style to train.py
This commit is contained in:
parent
6c40a08b41
commit
21675d7755
31
sample.py
31
sample.py
@ -9,7 +9,8 @@ import tiktoken
|
||||
from model import GPTConfig, GPT
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
out_dir = 'out'
|
||||
init_from = 'resume' # either 'resume' (from an out_dir) or a gpt2 variant (e.g. 'gpt2-xl')
|
||||
out_dir = 'out' # ignored if init_from is not 'resume'
|
||||
start = "\n" # or "<|endoftext|>" or whatever you like
|
||||
num_samples = 10 # number of samples to draw
|
||||
max_new_tokens = 500 # number of tokens generated in each sample
|
||||
@ -31,16 +32,22 @@ ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torc
|
||||
ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
|
||||
|
||||
# model
|
||||
ckpt_path = os.path.join(out_dir, 'ckpt.pt')
|
||||
checkpoint = torch.load(ckpt_path, map_location=device)
|
||||
gptconf = GPTConfig(**checkpoint['model_args'])
|
||||
model = GPT(gptconf)
|
||||
state_dict = checkpoint['model']
|
||||
unwanted_prefix = '_orig_mod.'
|
||||
for k,v in list(state_dict.items()):
|
||||
if k.startswith(unwanted_prefix):
|
||||
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
|
||||
model.load_state_dict(state_dict)
|
||||
if init_from == 'resume':
|
||||
# init from a model saved in a specific directory
|
||||
ckpt_path = os.path.join(out_dir, 'ckpt.pt')
|
||||
checkpoint = torch.load(ckpt_path, map_location=device)
|
||||
gptconf = GPTConfig(**checkpoint['model_args'])
|
||||
model = GPT(gptconf)
|
||||
state_dict = checkpoint['model']
|
||||
unwanted_prefix = '_orig_mod.'
|
||||
for k,v in list(state_dict.items()):
|
||||
if k.startswith(unwanted_prefix):
|
||||
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
|
||||
model.load_state_dict(state_dict)
|
||||
elif init_from.startswith('gpt2'):
|
||||
# init from a given GPT-2 model
|
||||
model = GPT.from_pretrained(init_from, dict(dropout=0.0))
|
||||
|
||||
model.eval()
|
||||
model.to(device)
|
||||
if compile:
|
||||
@ -48,7 +55,7 @@ if compile:
|
||||
|
||||
# look for the meta pickle in case it is available in the dataset folder
|
||||
load_meta = False
|
||||
if 'config' in checkpoint and 'dataset' in checkpoint['config']: # older checkpoints might not have these...
|
||||
if init_from == 'resume' and 'config' in checkpoint and 'dataset' in checkpoint['config']: # older checkpoints might not have these...
|
||||
meta_path = os.path.join('data', checkpoint['config']['dataset'], 'meta.pkl')
|
||||
load_meta = os.path.exists(meta_path)
|
||||
if load_meta:
|
||||
|
Loading…
Reference in New Issue
Block a user