GPT from “scratch” in Flax

# A sample generation
idx = bglm.apply(state.params, 600, method='generate')
generation = ts.decode(idx.tolist()[0])
print(generation.strip())

idx = bglm.apply(initial_params, 600, method='generate')
print(ts.decode(idx.tolist()[0]).strip())

import logging
import requests
from typing import Sequence

import matplotlib.pyplot as plt
import numpy as np

from tqdm.auto import tqdm

import jax
import jax.numpy as jnp
import flax
from flax import linen as nn
from flax.training import train_state, checkpoints
import optax

SEED = 1337

def create_logger():
    """
    Create and configure a logger for notebook use with a console handler.

    Returns:
        logging.Logger: Configured logger with a console handler.
    """
    logger = logging.getLogger("notebook")
    logger.setLevel(logging.INFO)

    if not logger.hasHandlers():
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.DEBUG)
        formatter = logging.Formatter("%(levelname)s - %(message)s")
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)

    return logger

logger = create_logger()

# Checking if Cuda is loaded
devices = jax.devices()
print(devices)

import abc

class BaseTextProcessor:
    """
    Base class for text processing which includes methods to encode, decode,
    and split text data for training and testing.

    Attributes:
        block_size (int): Size of a block of text to process.
        batch_size (int): Number of text blocks in a batch.
    """
    def __init__(self, block_size: int, batch_size: int):
        """
        Initializes the text processor with block size and batch size.
        
        Parameters:
            block_size (int): The size of a block of text.
            batch_size (int): The number of text blocks in each batch.
        """
        self.block_size = block_size
        self.batch_size = batch_size
        self._data = None

        self.text = self.set_text()
        self.chars = sorted(list(set(self.text)))
        self.vocab_size = len(self.chars)
        self.stoi = {ch: i for i, ch in enumerate(self.chars)}
        self.itos = {i: ch for i, ch in enumerate(self.chars)}

    @abc.abstractmethod
    def set_text(self) -> str:
        """
        Abstract method to set the text corpus to be used by the processor.
        
        Returns:
            str: The text corpus.
        """
        pass

    def encode(self, input_string: str) -> list:
        """
        Encodes the input string into its integer representation.
        
        Parameters:
            input_string (str): The string to encode.
        
        Returns:
            list: List of integers representing the encoded string.
        """
        return [self.stoi[c] for c in input_string]

    def decode(self, token_iter: iter) -> str:
        """
        Decodes a list of tokens into the corresponding string.
        
        Parameters:
            token_iter (iter): An iterable of integer tokens to decode.
        
        Returns:
            str: The decoded string.
        """
        return "".join([self.itos[int(c)] for c in token_iter])

    def batch_decoder(self, tokens_lists: list) -> list:
        """
        Decodes a list of token lists to their corresponding strings.
        
        Parameters:
            tokens_lists (list): A list of lists containing tokenized data.
        
        Returns:
            list: A list of decoded strings.
        """
        return [self.decode(tokens) for tokens in tokens_lists]

    @property
    def data(self):
        """
        Lazily encodes the text data and returns it. Ensures data is only encoded once.
        
        Returns:
            jnp.array: The encoded data as a JAX array.
        """
        if self._data is None:
            self._data = jnp.array(self.encode(self.text))
        return self._data

    def train_test_split(self, split: float = 0.9):
        """
        Splits the data into training and validation sets.
        
        Parameters:
            split (float): Fraction of data to be used as training data.
        
        Returns:
            tuple: A tuple containing the training data and validation data.
        """
        n = int(split * len(self.data))
        train_data = self.data[:n]
        val_data = self.data[n:]
        return train_data, val_data

    def get_batch(self, key):
        """
        Generates a batch of training examples.
        
        Parameters:
            key: JAX PRNG key for random number generation.
        
        Returns:
            tuple: A tuple containing the input examples (x), the target array (y),
                   and the indices (ix) of the starting positions in the data.
        """
        ix = jax.random.randint(key=key, minval=0, maxval=len(self.data) - self.block_size,
                                shape=(self.batch_size,))
        x = jnp.stack([self.data[i:i + self.block_size] for i in ix])
        y = jnp.stack([self.data[i + 1:i + self.block_size + 1] for i in ix])
        return x, y, ix

