From 57735f532df142cf9d1012061b93ab53e58c462b Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sat, 14 Jan 2023 02:26:44 +0000 Subject: [PATCH] correctly propagate the vocab_size from the rendered dataset into the model args --- train.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/train.py b/train.py index b888310..90482d0 100644 --- a/train.py +++ b/train.py @@ -12,6 +12,7 @@ $ torchrun --standalone --nproc_per_node=4 train.py import os import time import math +import pickle from contextlib import nullcontext import numpy as np @@ -102,8 +103,19 @@ def get_batch(split): iter_num = 0 best_val_loss = 1e9 -# model init. TODO: fix bug we should also propagate the correct vocab_size to the model_args -model_args = dict(n_layer = n_layer, n_head = n_head, n_embd = n_embd, block_size = block_size, dropout = dropout) +# attempt to derive vocab_size from the dataset +meta_path = os.path.join(data_dir, 'meta.pkl') +if os.path.exists(meta_path): + with open(meta_path, 'rb') as f: + meta = pickle.load(f) + vocab_size = meta['vocab_size'] + print(f"vocab_size = {vocab_size} (from {meta_path})") +else: + print(f"vocab_size not found in {meta_path}, using GPT-2 default of 50257") + vocab_size = 50257 + +# model init +model_args = dict(n_layer = n_layer, n_head = n_head, n_embd = n_embd, block_size = block_size, dropout = dropout, vocab_size = vocab_size) if init_from == 'scratch': # init a new model from scratch print("Initializing a new model from scratch")