Restructure: nanoGPT at root, custom code in src/

- Move model.py, train.py, configurator.py to root for nanoGPT compatibility
- data/ and config/ directories at root with Shakespeare dataset prep scripts
- src/jlens.py updated to import model from project root
- Cleaned up stale src/config/ and duplicate src/ files
- Fixed .gitignore: exclude out-shakespeare-char/ instead of raw data dirs

25 files changed, 265 insertions, 123 deletions

diff --git a/.gitignore b/.gitignore
index 91b3024..dc2eed8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,8 +5,7 @@ __pycache__/
 *.bin
 *.safetensors
 outputs/
-data/shakespeare/
-data/shakespeare_char/
+out-shakespeare-char/
 *.egg-info/
 .venv/
 .ipynb_checkpoints/
diff --git a/src/config/eval_gpt2.py b/config/eval_gpt2.py
index 53978cb..53978cb 100644
--- a/src/config/eval_gpt2.py
+++ b/config/eval_gpt2.py
diff --git a/src/config/eval_gpt2_large.py b/config/eval_gpt2_large.py
index 4cbeaef..4cbeaef 100644
--- a/src/config/eval_gpt2_large.py
+++ b/config/eval_gpt2_large.py
diff --git a/src/config/eval_gpt2_medium.py b/config/eval_gpt2_medium.py
index 9d0db11..9d0db11 100644
--- a/src/config/eval_gpt2_medium.py
+++ b/config/eval_gpt2_medium.py
diff --git a/src/config/eval_gpt2_xl.py b/config/eval_gpt2_xl.py
index 1bae34f..1bae34f 100644
--- a/src/config/eval_gpt2_xl.py
+++ b/config/eval_gpt2_xl.py
diff --git a/src/config/finetune_shakespeare.py b/config/finetune_shakespeare.py
index 148a4c4..148a4c4 100644
--- a/src/config/finetune_shakespeare.py
+++ b/config/finetune_shakespeare.py
diff --git a/src/config/train_gpt2.py b/config/train_gpt2.py
index 8f19273..8f19273 100644
--- a/src/config/train_gpt2.py
+++ b/config/train_gpt2.py
diff --git a/src/config/train_shakespeare_char.py b/config/train_shakespeare_char.py
index 41c81df..41c81df 100644
--- a/src/config/train_shakespeare_char.py
+++ b/config/train_shakespeare_char.py
diff --git a/configurator.py b/configurator.py
new file mode 100644
index 0000000..a8bba95
--- /dev/null
+++ b/configurator.py
@@ -0,0 +1,47 @@
+"""
+Poor Man's Configurator. Probably a terrible idea. Example usage:
+$ python train.py config/override_file.py --batch_size=32
+this will first run config/override_file.py, then override batch_size to 32
+
+The code in this file will be run as follows from e.g. train.py:
+>>> exec(open('configurator.py').read())
+
+So it's not a Python module, it's just shuttling this code away from train.py
+The code in this script then overrides the globals()
+
+I know people are not going to love this, I just really dislike configuration
+complexity and having to prepend config. to every single variable. If someone
+comes up with a better simple Python solution I am all ears.
+"""
+
+import sys
+from ast import literal_eval
+
+for arg in sys.argv[1:]:
+    if '=' not in arg:
+        # assume it's the name of a config file
+        assert not arg.startswith('--')
+        config_file = arg
+        print(f"Overriding config with {config_file}:")
+        with open(config_file) as f:
+            print(f.read())
+        exec(open(config_file).read())
+    else:
+        # assume it's a --key=value argument
+        assert arg.startswith('--')
+        key, val = arg.split('=')
+        key = key[2:]
+        if key in globals():
+            try:
+                # attempt to eval it it (e.g. if bool, number, or etc)
+                attempt = literal_eval(val)
+            except (SyntaxError, ValueError):
+                # if that goes wrong, just use the string
+                attempt = val
+            # ensure the types match ok
+            assert type(attempt) == type(globals()[key])
+            # cross fingers
+            print(f"Overriding: {key} = {attempt}")
+            globals()[key] = attempt
+        else:
+            raise ValueError(f"Unknown config key: {key}")
diff --git a/data/openwebtext/prepare.py b/data/openwebtext/prepare.py
new file mode 100644
index 0000000..2a9b975
--- /dev/null
+++ b/data/openwebtext/prepare.py
@@ -0,0 +1,81 @@
+# saves the openwebtext dataset to a binary file for training. following was helpful:
+# https://github.com/HazyResearch/flash-attention/blob/main/training/src/datamodules/language_modeling_hf.py
+
+import os
+from tqdm import tqdm
+import numpy as np
+import tiktoken
+from datasets import load_dataset # huggingface datasets
+
+# number of workers in .map() call
+# good number to use is ~order number of cpu cores // 2
+num_proc = 8
+
+# number of workers in load_dataset() call
+# best number might be different from num_proc above as it also depends on NW speed.
+# it is better than 1 usually though
+num_proc_load_dataset = num_proc
+
+enc = tiktoken.get_encoding("gpt2")
+
+if __name__ == '__main__':
+    # takes 54GB in huggingface .cache dir, about 8M documents (8,013,769)
+    dataset = load_dataset("openwebtext", num_proc=num_proc_load_dataset)
+
+    # owt by default only contains the 'train' split, so create a test split
+    split_dataset = dataset["train"].train_test_split(test_size=0.0005, seed=2357, shuffle=True)
+    split_dataset['val'] = split_dataset.pop('test') # rename the test split to val
+
+    # this results in:
+    # >>> split_dataset
+    # DatasetDict({
+    #     train: Dataset({
+    #         features: ['text'],
+    #         num_rows: 8009762
+    #     })
+    #     val: Dataset({
+    #         features: ['text'],
+    #         num_rows: 4007
+    #     })
+    # })
+
+    # we now want to tokenize the dataset. first define the encoding function (gpt2 bpe)
+    def process(example):
+        ids = enc.encode_ordinary(example['text']) # encode_ordinary ignores any special tokens
+        ids.append(enc.eot_token) # add the end of text token, e.g. 50256 for gpt2 bpe
+        # note: I think eot should be prepended not appended... hmm. it's called "eot" though...
+        out = {'ids': ids, 'len': len(ids)}
+        return out
+
+    # tokenize the dataset
+    tokenized = split_dataset.map(
+        process,
+        remove_columns=['text'],
+        desc="tokenizing the splits",
+        num_proc=num_proc,
+    )
+
+    # concatenate all the ids in each dataset into one large file we can use for training
+    for split, dset in tokenized.items():
+        arr_len = np.sum(dset['len'], dtype=np.uint64)
+        filename = os.path.join(os.path.dirname(__file__), f'{split}.bin')
+        dtype = np.uint16 # (can do since enc.max_token_value == 50256 is < 2**16)
+        arr = np.memmap(filename, dtype=dtype, mode='w+', shape=(arr_len,))
+        total_batches = 1024
+
+        idx = 0
+        for batch_idx in tqdm(range(total_batches), desc=f'writing {filename}'):
+            # Batch together samples for faster write
+            batch = dset.shard(num_shards=total_batches, index=batch_idx, contiguous=True).with_format('numpy')
+            arr_batch = np.concatenate(batch['ids'])
+            # Write into mmap
+            arr[idx : idx + len(arr_batch)] = arr_batch
+            idx += len(arr_batch)
+        arr.flush()
+
+    # train.bin is ~17GB, val.bin ~8.5MB
+    # train has ~9B tokens (9,035,582,198)
+    # val has ~4M tokens (4,434,897)
+
+    # to read the bin files later, e.g. with numpy:
+    # m = np.memmap('train.bin', dtype=np.uint16, mode='r')
diff --git a/data/openwebtext/readme.md b/data/openwebtext/readme.md
new file mode 100644
index 0000000..95eb1bf
--- /dev/null
+++ b/data/openwebtext/readme.md
@@ -0,0 +1,15 @@
+
+## openwebtext dataset
+
+after running `prepare.py` (preprocess) we get:
+
+- train.bin is ~17GB, val.bin ~8.5MB
+- train has ~9B tokens (9,035,582,198)
+- val has ~4M tokens (4,434,897)
+
+this came from 8,013,769 documents in total.
+
+references:
+
+- OpenAI's WebText dataset is discussed in [GPT-2 paper](https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf)
+- [OpenWebText](https://skylion007.github.io/OpenWebTextCorpus/) dataset
diff --git a/data/shakespeare/prepare.py b/data/shakespeare/prepare.py
new file mode 100644
index 0000000..bda25b1
--- /dev/null
+++ b/data/shakespeare/prepare.py
@@ -0,0 +1,33 @@
+import os
+import requests
+import tiktoken
+import numpy as np
+
+# download the tiny shakespeare dataset
+input_file_path = os.path.join(os.path.dirname(__file__), 'input.txt')
+if not os.path.exists(input_file_path):
+    data_url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
+    with open(input_file_path, 'w', encoding='utf-8') as f:
+        f.write(requests.get(data_url).text)
+
+with open(input_file_path, 'r', encoding='utf-8') as f:
+    data = f.read()
+n = len(data)
+train_data = data[:int(n*0.9)]
+val_data = data[int(n*0.9):]
+
+# encode with tiktoken gpt2 bpe
+enc = tiktoken.get_encoding("gpt2")
+train_ids = enc.encode_ordinary(train_data)
+val_ids = enc.encode_ordinary(val_data)
+print(f"train has {len(train_ids):,} tokens")
+print(f"val has {len(val_ids):,} tokens")
+
+# export to bin files
+train_ids = np.array(train_ids, dtype=np.uint16)
+val_ids = np.array(val_ids, dtype=np.uint16)
+train_ids.tofile(os.path.join(os.path.dirname(__file__), 'train.bin'))
+val_ids.tofile(os.path.join(os.path.dirname(__file__), 'val.bin'))
+
+# train.bin has 301,966 tokens
+# val.bin has 36,059 tokens
diff --git a/data/shakespeare/readme.md b/data/shakespeare/readme.md
new file mode 100644
index 0000000..1e6c457
--- /dev/null
+++ b/data/shakespeare/readme.md
@@ -0,0 +1,9 @@
+
+# tiny shakespeare
+
+Tiny shakespeare, of the good old char-rnn fame :)
+
+After running `prepare.py`:
+
+- train.bin has 301,966 tokens
+- val.bin has 36,059 tokens
diff --git a/data/shakespeare_char/prepare.py b/data/shakespeare_char/prepare.py
new file mode 100644
index 0000000..9fd1621
--- /dev/null
+++ b/data/shakespeare_char/prepare.py
@@ -0,0 +1,68 @@
+"""
+Prepare the Shakespeare dataset for character-level language modeling.
+So instead of encoding with GPT-2 BPE tokens, we just map characters to ints.
+Will save train.bin, val.bin containing the ids, and meta.pkl containing the
+encoder and decoder and some other related info.
+"""
+import os
+import pickle
+import requests
+import numpy as np
+
+# download the tiny shakespeare dataset
+input_file_path = os.path.join(os.path.dirname(__file__), 'input.txt')
+if not os.path.exists(input_file_path):
+    data_url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
+    with open(input_file_path, 'w') as f:
+        f.write(requests.get(data_url).text)
+
+with open(input_file_path, 'r') as f:
+    data = f.read()
+print(f"length of dataset in characters: {len(data):,}")
+
+# get all the unique characters that occur in this text
+chars = sorted(list(set(data)))
+vocab_size = len(chars)
+print("all the unique characters:", ''.join(chars))
+print(f"vocab size: {vocab_size:,}")
+
+# create a mapping from characters to integers
+stoi = { ch:i for i,ch in enumerate(chars) }
+itos = { i:ch for i,ch in enumerate(chars) }
+def encode(s):
+    return [stoi[c] for c in s] # encoder: take a string, output a list of integers
+def decode(l):
+    return ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string
+
+# create the train and test splits
+n = len(data)
+train_data = data[:int(n*0.9)]
+val_data = data[int(n*0.9):]
+
+# encode both to integers
+train_ids = encode(train_data)
+val_ids = encode(val_data)
+print(f"train has {len(train_ids):,} tokens")
+print(f"val has {len(val_ids):,} tokens")
+
+# export to bin files
+train_ids = np.array(train_ids, dtype=np.uint16)
+val_ids = np.array(val_ids, dtype=np.uint16)
+train_ids.tofile(os.path.join(os.path.dirname(__file__), 'train.bin'))
+val_ids.tofile(os.path.join(os.path.dirname(__file__), 'val.bin'))
+
+# save the meta information as well, to help us encode/decode later
+meta = {
+    'vocab_size': vocab_size,
+    'itos': itos,
+    'stoi': stoi,
+}
+with open(os.path.join(os.path.dirname(__file__), 'meta.pkl'), 'wb') as f:
+    pickle.dump(meta, f)
+
+# length of dataset in characters:  1115394
+# all the unique characters:
+#  !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz
+# vocab size: 65
+# train has 1003854 tokens
+# val has 111540 tokens
diff --git a/data/shakespeare_char/readme.md b/data/shakespeare_char/readme.md
new file mode 100644
index 0000000..d597b79
--- /dev/null
+++ b/data/shakespeare_char/readme.md
@@ -0,0 +1,9 @@
+
+# tiny shakespeare, character-level
+
+Tiny shakespeare, of the good old char-rnn fame :) Treated on character-level.
+
+After running `prepare.py`:
+
+- train.bin has 1,003,854 tokens
+- val.bin has 111,540 tokens
diff --git a/src/model.py b/model.py
index c698f8b..c698f8b 100644
--- a/src/model.py
+++ b/model.py
diff --git a/src/eval_gpt2.py b/src/eval_gpt2.py
deleted file mode 100644
index 53978cb..0000000
--- a/src/eval_gpt2.py
+++ /dev/null
@@ -1,8 +0,0 @@
-# evaluate the base gpt2
-# n_layer=12, n_head=12, n_embd=768
-# 124M parameters
-batch_size = 8
-eval_iters = 500 # use more iterations to get good estimate
-eval_only = True
-wandb_log = False
-init_from = 'gpt2'
diff --git a/src/eval_gpt2_large.py b/src/eval_gpt2_large.py
deleted file mode 100644
index 4cbeaef..0000000
--- a/src/eval_gpt2_large.py
+++ /dev/null
@@ -1,8 +0,0 @@
-# evaluate the base gpt2
-# n_layer=36, n_head=20, n_embd=1280
-# 774M parameters
-batch_size = 8
-eval_iters = 500 # use more iterations to get good estimate
-eval_only = True
-wandb_log = False
-init_from = 'gpt2-large'
diff --git a/src/eval_gpt2_medium.py b/src/eval_gpt2_medium.py
deleted file mode 100644
index 9d0db11..0000000
--- a/src/eval_gpt2_medium.py
+++ /dev/null
@@ -1,8 +0,0 @@
-# evaluate the base gpt2
-# n_layer=24, n_head=16, n_embd=1024
-# 350M parameters
-batch_size = 8
-eval_iters = 500 # use more iterations to get good estimate
-eval_only = True
-wandb_log = False
-init_from = 'gpt2-medium'
diff --git a/src/eval_gpt2_xl.py b/src/eval_gpt2_xl.py
deleted file mode 100644
index 1bae34f..0000000
--- a/src/eval_gpt2_xl.py
+++ /dev/null
@@ -1,8 +0,0 @@
-# evaluate the base gpt2
-# n_layer=48, n_head=25, n_embd=1600
-# 1558M parameters
-batch_size = 8
-eval_iters = 500 # use more iterations to get good estimate
-eval_only = True
-wandb_log = False
-init_from = 'gpt2-xl'
diff --git a/src/finetune_shakespeare.py b/src/finetune_shakespeare.py
deleted file mode 100644
index 148a4c4..0000000
--- a/src/finetune_shakespeare.py
+++ /dev/null
@@ -1,25 +0,0 @@
-import time
-
-out_dir = 'out-shakespeare'
-eval_interval = 5
-eval_iters = 40
-wandb_log = False # feel free to turn on
-wandb_project = 'shakespeare'
-wandb_run_name = 'ft-' + str(time.time())
-
-dataset = 'shakespeare'
-init_from = 'gpt2-xl' # this is the largest GPT-2 model
-
-# only save checkpoints if the validation loss improves
-always_save_checkpoint = False
-
-# the number of examples per iter:
-# 1 batch_size * 32 grad_accum * 1024 tokens = 32,768 tokens/iter
-# shakespeare has 301,966 tokens, so 1 epoch ~= 9.2 iters
-batch_size = 1
-gradient_accumulation_steps = 32
-max_iters = 20
-
-# finetune at constant LR
-learning_rate = 3e-5
-decay_lr = False
diff --git a/src/jlens.py b/src/jlens.py
index 515dc08..16ff7ab 100644
--- a/src/jlens.py
+++ b/src/jlens.py
@@ -346,9 +346,9 @@ if __name__ == '__main__':
     
     args = parser.parse_args()
     