class TinyShakespeare(BaseTextProcessor):
    """
    Processor for handling text data from the Tiny Shakespeare dataset.
    Inherits from BaseTextProcessor.
    """
    def set_text(self) -> str:
        """
        Fetches and sets the text data from the Tiny Shakespeare dataset.

        Returns:
            str: The text data fetched from the dataset.
        """
        text = requests.get(
            "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
        ).text
        return text

class QuickBrownFox(BaseTextProcessor):
    """
    Processor for handling a predefined text string.
    Inherits from BaseTextProcessor.
    """
    def set_text(self) -> str:
        """
        Sets a predefined text string for processing.

        Returns:
            str: A preset text.
        """
        return "The quick brown fox jumped over the lazy dog"

class Alphabet(BaseTextProcessor):
    """
    Processor for handling a simple alphabet string.
    Inherits from BaseTextProcessor.
    """
    def set_text(self) -> str:
        """
        Sets a simple alphabet string for processing.

        Returns:
            str: A simple alphabet string.
        """
        return "abcdefg"

# Initialize the random number generator
ROOT_KEY = jax.random.PRNGKey(seed=0)

# Split the root key into three separate keys
MAIN_KEY, PARAMS_KEY, DROPOUT_KEY = jax.random.split(ROOT_KEY, num=3)

BATCH_SIZE = 128  # Number of independent sequences processed in parallel
BLOCK_SIZE = 64   # Maximum context length for predictions

# Create an instance of TinyShakespeare with specified block size and batch size
ts = TinyShakespeare(block_size=BLOCK_SIZE, batch_size=BATCH_SIZE)

# Uncomment the appropriate line below to use a different text processor
# ts = utils.QuickBrownFox(block_size=BLOCK_SIZE, batch_size=BATCH_SIZE)
# ts = utils.Alphabet(block_size=BLOCK_SIZE, batch_size=BATCH_SIZE)

# Fetch a batch of data from the text processor
xb, yb, ix = ts.get_batch(MAIN_KEY)

# Display the first 50 characters of the text
ts.text[:50]

ts.vocab_size # Display the vocabulary size

xb[:2] # View the first two elements of the batch

ts.batch_decoder(xb[:2]) # Decode the first two batches for human-readable output

# Model Parameters Configuration
VOCAB_SIZE = ts.vocab_size  # Vocabulary size from the text processor instance
N_EMBD = 120  # Number of embedding dimensions
N_HEAD = 6    # Number of attention heads
N_LAYER = 6   # Number of transformer layers
DROPOUT_RATE = 0.4  # Dropout rate for training

# Alternative model configuration
# N_EMBD = 300
# N_HEAD = 6
# N_LAYER = 6
# DROPOUT_RATE = 0.2

global_mask = nn.make_causal_mask(xb) # Generate a causal mask

global_mask.shape # Display the shape of the global causal mask

class MultiHeadAttention(nn.Module):
    """
    A module to perform multi-head attention using Flax's linen library.
    This combines multiple attention heads into a single operation.
    """
    num_heads: int
    n_embd: int

    @nn.compact
    def __call__(self, x, training):
        """
        Apply multi-head attention to the input tensor.

        Parameters:
            x (tensor): Input tensor.
            training (bool): Flag to indicate if the model is training (affects dropout).

        Returns:
            tensor: Output tensor after applying multi-head attention and a dense layer.
        """
        x = nn.MultiHeadDotProductAttention(
            num_heads=self.num_heads,
            dropout_rate=DROPOUT_RATE,
            deterministic=not training
        )(x, mask=global_mask)
        x = nn.Dense(self.n_embd)(x)
        return x

