import requests

# Download the tiny shakespeare dataset
!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt

# Read the file
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

print("Length of the dataset in chracters: ", len(text))

print(text[:500])

# Check out the unique characters that occur in this text dataset
chars = sorted({char for char in text})
vocab_size = len(chars)
print(''.join(chars))
print(f"Unique characters: {vocab_size}")

# Create mappings from characters to integers and vice versa
char_to_int = {char: index for index, char in enumerate(chars)}
int_to_char = {index: char for index, char in enumerate(chars)}

# Define encoding and decoding functions
def encode(text):
  return [char_to_int[char] for char in text]

def decode(encoded):
  return ''.join([int_to_char[index] for index in encoded])

# Test the encoding and decoding functions
encoded = encode("hello there")
print(encoded)
print(decode(encoded))

# Encode the entire text dataset and save it into a torch.Tensor

import torch # PyTorch: https://pytorch.org

data = torch.tensor(encode(text), dtype=torch.long)
print(f"Shape of data tensor: {data.shape}")
print(f"Data type of data tensor: {data.dtype}")
print(f"First 500 characters in encoded form: {data[:500]}")

# Split the data into training and validation sets

train_ratio = 0.8
train_data_length = int(len(data) * train_ratio)
train_data = data[:train_data_length]
val_data = data[train_data_length:]

print(f"Length of training data: {len(train_data)}")
print(f"Length of validation data: {len(val_data)}")

block_size = 8
print("First", block_size + 1, "characters in training data:")
print(decode(train_data[:block_size+1].tolist()))

x = train_data[:block_size]
y = train_data[1:block_size+1]
for t in range(block_size):
    context = x[:t+1]
    target = y[t]
    print(f"when input is {context} the target: {target}")

torch.manual_seed(1337)
batch_size = 4 # The number of independent sequences that will be processed in parallel
block_size = 8 # Max context length for predictions


def get_batch(split):
    data = train_data if split == 'train' else val_data
    ix = torch.randint(0, len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    return x, y

xb, yb = get_batch('train')
print('Inputs:')
print(xb.shape)
print(xb)
print('Targets:')
print(yb.shape)
print(yb)

print('----')

for b in range(batch_size):
    for t in range(block_size):
        context = xb[b, :t+1]
        target = yb[b, t]
        print(f"When input is {context.tolist()}, the target is: {target}")

print(xb) # My input to the transformer

import torch
import torch.nn as nn
from torch.nn import functional as F
torch.manual_seed(1337)

class BigramLanguageModel(nn.Module):

    def __init__(self, vocab_size):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):
        logits = self.embedding(idx)
        B, T, C = logits.shape
        logits = logits.reshape(B * T, C)
        if targets is None:
            return logits, None
        targets = targets.reshape(-1)
        loss = F.cross_entropy(logits, targets)
        return logits, loss
    
    def generate(self, idx, max_new_tokens):
        for i in range(max_new_tokens):
            logits = self.embedding(idx)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, next_token), dim=1)
        return idx

m = BigramLanguageModel(vocab_size)
logits, loss = m(xb, yb)
print(logits.shape)
print(loss)

print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=100)[0].tolist()))

from torch.utils.tensorboard import SummaryWriter

batch_size = 32
max_steps = 100 # Increase number of steps for better results... 
early_stop_steps = 10

# Create a TensorBoard writer
writer = SummaryWriter()

# Keep track of the best validation loss
best_val_loss = float("inf")

# Early stopping counter
early_stop_counter = 0

for steps in range(max_steps):
    # Obtain a batch of data as sample
    xb, yb = get_batch('train')

    # Loss evaluation
    logits, loss = m(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

    if steps % 10 == 0:
        print("Step {}: Train Loss: {}".format(steps, loss.item()))
        writer.add_scalar('train_loss', loss.item(), steps)
    
    # Evaluate the model on the validation set
    with torch.no_grad():
        xb, yb = get_batch('val')
        logits, val_loss = m(xb, yb)
        writer.add_scalar('val_loss', val_loss.item(), steps)
        
        if val_loss.item() < best_val_loss:
            best_val_loss = val_loss.item()
            early_stop_counter = 0
        else:
            early_stop_counter += 1
        
        if early_stop_counter >= early_stop_steps:
            print("Early stopping at step {}".format(steps))
            break

writer.close()
print("Final Train Loss: {}".format(loss.item()))

print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=500)[0].tolist()))

