Piousbox

The below code snippet uses https://github.com/karpathy/minGPT and trains & evaluates the model on a concrete task: sorting an array of integers.

import torch
from torch.utils.data import Dataset
from torch.utils.data.dataloader import DataLoader
import pickle
from minigpt.utils import set_seed
set_seed(3407)
class SortDataset(Dataset):
    """
    Dataset for the Sort problem. E.g. for problem length 6:
    Input: 0 0 2 1 0 1 -> Output: 0 0 0 1 1 2
    Which will feed into the transformer concatenated as:
    input:  0 0 2 1 0 1 0 0 0 1 1
    output: I I I I I 0 0 0 1 1 2
    where I is "ignore", as the transformer is reading the input sequence
    """
    def __init__(self, split, length=6, num_digits=3):
        assert split in {'train', 'test'}
        self.split = split
        self.length = length
        self.num_digits = num_digits
    def __len__(self):
        return 10000 # ...
    def get_vocab_size(self):
        return self.num_digits
    def get_block_size(self):
        # the length of the sequence that will feed into transformer,
        # containing concatenated input and the output, but -1 because
        # the transformer starts making predictions at the last input element
        return self.length * 2 - 1
    def __getitem__(self, idx):
        # use rejection sampling to generate an input example from the desired split
        while True:
            # generate some random integers
            inp = torch.randint(self.num_digits, size=(self.length,), dtype=torch.long)
            # half of the time let's try to boost the number of examples that
            # have a large number of repeats, as this is what the model seems to struggle
            # with later in training, and they are kind of rare
            if torch.rand(1).item() < 0.5:
                if inp.unique().nelement() > self.length // 2:
                    # too many unqiue digits, re-sample
                    continue
            # figure out if this generated example is train or test based on its hash
            h = hash(pickle.dumps(inp.tolist()))
            inp_split = 'test' if h % 4 == 0 else 'train' # designate 25% of examples as test
            if inp_split == self.split:
                break # ok
        # solve the task: i.e. sort
        sol = torch.sort(inp)[0]
        # concatenate the problem specification and the solution
        cat = torch.cat((inp, sol), dim=0)
        # the inputs to the transformer will be the offset sequence
        x = cat[:-1].clone()
        y = cat[1:].clone()
        # we only want to predict at output locations, mask out the loss at the input locations
        y[:self.length-1] = -1
        return x, y
# print an example instance of the dataset
train_dataset = SortDataset('train', length=6, num_digits=10)
test_dataset = SortDataset('test', length=6, num_digits=10)
x, y = train_dataset[0]
for a, b in zip(x,y):
    print(int(a),int(b))
# create a GPT instance
from minigpt.model import GPT
model_path   = "out/model-1.pickle"
model_config = GPT.get_default_config()
model_config.model_type = 'gpt-nano'
model_config.vocab_size = train_dataset.get_vocab_size()
model_config.block_size = train_dataset.get_block_size()
model = GPT(model_config)
# create a Trainer object
from minigpt.trainer import Trainer
train_config = Trainer.get_default_config()
train_config.learning_rate = 5e-4 # the model we're using is so small that we can go a bit faster
train_config.max_iters = 2000
train_config.num_workers = 0
trainer = Trainer(train_config, model, train_dataset)
def batch_end_callback(trainer):
    if trainer.iter_num % 100 == 0:
        print(f"iter_dt {trainer.iter_dt * 1000:.2f}ms; iter {trainer.iter_num}: train loss {trainer.loss.item():.5f}")
trainer.set_callback('on_batch_end', batch_end_callback)
# trainer.run()
# torch.save(model.state_dict(), model_path)
model.load_state_dict(torch.load(model_path))
model.eval()
def eval_split(trainer, split, max_batches):
    dataset = {'train':train_dataset, 'test':test_dataset}[split]
    n = train_dataset.length # naugy direct access shrug
    results = []
    mistakes_printed_already = 0
    loader = DataLoader(dataset, batch_size=100, num_workers=0, drop_last=False)
    for b, (x, y) in enumerate(loader):
        x = x.to(trainer.device)
        y = y.to(trainer.device)
        # isolate the input pattern alone
        inp = x[:, :n]
        sol = y[:, -n:]
        # let the model sample the rest of the sequence
        cat = model.generate(inp, n, do_sample=False) # using greedy argmax, not sampling
        sol_candidate = cat[:, n:] # isolate the filled in sequence
        # compare the predicted sequence to the true sequence
        correct = (sol == sol_candidate).all(1).cpu() # Software 1.0 vs. Software 2.0 fight RIGHT on this line haha
        for i in range(x.size(0)):
            results.append(int(correct[i]))
            if not correct[i] and mistakes_printed_already < 3: # only print up to 5 mistakes to get a sense
                mistakes_printed_already += 1
                print("GPT claims that %s sorted is %s but gt is %s" % (inp[i].tolist(), sol_candidate[i].tolist(), sol[i].tolist()))
        if max_batches is not None and b+1 >= max_batches:
            break
    rt = torch.tensor(results, dtype=torch.float)
    print("%s final score: %d/%d = %.2f%% correct" % (split, rt.sum(), len(results), 100*rt.mean()))
    return rt.sum()
## run a lot of examples from both train and test through the model and verify the output correctness
with torch.no_grad():
    train_score = eval_split(trainer, 'train', max_batches=50)
    test_score  = eval_split(trainer, 'test',  max_batches=50)
## let's run a random given sequence through the model as well
n = train_dataset.length # naugy direct access shrug
inp = torch.tensor([[5, 4, 3, 2, 1, 0]], dtype=torch.long).to(trainer.device)
assert inp[0].nelement() == n
with torch.no_grad():
    cat = model.generate(inp, n, do_sample=False)
sol = torch.sort(inp[0])[0]
sol_candidate = cat[:, n:]
print('input sequence  :', inp.tolist())
print('predicted sorted:', sol_candidate.tolist())
print('gt sort         :', sol.tolist())
print('matches         :', bool((sol == sol_candidate).all()))

Outcome:

(zenvmac) piousbox@mac:~/projects/python/simple_ai_1/202312_sort_gpt$ python main.py
9 -1
8 -1
1 -1
9 -1
1 -1
0 0
0 1
1 1
1 8
8 9
9 9
number of parameters: 0.09M
running on device cpu
GPT claims that [8, 8, 9, 9, 8, 9] sorted is [8, 8, 8, 8, 9, 9] but gt is [8, 8, 8, 9, 9, 9]
GPT claims that [7, 8, 8, 8, 8, 6] sorted is [7, 7, 8, 8, 8, 8] but gt is [6, 7, 8, 8, 8, 8]
train final score: 4998/5000 = 99.96% correct
GPT claims that [7, 8, 8, 8, 6, 7] sorted is [7, 7, 7, 8, 8, 8] but gt is [6, 7, 7, 8, 8, 8]
GPT claims that [7, 9, 7, 7, 6, 7] sorted is [7, 7, 7, 7, 7, 9] but gt is [6, 7, 7, 7, 7, 9]
GPT claims that [6, 8, 8, 8, 8, 8] sorted is [8, 8, 8, 8, 8, 8] but gt is [6, 8, 8, 8, 8, 8]
test final score: 4997/5000 = 99.94% correct
input sequence  : [[5, 4, 3, 2, 1, 0]]
predicted sorted: [[0, 1, 2, 3, 4, 5]]
gt sort         : [0, 1, 2, 3, 4, 5]
matches         : True
(zenvmac) piousbox@mac:~/projects/python/simple_ai_1/202312_sort_gpt$