Implement simple character level GPT and the trainer for it

[Aisuko]

class SimpleGPT(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        # each token directly reads off the logits for the next token from the lookup table
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    
    def forward(self, idx, targets=None):
        # idx and targets are both(B,T) tensors of integers
        logits=self.token_embedding_table(idx) # (B,T,C)

        if targets is None:
            loss=None
        else:
            B,T,C=logits.shape
            logits=logits.view(B*T, C)
            targets=targets.view(B*T) # B*T it also ok here
            loss=F.cross_entropy(logits, targets)
        return logits, loss
    
    def generate(self, idx, max_new_tokens):
        # idx is (B,T) array of indices in the current context
        for _ in range(max_new_tokens):
            # get the predictions
            logits, loss=self(idx) # call forward automatically
            # focus only on the last time step
            logits=logits[:,-1,:] # becomes (B,C)
            # apply softmax to get probabilities
            probs=F.softmax(logits, dim=-1) # (B,C)
            # sample from the distribution
            idx_next=torch.multinomial(probs, num_samples=1) #(B,1)
            # append sampled index to the running sequence
            idx=torch.cat((idx, idx_next), dim=1) # (B,T+1)
        return idx