Upload folder using huggingface_hub

cac4140 verified 9 months ago

16 kB

import tiktoken
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
import numpy as np




class GPTDatasetV1(Dataset):
    def __init__(self, txt, tokenizer, max_length, stride):
        self.input_ids = []
        self.target_ids = []

        # Tokenize the entire text
        token_ids = tokenizer.encode(txt, allowed_special={"<|endoftext|>"})

        # Use a sliding window to chunk the book into overlapping sequences of max_length
        for i in range(0, len(token_ids) - max_length, stride):
            input_chunk = token_ids[i:i + max_length]
            target_chunk = token_ids[i + 1: i + max_length + 1]
            self.input_ids.append(torch.tensor(input_chunk))
            self.target_ids.append(torch.tensor(target_chunk))

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, idx):
        return self.input_ids[idx], self.target_ids[idx]


def create_dataloader_v1(txt, batch_size=4, max_length=256,
                         stride=128, shuffle=True, drop_last=True, num_workers=0):
    # Initialize the tokenizer
    tokenizer = tiktoken.get_encoding("gpt2")

    # Create dataset
    dataset = GPTDatasetV1(txt, tokenizer, max_length, stride)

    # Create dataloader
    dataloader = DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers)

    return dataloader



class MultiHeadAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        assert d_out % num_heads == 0, "d_out must be divisible by num_heads"

        self.d_out = d_out
        self.num_heads = num_heads
        self.head_dim = d_out // num_heads  # Reduce the projection dim to match desired output dim

        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
        self.dropout = nn.Dropout(dropout)
        self.register_buffer("mask", torch.triu(torch.ones(context_length, context_length), diagonal=1))

    def forward(self, x):
        b, num_tokens, d_in = x.shape

        keys = self.W_key(x)  # Shape: (b, num_tokens, d_out)
        queries = self.W_query(x)
        values = self.W_value(x)

        # We implicitly split the matrix by adding a `num_heads` dimension
        # Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
        values = values.view(b, num_tokens, self.num_heads, self.head_dim)
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

        # Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
        keys = keys.transpose(1, 2)
        queries = queries.transpose(1, 2)
        values = values.transpose(1, 2)

        # Compute scaled dot-product attention (aka self-attention) with a causal mask
        attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head

        # Original mask truncated to the number of tokens and converted to boolean
        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]

        # Use the mask to fill attention scores
        attn_scores.masked_fill_(mask_bool, -torch.inf)

        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # Shape: (b, num_tokens, num_heads, head_dim)
        context_vec = (attn_weights @ values).transpose(1, 2)

        # Combine heads, where self.d_out = self.num_heads * self.head_dim
        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
        context_vec = self.out_proj(context_vec)  # optional projection

        return context_vec


class LayerNorm(nn.Module):
    def __init__(self, emb_dim):
        super().__init__()
        self.eps = 1e-5
        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * norm_x + self.shift


class GELU(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(
            torch.sqrt(torch.tensor(2.0 / torch.pi)) *
            (x + 0.044715 * torch.pow(x, 3))
        ))


class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
            GELU(),
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
        )

    def forward(self, x):
        return self.layers(x)


class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=cfg["emb_dim"],
            d_out=cfg["emb_dim"],
            context_length=cfg["context_length"],
            num_heads=cfg["n_heads"],
            dropout=cfg["drop_rate"],
            qkv_bias=cfg["qkv_bias"])
        self.ff = FeedForward(cfg)
        self.norm1 = LayerNorm(cfg["emb_dim"])
        self.norm2 = LayerNorm(cfg["emb_dim"])
        self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

    def forward(self, x):
        # Shortcut connection for attention block
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)   # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        # Shortcut connection for feed-forward block
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        return x


class GPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate"])

        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits


def generate_text_simple(model, idx, max_new_tokens, context_size):
    # idx is (B, T) array of indices in the current context
    for _ in range(max_new_tokens):

        # Crop current context if it exceeds the supported context size
        # E.g., if LLM supports only 5 tokens, and the context size is 10
        # then only the last 5 tokens are used as context
        idx_cond = idx[:, -context_size:]

        # Get the predictions
        with torch.no_grad():
            logits = model(idx_cond)

        # Focus only on the last time step
        # (batch, n_token, vocab_size) becomes (batch, vocab_size)
        logits = logits[:, -1, :]

        # Get the idx of the vocab entry with the highest logits value
        idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # (batch, 1)

        # Append sampled index to the running sequence
        idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)

    return idx


