play_char.ipynb

Kernel: Python 3

Train a character-level GPT on some text data

The inputs here are simple text files, which we chop up to individual characters and then train GPT on. So you could say this is a char-transformer instead of a char-rnn. Doesn't quite roll off the tongue as well. In this example we will feed it some Shakespeare, which we'll get it to predict character-level.

# 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 torch.nn as nn
from torch.nn import functional as F

import math
from torch.utils.data import Dataset

class CharDataset(Dataset):

    def __init__(self, data, block_size):
        chars = sorted(list(set(data)))
        data_size, vocab_size = len(data), len(chars)
        print('data has %d characters, %d unique.' % (data_size, vocab_size))
        
        self.stoi = { ch:i for i,ch in enumerate(chars) }
        self.itos = { i:ch for i,ch in enumerate(chars) }
        self.block_size = block_size
        self.vocab_size = vocab_size
        self.data = data
    
    def __len__(self):
        return len(self.data) - self.block_size

    def __getitem__(self, idx):
        # grab a chunk of (block_size + 1) characters from the data
        chunk = self.data[idx:idx + self.block_size + 1]
        # encode every character to an integer
        dix = [self.stoi[s] for s in chunk]
        """
        arrange data and targets so that the first i elements of x
        will be asked to predict the i-th element of y. Notice that
        the eventual language model will actually make block_size
        individual predictions at the same time based on this data,
        so we are being clever and amortizing the cost of the forward
        pass of the network. So for example if block_size is 4, then
        we could e.g. sample a chunk of text "hello", the integers in
        x will correspond to "hell" and in y will be "ello". This will
        then actually "multitask" 4 separate examples at the same time
        in the language model:
        - given just "h", please predict "e" as next
        - given "he" please predict "l" next
        - given "hel" predict "l" next
        - given "hell" predict "o" next
        
        In addition, because the DataLoader will create batches of examples,
        every forward/backward pass during traning will simultaneously train
        a LOT of predictions, amortizing a lot of computation. In particular,
        for a batched input of integers X (B, T) where B is batch size and
        T is block_size and Y (B, T), the network will during training be
        simultaneously training to make B*T predictions, all at once! Of course,
        at test time we can paralellize across batch B, but unlike during training
        we cannot parallelize across the time dimension T - we have to run
        a forward pass of the network to recover the next single character of the 
        sequence along each batch dimension, and repeatedly always feed in a next
        character to get the next one.
        
        So yes there is a big asymmetry between train/test time of autoregressive
        models. During training we can go B*T at a time with every forward pass,
        but during test time we can only go B at a time, T times, with T forward 
        passes.
        """
        x = torch.tensor(dix[:-1], dtype=torch.long)
        y = torch.tensor(dix[1:], dtype=torch.long)
        return x, y

block_size = 128 # spatial extent of the model for its context

# you can download this file at https://github.com/karpathy/char-rnn/blob/master/data/tinyshakespeare/input.txt
text = open('input.txt', 'r').read() # don't worry we won't run out of file handles
train_dataset = CharDataset(text, block_size) # one line of poem is roughly 50 characters

data has 1115394 characters, 65 unique.

from mingpt.model import GPT, GPTConfig
mconf = GPTConfig(train_dataset.vocab_size, train_dataset.block_size,
                  n_layer=8, n_head=8, n_embd=512)
model = GPT(mconf)

08/24/2020 22:59:58 - INFO - mingpt.model -   number of parameters: 2.535219e+07

from mingpt.trainer import Trainer, TrainerConfig

# initialize a trainer instance and kick off training
tconf = TrainerConfig(max_epochs=2, batch_size=512, learning_rate=6e-4,
                      lr_decay=True, warmup_tokens=512*20, final_tokens=2*len(train_dataset)*block_size,
                      num_workers=4)
trainer = Trainer(model, train_dataset, None, tconf)
trainer.train()

  0%|          | 0/2179 [00:00<?, ?it/s]/apcv/shared/conda-envs/apcv-6244e1d-566/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:61: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.
  warnings.warn('Was asked to gather along dimension 0, but all '
epoch 1 iter 2178: train loss 0.33318. lr 3.000169e-04: 100%|██████████| 2179/2179 [09:22<00:00,  3.88it/s]
epoch 2 iter 2178: train loss 0.17447. lr 6.000000e-05: 100%|██████████| 2179/2179 [08:45<00:00,  4.15it/s]

# alright, let's sample some character-level Shakespeare
from mingpt.utils import sample

context = "O God, O God!"
x = torch.tensor([train_dataset.stoi[s] for s in context], dtype=torch.long)[None,...].to(trainer.device)
y = sample(model, x, 2000, temperature=1.0, sample=True, top_k=10)[0]
completion = ''.join([train_dataset.itos[int(i)] for i in y])
print(completion)

O God, O God! that e'er this tongue of mine,
That laid the sentence of dread banishment
On yon proud man, should take it off again
With words of sooth! O that I were as great
As is my grief, or lesser than my name!
Or that I could forget
With Richmond, I'll tell you what I am,
The Lord Aumerle, .

CLAUDIO:
The prenzie Angelo!

ISABELLA:
O, 'tis the cunning livery of hell,
The damned'st body to invest and cover
In prenzie guards! Dost thou think, Claudio?
If I would yield him my virginity,
Thou mightst be freed.

CLAUDIO:
O heavens! it cannot be.

ISABELLA:
Yes, he would give't thee, from this rank offence,
So to offend him still. This night's the time
That I should do what I abhor to name,
Or else thou diest to-morrow.

CLAUDIO:
Thou shalt not do't.

ISABELLA:
O, were it but my life,
I'ld throw it down for your deliverance
As frankly as a pin.

CLAUDIO:
Thanks, dear Isabel.

ISABELLA:
Be ready, Claudio, for your death tomorrow.

CLAUDIO:
Yes. Has he affections
That profit us.

DUKE VINCENTIO:
By the world they see the word in 's doom.

ANGELO:
Thou art the like, thus hate the course in heaven.
What foul hath bled the wheel or at wild,
And wash him fresh again with true-love tears.
Ah, thou, the model where old Troy did stand,
Thou map of honour, thou King Richard's tomb,
And not King Richard; thou most beauteous inn,
Why should hard-favour'd grief be lodged in thee,
When triumph is become an alehouse guest?

KING RICHARD II:
Join not with grief, fair woman, do not so,
To make my end too sudden: learn, good soul,
To think our former state a happy dream;
From which awaked, the truth of what we are
Shows us but this: I am sworn brother, sweet,
To grim Necessity, and he and I
Will keep a league till death. Hie thee to France
And cloister thee in some religious house:
Our holy lives must win a new world's crown,
Which our profane hours here have stricken down.

QUEEN:
What, is my Richard both in shape and mind
Transform'd and weaken'd? hath Bolingbroke deposed
Thine intellect? hath h

# well that was fun