# Set random seed for reproducibility
torch.manual_seed(42)

# Create a lower triangular matrix of ones
a = torch.tril(torch.ones(3, 3))

# Normalize the lower triangular matrix along the rows
a = a / torch.sum(a, 1, keepdim=True)

# Generate random weights from a normal distribution
b = torch.randn(3, 2).float()

# Compute the weighted aggregation using matrix multiplication
c = a @ b

# Print results
print('a=')
print(a)
print('--')
print('b=')
print(b)
print('--')
print('c=')
print(c)

# Set random seed for reproducibility
torch.manual_seed(1337)

# Define the batch size (B), time steps (T), and number of channels (C)
B, T, C = 4, 8, 2

# Generate a tensor with random values, with shape (B, T, C)
x = torch.randn(B, T, C)

# Print the shape of the tensor
print("Shape of the tensor:", x.shape)

wei = torch.tril(torch.ones(T, T))
wei = wei / wei.sum(1, keepdim=True)
xbow2 = wei @ x # (B, T, T) @ (B, T, C) ----> (B, T, C)
torch.allclose(xbow, xbow2)

# Create a mask for the lower triangular matrix of shape (T, T)
mask = torch.tril(torch.ones(T, T)).bool()

# Set the values outside the lower triangle to -inf to ensure zero weights
mask = mask.float().masked_fill(~mask, float('-inf'))

# Apply softmax to the masked values along the last dimension to obtain the weights
weights = F.softmax(mask, dim=-1)

# Obtain the weighted average of x using the computed weights
xbow3 = torch.matmul(weights, x)

# Check if xbow and xbow3 are equal
torch.allclose(xbow, xbow3)

# Set the random seed to a fixed value for reproducibility
torch.manual_seed(1337)

# Define the batch size (B), number of time steps (T), and number of channels (C)
B,T,C = 4,8,32

# Generate a tensor with random values of shape (B, T, C)
x = torch.randn(B,T,C)

# Define the size of each head in the attention mechanism
head_size = 16

# Define the linear layers for the keys, queries, and values
key = nn.Linear(C, head_size, bias=False)
query = nn.Linear(C, head_size, bias=False)
value = nn.Linear(C, head_size, bias=False)

# Apply the linear layer for the keys to the input tensor to get k of shape (B, T, head_size)
k = key(x)

# Apply the linear layer for the queries to the input tensor to get q of shape (B, T, head_size)
q = query(x)

# Calculate the attention weights by multiplying q and k transposed together. This produces a tensor wei of shape (B, T, T).
wei = q @ k.transpose(-2, -1) # (B, T, 16) @ (B, 16, T) ---> (B, T, T)

# Create a lower triangular matrix with ones and zero out the upper triangular part
tril = torch.tril(torch.ones(T, T))

# Replace the 0 elements in the triangular part of the attention weights with negative infinity
wei = wei.masked_fill(tril == 0, float('-inf'))

# Apply the softmax function along the last dimension of wei to get the final attention weights
wei = F.softmax(wei, dim=-1)

# Apply the linear layer for the values to the input tensor to get v of shape (B, T, head_size)
v = value(x)

# Compute the weighted sum of the values with the attention weights to get the output tensor of shape (B, T, head_size)
out = wei @ v

# Print the shape of the output tensor
print(out.shape)

print(wei[0])

k = torch.randn(B, T, head_size)
q = torch.randn(B, T, head_size)
scale = torch.rsqrt(torch.tensor(head_size, dtype=torch.float32))
wei = q @ k.transpose(-2, -1) * scale

print("variance of k:", k.var())
print("variance of q:", q.var())
print("variance of wei:", wei.var())

# softmax over last dimension
softmaxed_wei = torch.softmax(wei, dim=-1)