def generate(model, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None):

    # For-loop is the same as before: Get logits, and only focus on last time step
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -context_size:]
        with torch.no_grad():
            logits = model(idx_cond)
        logits = logits[:, -1, :]

        # New: Filter logits with top_k sampling
        if top_k is not None:
            # Keep only top_k values
            top_logits, _ = torch.topk(logits, top_k)
            min_val = top_logits[:, -1]
            logits = torch.where(logits < min_val, torch.tensor(float("-inf")).to(logits.device), logits)

        # New: Apply temperature scaling
        if temperature > 0.0:
            logits = logits / temperature

            # New (not in book): numerical stability tip to get equivalent results on mps device
            # subtract rowwise max before softmax
            logits = logits - logits.max(dim=-1, keepdim=True).values

            # Apply softmax to get probabilities
            probs = torch.softmax(logits, dim=-1)  # (batch_size, context_len)

            # Sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1)  # (batch_size, 1)

        # Otherwise same as before: get idx of the vocab entry with the highest logits value
        else:
            idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # (batch_size, 1)

        if idx_next == eos_id:  # Stop generating early if end-of-sequence token is encountered and eos_id is specified
            break

        # Same as before: append sampled index to the running sequence
        idx = torch.cat((idx, idx_next), dim=1)  # (batch_size, num_tokens+1)

    return idx


