Spaces:
Sleeping
Sleeping
File size: 3,658 Bytes
1ca3221 |
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 95 96 97 98 99 100 101 102 103 104 |
"""
Sampling script for the nano-coder model.
This script loads a trained nano-coder model and generates Python code completions.
"""
import os
import pickle
import torch
import torch.nn.functional as F
from model import GPTConfig, GPT
# Configuration
out_dir = 'out-nano-coder'
start = "def fibonacci(n):\n " # or start with any Python code
num_samples = 5 # number of samples to generate
max_new_tokens = 500 # number of tokens generated in each sample
temperature = 0.8 # 1.0 = no change, lower values make output more focused
top_k = 200 # retain only the top_k most likely tokens, clamp others to have 0 probability
seed = 1337
device = 'cuda' if torch.cuda.is_available() else 'cpu'
dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16'
# Load the model
def load_model():
"""Load the trained nano-coder model."""
# Load the checkpoint
ckpt_path = os.path.join(out_dir, 'ckpt.pt')
if not os.path.exists(ckpt_path):
raise FileNotFoundError(f"Checkpoint not found at {ckpt_path}. Please train the model first.")
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)
model.eval()
model.to(device)
return model, checkpoint
def load_vocab():
"""Load the vocabulary from the dataset."""
data_dir = os.path.join('data', 'python-codes-25k')
meta_path = os.path.join(data_dir, 'meta.pkl')
if not os.path.exists(meta_path):
raise FileNotFoundError(f"Vocabulary not found at {meta_path}. Please run prepare_code_dataset.py first.")
with open(meta_path, 'rb') as f:
meta = pickle.load(f)
return meta['stoi'], meta['itos']
def encode(text, stoi):
"""Encode text to token ids."""
return [stoi[c] for c in text]
def decode(ids, itos):
"""Decode token ids to text."""
return ''.join([itos[i] for i in ids])
def generate_code(model, stoi, itos, start_text, max_new_tokens, temperature, top_k):
"""Generate code completion."""
# Encode the start text
start_ids = encode(start_text, stoi)
x = torch.tensor(start_ids, dtype=torch.long, device=device)[None, ...]
# Generate tokens
with torch.no_grad():
with torch.amp.autocast(device_type='cuda' if device == 'cuda' else 'cpu', dtype=torch.bfloat16 if dtype == 'bfloat16' else torch.float16):
y = model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)
completion = decode(y[0].tolist(), itos)
return completion
def main():
print("Loading nano-coder model...")
model, checkpoint = load_model()
stoi, itos = load_vocab()
print(f"Model loaded successfully!")
print(f"Vocabulary size: {len(stoi)}")
print(f"Model parameters: {model.get_num_params()/1e6:.2f}M")
print(f"Context length: {model.config.block_size}")
print(f"Generating {num_samples} samples...")
print(f"Start text: {repr(start)}")
print("-" * 80)
# Set random seed for reproducibility
torch.manual_seed(seed)
# Generate samples
for i in range(num_samples):
print(f"\n--- Sample {i+1} ---")
completion = generate_code(model, stoi, itos, start, max_new_tokens, temperature, top_k)
print(completion)
print("-" * 80)
if __name__ == '__main__':
main() |