[0,1,0] ---> GPT ---> [P (0) = 20%, P (1) = 80%]

# hyperparameters for our GPT
# vocab size is 2, so we only have two possible tokens: 0,1
vocab_size = 2
# context length is 3, so we take 3 bits to predict the next bit probability
context_length = 3

print ('state space (for this exercise) = ', vocab_size ** context_length)
# state space (for this exercise) = 8

print ('actual state space (in reality) = ', sum (vocab_size ** i for i in range (1, context_length+1)))
# actual state space (in reality) = 14

config = GPTConfig (
    block_size = context_length,
    vocab_size = vocab_size,
    n_layer = 4,
    n_head = 4,
    n_embd = 16,
    bias = False,
)
gpt = GPT (config)

def all_possible (n, k):
    # return all possible lists of k elements, each in range of [0,n)
    if k == 0:
        yield []
    else:
        for i in range (n):
            for c in all_possible (n, k - 1):
                yield [i] + c
list (all_possible (vocab_size, context_length))

[[0, 0, 0],
 [0, 0, 1],
 [0, 1, 0],
 [0, 1, 1],
 [1, 0, 0],
 [1, 0, 1],
 [1, 1, 0],
 [1, 1, 1]]

# we'll use graphviz for pretty plotting the current state of the GPT
from graphviz import Digraph

def plot_model ():
    dot = Digraph (comment='Baby GPT', engine='circo')

    for xi in all_possible (gpt.config.vocab_size, gpt.config.block_size):

        # forward the GPT and get probabilities for next token
        x = torch.tensor (xi, dtype=torch.long)[None, ...] # turn the list into a torch tensor and add a batch dimension
        logits = gpt (x) # forward the gpt neural net
        probs = nn.functional.softmax (logits, dim=-1) # get the probabilities
        y = probs [0].tolist () # remove the batch dimension and unpack the tensor into simple list
        print (f"input {xi} ---> {y}")

        # also build up the transition graph for plotting later
        current_node_signature = "".join (str (d) for d in xi)
        dot.node (current_node_signature)
        for t in range (gpt.config.vocab_size):
            next_node = xi [1:] + [t] # crop the context and append the next character
            next_node_signature = "".join (str (d) for d in next_node)
            p = y [t]
            label=f"{t}({p*100:.0f}%)"
            dot.edge (current_node_signature, next_node_signature, label=label)

    return dot

plot_model ()

input [0, 0, 0] ---> [0.4963349997997284, 0.5036649107933044] 
input [0, 0, 1] ---> [0.4515703618526459, 0.5484296679496765] 
input [0, 1, 0] ---> [0.49648362398147583, 0.5035163760185242] 
input [0, 1, 1] ---> [0.45181113481521606, 0.5481888651847839] 
input [1, 0, 0] ---> [0.4961162209510803, 0.5038837194442749] 
input [1, 0, 1] ---> [0.4517717957496643, 0.5482282042503357] 
input [1, 1, 0] ---> [0.4962802827358246, 0.5037197470664978] 
input [1, 1, 1] ---> [0.4520467519760132, 0.5479532480239868]

# let's train our baby GPT on this sequence
seq = list (map (int, "111101111011110"))
seq

[1, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0]

# convert the sequence to a tensor holding all the individual examples in that sequence
X, Y = [], []
# iterate over the sequence and grab every consecutive 3 bits
# the correct label for what's next is the next bit at each position
for i in range (len (seq) - context_length):
    X.append (seq [i:i+context_length])
    Y.append (seq [i+context_length])
    print (f"example {i+1:2d}: {X [-1]} --> {Y [-1]}")
X = torch.tensor (X, dtype=torch.long)
Y = torch.tensor (Y, dtype=torch.long)
print (X.shape, Y.shape)

# init a GPT and the optimizer
torch.manual_seed (1337)
gpt = GPT (config)
optimizer = torch.optim.AdamW (gpt.parameters (), lr=1e-3, weight_decay=1e-1)

# train the GPT for some number of iterations
for i in range (50):
    logits = gpt (X)
    loss = F.cross_entropy (logits, Y)
    loss.backward ()
    optimizer.step ()
    optimizer.zero_grad ()
    print (i, loss.item ())

print ("Training data sequence, as a reminder:", seq)
plot_model ()