# Create an instance of the MultiHeadAttention module
mha = MultiHeadAttention(
    # head_size=head_size,  # Uncomment and define head_size if needed
    num_heads=N_HEAD,
    n_embd=N_EMBD,
)

# Initialize the input tensor
input_x = jnp.ones((BATCH_SIZE, BLOCK_SIZE, N_EMBD))
logger.debug(f"Input tensor shape: {input_x.shape}")

# Initialize the parameters and apply the model
params = mha.init({'params': PARAMS_KEY, "dropout": DROPOUT_KEY}, input_x, training=True)
print(f"Output shape after applying MultiHeadAttention: {mha.apply(params, input_x, training=False).shape}")

# Example of initializing and applying with a different key
input_x = jnp.ones((BATCH_SIZE, BLOCK_SIZE, N_EMBD))
params = mha.init(ROOT_KEY, input_x, training=False)
mha.apply(params, input_x, training=False)

class FeedForward(nn.Module):
    """
    A feedforward neural network module using Flax's Linen API with two dense layers
    and a dropout layer for regularization.
    """
    
    @nn.compact      
    def __call__(self, x, training):
        """
        Applies a sequence of layers to the input tensor.

        Parameters:
            x (tensor): Input tensor to the feedforward network.
            training (bool): Flag to indicate if the model is training.

        Returns:
            tensor: The output tensor after processing through dense and dropout layers.
        """
        x = nn.Sequential([
            nn.Dense(4 * N_EMBD),
            nn.relu,
            nn.Dense(N_EMBD),
            nn.Dropout(DROPOUT_RATE, deterministic=not training)
        ])(x)
        return x

# Create an instance of the FeedForward module
ff = FeedForward()

# Prepare input tensor
input_x = jnp.ones((BATCH_SIZE, BLOCK_SIZE, N_EMBD))
logger.debug(f"Input tensor shape: {input_x.shape}")

# Initialize the model parameters and apply the model
params = ff.init({'params': PARAMS_KEY, "dropout": DROPOUT_KEY}, input_x, training=True)
print(f"Output shape after applying FeedForward: {ff.apply(params, input_x, training=False).shape}")

print(ff.tabulate(ROOT_KEY, input_x, training=False,
                  console_kwargs={'force_terminal': False, 'force_jupyter': True}))

class Block(nn.Module):
    """
    A transformer block module using Flax's linen API, which integrates multi-head attention
    and feedforward neural network layers.
    """
    @nn.compact         
    def __call__(self, x, training):
        """
        Process the input tensor through the transformer block.

        Parameters:
            x (tensor): Input tensor.
            training (bool): Whether the model is in training mode.

        Returns:
            dict: A dictionary containing the output tensor and the training status.
        """
        # Initialize the MultiHeadAttention and FeedForward modules
        sa = MultiHeadAttention(n_embd=N_EMBD, num_heads=N_HEAD)
        ff = FeedForward()
        
        # Apply self-attention and residual connection followed by layer normalization
        x = x + sa(nn.LayerNorm(N_EMBD)(x), training=training)
        
        # Apply feedforward network and residual connection followed by layer normalization
        x = x + ff(nn.LayerNorm(N_EMBD)(x), training=training)
        
        return dict(x=x, training=training)

# Instantiate the Block module
block = Block()

# Prepare the input tensor
input_x = jnp.ones((BATCH_SIZE, BLOCK_SIZE, N_EMBD))
logger.debug(f"Input tensor shape: {input_x.shape=}")

# Initialize the model parameters and apply the model
block_params = block.init({'params': PARAMS_KEY, "dropout": DROPOUT_KEY}, input_x, training=True)
logger.info(f"Output tensor shape after block application: {block.apply(block_params, input_x, training=False)['x'].shape=}")

