<jupyter_start><jupyter_text>Train GPT on additionTrain a GPT model on a dedicated addition dataset to see if a Transformer can learn to add.<jupyter_code># set up logging
import logging
logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
)
# make deterministic
from mingpt.utils import set_seed
set_seed(42)
import numpy as np
import torch
import string
import os
from tqdm import tqdm
import torch.nn as nn
from torch.nn import functional as F
from mingpt.md import MemData
from mingpt.marker_dataset import MarkerDataset
from mingpt.math_dataset import MathDataset
%load_ext autoreload
%autoreload 2
# create a dataset 
easy = 'run/numbers__place_value.txt'
medium = 'run/numbers__is_prime.txt'
hard = 'run/numbers__list_prime_factors.txt'
!rm -rf run
!cp -r data run
memory_slots = 7
MD = MemData(memory_slots)
MD.initiate_mem_slot_data(hard)
# create a dataset 
easy_test = 'run/test_numbers__place_value.txt'
medium_test = 'run/test_numbers__is_prime.txt'
hard_test = 'run/test_numbers__list_prime_factors.txt'
easy_train = 'run/train_buffer_numbers__place_value.txt'
medium_train = 'run/train_buffer_numbers__is_prime.txt'
hard_train = 'run/train_buffer_numbers__list_prime_factors.txt'
train_dataset = MathDataset(fname=hard_train, MD=MD)
test_dataset = MathDataset(fname=hard_test, MD=MD)
MD.block_size
train_dataset[1000]
from mingpt.model import GPT, GPTConfig, GPT1Config

# initialize a baby GPT model
mconf = GPTConfig(MD.vocab_size, MD.block_size, 
                  n_layer=2, n_head=4, n_embd=128)
model = GPT(mconf)
from mingpt.trainer import Trainer, TrainerConfig

# initialize a trainer instance and kick off training
tconf = TrainerConfig(max_epochs=1, batch_size=512, learning_rate=6e-4,
                      lr_decay=True, warmup_tokens=1024, final_tokens=50*len(train_dataset)*(14+1),
                      num_workers=0)
trainer = Trainer(model, train_dataset, test_dataset, tconf)
trainer.train()
trainer.save_checkpoint()
# now let's give the trained model an addition exam
from torch.utils.data.dataloader import DataLoader
from mingpt.examiner import Examiner
examiner = Examiner(MD)
examiner.exam(hard_train, train_dataset, trainer)
# training set: how well did we memorize?
examples = give_exam(test_dataset, batch_size=1, max_batches=-1)
print("Q: %s\nX:%s\nO:%s\n" % (examples[0][0], examples[0][1] , examples[0][2]))
for item in examples:
    print("Question:", item[0])
    print("X:", item[1])
    print("Out:", item[2])


# test set: how well did we generalize?
give_exam(train_dataset, batch_size=1024, max_batches=1)
# well that's amusing... our model learned everything except 55 + 45
import itertools as it
f = ['-1', '-1', '2', '1', '1']

it.takewhile(lambda x: x!='2', f)
f<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1106830122591944, 'fraction_numerical': 0.019614711033274956, 'mean_word_length': 5.24945295404814, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '54465', 'license': 'MIT', 'n_tokens_mistral': 1056, 'n_tokens_neox': 995, 'n_words': 255, 'path': 'archive/play_math.ipynb', 'repo_name': 'emilwallner/minGPT', 'size': '2855'}