-    # Import model from local src
+    # Import model from local dir
     import sys
-    sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
+    sys.path.insert(0, os.path.dirname(__file__))
     from model import GPT, GPTConfig
     
     # Load model
diff --git a/src/train_gpt2.py b/src/train_gpt2.py
deleted file mode 100644
index 8f19273..0000000
--- a/src/train_gpt2.py
+++ /dev/null
@@ -1,25 +0,0 @@
-# config for training GPT-2 (124M) down to very nice loss of ~2.85 on 1 node of 8X A100 40GB
-# launch as the following (e.g. in a screen session) and wait ~5 days:
-# $ torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py
-
-wandb_log = True
-wandb_project = 'owt'
-wandb_run_name='gpt2-124M'
-
-# these make the total batch size be ~0.5M
-# 12 batch size * 1024 block size * 5 gradaccum * 8 GPUs = 491,520
-batch_size = 12
-block_size = 1024
-gradient_accumulation_steps = 5 * 8
-
-# this makes total number of tokens be 300B
-max_iters = 600000
-lr_decay_iters = 600000
-
-# eval stuff
-eval_interval = 1000
-eval_iters = 200
-log_interval = 10
-
-# weight decay
-weight_decay = 1e-1
diff --git a/src/train_shakespeare_char.py b/src/train_shakespeare_char.py
deleted file mode 100644
index 41c81df..0000000
--- a/src/train_shakespeare_char.py
+++ /dev/null
@@ -1,37 +0,0 @@
-# train a miniature character-level shakespeare model
-# good for debugging and playing on macbooks and such
-
-out_dir = 'out-shakespeare-char'
-eval_interval = 250 # keep frequent because we'll overfit
-eval_iters = 200
-log_interval = 10 # don't print too too often
-
-# we expect to overfit on this small dataset, so only save when val improves
-always_save_checkpoint = False
-
-wandb_log = False # override via command line if you like
-wandb_project = 'shakespeare-char'
-wandb_run_name = 'mini-gpt'
-
-dataset = 'shakespeare_char'
-gradient_accumulation_steps = 1
-batch_size = 64
-block_size = 256 # context of up to 256 previous characters
-
-# baby GPT model :)
-n_layer = 6
-n_head = 6
-n_embd = 384
-dropout = 0.2
-
-learning_rate = 1e-3 # with baby networks can afford to go a bit higher
-max_iters = 5000
-lr_decay_iters = 5000 # make equal to max_iters usually
-min_lr = 1e-4 # learning_rate / 10 usually
-beta2 = 0.99 # make a bit bigger because number of tokens per iter is small
-
-warmup_iters = 100 # not super necessary potentially
-
-# on macbook also add
-# device = 'cpu'  # run on cpu only
-# compile = False # do not torch compile the model
diff --git a/src/train.py b/train.py
index de57850..de57850 100644
--- a/src/train.py
+++ b/train.py