print(block.tabulate(ROOT_KEY, input_x, training=False,
                     console_kwargs={'force_terminal': True, 'force_jupyter': True},
                     column_kwargs={'width': 400}))

vocab_size = ts.vocab_size  # Vocabulary size from the text processor instance

class BigramLanguageModel(nn.Module):
    """
    A Bigram Language Model using a neural network approach with Flax's linen module.
    This model uses embedding layers, a sequence of blocks, and a final dense layer
    to predict the next token in the sequence.
    """

    @nn.compact
    def __call__(self, idx, training):
        """
        Process the input indices through the model and compute logits for next token prediction.

        Parameters:
            idx (array): Input tensor of token indices.
            training (bool): Flag indicating whether the model is training.

        Returns:
            array: Logits representing the probability distribution of the next token.
        """
        logger.debug(f"Input shape in __call__: {idx.shape}")
        B, T = idx.shape

        tok_emb = nn.Embed(vocab_size, N_EMBD, name="TokenEmbedding")(idx)  # Token embeddings
        pos_emb = nn.Embed(BLOCK_SIZE, N_EMBD, name="Position Embedding")(jnp.arange(T))  # Positional embeddings

        x = tok_emb + pos_emb  # Combine embeddings
        x = nn.Sequential([Block() for _ in range(N_LAYER)])(x, training=training)["x"]
        x = nn.LayerNorm(N_EMBD, name="LayerNorm")(x)
        logits = nn.Dense(vocab_size, name="Final Dense")(x)

        return logits

    def generate(self, max_new_tokens):
        """
        Generate a sequence of tokens using the model.

        Parameters:
            max_new_tokens (int): Maximum number of new tokens to generate.

        Returns:
            array: Indices of the generated tokens.
        """
        idx = jnp.zeros((1, BLOCK_SIZE), dtype=jnp.int32) * 4
        key = jax.random.PRNGKey(0)  # Key for randomness

        for i in range(max_new_tokens):
            logits = self.__call__(idx[:, -BLOCK_SIZE:], training=False)
            logits_last_t = logits[0, -1]

            key, subkey = jax.random.split(key)
            idx_next = jax.random.categorical(subkey, logits_last_t)
            idx = jnp.atleast_2d(jnp.append(idx, idx_next))

        return idx

# Create an instance of BigramLanguageModel and configure logger
bglm = BigramLanguageModel()
logger.setLevel(logging.DEBUG)

# Prepare input tensor
input_x = jnp.ones((BATCH_SIZE, BLOCK_SIZE), dtype=jnp.int16)
logger.debug(f"Input tensor shape: {input_x.shape}")

# Initialize the model
initial_params = bglm.init({'params':PARAMS_KEY, "dropout": DROPOUT_KEY}, input_x, training=True)
initial_params["params"].keys()

print(bglm.tabulate({'params': PARAMS_KEY, "dropout": DROPOUT_KEY}, input_x, training=False,
                    console_kwargs={'force_terminal': False, 'force_jupyter': True},
                    column_kwargs={'width': 400}))

# Configure logger to not propagate messages to higher level loggers
logger.propagate = False

# Create an instance of BigramLanguageModel
bglm = BigramLanguageModel()

# Generate a sequence of tokens using the model and initial parameters
idx = bglm.apply(initial_params, 50, method='generate')

# Decode the generated indices back to text using the text processor's decode method
decoded_text = ts.decode(idx.tolist()[0])

# Optionally, print or process the decoded text further
print(decoded_text)

# Set the learning rate for the optimizer
LEARNING_RATE = 1e-2

class TrainState(train_state.TrainState):
    """
    Extends the flax TrainState to include the dropout key in the training state.
    """
    key: jax.Array

# Initialize the training state with the model's apply function, parameters, and an optimizer
state = TrainState.create(
    apply_fn=bglm.apply,
    params=initial_params,
    key=DROPOUT_KEY,
    tx=optax.adam(LEARNING_RATE)
)