# a different softmax that reduces the "peakiness" of the attention distribution
sm = torch.nn.Softmax(dim=-1)
scaled_sm_wei = sm(wei * 8)

torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5]), dim=-1)

torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5])*8, dim=-1)

class LayerNorm1d:
    def __init__(self, dim, eps=1e-5, momentum=0.1):
        self.eps = eps
        self.gamma = torch.ones(dim)
        self.beta = torch.zeros(dim)
        
    def __call__(self, x):
        # Calculate the forward pass
        xmean = x.mean(1, keepdim=True)  # Batch mean
        xvar = x.var(1, keepdim=True)  # Batch variance
        xhat = (x - xmean) / torch.sqrt(xvar + self.eps)  # Normalize to unit variance
        self.out = self.gamma * xhat + self.beta
        return self.out
    
    def parameters(self):
        return [self.gamma, self.beta]

# Instantiate the module
torch.manual_seed(1337)
module = LayerNorm1d(100)

# Generate a batch of size 32 with 100-dimensional vectors
x = torch.randn(32, 100)

# Apply the layer normalization to the batch
x = module(x)

# Print the shape of the output tensor
print(f"Output shape: {x.shape}")

# Compute the mean and standard deviation of the first feature across all inputs in the batch
batch_feature_mean, batch_feature_std = x[:, 0].mean(), x[:, 0].std()
print(f"Batch feature mean: {batch_feature_mean}, Batch feature std: {batch_feature_std}")

# Compute the mean and standard deviation of all features for a single input in the batch
single_input_mean, single_input_std = x[0, :].mean(), x[0, :].std()
print(f"Single input mean: {single_input_mean}, Single input std: {single_input_std}")

import torch
import torch.nn as nn
from torch.nn import functional as F

# Hyperparameters
batch_size = 16 
block_size = 32 
max_iters = 5000
eval_interval = 100
learning_rate = 1e-3
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 64
n_head = 4
n_layer = 4
dropout = 0.0

torch.manual_seed(1337)

# wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

# Create a sorted list of unique characters in the text
chars = sorted(list(set(text)))

# Get the number of unique characters
vocab_size = len(chars)

# Create a dictionary mapping each character to an integer index
stoi = {ch: i for i, ch in enumerate(chars)}

# Create a dictionary mapping each integer index to a character
itos = {i: ch for i, ch in enumerate(chars)}

# Define a lambda function to encode a string as a list of integers
encode = lambda s: [stoi[c] for c in s]

# Define a lambda function to decode a list of integers as a string
decode = lambda l: ''.join([itos[i] for i in l])

# Convert the text to a tensor of integers
data = torch.tensor(encode(text), dtype=torch.long)

# Split the data into train and validation sets
n = int(0.9 * len(data))  # use first 90% for training, last 10% for validation
train_data = data[:n]
val_data = data[n:]

# Define a function to generate a small batch of input-target pairs
def get_batch(split):
    # Select either the training or validation set
    data = train_data if split == 'train' else val_data
    
    # Generate random indices to start each block of input
    ix = torch.randint(len(data) - block_size, (batch_size,))
    
    # Select block_size characters starting at each index for input
    x = torch.stack([data[i:i + block_size] for i in ix])
    
    # Select block_size characters starting at each index + 1 for target
    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
    
    # Send tensors to GPU if available
    x, y = x.to(device), y.to(device)
    return x, y

# Define a function to estimate the model's loss on the train and validation sets
@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out


class Head(nn.Module):
    """ one head of self-attention """

    def __init__(self, head_size):
        super().__init__()

        # Linear transformations for key, query, and value
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)

        # Lower triangular matrix to mask future values
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

        # Dropout layer
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # Get batch size, sequence length, and number of features
        B, T, C = x.shape

        # Linear transformations of key and query
        k = self.key(x)   # (B,T,C)
        q = self.query(x) # (B,T,C)

        # Compute attention scores ("affinities")
        wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)

        # Mask future values
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)

        # Apply softmax to get attention weights
        wei = F.softmax(wei, dim=-1) # (B, T, T)

        # Apply dropout to attention weights
        wei = self.dropout(wei)

        # Linear transformation of value
        v = self.value(x) # (B,T,C)

        # Weighted sum of values using attention weights
        out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)

        return out