def train_model_simple(model, train_loader, val_loader, optimizer, device, num_epochs,
                       eval_freq, eval_iter, start_context, tokenizer):
    # Initialize lists to track losses and tokens seen
    train_losses, val_losses, track_tokens_seen = [], [], []
    tokens_seen, global_step = 0, -1

    # Main training loop
    for epoch in range(num_epochs):
        model.train()  # Set model to training mode

        for input_batch, target_batch in train_loader:
            optimizer.zero_grad()  # Reset loss gradients from previous batch iteration
            loss = calc_loss_batch(input_batch, target_batch, model, device)
            loss.backward()  # Calculate loss gradients
            optimizer.step()  # Update model weights using loss gradients
            tokens_seen += input_batch.numel()
            global_step += 1

            # Optional evaluation step
            if global_step % eval_freq == 0:
                train_loss, val_loss = evaluate_model(
                    model, train_loader, val_loader, device, eval_iter)
                train_losses.append(train_loss)
                val_losses.append(val_loss)
                track_tokens_seen.append(tokens_seen)
                print(f"Ep {epoch+1} (Step {global_step:06d}): "
                      f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")

        # Print a sample text after each epoch
        generate_and_print_sample(
            model, tokenizer, device, start_context
        )

    return train_losses, val_losses, track_tokens_seen


def evaluate_model(model, train_loader, val_loader, device, eval_iter):
    model.eval()
    with torch.no_grad():
        train_loss = calc_loss_loader(train_loader, model, device, num_batches=eval_iter)
        val_loss = calc_loss_loader(val_loader, model, device, num_batches=eval_iter)
    model.train()
    return train_loss, val_loss


def generate_and_print_sample(model, tokenizer, device, start_context):
    model.eval()
    context_size = model.pos_emb.weight.shape[0]
    encoded = text_to_token_ids(start_context, tokenizer).to(device)
    with torch.no_grad():
        token_ids = generate_text_simple(
            model=model, idx=encoded,
            max_new_tokens=50, context_size=context_size
        )
        decoded_text = token_ids_to_text(token_ids, tokenizer)
        print(decoded_text.replace("\n", " "))  # Compact print format
    model.train()


def assign(left, right):
    if left.shape != right.shape:
        raise ValueError(f"Shape mismatch. Left: {left.shape}, Right: {right.shape}")
    return torch.nn.Parameter(torch.tensor(right))


def text_to_token_ids(text, tokenizer):
    encoded = tokenizer.encode(text, allowed_special={"<|endoftext|>"})
    encoded_tensor = torch.tensor(encoded).unsqueeze(0)  # add batch dimension
    return encoded_tensor


def token_ids_to_text(token_ids, tokenizer):
    flat = token_ids.squeeze(0)  # remove batch dimension
    return tokenizer.decode(flat.tolist())


def calc_loss_batch(input_batch, target_batch, model, device):
    input_batch, target_batch = input_batch.to(device), target_batch.to(device)
    logits = model(input_batch)
    loss = torch.nn.functional.cross_entropy(logits.flatten(0, 1), target_batch.flatten())
    return loss


def calc_loss_loader(data_loader, model, device, num_batches=None):
    total_loss = 0.
    if len(data_loader) == 0:
        return float("nan")
    elif num_batches is None:
        num_batches = len(data_loader)
    else:
        # Reduce the number of batches to match the total number of batches in the data loader
        # if num_batches exceeds the number of batches in the data loader
        num_batches = min(num_batches, len(data_loader))
    for i, (input_batch, target_batch) in enumerate(data_loader):
        if i < num_batches:
            loss = calc_loss_batch(input_batch, target_batch, model, device)
            total_loss += loss.item()
        else:
            break
    return total_loss / num_batches


def plot_losses(epochs_seen, tokens_seen, train_losses, val_losses):
    fig, ax1 = plt.subplots(figsize=(5, 3))

    # Plot training and validation loss against epochs
    ax1.plot(epochs_seen, train_losses, label="Training loss")
    ax1.plot(epochs_seen, val_losses, linestyle="-.", label="Validation loss")
    ax1.set_xlabel("Epochs")
    ax1.set_ylabel("Loss")
    ax1.legend(loc="upper right")
    ax1.xaxis.set_major_locator(MaxNLocator(integer=True))  # only show integer labels on x-axis

    # Create a second x-axis for tokens seen
    ax2 = ax1.twiny()  # Create a second x-axis that shares the same y-axis
    ax2.plot(tokens_seen, train_losses, alpha=0)  # Invisible plot for aligning ticks
    ax2.set_xlabel("Tokens seen")

    fig.tight_layout()  # Adjust layout to make room
    plt.savefig("loss-plot.pdf")
    plt.show()

def main():
    GPT_CONFIG_124M = {
        "vocab_size": 50257,     # Vocabulary size
        "context_length": 1024,  # Context length
        "emb_dim": 768,          # Embedding dimension
        "n_heads": 12,           # Number of attention heads
        "n_layers": 12,          # Number of layers
        "drop_rate": 0.1,        # Dropout rate
        "qkv_bias": False        # Query-Key-Value bias
    }

    torch.manual_seed(123)
    model = GPTModel(GPT_CONFIG_124M)
    model.eval()  # disable dropout

    start_context = "Hi, there"

    tokenizer = tiktoken.get_encoding("gpt2")
    encoded = tokenizer.encode(start_context)
    encoded_tensor = torch.tensor(encoded).unsqueeze(0)

    print(f"\n{50*'='}\n{22*' '}IN\n{50*'='}")
    print("\nInput text:", start_context)
    print("Encoded input text:", encoded)
    print("encoded_tensor.shape:", encoded_tensor.shape)

    out = generate_text_simple(
        model=model,
        idx=encoded_tensor,
        max_new_tokens=10,
        context_size=GPT_CONFIG_124M["context_length"]
    )
    decoded_text = tokenizer.decode(out.squeeze(0).tolist())

    print(f"\n\n{50*'='}\n{22*' '}OUT\n{50*'='}")
    print("\nOutput:", out)
    print("Output length:", len(out[0]))
    print("Output text:", decoded_text)


if __name__ == "__main__":
    main()

	import tiktoken
	import torch
	import torch.nn as nn
	from torch.utils.data import Dataset, DataLoader

	import matplotlib.pyplot as plt
	from matplotlib.ticker import MaxNLocator
	import numpy as np




	class GPTDatasetV1(Dataset):
	def __init__(self, txt, tokenizer, max_length, stride):
	self.input_ids = []
	self.target_ids = []

	# Tokenize the entire text
	token_ids = tokenizer.encode(txt, allowed_special={"<\|endoftext\|>"})

	# Use a sliding window to chunk the book into overlapping sequences of max_length
	for i in range(0, len(token_ids) - max_length, stride):
	input_chunk = token_ids[i:i + max_length]
	target_chunk = token_ids[i + 1: i + max_length + 1]
	self.input_ids.append(torch.tensor(input_chunk))
	self.target_ids.append(torch.tensor(target_chunk))

	def __len__(self):
	return len(self.input_ids)

	def __getitem__(self, idx):
	return self.input_ids[idx], self.target_ids[idx]


	def create_dataloader_v1(txt, batch_size=4, max_length=256,
	stride=128, shuffle=True, drop_last=True, num_workers=0):
	# Initialize the tokenizer
	tokenizer = tiktoken.get_encoding("gpt2")

	# Create dataset
	dataset = GPTDatasetV1(txt, tokenizer, max_length, stride)

	# Create dataloader
	dataloader = DataLoader(
	dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers)

	return dataloader



	class MultiHeadAttention(nn.Module):
	def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
	super().__init__()
	assert d_out % num_heads == 0, "d_out must be divisible by num_heads"

	self.d_out = d_out
	self.num_heads = num_heads
	self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim

	self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
	self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
	self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
	self.out_proj = nn.Linear(d_out, d_out) # Linear layer to combine head outputs
	self.dropout = nn.Dropout(dropout)
	self.register_buffer("mask", torch.triu(torch.ones(context_length, context_length), diagonal=1))

	def forward(self, x):
	b, num_tokens, d_in = x.shape

	keys = self.W_key(x) # Shape: (b, num_tokens, d_out)
	queries = self.W_query(x)
	values = self.W_value(x)

	# We implicitly split the matrix by adding a `num_heads` dimension
	# Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
	keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
	values = values.view(b, num_tokens, self.num_heads, self.head_dim)
	queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

	# Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
	keys = keys.transpose(1, 2)
	queries = queries.transpose(1, 2)
	values = values.transpose(1, 2)

	# Compute scaled dot-product attention (aka self-attention) with a causal mask
	attn_scores = queries @ keys.transpose(2, 3) # Dot product for each head

	# Original mask truncated to the number of tokens and converted to boolean
	mask_bool = self.mask.bool()[:num_tokens, :num_tokens]

	# Use the mask to fill attention scores
	attn_scores.masked_fill_(mask_bool, -torch.inf)

	attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
	attn_weights = self.dropout(attn_weights)

	# Shape: (b, num_tokens, num_heads, head_dim)
	context_vec = (attn_weights @ values).transpose(1, 2)

	# Combine heads, where self.d_out = self.num_heads * self.head_dim
	context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
	context_vec = self.out_proj(context_vec) # optional projection

	return context_vec


	class LayerNorm(nn.Module):
	def __init__(self, emb_dim):
	super().__init__()
	self.eps = 1e-5
	self.scale = nn.Parameter(torch.ones(emb_dim))
	self.shift = nn.Parameter(torch.zeros(emb_dim))

	def forward(self, x):
	mean = x.mean(dim=-1, keepdim=True)
	var = x.var(dim=-1, keepdim=True, unbiased=False)
	norm_x = (x - mean) / torch.sqrt(var + self.eps)
	return self.scale * norm_x + self.shift


	class GELU(nn.Module):
	def __init__(self):
	super().__init__()

	def forward(self, x):
	return 0.5 * x * (1 + torch.tanh(
	torch.sqrt(torch.tensor(2.0 / torch.pi)) *
	(x + 0.044715 * torch.pow(x, 3))
	))


	class FeedForward(nn.Module):
	def __init__(self, cfg):
	super().__init__()
	self.layers = nn.Sequential(
	nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
	GELU(),
	nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
	)

	def forward(self, x):
	return self.layers(x)


	class TransformerBlock(nn.Module):
	def __init__(self, cfg):
	super().__init__()
	self.att = MultiHeadAttention(
	d_in=cfg["emb_dim"],
	d_out=cfg["emb_dim"],
	context_length=cfg["context_length"],
	num_heads=cfg["n_heads"],
	dropout=cfg["drop_rate"],
	qkv_bias=cfg["qkv_bias"])
	self.ff = FeedForward(cfg)
	self.norm1 = LayerNorm(cfg["emb_dim"])
	self.norm2 = LayerNorm(cfg["emb_dim"])
	self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

	def forward(self, x):
	# Shortcut connection for attention block
	shortcut = x
	x = self.norm1(x)
	x = self.att(x) # Shape [batch_size, num_tokens, emb_size]
	x = self.drop_shortcut(x)
	x = x + shortcut # Add the original input back

	# Shortcut connection for feed-forward block
	shortcut = x
	x = self.norm2(x)
	x = self.ff(x)
	x = self.drop_shortcut(x)
	x = x + shortcut # Add the original input back

	return x


	class GPTModel(nn.Module):
	def __init__(self, cfg):
	super().__init__()
	self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
	self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
	self.drop_emb = nn.Dropout(cfg["drop_rate"])

	self.trf_blocks = nn.Sequential(
	*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

	self.final_norm = LayerNorm(cfg["emb_dim"])
	self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

	def forward(self, in_idx):
	batch_size, seq_len = in_idx.shape
	tok_embeds = self.tok_emb(in_idx)
	pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
	x = tok_embeds + pos_embeds # Shape [batch_size, num_tokens, emb_size]
	x = self.drop_emb(x)
	x = self.trf_blocks(x)
	x = self.final_norm(x)
	logits = self.out_head(x)
	return logits


	def generate_text_simple(model, idx, max_new_tokens, context_size):
	# idx is (B, T) array of indices in the current context
	for _ in range(max_new_tokens):

	# Crop current context if it exceeds the supported context size
	# E.g., if LLM supports only 5 tokens, and the context size is 10
	# then only the last 5 tokens are used as context
	idx_cond = idx[:, -context_size:]

	# Get the predictions
	with torch.no_grad():
	logits = model(idx_cond)

	# Focus only on the last time step
	# (batch, n_token, vocab_size) becomes (batch, vocab_size)
	logits = logits[:, -1, :]

	# Get the idx of the vocab entry with the highest logits value
	idx_next = torch.argmax(logits, dim=-1, keepdim=True) # (batch, 1)

	# Append sampled index to the running sequence
	idx = torch.cat((idx, idx_next), dim=1) # (batch, n_tokens+1)

	return idx


	def generate(model, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None):

	# For-loop is the same as before: Get logits, and only focus on last time step
	for _ in range(max_new_tokens):
	idx_cond = idx[:, -context_size:]
	with torch.no_grad():
	logits = model(idx_cond)
	logits = logits[:, -1, :]

	# New: Filter logits with top_k sampling
	if top_k is not None:
	# Keep only top_k values
	top_logits, _ = torch.topk(logits, top_k)
	min_val = top_logits[:, -1]
	logits = torch.where(logits < min_val, torch.tensor(float("-inf")).to(logits.device), logits)

	# New: Apply temperature scaling
	if temperature > 0.0:
	logits = logits / temperature

	# New (not in book): numerical stability tip to get equivalent results on mps device
	# subtract rowwise max before softmax
	logits = logits - logits.max(dim=-1, keepdim=True).values

	# Apply softmax to get probabilities
	probs = torch.softmax(logits, dim=-1) # (batch_size, context_len)

	# Sample from the distribution
	idx_next = torch.multinomial(probs, num_samples=1) # (batch_size, 1)

	# Otherwise same as before: get idx of the vocab entry with the highest logits value
	else:
	idx_next = torch.argmax(logits, dim=-1, keepdim=True) # (batch_size, 1)

	if idx_next == eos_id: # Stop generating early if end-of-sequence token is encountered and eos_id is specified
	break

	# Same as before: append sampled index to the running sequence
	idx = torch.cat((idx, idx_next), dim=1) # (batch_size, num_tokens+1)

	return idx


	def train_model_simple(model, train_loader, val_loader, optimizer, device, num_epochs,
	eval_freq, eval_iter, start_context, tokenizer):
	# Initialize lists to track losses and tokens seen
	train_losses, val_losses, track_tokens_seen = [], [], []
	tokens_seen, global_step = 0, -1

	# Main training loop
	for epoch in range(num_epochs):
	model.train() # Set model to training mode

	for input_batch, target_batch in train_loader:
	optimizer.zero_grad() # Reset loss gradients from previous batch iteration
	loss = calc_loss_batch(input_batch, target_batch, model, device)
	loss.backward() # Calculate loss gradients
	optimizer.step() # Update model weights using loss gradients
	tokens_seen += input_batch.numel()
	global_step += 1

	# Optional evaluation step
	if global_step % eval_freq == 0:
	train_loss, val_loss = evaluate_model(
	model, train_loader, val_loader, device, eval_iter)
	train_losses.append(train_loss)
	val_losses.append(val_loss)
	track_tokens_seen.append(tokens_seen)
	print(f"Ep {epoch+1} (Step {global_step:06d}): "
	f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")

	# Print a sample text after each epoch
	generate_and_print_sample(
	model, tokenizer, device, start_context
	)

	return train_losses, val_losses, track_tokens_seen


	def evaluate_model(model, train_loader, val_loader, device, eval_iter):
	model.eval()
	with torch.no_grad():
	train_loss = calc_loss_loader(train_loader, model, device, num_batches=eval_iter)
	val_loss = calc_loss_loader(val_loader, model, device, num_batches=eval_iter)
	model.train()
	return train_loss, val_loss


	def generate_and_print_sample(model, tokenizer, device, start_context):
	model.eval()
	context_size = model.pos_emb.weight.shape[0]
	encoded = text_to_token_ids(start_context, tokenizer).to(device)
	with torch.no_grad():
	token_ids = generate_text_simple(
	model=model, idx=encoded,
	max_new_tokens=50, context_size=context_size
	)
	decoded_text = token_ids_to_text(token_ids, tokenizer)
	print(decoded_text.replace("\n", " ")) # Compact print format
	model.train()


	def assign(left, right):
	if left.shape != right.shape:
	raise ValueError(f"Shape mismatch. Left: {left.shape}, Right: {right.shape}")
	return torch.nn.Parameter(torch.tensor(right))


	def text_to_token_ids(text, tokenizer):
	encoded = tokenizer.encode(text, allowed_special={"<\|endoftext\|>"})
	encoded_tensor = torch.tensor(encoded).unsqueeze(0) # add batch dimension
	return encoded_tensor


	def token_ids_to_text(token_ids, tokenizer):
	flat = token_ids.squeeze(0) # remove batch dimension
	return tokenizer.decode(flat.tolist())


	def calc_loss_batch(input_batch, target_batch, model, device):
	input_batch, target_batch = input_batch.to(device), target_batch.to(device)
	logits = model(input_batch)
	loss = torch.nn.functional.cross_entropy(logits.flatten(0, 1), target_batch.flatten())
	return loss


	def calc_loss_loader(data_loader, model, device, num_batches=None):
	total_loss = 0.
	if len(data_loader) == 0:
	return float("nan")
	elif num_batches is None:
	num_batches = len(data_loader)
	else:
	# Reduce the number of batches to match the total number of batches in the data loader
	# if num_batches exceeds the number of batches in the data loader
	num_batches = min(num_batches, len(data_loader))
	for i, (input_batch, target_batch) in enumerate(data_loader):
	if i < num_batches:
	loss = calc_loss_batch(input_batch, target_batch, model, device)
	total_loss += loss.item()
	else:
	break
	return total_loss / num_batches


	def plot_losses(epochs_seen, tokens_seen, train_losses, val_losses):
	fig, ax1 = plt.subplots(figsize=(5, 3))

	# Plot training and validation loss against epochs
	ax1.plot(epochs_seen, train_losses, label="Training loss")
	ax1.plot(epochs_seen, val_losses, linestyle="-.", label="Validation loss")
	ax1.set_xlabel("Epochs")
	ax1.set_ylabel("Loss")
	ax1.legend(loc="upper right")
	ax1.xaxis.set_major_locator(MaxNLocator(integer=True)) # only show integer labels on x-axis

	# Create a second x-axis for tokens seen
	ax2 = ax1.twiny() # Create a second x-axis that shares the same y-axis
	ax2.plot(tokens_seen, train_losses, alpha=0) # Invisible plot for aligning ticks
	ax2.set_xlabel("Tokens seen")

	fig.tight_layout() # Adjust layout to make room
	plt.savefig("loss-plot.pdf")
	plt.show()

	def main():
	GPT_CONFIG_124M = {
	"vocab_size": 50257, # Vocabulary size
	"context_length": 1024, # Context length
	"emb_dim": 768, # Embedding dimension
	"n_heads": 12, # Number of attention heads
	"n_layers": 12, # Number of layers
	"drop_rate": 0.1, # Dropout rate
	"qkv_bias": False # Query-Key-Value bias
	}

	torch.manual_seed(123)
	model = GPTModel(GPT_CONFIG_124M)
	model.eval() # disable dropout

	start_context = "Hi, there"

	tokenizer = tiktoken.get_encoding("gpt2")
	encoded = tokenizer.encode(start_context)
	encoded_tensor = torch.tensor(encoded).unsqueeze(0)

	print(f"\n{50'='}\n{22' '}IN\n{50*'='}")
	print("\nInput text:", start_context)
	print("Encoded input text:", encoded)
	print("encoded_tensor.shape:", encoded_tensor.shape)

	out = generate_text_simple(
	model=model,
	idx=encoded_tensor,
	max_new_tokens=10,
	context_size=GPT_CONFIG_124M["context_length"]
	)
	decoded_text = tokenizer.decode(out.squeeze(0).tolist())

	print(f"\n\n{50'='}\n{22' '}OUT\n{50*'='}")
	print("\nOutput:", out)
	print("Output length:", len(out[0]))
	print("Output text:", decoded_text)


	if __name__ == "__main__":
	main()