@jax.jit
def train_step(state: TrainState, inputs, labels):
    """
    Performs a single training step.

    Parameters:
        state (TrainState): Current state of the training including parameters.
        inputs (array): Input data to the model.
        labels (array): Target labels corresponding to the input data.

    Returns:
        tuple: Updated training state and the computed loss for the step.
    """
    # Update the dropout key based on the current training step
    dropout_train_key = jax.random.fold_in(key=state.key, data=state.step)

    def cross_entropy_loss(params):
        """
        Computes the cross entropy loss for a set of parameters.

        Parameters:
            params (dict): Model parameters.

        Returns:
            tuple: Loss value and logits from the model's forward pass.
        """
        logits = bglm.apply(params, inputs, training=True, rngs={'dropout': dropout_train_key})
        logger.debug(f"Logits shape: {logits.shape}")
        loss = jnp.mean(optax.softmax_cross_entropy_with_integer_labels(logits=logits, labels=labels))
        return loss, logits
    
    grad_fn = jax.value_and_grad(cross_entropy_loss, has_aux=True)
    (loss, logits), grads = grad_fn(state.params)
    state = state.apply_gradients(grads=grads)
    return state, loss

# Debugging output to inspect parameters and step count
print(state.params["params"].keys())

print(f"Training step count: {state.step}")

# Set the logger to debug level
logger.setLevel(logging.DEBUG)

# Training configuration
EVAL_INTERVAL = 100
CKPT_DIR = 'ckpts' # Set an absolute path for the checkpoint directory
EPOCHS = 10000
losses = []  # List to store loss values

# Split the random key for training and dropout
train_key, dropout_key = jax.random.split(ROOT_KEY, num=2)

# Main training loop
for epoch in tqdm(range(EPOCHS)):
    # Generate a new random key for the current training step
    train_key = jax.random.fold_in(key=ROOT_KEY, data=state.step)

    # Fetch a new batch of data
    xb, yb, ix = ts.get_batch(train_key)
    
    # Perform a training step and capture the loss
    state, loss = train_step(state, xb, yb)
    losses.append(loss)
    
    # Periodically evaluate and report the loss
    if epoch % EVAL_INTERVAL == 0 or epoch == EPOCHS - 1:
        print(f"Epoch {epoch}: Train Loss {loss:.4f}")
    
    # Update model parameters 
    checkpoints.save_checkpoint(ckpt_dir=CKPT_DIR,
                                target=state,
                                keep_every_n_steps=100,
                                prefix='nn_compact',
                                overwrite=True,
                                step=epoch)

fig, ax = plt.subplots()
ax.plot(np.arange(EPOCHS), losses)

ax.set_xlabel("Training Step")
ax.set_ylabel("Training Loss")

# Configure the logger to information level
logger.setLevel(logging.INFO)

# Generate text using the current state parameters
idx = bglm.apply(state.params, 600, method='generate')
generation = ts.decode(idx.tolist()[0])
print(generation.strip())  # Print the generated text, stripped of leading/trailing whitespace

ckpt_301 = checkpoints.restore_checkpoint(ckpt_dir=f'{CKPT_DIR}/nn_compact301', target=None) # Restore model from a specific checkpoint

# Generate text using the restored model parameters
idx = bglm.apply(ckpt_301["params"], 600, method='generate')
generation = ts.decode(idx.tolist()[0])
print(generation.strip())

ckpt_1101 = checkpoints.restore_checkpoint(ckpt_dir=f'{CKPT_DIR}/nn_compact1101', target=None)

idx = bglm.apply(ckpt_1101["params"], 600, method='generate')
generation = ts.decode(idx.tolist()[0])
print(generation.strip())

ckpt_5901 = checkpoints.restore_checkpoint(ckpt_dir=f'{CKPT_DIR}/nn_compact5901', target=None)

idx = bglm.apply(ckpt_5901["params"], 600, method='generate')
generation = ts.decode(idx.tolist()[0])
print(generation.strip())