class MultiHeadAttention(nn.Module):
    """ multiple heads of self-attention in parallel """

    def __init__(self, num_heads, head_size):
        super().__init__()  # Initialize the superclass (nn.Module)
        # Instantiate a list of head modules, and assign it to self.heads
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        # A linear transformation to project the concatenated attention heads back to the input dimension
        self.proj = nn.Linear(n_embd, n_embd)
        # Dropout layer to avoid overfitting
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # Apply the self-attention mechanism by passing the input through each attention head and concatenate the results along the feature dimension
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        # Apply the projection layer and the dropout layer to the concatenated output
        out = self.dropout(self.proj(out))
        return out


class FeedFoward(nn.Module):
    """ a simple linear layer followed by a non-linearity """

    def __init__(self, n_embd):
        super().__init__()

        # Define a sequential neural network composed of two linear layers and a ReLU activation function
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        # Forward pass through the neural network
        return self.net(x)


class Block(nn.Module):
    """ Transformer block: communication followed by computation """

    def __init__(self, n_embd, n_head):
        # n_embd: embedding dimension, n_head: the number of heads we'd like
        super().__init__()
        head_size = n_embd // n_head
        
        # Multi-head self-attention layer
        self.sa = MultiHeadAttention(n_head, head_size)
        
        # Feedforward layer
        self.ffwd = FeedFoward(n_embd)
        
        # Layer normalization layer 1
        self.ln1 = nn.LayerNorm(n_embd)
        
        # Layer normalization layer 2
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        # Communication followed by computation
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x


# super simple bigram model
class BigramLanguageModel(nn.Module):

    def __init__(self):
        super().__init__()
        # Define the model architecture using embedding, multi-head attention, and linear layers
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)  # Lookup table for token embedding
        self.position_embedding_table = nn.Embedding(block_size, n_embd)  # Lookup table for position embedding
        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])  # A sequence of transformer blocks
        self.ln_f = nn.LayerNorm(n_embd)  # Layer normalization
        self.lm_head = nn.Linear(n_embd, vocab_size)  # Linear layer to get logits

    def forward(self, idx, targets=None):
        B, T = idx.shape

        # Embed tokens and positions
        tok_emb = self.token_embedding_table(idx)  # (B,T,C)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))  # (T,C)
        x = tok_emb + pos_emb  # (B,T,C)

        # Pass through the transformer blocks
        x = self.blocks(x)  # (B,T,C)

        # Apply layer normalization and linear layer
        x = self.ln_f(x)  # (B,T,C)
        logits = self.lm_head(x)  # (B,T,vocab_size)

        # Compute the cross-entropy loss if targets are provided
        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        # Generate new text by sampling from the learned distribution
        for _ in range(max_new_tokens):
            # Crop idx to the last block_size tokens
            idx_cond = idx[:, -block_size:]
            # Get the predictions
            logits, loss = self(idx_cond)
            # 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 the sampled index to the running sequence
            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)
        return idx

# Instantiate the bigram language model
model = BigramLanguageModel()  

# Move the model to the specified device (CPU or GPU)
m = model.to(device)  

# Print the number of parameters in the model
print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')  

# Create a PyTorch optimizer with the AdamW algorithm
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# Loop over training iterations
for iter in range(max_iters):  

    # Every once in a while evaluate the loss on train and val sets
    if iter % eval_interval == 0 or iter == max_iters - 1:
        losses = estimate_loss()
        print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

    # Sample a batch of data (inputs and targets)
    xb, yb = get_batch('train')

    # Evaluate the loss and update the model parameters
    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

# Generate text from the model

# Initialize the context with a zero tensor
context = torch.zeros((1, 1), dtype=torch.long, device=device)  

# Generate a sequence of tokens using the model
generated_sequence = m.generate(context, max_new_tokens=2000)

# Decode the generated sequence into a string
decoded_sequence = decode(generated_sequence[0].tolist())  

# Print the generated text
print(decoded_sequence)