%%capture
!pip install torchinfo datasets transformers scikit-learn matplotlib watermark

%load_ext watermark
%watermark -a 'Chizoba Obasi' -v -p torch

Author: Chizoba Obasi

Python implementation: CPython
Python version       : 3.12.13
IPython version      : 7.34.0

torch: 2.10.0+cu128

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
from datasets import load_dataset

from sklearn.metrics import (
    accuracy_score, f1_score, precision_score, recall_score,
    classification_report, confusion_matrix, ConfusionMatrixDisplay
)

from torch.amp import GradScaler, autocast
from torchinfo import summary

import seaborn as sns
import matplotlib.pyplot as plt
# %matplotlib inline

import numpy as np
import math
import time
import os
import random
import re
import json

from collections import Counter

if torch.cuda.is_available():
    torch.backends.cudnn.deterministic = True

def set_all_seeds(seed):
    os.environ["PL_GLOBAL_SEED"] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

##########################
### SETTINGS
##########################

RANDOM_SEED = 123
set_all_seeds(RANDOM_SEED)

# Vocabulary (reuse from pre-training)
VOCABULARY_SIZE = 30000

# GPT-1 Architecture (must match pre-trained model)
D_MODEL = 768
N_LAYERS = 12
N_HEADS = 12
D_FF = 4 * D_MODEL
DROPOUT = 0.1
MAX_LEN = 64  # ✅ Optimized for MRPC: shorter sentences, faster attention

# Pre-training (Autoregressive LM on WikiText-2)
PRETRAIN_EPOCHS = 10  # Reduced for Colab time limits
PRETRAIN_BATCH_SIZE = 64
PRETRAIN_LR = 2.5e-4  # GPT-1 default
GRAD_ACCUM_STEPS = 8  # Effective batch = 512

# Fine-tuning (MRPC) — ✅ Use 100% of train data
FINETUNE_EPOCHS = 10  # MRPC is small; train longer with early stopping
FINETUNE_BATCH_SIZE = 16  # Memory-safe for Colab T4
FINETUNE_GRAD_ACCUM_STEPS = 4  # Effective batch = 64
FINETUNE_LR = 6.25e-5  # GPT-1 default for fine-tuning
AUX_LM_WEIGHT = 0.5  # λ = 0.5 per GPT-1 paper
PATIENCE = 3  # Early stopping patience

# MRPC-specific
NUM_CLASSES = 2  # Binary: Paraphrase (1) / Not Paraphrase (0)
LABEL_NAMES = ['Not Paraphrase', 'Paraphrase']

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {DEVICE}")
print(f"MRPC Classes: {NUM_CLASSES} → {LABEL_NAMES}")
print(f"MAX_LEN: {MAX_LEN} (optimized for MRPC sentence lengths)")

Device: cuda
MRPC Classes: 2 → ['Not Paraphrase', 'Paraphrase']
MAX_LEN: 64 (optimized for MRPC sentence lengths)

# ==========================================
# 1a. Load Datasets
# ==========================================

# Pre-training corpus (WikiText-2 as BookCorpus proxy)
print("Loading WikiText-2 for pre-training...")
wiki_dataset = load_dataset("wikitext", "wikitext-2-raw-v1")

print(f"WikiText-2 samples - train: {len(wiki_dataset['train'])}, "
      f"val: {len(wiki_dataset['validation'])}, "
      f"test: {len(wiki_dataset['test'])}")

# MRPC dataset for fine-tuning
print("\nLoading MRPC dataset...")
mrpc_dataset = load_dataset("glue", "mrpc")

print(f"MRPC samples - train: {len(mrpc_dataset['train'])}, "
      f"validation: {len(mrpc_dataset['validation'])}, "
      f"test: {len(mrpc_dataset['test'])}")

# Label distribution check
train_labels = [ex['label'] for ex in mrpc_dataset['train']]
label_counts = Counter(train_labels)
print(f"\nMRPC Train label distribution: {dict(label_counts)}")
print(f"Paraphrase ratio: {label_counts[1]/len(train_labels)*100:.1f}%")

Loading WikiText-2 for pre-training...

/usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_auth.py:104: UserWarning: 
Error while fetching `HF_TOKEN` secret value from your vault: 'Requesting secret HF_TOKEN timed out. Secrets can only be fetched when running from the Colab UI.'.
You are not authenticated with the Hugging Face Hub in this notebook.
If the error persists, please let us know by opening an issue on GitHub (https://github.com/huggingface/huggingface_hub/issues/new).
  warnings.warn(

Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
WARNING:huggingface_hub.utils._http:Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.

WikiText-2 samples - train: 36718, val: 3760, test: 4358

Loading MRPC dataset...

MRPC samples - train: 3668, validation: 408, test: 1725

MRPC Train label distribution: {1: 2474, 0: 1194}
Paraphrase ratio: 67.4%

# ==========================================
# 1b. Tokenizer & Vocabulary
# ==========================================

def tokenizer(text):
    """Simple word-level tokenizer with punctuation splitting."""
    if not isinstance(text, str):
        text = str(text)
    text = text.lower()
    text = re.sub(r'([.,!?;:\-])', r' \1 ', text)
    return text.split()

# GPT-1 (MRPC) Special Tokens
SPECIAL_TOKENS = ['[PAD]', '[UNK]', '[BOS]', '[EOS]', '[CLF]', '[SEP]']
PAD_IDX, UNK_IDX, BOS_IDX, EOS_IDX, CLF_IDX, SEP_IDX = range(len(SPECIAL_TOKENS))

# Build vocab from WikiText-2 (same as pre-training)
print("\nBuilding vocabulary from WikiText-2...")
word_counts = Counter()
for item in wiki_dataset['train']:
    text = item['text'].strip()
    if len(text) > 0 and isinstance(text, str):
        word_counts.update(tokenizer(text))

vocab = {word: i for i, word in enumerate(SPECIAL_TOKENS)}
for word, _ in word_counts.most_common(VOCABULARY_SIZE - len(SPECIAL_TOKENS)):
    if word not in vocab:
        vocab[word] = len(vocab)

inv_vocab = {v: k for k, v in vocab.items()}
ACTUAL_VOCAB_SIZE = len(vocab)

print(f"Vocabulary size: {ACTUAL_VOCAB_SIZE}")
print(f"Special tokens: {SPECIAL_TOKENS}")
print(f"Sample tokens: {list(vocab.keys())[5:20]}")

Building vocabulary from WikiText-2...
Vocabulary size: 30000
Special tokens: ['[PAD]', '[UNK]', '[BOS]', '[EOS]', '[CLF]', '[SEP]']
Sample tokens: ['[SEP]', 'the', ',', '.', 'of', 'and', '@', 'in', 'to', 'a', '=', '"', 'was', '-', 'on']

# ==========================================
# 1c. Pre-training Dataset (Causal LM)
# ==========================================

class GPTPretrainDataset(Dataset):
    """
    Autoregressive LM dataset for GPT-1 pre-training.
    Chunks text into fixed-length sequences with [BOS] and [EOS].
    Input:  [BOS] t1 t2 ... t_{n-1}
    Target:  t1  t2 t3 ... t_n  (shifted by 1)
    """
    def __init__(self, text_dataset, vocab, max_len=128):
        super().__init__()
        self.max_len = max_len
        self.vocab = vocab

        # Concatenate all text into a single token stream
        all_tokens = []
        for item in text_dataset:
            text = item['text'].strip()
            if len(text) > 10:
                tokens = [vocab.get(t, UNK_IDX) for t in tokenizer(text)]
                all_tokens.extend(tokens)

        # Chunk into sequences of max_len (with room for [BOS] and [EOS])
        chunk_size = max_len - 2  # Reserve space for [BOS] and [EOS]
        self.samples = []
        for i in range(0, len(all_tokens) - chunk_size, chunk_size):
            chunk = all_tokens[i:i + chunk_size]
            # Input:  [BOS] chunk
            # Target: chunk [EOS]
            input_ids = [BOS_IDX] + chunk
            target_ids = chunk + [EOS_IDX]
            self.samples.append((input_ids, target_ids))

        print(f"  GPTPretrainDataset: {len(self.samples)} samples from {len(all_tokens):,} tokens.")

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

    def __getitem__(self, idx):
        inp, tgt = self.samples[idx]
        return torch.tensor(inp, dtype=torch.long), torch.tensor(tgt, dtype=torch.long)


def pretrain_collate_fn(batch):
    inputs, targets = zip(*batch)
    inputs = pad_sequence(inputs, batch_first=True, padding_value=PAD_IDX)
    targets = pad_sequence(targets, batch_first=True, padding_value=-100)  # ignore padding in loss
    return inputs, targets

# Build Pre-training DataLoaders
pretrain_train_dataset = GPTPretrainDataset(wiki_dataset['train'], vocab, max_len=MAX_LEN)
pretrain_val_dataset   = GPTPretrainDataset(wiki_dataset['validation'], vocab, max_len=MAX_LEN)

pretrain_train_loader = DataLoader(
    pretrain_train_dataset, batch_size=PRETRAIN_BATCH_SIZE,
    shuffle=True, collate_fn=pretrain_collate_fn, num_workers=0
)
pretrain_val_loader = DataLoader(
    pretrain_val_dataset, batch_size=PRETRAIN_BATCH_SIZE,
    shuffle=False, collate_fn=pretrain_collate_fn, num_workers=0
)

# Sanity check
sample_in, sample_tgt = next(iter(pretrain_train_loader))
print(f"Batch shapes — input: {sample_in.shape}, target: {sample_tgt.shape}")
print(f"Train batches: {len(pretrain_train_loader)}, Val batches: {len(pretrain_val_loader)}")

# Show example
print(f"\nExample input:  {[inv_vocab.get(t.item(), '?') for t in sample_in[0][:10]]}")
print(f"Example target: {[inv_vocab.get(t.item(), '?') for t in sample_tgt[0][:10] if t.item() != -100]}")

  GPTPretrainDataset: 33964 samples from 2,105,796 tokens.
  GPTPretrainDataset: 3547 samples from 219,976 tokens.
Batch shapes — input: torch.Size([64, 63]), target: torch.Size([64, 63])
Train batches: 531, Val batches: 56

Example input:  ['[BOS]', 'his', '[UNK]', 'in', 'the', 'times', 'described', 'him', 'as', '"']
Example target: ['his', '[UNK]', 'in', 'the', 'times', 'described', 'him', 'as', '"', 'a']

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        # Fused Q, K, V projection: [d_model -> 3 * d_model]
        self.c_attn = nn.Linear(d_model, 3 * d_model)

        # Output projection
        self.c_proj = nn.Linear(d_model, d_model)
        self.c_proj.NANOGPT_SCALE_INIT = True  # Scaling flag for residual paths

        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)

        self.scale = 1.0 / math.sqrt(self.d_k)

    def forward(self, x, mask=None):
        batch_size, seq_len, d_model = x.size()

        # 1. Linear projection and split into Q, K, V
        qkv = self.c_attn(x)
        q, k, v = qkv.split(self.d_model, dim=2)

        # 2. Reshape for Multi-Head: [batch, n_heads, seq_len, d_k]
        q = q.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # 3. Scaled Dot-Product Attention
        att = (q @ k.transpose(-2, -1)) * self.scale

        if mask is not None:
            att = att.masked_fill(mask == 0, float('-inf'))

        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)

        # 4. Combine heads: [batch, seq_len, d_model]
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)

        # 5. Output projection and residual dropout
        y = self.resid_dropout(self.c_proj(y))

        return y

# ==========================================
# 2b. Optimized Feed-Forward Network
# ==========================================

class PositionwiseFeedForward(nn.Module):
    """GPT-1 Feed-Forward: Linear → GELU → Dropout → Linear → Dropout"""
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
        self.fc2.NANOGPT_SCALE_INIT = True

        self.activation_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = self.fc1(x)
        x = F.gelu(x)
        x = self.activation_dropout(x)
        x = self.fc2(x)
        return self.resid_dropout(x)

# ==========================================
# 2c. Decoder Layer (Causal Self-Attention Only)
# ==========================================

class DecoderLayer(nn.Module):
    """
    GPT-1 Decoder Layer with Pre-LayerNorm.

    Structure:
      x → [Self-Attn + Dropout] → +x → LayerNorm → [FFN + Dropout] → +x → LayerNorm → output

    Key differences from Transformer decoder:
    - NO cross-attention (decoder-only architecture)
    - Causal mask prevents attending to future tokens
    - Pre-LayerNorm (modern standard for stability)
    """
    def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout=dropout)
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout=dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        # Pre-Norm Self-Attention + Residual
        attn_out = self.self_attn(x, mask)
        x = x + self.norm1(attn_out)

        # Pre-Norm Feed-Forward + Residual
        ff_out = self.feed_forward(x)
        x = x + self.norm2(ff_out)

        return x

# ==========================================
# 2d. GPT-1 Model
# ==========================================

class GPT(nn.Module):
    """
    GPT-1: Generative Pre-trained Transformer (decoder-only).

    Key differences from BERT:
    - Causal (lower-triangular) mask for autoregressive prediction
    - No segment embeddings (single contiguous sequence)
    - LM head with weight tying to token embedding
    - Output at LAST token for classification (not [CLS] at position 0)
    """
    def __init__(self, vocab_size, d_model, n_layers, n_heads, d_ff,
                 dropout, max_len, pad_idx, device):
        super().__init__()
        self.n_layers = n_layers
        self.pad_idx = pad_idx
        self.d_model = d_model
        self.device = device

        # --- GPT Embeddings (Learned positional, NO segment) ---
        self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=pad_idx)
        self.position_embedding = nn.Embedding(max_len, d_model)
        self.embed_dropout = nn.Dropout(dropout)

        # --- Decoder Stack (causal self-attention only) ---
        self.decoder_layers = nn.ModuleList([
            DecoderLayer(d_model, n_heads, d_ff, dropout)
            for _ in range(n_layers)
        ])

        self.final_norm = nn.LayerNorm(d_model)

        # --- LM Head (weight-tied with token embedding) ---
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.lm_head.weight = self.token_embedding.weight  # Weight tying

        self.apply(self._init_weights)

    def _init_weights(self, module):
        """
        Initialize weights with GPT-1 / nanoGPT conventions.

        - Linear/Embedding: N(0, 0.02)
        - Linear biases: zeros
        - Residual projections: scaled by 1/sqrt(2*N_layers)
        - LayerNorm: ones (weight), zeros (bias)
        - Padding embeddings: zeros
        """
        if isinstance(module, nn.Linear):
            std = 0.02
            if hasattr(module, 'NANOGPT_SCALE_INIT'):
                std *= (2 * self.n_layers) ** -0.5
            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.padding_idx is not None:
                with torch.no_grad():
                    module.weight[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            torch.nn.init.ones_(module.weight)
            torch.nn.init.zeros_(module.bias)

    def make_causal_mask(self, input_ids):
        """
        Combined padding + causal mask.
        Causal: each position can only attend to itself and earlier positions.
        """
        batch_size, seq_len = input_ids.size()

        # Padding mask: [batch, 1, 1, seq_len]
        pad_mask = (input_ids != self.pad_idx).unsqueeze(1).unsqueeze(2)

        # Causal mask: [1, 1, seq_len, seq_len]
        causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=input_ids.device)).bool()
        causal_mask = causal_mask.unsqueeze(0).unsqueeze(0)

        return pad_mask & causal_mask

    def get_embeddings(self, input_ids):
        """GPT embedding = token + position (no segment), then dropout."""
        seq_len = input_ids.size(1)
        positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
        tok_emb = self.token_embedding(input_ids) * math.sqrt(self.d_model)
        pos_emb = self.position_embedding(positions)
        return self.embed_dropout(tok_emb + pos_emb)

    def forward(self, input_ids):
        """Forward pass. Returns lm_logits [batch, seq_len, vocab_size]."""
        mask = self.make_causal_mask(input_ids)
        x = self.get_embeddings(input_ids)
        for layer in self.decoder_layers:
            x = layer(x, mask)
        x = self.final_norm(x)
        return self.lm_head(x)  # [batch, seq_len, vocab_size]

    def get_hidden_states(self, input_ids):
        """Returns hidden states (for fine-tuning heads)."""
        mask = self.make_causal_mask(input_ids)
        x = self.get_embeddings(input_ids)
        for layer in self.decoder_layers:
            x = layer(x, mask)
        return self.final_norm(x)

def count_parameters(model):
    """Parameter breakdown for GPT-1."""
    def count(m):
      return sum(p.numel() for p in m.parameters() if p.requires_grad)
    emb = count(model.token_embedding) + count(model.position_embedding)
    dec = count(model.decoder_layers)
    norm = count(model.final_norm)
    # LM head is weight-tied, so 0 extra params
    total = count(model)

    print(f"{'Component':<25} | {'Parameters':<12}")
    print("-" * 42)
    print(f"{'Embeddings (tok+pos)':<25} | {emb:,}")
    print(f"{'Decoder Stack (N Layers)':<25} | {dec:,}")
    print(f"{'Final LayerNorm':<25} | {norm:,}")
    print(f"{'LM Head (weight-tied)':<25} | 0 (shared)")
    print("-" * 42)
    print(f"{'TOTAL TRAINABLE':<25} | {total:,}")

# ==========================================
# 3a. Pre-training Evaluation + Training Loop
# ==========================================

def evaluate_lm(model, data_loader, criterion, device):
    """Evaluate causal LM. Returns avg loss and perplexity."""
    model.eval()
    total_loss = 0.0
    total_tokens = 0

    with torch.no_grad():
        for inputs, targets in data_loader:
            inputs = inputs.to(device)
            targets = targets.to(device)
            logits = model(inputs)
            loss = criterion(logits.view(-1, logits.size(-1)), targets.view(-1))
            n_tokens = (targets != -100).sum().item()
            total_loss += loss.item() * n_tokens
            total_tokens += n_tokens

    avg_loss = total_loss / max(total_tokens, 1)
    ppl = math.exp(min(avg_loss, 100))
    return avg_loss, ppl


def pretrain_gpt(model, train_loader, val_loader, num_epochs, lr, device, patience=5, grad_accum_steps=8):
    """GPT-1 pre-training loop with causal LM objective."""
    optimizer = optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.999), weight_decay=0.01)

    warmup_steps = 2000
    total_steps = len(train_loader) * num_epochs
    scheduler = torch.optim.lr_scheduler.SequentialLR(
        optimizer,
        schedulers=[
            torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps),
            torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps - warmup_steps)
        ],
        milestones=[warmup_steps]
    )

    criterion = nn.CrossEntropyLoss(ignore_index=-100)
    use_amp = (device.type == 'cuda')
    scaler = GradScaler('cuda', enabled=use_amp)

    best_val_loss = float('inf')
    epochs_no_improve = 0
    start_time = time.time()

    history = {'train_loss': [], 'val_loss': [], 'train_ppl': [], 'val_ppl': [], 'lrs': []}

    print(f"Starting GPT-1 pre-training on {device}...")
    print(f"Epochs: {num_epochs} | Batches/epoch: {len(train_loader)} | LR: {lr} | Patience: {patience}")
    print("-" * 100)

    for epoch in range(num_epochs):
        model.train()
        epoch_loss = 0.0
        epoch_tokens = 0

        for step, (inputs, targets) in enumerate(train_loader):
            inputs = inputs.to(device, non_blocking=True)
            targets = targets.to(device, non_blocking=True)

            if step % grad_accum_steps == 0:
                optimizer.zero_grad(set_to_none=True)

            with autocast('cuda', enabled=use_amp):
                logits = model(inputs)
                loss = criterion(logits.view(-1, logits.size(-1)), targets.view(-1))
                loss = loss / grad_accum_steps

            scaler.scale(loss).backward()

            if (step + 1) % grad_accum_steps == 0:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                scaler.step(optimizer)
                scaler.update()
                scheduler.step()

            n_tokens = (targets != -100).sum().item()
            epoch_loss += loss.item() * grad_accum_steps * n_tokens
            epoch_tokens += n_tokens

        # --- Epoch stats ---
        t_loss = epoch_loss / max(epoch_tokens, 1)
        t_ppl = math.exp(min(t_loss, 100))
        v_loss, v_ppl = evaluate_lm(model, val_loader, criterion, device)

        # --- Early stopping ---
        if v_loss < best_val_loss:
            best_val_loss = v_loss
            torch.save(model.state_dict(), 'gpt1_pretrained.pth')
            epochs_no_improve = 0
            status = " * (Saved Best)"
        else:
            epochs_no_improve += 1
            status = f" (No improve: {epochs_no_improve}/{patience})"

        current_lr = optimizer.param_groups[0]['lr']
        elapsed = time.time() - start_time
        tok_per_sec = epoch_tokens / elapsed

        history['train_loss'].append(t_loss)
        history['val_loss'].append(v_loss)
        history['train_ppl'].append(t_ppl)
        history['val_ppl'].append(v_ppl)
        history['lrs'].append(current_lr)

        print(f"Epoch {epoch+1:03d}/{num_epochs:03d} | "
              f"Loss (T/V): {t_loss:.4f} / {v_loss:.4f} | "
              f"PPL (T/V): {t_ppl:.2f} / {v_ppl:.2f} | "
              f"tok/s: {tok_per_sec:07.0f} | "
              f"LR: {current_lr:.6f}{status}")

        if epochs_no_improve >= patience:
            print(f"\nEarly stopping triggered.")
            break

    print(f'\nDone! Total Time: {(time.time() - start_time) / 60:.2f} min')
    model.load_state_dict(torch.load('gpt1_pretrained.pth', map_location=device))
    return history

# ==========================================
# 3b. Initialize & Pre-train GPT-1
# ==========================================

set_all_seeds(RANDOM_SEED)

print("\n" + "="*70)
print("PHASE 1: Initializing GPT-1 for Pre-training")
print("="*70)

gpt_model = GPT(
    vocab_size=ACTUAL_VOCAB_SIZE,
    d_model=D_MODEL,
    n_layers=N_LAYERS,
    n_heads=N_HEADS,
    d_ff=D_FF,
    dropout=DROPOUT,
    max_len=MAX_LEN,
    pad_idx=PAD_IDX,
    device=DEVICE
).to(DEVICE)

count_parameters(gpt_model)

======================================================================
PHASE 1: Initializing GPT-1 for Pre-training
======================================================================
Component                 | Parameters  
------------------------------------------
Embeddings (tok+pos)      | 23,089,152
Decoder Stack (N Layers)  | 85,054,464
Final LayerNorm           | 1,536
LM Head (weight-tied)     | 0 (shared)
------------------------------------------
TOTAL TRAINABLE           | 108,145,152

# Run Pre-training (Phase 1)
print("\n" + "="*70)
print("PHASE 1: Starting Pre-training on WikiText-2")
print("="*70)

pretrain_history = pretrain_gpt(
    model=gpt_model,
    train_loader=pretrain_train_loader,
    val_loader=pretrain_val_loader,
    num_epochs=PRETRAIN_EPOCHS,
    lr=PRETRAIN_LR,
    device=DEVICE,
    patience=PATIENCE,
    grad_accum_steps=GRAD_ACCUM_STEPS
)

======================================================================
PHASE 1: Starting Pre-training on WikiText-2
======================================================================
Starting GPT-1 pre-training on cuda...
Epochs: 10 | Batches/epoch: 531 | LR: 0.00025 | Patience: 3
----------------------------------------------------------------------------------------------------
Epoch 001/010 | Loss (T/V): 9.2494 / 8.4147 | PPL (T/V): 10398.02 / 4512.87 | tok/s: 0012869 | LR: 0.000008 * (Saved Best)
Epoch 002/010 | Loss (T/V): 8.0756 / 7.3275 | PPL (T/V): 3215.21 / 1521.63 | tok/s: 0006132 | LR: 0.000017 * (Saved Best)
Epoch 003/010 | Loss (T/V): 7.1481 / 6.6281 | PPL (T/V): 1271.72 / 756.05 | tok/s: 0004023 | LR: 0.000025 * (Saved Best)
Epoch 004/010 | Loss (T/V): 6.5819 / 6.1971 | PPL (T/V): 721.93 / 491.30 | tok/s: 0002995 | LR: 0.000033 * (Saved Best)
Epoch 005/010 | Loss (T/V): 6.2815 / 5.9776 | PPL (T/V): 534.57 / 394.47 | tok/s: 0002382 | LR: 0.000041 * (Saved Best)
Epoch 006/010 | Loss (T/V): 6.0635 / 5.7652 | PPL (T/V): 429.88 / 319.00 | tok/s: 0001979 | LR: 0.000050 * (Saved Best)
Epoch 007/010 | Loss (T/V): 5.8419 / 5.5905 | PPL (T/V): 344.43 / 267.88 | tok/s: 0001693 | LR: 0.000058 * (Saved Best)
Epoch 008/010 | Loss (T/V): 5.6380 / 5.4334 | PPL (T/V): 280.89 / 228.94 | tok/s: 0001478 | LR: 0.000066 * (Saved Best)
Epoch 009/010 | Loss (T/V): 5.4565 / 5.2993 | PPL (T/V): 234.27 / 200.19 | tok/s: 0001312 | LR: 0.000074 * (Saved Best)
Epoch 010/010 | Loss (T/V): 5.2932 / 5.1836 | PPL (T/V): 198.97 / 178.32 | tok/s: 0001180 | LR: 0.000083 * (Saved Best)

Done! Total Time: 30.21 min

# ==========================================
# 3c. Plot Pre-training Curves
# ==========================================

fig, axes = plt.subplots(1, 3, figsize=(16, 4))
epochs = range(1, len(pretrain_history['train_loss']) + 1)

# Loss
axes[0].plot(epochs, pretrain_history['train_loss'], 'o-', label='Train', color='#1f77b4')
axes[0].plot(epochs, pretrain_history['val_loss'], 's--', label='Val', color='#ff7f0e')
axes[0].set_title('Causal LM Loss')
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Loss')
axes[0].legend()
axes[0].grid(True, alpha=0.3)

# Perplexity
axes[1].plot(epochs, pretrain_history['train_ppl'], 'o-', label='Train PPL', color='#2ca02c')
axes[1].plot(epochs, pretrain_history['val_ppl'], 's--', label='Val PPL', color='#d62728')
axes[1].set_title('Perplexity')
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('PPL')
axes[1].legend()
axes[1].grid(True, alpha=0.3)

# LR
axes[2].plot(epochs, pretrain_history['lrs'], label='LR', color='black', linewidth=2)
axes[2].set_title('LR Schedule')
axes[2].set_xlabel('Epoch')
axes[2].set_ylabel('Rate')
axes[2].ticklabel_format(style='sci', axis='y', scilimits=(0,0))
axes[2].grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('gpt1_pretraining_curves.png', dpi=150, bbox_inches='tight')
plt.show()

print("✓ Pre-training curves saved to 'gpt1_pretraining_curves.png'")

✓ Pre-training curves saved to 'gpt1_pretraining_curves.png'

# ==========================================
# 4a. MRPC Dataset Preparation for GPT-1 (Colab-Optimized)
# ==========================================

"""
Input format for GPT-1 MRPC fine-tuning:

    [BOS] sentence₁ [SEP] sentence₂ [CLF]

Key design choices:
- [SEP] clearly demarcates sentence₁ vs sentence₂
- [CLF] at the END captures full context (GPT reads left-to-right)
- LM targets: shifted input for auxiliary language modeling loss
- Labels: 0=Not Paraphrase, 1=Paraphrase

🚀 COLAB OPTIMIZATIONS:
- ✅ Use 100% of MRPC train data (~3.7k examples — already small!)
- MAX_LEN = 64 (fits ~95% of MRPC examples, 4x faster attention)
- Gradient accumulation (maintain effective batch size)
- Checkpointing every epoch (survive Colab disconnects)
- Progress tracking (batches/sec, ETA)
"""

# ==========================================
# 4a-i. Verify Sequence Length Distribution for MRPC
# ==========================================

def estimate_seq_length(example):
    """Estimate token length for an MRPC example."""
    sent1_len = len(tokenizer(example['sentence1']))
    sent2_len = len(tokenizer(example['sentence2']))
    return sent1_len + sent2_len + 4  # +4 for [BOS], [SEP], [CLF]

# Sample all examples to check length distribution (MRPC is small)
length_samples = [estimate_seq_length(ex) for ex in mrpc_dataset['train']]
avg_len = sum(length_samples) / len(length_samples)
pct_under_64 = sum(1 for l in length_samples if l <= 64) / len(length_samples) * 100
pct_under_128 = sum(1 for l in length_samples if l <= 128) / len(length_samples) * 100

print(f"\n📊 MRPC Sequence Length Analysis (full train set):")
print(f"  Total examples: {len(length_samples):,}")
print(f"  Average length: {avg_len:.1f} tokens")
print(f"  Fit in MAX_LEN=64:  {pct_under_64:.1f}% ✅")
print(f"  Fit in MAX_LEN=128: {pct_under_128:.1f}%")

if pct_under_64 < 90:
    print(f"  ⚠️  Warning: Some examples will be truncated at MAX_LEN=64")

📊 MRPC Sequence Length Analysis (full train set):
  Total examples: 3,668
  Average length: 50.6 tokens
  Fit in MAX_LEN=64:  88.2% ✅
  Fit in MAX_LEN=128: 100.0%
  ⚠️  Warning: Some examples will be truncated at MAX_LEN=64

# ==========================================
# 4a-ii. MRPC Dataset Class for GPT-1
# ==========================================

class MRPCDataset(Dataset):
    """
    MRPC dataset for GPT-1 fine-tuning.

    Input format: [BOS] sentence₁ [SEP] sentence₂ [CLF]

    Key features:
    - [CLF] token at END provides representation for binary classification
    - LM targets: shifted input for auxiliary language modeling loss
    - Dynamic truncation to fit max_len while preserving both sentences

    Returns:
        input_ids:     [BOS] sent1 [SEP] sent2 [CLF]
        lm_targets:    sent1 [SEP] sent2 [CLF] + padding (-100)
        label:         0=Not Paraphrase, 1=Paraphrase
        clf_position:  index of [CLF] token (for extracting hidden state)
    """
    def __init__(self, hf_dataset, vocab, tokenizer, max_len=64):
        self.data = hf_dataset
        self.vocab = vocab
        self.tokenizer = tokenizer
        self.max_len = max_len

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

    def __getitem__(self, idx):
        item = self.data[idx]
        sentence1 = item['sentence1']
        sentence2 = item['sentence2']
        label = item['label']  # 0=not paraphrase, 1=paraphrase

        # Tokenize sentences
        sent1_tokens = self.tokenizer(sentence1)
        sent2_tokens = self.tokenizer(sentence2)

        # Convert to IDs
        sent1_ids = [self.vocab.get(t, UNK_IDX) for t in sent1_tokens]
        sent2_ids = [self.vocab.get(t, UNK_IDX) for t in sent2_tokens]

        # Truncate to fit max_len: [BOS] + sent1 + [SEP] + sent2 + [CLF]
        max_content = self.max_len - 4  # Reserve 4 slots for special tokens

        # Dynamic truncation: allocate space proportionally
        sent1_max = max_content // 2
        sent2_max = max_content - sent1_max

        sent1_ids = sent1_ids[:sent1_max]
        sent2_ids = sent2_ids[:sent2_max]

        # Build input: [BOS] sentence₁ [SEP] sentence₂ [CLF]
        input_ids = [BOS_IDX] + sent1_ids + [SEP_IDX] + sent2_ids + [CLF_IDX]

        # LM targets: shifted by 1 (autoregressive prediction)
        lm_targets = input_ids[1:]  # Remove BOS from start

        # Pad lm_targets to match input_ids length (for batch collation)
        while len(lm_targets) < len(input_ids):
            lm_targets.append(-100)  # -100 = ignore in CrossEntropyLoss
        lm_targets = lm_targets[:len(input_ids)]

        # Position of [CLF] token (last position) for classification head
        clf_position = len(input_ids) - 1

        return (
            torch.tensor(input_ids, dtype=torch.long),
            torch.tensor(lm_targets, dtype=torch.long),
            torch.tensor(label, dtype=torch.long),
            torch.tensor(clf_position, dtype=torch.long),
        )


# ==========================================
# 4a-iii. Collate Function for Batching
# ==========================================

def mrpc_collate_fn(batch):
    """Collate function for MRPC batches with padding."""
    input_ids, lm_targets, labels, clf_positions = zip(*batch)
    input_ids = pad_sequence(input_ids, batch_first=True, padding_value=PAD_IDX)
    lm_targets = pad_sequence(lm_targets, batch_first=True, padding_value=-100)
    labels = torch.stack(labels)
    clf_positions = torch.stack(clf_positions)
    return input_ids, lm_targets, labels, clf_positions

# ==========================================
# 4a-iv. Build MRPC DataLoaders
# ==========================================

print(f"\n{'='*70}")
print("Building MRPC DataLoaders (100% of train data)")
print(f"{'='*70}")
print(f"  Train: {len(mrpc_dataset['train']):,} samples (100%)")
print(f"  Val:   {len(mrpc_dataset['validation']):,} samples (100%)")
print(f"  Test:  {len(mrpc_dataset['test']):,} samples (100%)")
print(f"  MAX_LEN: {MAX_LEN} (optimized for MRPC)")
print(f"  Batch size: {FINETUNE_BATCH_SIZE}")
print(f"  Gradient accumulation: {FINETUNE_GRAD_ACCUM_STEPS}")
print(f"  Effective batch size: {FINETUNE_BATCH_SIZE * FINETUNE_GRAD_ACCUM_STEPS}")
print(f"{'='*70}\n")

mrpc_train_ds = MRPCDataset(mrpc_dataset['train'], vocab, tokenizer, max_len=MAX_LEN)
mrpc_val_ds = MRPCDataset(mrpc_dataset['validation'], vocab, tokenizer, max_len=MAX_LEN)
mrpc_test_ds = MRPCDataset(mrpc_dataset['test'], vocab, tokenizer, max_len=MAX_LEN)

mrpc_train_loader = DataLoader(
    mrpc_train_ds, batch_size=FINETUNE_BATCH_SIZE,
    shuffle=True, collate_fn=mrpc_collate_fn, num_workers=0, pin_memory=True
)
mrpc_val_loader = DataLoader(
    mrpc_val_ds, batch_size=FINETUNE_BATCH_SIZE,
    shuffle=False, collate_fn=mrpc_collate_fn, num_workers=0, pin_memory=True
)
mrpc_test_loader = DataLoader(
    mrpc_test_ds, batch_size=FINETUNE_BATCH_SIZE,
    shuffle=False, collate_fn=mrpc_collate_fn, num_workers=0, pin_memory=True
)

print(f"✓ DataLoaders created:")
print(f"  Train: {len(mrpc_train_loader)} batches × {FINETUNE_BATCH_SIZE} = ~{len(mrpc_train_ds):,} samples")
print(f"  Val:   {len(mrpc_val_loader)} batches × {FINETUNE_BATCH_SIZE} = ~{len(mrpc_val_ds):,} samples")
print(f"  Test:  {len(mrpc_test_loader)} batches × {FINETUNE_BATCH_SIZE} = ~{len(mrpc_test_ds):,} samples")

======================================================================
Building MRPC DataLoaders (100% of train data)
======================================================================
  Train: 3,668 samples (100%)
  Val:   408 samples (100%)
  Test:  1,725 samples (100%)
  MAX_LEN: 64 (optimized for MRPC)
  Batch size: 16
  Gradient accumulation: 4
  Effective batch size: 64
======================================================================

✓ DataLoaders created:
  Train: 230 batches × 16 = ~3,668 samples
  Val:   26 batches × 16 = ~408 samples
  Test:  108 batches × 16 = ~1,725 samples

# ==========================================
# 4a-v. Sanity Check & Debug Output
# ==========================================

print(f"\n{'='*70}")
print("SANITY CHECK: First MRPC Batch")
print(f"{'='*70}")

sample_in, sample_tgt, sample_lbl, sample_pos = next(iter(mrpc_train_loader))

print(f"Batch shapes:")
print(f"  input_ids:     {sample_in.shape}     [batch, seq_len]")
print(f"  lm_targets:    {sample_tgt.shape}    [batch, seq_len]")
print(f"  labels:        {sample_lbl.shape}    [batch]")
print(f"  clf_positions: {sample_pos.shape}    [batch]")

print(f"\nFirst example decoded:")
print(f"  Input tokens:  {[inv_vocab.get(t.item(), '?') for t in sample_in[0][:20]]}")
print(f"  LM targets:    {[inv_vocab.get(t.item(), '?') if t.item()!=-100 else '[IGN]' for t in sample_tgt[0][:20]]}")
print(f"  Label:         {sample_lbl[0].item()} → {LABEL_NAMES[sample_lbl[0].item()]}")
print(f"  CLF position:  {sample_pos[0].item()}")

# Verify [CLF] is at expected position
clf_token = sample_in[0, sample_pos[0].item()].item()
print(f"  Token at CLF pos: {inv_vocab.get(clf_token, '?')} (expected: [CLF])")

# Verify LM target alignment
print(f"\nLM target alignment check:")
inp_trim = sample_in[0, 1:]
tgt_trim = sample_tgt[0, :-1]
mask = (tgt_trim != -100)
aligned = torch.all(inp_trim[mask] == tgt_trim[mask])
print(f"  input_ids[1:] == lm_targets[:-1] (ignoring padding): {'✓ PASS' if aligned else '✗ FAIL'}")

print(f"\n{'='*70}")

======================================================================
SANITY CHECK: First MRPC Batch
======================================================================
Batch shapes:
  input_ids:     torch.Size([16, 63])     [batch, seq_len]
  lm_targets:    torch.Size([16, 63])    [batch, seq_len]
  labels:        torch.Size([16])    [batch]
  clf_positions: torch.Size([16])    [batch]

First example decoded:
  Input tokens:  ['[BOS]', 'the', 'government', 'defeated', 'the', 'rebel', 'motion', 'by', '[UNK]', 'votes', 'to', '117', 'in', 'the', '[UNK]', '-', 'seat', 'house', 'of', 'commons']
  LM targets:    ['the', 'government', 'defeated', 'the', 'rebel', 'motion', 'by', '[UNK]', 'votes', 'to', '117', 'in', 'the', '[UNK]', '-', 'seat', 'house', 'of', 'commons', '.']
  Label:         0 → Not Paraphrase
  CLF position:  37
  Token at CLF pos: [CLF] (expected: [CLF])

LM target alignment check:
  input_ids[1:] == lm_targets[:-1] (ignoring padding): ✓ PASS

======================================================================

# ==========================================
# 4b. GPT-1 Classifier Wrapper for MRPC (Binary)
# ==========================================

class GPTClassifier(nn.Module):
    """
    GPT-1 Classifier for MRPC (binary paraphrase detection).
    Uses [CLF] token (last position) hidden state for 2-way classification.
    Also returns LM logits for the auxiliary LM loss.
    """
    def __init__(self, gpt_model, d_model, num_classes, dropout=0.1):
        super().__init__()
        self.gpt = gpt_model
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(d_model, num_classes)  # ✅ Binary: 2 output classes
        )

    def forward(self, input_ids, clf_positions):
        # Get hidden states from the full decoder
        hidden = self.gpt.get_hidden_states(input_ids)  # [batch, seq_len, d_model]

        # LM logits (for auxiliary loss)
        lm_logits = self.gpt.lm_head(hidden)  # [batch, seq_len, vocab_size]

        # Classification: extract [CLF] token hidden state per sample
        batch_size = input_ids.size(0)
        clf_hidden = hidden[torch.arange(batch_size, device=input_ids.device), clf_positions]
        cls_logits = self.classifier(clf_hidden)  # [batch, num_classes]

        return cls_logits, lm_logits

# ==========================================
# 4c. Fine-tuning Training & Evaluation Loop (MRPC)
# ==========================================

def evaluate_mrpc(model, data_loader, task_criterion, lm_criterion, lm_weight, device, num_classes):
    """Evaluate MRPC fine-tuning. Returns loss, accuracy, F1, predictions."""
    model.eval()
    total_loss = 0.0
    all_preds, all_labels = [], []

    with torch.no_grad():
        for input_ids, lm_targets, labels, clf_positions in data_loader:
            input_ids = input_ids.to(device)
            lm_targets = lm_targets.to(device)
            labels = labels.to(device)
            clf_positions = clf_positions.to(device)

            cls_logits, lm_logits = model(input_ids, clf_positions)

            task_loss = task_criterion(cls_logits, labels)
            lm_loss = lm_criterion(lm_logits.view(-1, lm_logits.size(-1)), lm_targets.view(-1))
            loss = task_loss + lm_weight * lm_loss

            total_loss += loss.item() * labels.size(0)
            preds = cls_logits.argmax(dim=-1)
            all_preds.extend(preds.cpu().tolist())
            all_labels.extend(labels.cpu().tolist())

    avg_loss = total_loss / len(all_labels)
    acc = accuracy_score(all_labels, all_preds)
    f1_binary = f1_score(all_labels, all_preds, average='binary')  # ✅ MRPC standard
    precision = precision_score(all_labels, all_preds, average='binary')
    recall = recall_score(all_labels, all_preds, average='binary')

    return avg_loss, acc, f1_binary, precision, recall, all_preds, all_labels


def finetune_mrpc(model, train_loader, val_loader, num_epochs, lr, lm_weight,
                  device, patience=3, num_classes=2, grad_accum_steps=4,
                  checkpoint_path='gpt1_mrpc_checkpoint.pth'):
    """
    🚀 Optimized GPT-1 fine-tuning for MRPC with auxiliary LM loss.

    L_total = L_task + λ * L_LM (λ = 0.5 per the paper).
    """
    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)

    warmup_steps = min(500, len(train_loader) // 2)  # Cap warmup for small MRPC dataset
    total_steps = len(train_loader) * num_epochs
    scheduler = torch.optim.lr_scheduler.SequentialLR(
        optimizer,
        schedulers=[
            torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps),
            torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps - warmup_steps)
        ],
        milestones=[warmup_steps]
    )

    task_criterion = nn.CrossEntropyLoss()
    lm_criterion = nn.CrossEntropyLoss(ignore_index=-100)

    use_amp = (device.type == 'cuda')
    scaler = GradScaler('cuda', enabled=use_amp)

    if use_amp:
        print(f"✓ AMP (Mixed Precision) enabled")

    history = {
        'train_loss': [], 'val_loss': [],
        'val_acc': [], 'val_f1': [], 'val_precision': [], 'val_recall': [],
        'lrs': [], 'epoch_time': []
    }
    best_val_f1 = 0.0
    epochs_no_improve = 0
    start_time = time.time()

    print(f"\n{'='*90}")
    print(f"Fine-tuning GPT-1 on MRPC ({device})...")
    print(f"{'='*90}")
    print(f"  Epochs: {num_epochs} | LR: {lr} | Aux LM weight: {lm_weight} | Patience: {patience}")
    print(f"  Gradient accumulation: {grad_accum_steps} steps")
    print(f"  Checkpoint path: {checkpoint_path}")
    print(f"{'='*90}\n")

    for epoch in range(num_epochs):
        model.train()
        train_loss = 0.0
        epoch_start = time.time()
        batches_processed = 0
        optimizer.zero_grad(set_to_none=True)

        for step, (input_ids, lm_targets, labels, clf_positions) in enumerate(train_loader):
            input_ids = input_ids.to(device, non_blocking=True)
            lm_targets = lm_targets.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)
            clf_positions = clf_positions.to(device, non_blocking=True)

            with autocast('cuda', enabled=use_amp):
                cls_logits, lm_logits = model(input_ids, clf_positions)
                task_loss = task_criterion(cls_logits, labels)
                lm_loss = lm_criterion(lm_logits.view(-1, lm_logits.size(-1)), lm_targets.view(-1))
                loss = task_loss + lm_weight * lm_loss
                loss = loss / grad_accum_steps

            scaler.scale(loss).backward()

            if (step + 1) % grad_accum_steps == 0:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                scaler.step(optimizer)
                scaler.update()
                scheduler.step()
                optimizer.zero_grad(set_to_none=True)

            train_loss += loss.item() * grad_accum_steps
            batches_processed += 1

            # Progress tracking (every 50 batches for small MRPC dataset)
            if batches_processed % 50 == 0 and batches_processed > 0:
                elapsed = time.time() - epoch_start
                batches_per_sec = batches_processed / max(elapsed, 0.001)
                remaining_batches = len(train_loader) - batches_processed
                eta_min = remaining_batches / max(batches_per_sec, 0.001) / 60
                print(f"  [{batches_processed}/{len(train_loader)}] "
                      f"{batches_per_sec:.1f} batch/s | ETA: {eta_min:.1f} min")

        # Evaluate
        epoch_time = time.time() - epoch_start
        avg_train_loss = train_loss / len(train_loader)
        val_loss, val_acc, val_f1, val_prec, val_rec, _, _ = evaluate_mrpc(
            model, val_loader, task_criterion, lm_criterion, lm_weight, device, num_classes
        )

        current_lr = optimizer.param_groups[0]['lr']
        history['train_loss'].append(avg_train_loss)
        history['val_loss'].append(val_loss)
        history['val_acc'].append(val_acc)
        history['val_f1'].append(val_f1)
        history['val_precision'].append(val_prec)
        history['val_recall'].append(val_rec)
        history['lrs'].append(current_lr)
        history['epoch_time'].append(epoch_time)

        # Save checkpoint every epoch (survive Colab disconnects)
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict(),
            'val_f1': val_f1,
            'history': history,
        }
        torch.save(checkpoint, checkpoint_path)

        # Track best model by F1 (MRPC standard metric)
        if val_f1 > best_val_f1:
            best_val_f1 = val_f1
            torch.save(model.state_dict(), 'best_gpt1_mrpc.pth')
            epochs_no_improve = 0
            status = " * (Saved Best)"
        else:
            epochs_no_improve += 1
            status = f" (No improve: {epochs_no_improve}/{patience})"

        print(f"\nEpoch {epoch+1:02d}/{num_epochs} | "
              f"Time: {epoch_time/60:.1f} min | "
              f"Loss (T/V): {avg_train_loss:.4f} / {val_loss:.4f} | "
              f"Acc: {val_acc*100:.2f}% | F1: {val_f1*100:.2f}% | "
              f"P/R: {val_prec*100:.1f}%/{val_rec*100:.1f}% | "
              f"LR: {current_lr:.2e}{status}")
        print(f"  ✓ Checkpoint saved: {checkpoint_path}")

        if epochs_no_improve >= patience:
            print(f"\n⏹️  Early stopping triggered.")
            break

    total_time = time.time() - start_time
    print(f"\n{'='*90}")
    print(f"Fine-tuning complete! Total time: {total_time/60:.1f} min")
    print(f"Best Val F1 (MRPC metric): {best_val_f1*100:.2f}%")
    print(f"{'='*90}")

    # Load best model
    model.load_state_dict(torch.load('best_gpt1_mrpc.pth', map_location=device))
    return history


# ==========================================
# 4d. Resume from Checkpoint (Optional)
# ==========================================

def load_checkpoint(model, optimizer, scheduler, checkpoint_path, device):
    """Resume training from a saved checkpoint."""
    if not os.path.exists(checkpoint_path):
        print(f"⚠️  No checkpoint found at {checkpoint_path}")
        return 0, None

    checkpoint = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    scheduler.load_state_dict(checkpoint['scheduler_state_dict'])

    start_epoch = checkpoint['epoch'] + 1
    history = checkpoint['history']

    print(f"✓ Resumed from epoch {start_epoch} (val_acc: {checkpoint['val_acc']*100:.2f}%)")
    return start_epoch, history

# ==========================================
# 4d. Initialize & Fine-tune on MRPC (Phase 2)
# ==========================================

set_all_seeds(RANDOM_SEED)

# Load pre-trained weights from Phase 1
assert os.path.exists('gpt1_pretrained.pth'), "⚠ Phase 2 Error: Pre-trained weights not found!"
print("\n✓ Loading pre-trained GPT-1 weights from Phase 1...")
gpt_model.load_state_dict(torch.load('gpt1_pretrained.pth', map_location=DEVICE))

# Wrap with binary classification head for MRPC
classifier = GPTClassifier(
    gpt_model=gpt_model,
    d_model=D_MODEL,
    num_classes=NUM_CLASSES,  # ✅ Binary: 2 classes
    dropout=DROPOUT
).to(DEVICE)

print(f"Classifier parameters: {sum(p.numel() for p in classifier.parameters() if p.requires_grad):,}")
print(f"Classifier head: {sum(p.numel() for p in classifier.classifier.parameters()):,} (added on top of pre-trained GPT)")

✓ Loading pre-trained GPT-1 weights from Phase 1...
Classifier parameters: 108,146,690
Classifier head: 1,538 (added on top of pre-trained GPT)

# ==========================================
# Pre-Fine-Tuning Checklist (MRPC)
# ==========================================

print("\n" + "=" * 75)
print("🚀 GPT-1 MRPC FINE-TUNING: READY CHECKLIST")
print("=" * 75)

# Hardware
gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
print(f"✓ Device: {gpu_name}")

# Data
print(f"\n✓ Data splits (100% of MRPC):")
print(f"  └─ Train: {len(mrpc_dataset['train']):,} samples")
print(f"  └─ Val:   {len(mrpc_dataset['validation']):,} samples")
print(f"  └─ Test:  {len(mrpc_dataset['test']):,} samples")

# Architecture
print(f"\n✓ Model config:")
print(f"  └─ MAX_LEN: {MAX_LEN} (optimized for MRPC)")
print(f"  └─ Batch: {FINETUNE_BATCH_SIZE} (effective: {FINETUNE_BATCH_SIZE * FINETUNE_GRAD_ACCUM_STEPS})")
print(f"  └─ Epochs: {FINETUNE_EPOCHS} | Patience: {PATIENCE}")
print(f"  └─ Aux LM weight (λ): {AUX_LM_WEIGHT}")

# Time estimate (MRPC is small!)
TIME_PER_TRAIN_BATCH_SEC = 0.025  # ~25ms on Colab T4, MAX_LEN=64, small batches
TIME_PER_EVAL_BATCH_SEC = 0.010
train_batches = len(mrpc_dataset['train']) / FINETUNE_BATCH_SIZE
val_batches = len(mrpc_dataset['validation']) / FINETUNE_BATCH_SIZE
train_sec = train_batches * TIME_PER_TRAIN_BATCH_SEC * FINETUNE_EPOCHS
val_sec = val_batches * TIME_PER_EVAL_BATCH_SEC * FINETUNE_EPOCHS
total_min = (train_sec + val_sec) / 60

print(f"\n✓ Estimated fine-tuning time: ~{total_min:.1f} minutes")
print(f"  └─ Training: {train_sec/60:.1f} min | Validation: {val_sec/60:.1f} min")

# Files
print(f"\n✓ Checkpointing:")
print(f"  └─ Best model: best_gpt1_mrpc.pth")
print(f"  └─ Epoch checkpoints: gpt1_mrpc_checkpoint.pth")

print("\n" + "=" * 75)
print("✨ ALL SET — STARTING MRPC FINE-TUNING!")
print("=" * 75)

===========================================================================
🚀 GPT-1 MRPC FINE-TUNING: READY CHECKLIST
===========================================================================
✓ Device: Tesla T4

✓ Data splits (100% of MRPC):
  └─ Train: 3,668 samples
  └─ Val:   408 samples
  └─ Test:  1,725 samples

✓ Model config:
  └─ MAX_LEN: 64 (optimized for MRPC)
  └─ Batch: 16 (effective: 64)
  └─ Epochs: 10 | Patience: 3
  └─ Aux LM weight (λ): 0.5

✓ Estimated fine-tuning time: ~1.0 minutes
  └─ Training: 1.0 min | Validation: 0.0 min

✓ Checkpointing:
  └─ Best model: best_gpt1_mrpc.pth
  └─ Epoch checkpoints: gpt1_mrpc_checkpoint.pth

===========================================================================
✨ ALL SET — STARTING MRPC FINE-TUNING!
===========================================================================

# ==========================================
# Run Fine-tuning (Phase 2)
# ==========================================

CHECKPOINT_PATH = 'gpt1_mrpc_checkpoint.pth'

# Check for existing checkpoint (resume capability)
start_epoch = 0
finetune_history = None

if os.path.exists(CHECKPOINT_PATH):
    print(f"\n📁 Found existing checkpoint: {CHECKPOINT_PATH}")
    response = input("Resume from checkpoint? (y/n): ").strip().lower()
    if response == 'y':
        optimizer = optim.AdamW(classifier.parameters(), lr=FINETUNE_LR, weight_decay=0.01)
        warmup_steps = min(500, len(mrpc_train_loader) // 2)
        total_steps = len(mrpc_train_loader) * FINETUNE_EPOCHS
        scheduler = torch.optim.lr_scheduler.SequentialLR(
            optimizer,
            schedulers=[
                torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps),
                torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps - warmup_steps)
            ],
            milestones=[warmup_steps]
        )
        start_epoch, finetune_history = load_checkpoint(
            classifier, optimizer, scheduler, CHECKPOINT_PATH, DEVICE
        )


# Run Fine-tuning
if start_epoch == 0 or finetune_history is None:
  print(f"\n🚀 Starting MRPC fine-tuning run...")
  finetune_history = finetune_mrpc(
      classifier, mrpc_train_loader, mrpc_val_loader,
      num_epochs=FINETUNE_EPOCHS,
      lr=FINETUNE_LR,
      lm_weight=AUX_LM_WEIGHT,
      device=DEVICE,
      patience=PATIENCE,
      num_classes=NUM_CLASSES,
      grad_accum_steps=FINETUNE_GRAD_ACCUM_STEPS,
      checkpoint_path=CHECKPOINT_PATH
  )

🚀 Starting MRPC fine-tuning run...
✓ AMP (Mixed Precision) enabled

==========================================================================================
Fine-tuning GPT-1 on MRPC (cuda)...
==========================================================================================
  Epochs: 10 | LR: 6.25e-05 | Aux LM weight: 0.5 | Patience: 3
  Gradient accumulation: 4 steps
  Checkpoint path: gpt1_mrpc_checkpoint.pth
==========================================================================================

  [50/230] 9.9 batch/s | ETA: 0.3 min
  [100/230] 10.1 batch/s | ETA: 0.2 min
  [150/230] 9.0 batch/s | ETA: 0.1 min
  [200/230] 8.7 batch/s | ETA: 0.1 min

Epoch 01/10 | Time: 0.4 min | Loss (T/V): 3.3967 / 3.3367 | Acc: 69.12% | F1: 81.19% | P/R: 69.6%/97.5% | LR: 3.10e-05 * (Saved Best)
  ✓ Checkpoint saved: gpt1_mrpc_checkpoint.pth
  [50/230] 12.0 batch/s | ETA: 0.2 min
  [100/230] 11.9 batch/s | ETA: 0.2 min
  [150/230] 11.8 batch/s | ETA: 0.1 min
  [200/230] 11.7 batch/s | ETA: 0.0 min

Epoch 02/10 | Time: 0.3 min | Loss (T/V): 3.2941 / 3.3085 | Acc: 64.46% | F1: 74.25% | P/R: 73.6%/74.9% | LR: 6.20e-05 (No improve: 1/3)
  ✓ Checkpoint saved: gpt1_mrpc_checkpoint.pth
  [50/230] 11.2 batch/s | ETA: 0.3 min
  [100/230] 11.0 batch/s | ETA: 0.2 min
  [150/230] 10.6 batch/s | ETA: 0.1 min
  [200/230] 10.5 batch/s | ETA: 0.0 min

Epoch 03/10 | Time: 0.4 min | Loss (T/V): 3.1497 / 3.3409 | Acc: 63.24% | F1: 70.00% | P/R: 79.2%/62.7% | LR: 6.24e-05 (No improve: 2/3)
  ✓ Checkpoint saved: gpt1_mrpc_checkpoint.pth
  [50/230] 12.0 batch/s | ETA: 0.3 min
  [100/230] 11.9 batch/s | ETA: 0.2 min
  [150/230] 11.2 batch/s | ETA: 0.1 min
  [200/230] 11.3 batch/s | ETA: 0.0 min

Epoch 04/10 | Time: 0.3 min | Loss (T/V): 2.9203 / 3.6637 | Acc: 60.05% | F1: 68.47% | P/R: 74.4%/63.4% | LR: 6.21e-05 (No improve: 3/3)
  ✓ Checkpoint saved: gpt1_mrpc_checkpoint.pth

⏹️  Early stopping triggered.

==========================================================================================
Fine-tuning complete! Total time: 3.1 min
Best Val F1 (MRPC metric): 81.19%
==========================================================================================

# ==========================================
# 4e. Plot Fine-tuning Curves (MRPC)
# ==========================================

fig, axes = plt.subplots(2, 2, figsize=(14, 10))
epochs = range(1, len(finetune_history['train_loss']) + 1)

# Loss
axes[0, 0].plot(epochs, finetune_history['train_loss'], 'o-', label='Train', color='#1f77b4')
axes[0, 0].plot(epochs, finetune_history['val_loss'], 's--', label='Val', color='#ff7f0e')
axes[0, 0].set_title('Loss (Task + Aux LM)')
axes[0, 0].set_xlabel('Epoch')
axes[0, 0].set_ylabel('Loss')
axes[0, 0].legend()
axes[0, 0].grid(True, alpha=0.3)

# Accuracy
axes[0, 1].plot(epochs, [a*100 for a in finetune_history['val_acc']], 'o-', label='Val Acc', color='#2ca02c')
axes[0, 1].set_title('Validation Accuracy (MRPC)')
axes[0, 1].set_xlabel('Epoch')
axes[0, 1].set_ylabel('Accuracy (%)')
axes[0, 1].legend()
axes[0, 1].grid(True, alpha=0.3)

# F1-Binary (MRPC primary metric)
axes[1, 0].plot(epochs, [f*100 for f in finetune_history['val_f1']], 'D-', label='F1-binary', color='purple')
axes[1, 0].plot(epochs, [p*100 for p in finetune_history['val_precision']], '^--', label='Precision', color='brown')
axes[1, 0].plot(epochs, [r*100 for r in finetune_history['val_recall']], 's--', label='Recall', color='green')
axes[1, 0].set_title('Validation Metrics (MRPC Standard)')
axes[1, 0].set_xlabel('Epoch')
axes[1, 0].set_ylabel('Score (%)')
axes[1, 0].legend()
axes[1, 0].grid(True, alpha=0.3)

# LR
axes[1, 1].plot(epochs, finetune_history['lrs'], label='LR', color='black', linewidth=2)
axes[1, 1].set_title('LR Schedule')
axes[1, 1].set_xlabel('Epoch')
axes[1, 1].set_ylabel('Rate')
axes[1, 1].ticklabel_format(style='sci', axis='y', scilimits=(0,0))
axes[1, 1].grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('gpt1_mrpc_training_curves.png', dpi=150, bbox_inches='tight')
plt.show()

print("✓ Fine-tuning curves saved to 'gpt1_mrpc_training_curves.png'")

✓ Fine-tuning curves saved to 'gpt1_mrpc_training_curves.png'

# ==========================================
# 5a. Full Test Set Evaluation
# ==========================================

task_criterion = nn.CrossEntropyLoss()
lm_criterion = nn.CrossEntropyLoss(ignore_index=-100)

print("\n" + "="*70)
print("EVALUATING ON MRPC TEST SET")
print("="*70)

test_loss, test_acc, test_f1, test_prec, test_rec, test_preds, test_labels = evaluate_mrpc(
    classifier, mrpc_test_loader, task_criterion, lm_criterion, AUX_LM_WEIGHT, DEVICE, NUM_CLASSES
)

print(f"  Loss:              {test_loss:.4f}")
print(f"  Accuracy:          {test_acc*100:.2f}%")
print(f"  F1-Binary (MRPC):  {test_f1*100:.2f}% ← Primary metric")
print(f"  Precision:         {test_prec*100:.2f}%")
print(f"  Recall:            {test_rec*100:.2f}%")
print("="*70)

print("\nClassification Report:")
print(classification_report(test_labels, test_preds, target_names=LABEL_NAMES))

======================================================================
EVALUATING ON MRPC TEST SET
======================================================================
  Loss:              3.3382
  Accuracy:          69.62%
  F1-Binary (MRPC):  81.01% ← Primary metric
  Precision:         69.31%
  Recall:            97.47%
======================================================================

Classification Report:
                precision    recall  f1-score   support

Not Paraphrase       0.74      0.14      0.24       578
    Paraphrase       0.69      0.97      0.81      1147

      accuracy                           0.70      1725
     macro avg       0.72      0.56      0.53      1725
  weighted avg       0.71      0.70      0.62      1725

# ==========================================
# 5b. Confusion Matrix
# ==========================================

cm = confusion_matrix(test_labels, test_preds)
plt.figure(figsize=(6, 5))
ConfusionMatrixDisplay(cm, display_labels=LABEL_NAMES).plot(cmap='Blues', values_format='d')
plt.title('MRPC Confusion Matrix (GPT-1 Fine-tuned)')
plt.tight_layout()
plt.savefig('gpt1_mrpc_confusion_matrix.png', dpi=150, bbox_inches='tight')
plt.show()

<Figure size 600x500 with 0 Axes>

# ==========================================
# 5c. Per-Class Analysis
# ==========================================

print("\nPer-Class Breakdown:")
for i, label in enumerate(LABEL_NAMES):
    class_mask = [l == i for l in test_labels]
    class_correct = [p == l for p, l in zip(test_preds, test_labels) if l == i]
    if sum(class_mask) > 0:
        acc = sum(class_correct) / sum(class_mask) * 100
        print(f"  {label:12s}: {acc:5.2f}% accuracy ({sum(class_correct)}/{sum(class_mask)} correct)")

Per-Class Breakdown:
  Not Paraphrase: 14.36% accuracy (83/578 correct)
  Paraphrase  : 97.47% accuracy (1118/1147 correct)

# ==========================================
# 6a. Sample Misclassifications
# ==========================================

def analyze_mrpc_errors(model, data_loader, vocab, inv_vocab, device, num_samples=10):
    """Print MRPC examples where model prediction != true label."""
    model.eval()
    errors = []

    with torch.no_grad():
        for input_ids, lm_targets, labels, clf_positions in data_loader:
            input_ids = input_ids.to(device)
            labels = labels.to(device)
            clf_positions = clf_positions.to(device)

            cls_logits, _ = model(input_ids, clf_positions)
            preds = cls_logits.argmax(dim=-1)

            for i in range(len(labels)):
                if preds[i] != labels[i]:
                    # Decode input
                    tokens = [inv_vocab.get(t.item(), '?') for t in input_ids[i]
                             if t.item() not in [PAD_IDX, BOS_IDX, EOS_IDX, CLF_IDX, SEP_IDX]]
                    text = ' '.join(tokens).replace('[SEP]', ' | ')
                    errors.append({
                        'sentence_pair': text,
                        'true_label': labels[i].item(),
                        'pred_label': preds[i].item(),
                        'confidence': F.softmax(cls_logits[i], dim=-1)[preds[i]].item()
                    })
                    if len(errors) >= num_samples:
                        return errors
    return errors


print("=" * 80)
print("SAMPLE MISCLASSIFICATIONS (MRPC Error Analysis)")
print("=" * 80)

error_samples = analyze_mrpc_errors(classifier, mrpc_val_loader, vocab, inv_vocab, DEVICE, num_samples=10)

for i, err in enumerate(error_samples, 1):
    print(f"\n{i}. Input: {err['sentence_pair'][:150]}...")
    print(f"   True: {LABEL_NAMES[err['true_label']]} | Pred: {LABEL_NAMES[err['pred_label']]} (conf: {err['confidence']:.2%})")

print("\n" + "=" * 80)

================================================================================
SAMPLE MISCLASSIFICATIONS (MRPC Error Analysis)
================================================================================

1. Input: [UNK] said [UNK] hated the iraqi regime and looked forward to using his long years of training in the war . his wife said he was " 100 percent behind ...
   True: Not Paraphrase | Pred: Paraphrase (conf: 77.88%)

2. Input: the dollar was at 116 . 92 yen against the yen , flat on the session , and at 1 . [UNK] against the swiss [UNK] , also flat . the dollar was at 116 . ...
   True: Not Paraphrase | Pred: Paraphrase (conf: 58.55%)

3. Input: no dates have been set for the civil or the criminal trial . no dates have been set for the criminal or civil cases , but [UNK] has pleaded not guilty...
   True: Not Paraphrase | Pred: Paraphrase (conf: 68.03%)

4. Input: while [UNK] levels in the environment were up last year , they have dropped by 75 percent since the 1970s , said [UNK] . the institute said [UNK] leve...
   True: Not Paraphrase | Pred: Paraphrase (conf: 75.94%)

5. Input: " sanitation is poor . . . there could be [UNK] and cholera , " he said . " sanitation is poor , drinking water is generally left behind . . . there c...
   True: Not Paraphrase | Pred: Paraphrase (conf: 81.30%)

6. Input: the pound also made progress against the dollar , reached fresh three - year highs at $ 1 . [UNK] . the british pound [UNK] its muscle against the dol...
   True: Not Paraphrase | Pred: Paraphrase (conf: 60.68%)

7. Input: the driver , eugene rogers , helped to remove children from the bus , wood said . at the accident scene , the driver was " covered in blood " but help...
   True: Not Paraphrase | Pred: Paraphrase (conf: 70.36%)

8. Input: [UNK] said he expects muhammad will similarly be called as a witness at a [UNK] hearing for [UNK] . lee boyd [UNK] will be called as a witness wednesd...
   True: Not Paraphrase | Pred: Paraphrase (conf: 72.06%)

9. Input: with these assets , funny [UNK] has a solid chance to become the first triple crown winner since affirmed in 1978 . funny [UNK] is looking to become h...
   True: Not Paraphrase | Pred: Paraphrase (conf: 74.56%)

10. Input: bush wanted " to see an aircraft landing the same way that the pilots saw an aircraft landing , " white house press secretary ari [UNK] said yesterday...
   True: Not Paraphrase | Pred: Paraphrase (conf: 88.41%)

================================================================================

# ==========================================
# 6b. High-Confidence Correct Predictions (MRPC)
# ==========================================

print("\n" + "="*80)
print("✅ HIGH-CONFIDENCE CORRECT PREDICTIONS (MRPC)")
print("="*80)
print("Examples where model was correct AND confident (prob > 95%)\n")

classifier.eval()
shown = 0
MAX_SHOW = 10

with torch.no_grad():
    for input_ids, lm_targets, labels, clf_positions in mrpc_test_loader:
        input_ids = input_ids.to(DEVICE)
        labels = labels.to(DEVICE)
        clf_positions = clf_positions.to(DEVICE)

        cls_logits, _ = classifier(input_ids, clf_positions)
        preds = cls_logits.argmax(dim=-1)
        probs = F.softmax(cls_logits, dim=-1)

        for i in range(len(labels)):
            if preds[i] == labels[i] and probs[i][preds[i]] > 0.5:
                # Decode sentence pair (split at [SEP])
                tokens = [inv_vocab.get(t.item(), '?') for t in input_ids[i]
                          if t.item() not in [PAD_IDX, BOS_IDX, EOS_IDX, CLF_IDX]]
                text = ' '.join(tokens)
                parts = text.split('[SEP]')
                sent1 = parts[0].strip() if len(parts) > 0 else ""
                sent2 = parts[1].strip() if len(parts) > 1 else ""

                label_name = LABEL_NAMES[labels[i].item()]
                conf = probs[i][preds[i]].item()

                print(f"[{label_name}] conf: {conf:.2%}")
                print(f"  S1: {sent1[:80]}{'...' if len(sent1)>80 else ''}")
                print(f"  S2: {sent2[:80]}{'...' if len(sent2)>80 else ''}")
                print()

                shown += 1
                if shown >= MAX_SHOW:
                    break
        if shown >= MAX_SHOW:
            break

if shown == 0:
    print("⚠️  No high-confidence (≥95%) correct predictions found in this batch.")
print(f"{'='*80}\n")

================================================================================
✅ HIGH-CONFIDENCE CORRECT PREDICTIONS (MRPC)
================================================================================
Examples where model was correct AND confident (prob > 95%)

⚠️  No high-confidence (≥95%) correct predictions found in this batch.
================================================================================

# ==========================================
# 7a. MRPC Inference Function
# ==========================================

def predict_paraphrase(sentence1, sentence2, model, vocab, device, max_len=64):
    """
    Predict paraphrase label for a sentence pair.

    Returns:
        dict with 'prediction', 'probabilities', and metadata
    """
    model.eval()

    # Handle empty inputs
    if not sentence1.strip():
        sentence1 = "[UNK]"
    if not sentence2.strip():
        sentence2 = "[UNK]"

    # Tokenize
    sent1_tokens = tokenizer(sentence1)
    sent2_tokens = tokenizer(sentence2)

    sent1_ids = [vocab.get(t, UNK_IDX) for t in sent1_tokens]
    sent2_ids = [vocab.get(t, UNK_IDX) for t in sent2_tokens]

    # Truncate to fit max_len
    max_content = max_len - 4
    max_sent1_len = max_content // 2
    max_sent2_len = max_content - max_sent1_len

    sent1_ids = sent1_ids[:max_sent1_len]
    sent2_ids = sent2_ids[:max_sent2_len]

    # Construct input: [BOS] sent1 [SEP] sent2 [CLF]
    input_ids = [BOS_IDX] + sent1_ids + [SEP_IDX] + sent2_ids + [CLF_IDX]
    clf_position = len(input_ids) - 1

    # Convert to tensors
    input_tensor = torch.tensor([input_ids], dtype=torch.long, device=device)
    clf_positions = torch.tensor([clf_position], dtype=torch.long, device=device)

    # Predict
    with torch.no_grad():
        cls_logits, _ = model(input_tensor, clf_positions)
        probs = torch.softmax(cls_logits, dim=-1)
        pred = cls_logits.argmax(dim=-1).item()

    return {
        'prediction': LABEL_NAMES[pred],
        'predicted_label_id': pred,
        'probabilities': {
            LABEL_NAMES[0]: probs[0][0].item(),
            LABEL_NAMES[1]: probs[0][1].item()
        },
        'input_length': len(input_ids),
        'truncated': len(input_ids) == max_len
    }

# ==========================================
# 7b. Test Inference Examples
# ==========================================

examples = [
    {
        'sentence1': 'Amrozi accused his brother of distorting evidence.',
        'sentence2': 'Amrozi accused his brother, whom he called "the witness", of deliberately distorting his evidence.',
        'expected': 'Paraphrase'
    },
    {
        'sentence1': 'The cat is sleeping on the couch.',
        'sentence2': 'A dog is barking in the yard.',
        'expected': 'Not Paraphrase'
    },
    {
        'sentence1': 'She bought a new car yesterday.',
        'sentence2': 'Yesterday, she purchased a new automobile.',
        'expected': 'Paraphrase'
    },
    {
        'sentence1': 'The meeting was scheduled for 3 PM.',
        'sentence2': 'The conference call is at 15:00.',
        'expected': 'Paraphrase'  # Tricky: same time, different event type
    }
]

print(f"\n{'='*80}")
print("MRPC Inference Examples (GPT-1 Fine-tuned)")
print(f"{'='*80}")

for i, ex in enumerate(examples, 1):
    result = predict_paraphrase(ex['sentence1'], ex['sentence2'], classifier, vocab, DEVICE)

    print(f"\n[Example {i}] {'─'*60}")
    print(f"Sentence 1: \"{ex['sentence1']}\"")
    print(f"Sentence 2: \"{ex['sentence2']}\"")
    print(f"Expected:   {ex['expected']}")
    print(f"Predicted:  {result['prediction']}")

    # Probability bar chart
    print(f"Probabilities:")
    for label, prob in result['probabilities'].items():
        bar = '█' * int(prob * 40)
        print(f"  {label:20s} [{prob:6.4f}] {bar}")

    # Confidence + status
    confidence = result['probabilities'][result['prediction']]
    status = "✓" if result['prediction'] == ex['expected'] else "✗"
    print(f"Confidence: {confidence:.2%} | Status: {status}")
    if result['truncated']:
        print(f"⚠ Input was truncated (length={result['input_length']})")

print(f"\n{'='*80}")

================================================================================
MRPC Inference Examples (GPT-1 Fine-tuned)
================================================================================

[Example 1] ────────────────────────────────────────────────────────────
Sentence 1: "Amrozi accused his brother of distorting evidence."
Sentence 2: "Amrozi accused his brother, whom he called "the witness", of deliberately distorting his evidence."
Expected:   Paraphrase
Predicted:  Paraphrase
Probabilities:
  Not Paraphrase       [0.4167] ████████████████
  Paraphrase           [0.5833] ███████████████████████
Confidence: 58.33% | Status: ✓

[Example 2] ────────────────────────────────────────────────────────────
Sentence 1: "The cat is sleeping on the couch."
Sentence 2: "A dog is barking in the yard."
Expected:   Not Paraphrase
Predicted:  Paraphrase
Probabilities:
  Not Paraphrase       [0.3601] ██████████████
  Paraphrase           [0.6399] █████████████████████████
Confidence: 63.99% | Status: ✗

[Example 3] ────────────────────────────────────────────────────────────
Sentence 1: "She bought a new car yesterday."
Sentence 2: "Yesterday, she purchased a new automobile."
Expected:   Paraphrase
Predicted:  Not Paraphrase
Probabilities:
  Not Paraphrase       [0.5015] ████████████████████
  Paraphrase           [0.4985] ███████████████████
Confidence: 50.15% | Status: ✗

[Example 4] ────────────────────────────────────────────────────────────
Sentence 1: "The meeting was scheduled for 3 PM."
Sentence 2: "The conference call is at 15:00."
Expected:   Paraphrase
Predicted:  Paraphrase
Probabilities:
  Not Paraphrase       [0.4867] ███████████████████
  Paraphrase           [0.5133] ████████████████████
Confidence: 51.33% | Status: ✓

================================================================================

# ==========================================
# Save Final Model & Vocabulary
# ==========================================

print("\n" + "="*70)
print("SAVING FINAL ARTIFACTS")
print("="*70)

# Save fine-tuned classifier
torch.save(classifier.state_dict(), 'gpt1_mrpc_finetuned.pth')
print("✓ Saved: gpt1_mrpc_finetuned.pth")

# Save vocabulary for inference
with open('gpt1_vocab_mrpc.json', 'w') as f:
    json.dump({'vocab': vocab, 'inv_vocab': {str(k): v for k, v in inv_vocab.items()}}, f)
print("✓ Saved: gpt1_vocab_mrpc.json")

# Save hyperparameters config
config = {
    'model': 'GPT-1',
    'task': 'MRPC Paraphrase Detection',
    'd_model': D_MODEL,
    'n_layers': N_LAYERS,
    'n_heads': N_HEADS,
    'max_len': MAX_LEN,
    'finetune_lr': FINETUNE_LR,
    'aux_lm_weight': AUX_LM_WEIGHT,
    'num_classes': NUM_CLASSES,
    'label_names': LABEL_NAMES,
    'random_seed': RANDOM_SEED,
    'device': str(DEVICE)
}
with open('gpt1_mrpc_config.json', 'w') as f:
    json.dump(config, f, indent=2)
print("✓ Saved: gpt1_mrpc_config.json")

print("\n✨ All artifacts saved! Ready for inference or deployment.")
print("="*70)

======================================================================
SAVING FINAL ARTIFACTS
======================================================================
✓ Saved: gpt1_mrpc_finetuned.pth
✓ Saved: gpt1_vocab_mrpc.json
✓ Saved: gpt1_mrpc_config.json

✨ All artifacts saved! Ready for inference or deployment.
======================================================================

%watermark -iv

datasets  : 4.0.0
json      : 2.0.9
matplotlib: 3.10.0
numpy     : 2.0.2
re        : 2.2.1
seaborn   : 0.13.2
sklearn   : 1.6.1
torch     : 2.10.0+cu128
torchinfo : 1.8.0

Label	Relationship	Description
`0`	Not Paraphrase	The sentences express different meanings or information
`1`	Paraphrase	The sentences express the same meaning using different wording

Feature	SNLI (Task 2)	MRPC (Task 3)
Task type	3-class natural language inference	Binary paraphrase detection
Input structure	`[BOS] premise [SEP] hypothesis [CLF]`	`[BOS] sentence₁ [SEP] sentence₂ [CLF]`
Output classes	3 (entailment/neutral/contradiction)	2 (paraphrase/not-paraphrase)
Classification head	`Linear(d_model, 3)`	`Linear(d_model, 2)`
Evaluation metric	Accuracy + F1-macro	Accuracy + F1-binary (MRPC standard)
Dataset size	~550k examples	~3.7k examples (much smaller!)
Challenge	Cross-sentence reasoning	Lexical variation, syntactic reordering

Split	Examples	Avg. Sentence₁ Len	Avg. Sentence₂ Len	Paraphrase %
Train	3,668	~20 tokens	~20 tokens	~66.5%
Dev	408	~20 tokens	~20 tokens	~66.9%
Test	1,725	~20 tokens	~20 tokens	~66.4%

Model	MRPC Test Accuracy	MRPC Test F1
GPT-1 (fine-tuned)	~82-84%	~87-89%
BERT-base (for comparison)	~88.9%	~91.8%
Human performance	~90-92%	~93-95%

Metric	Expected Range	Notes
Test Accuracy	82-84%	Matches GPT-1 paper expectations
F1-Binary	87-89%	✅ Primary MRPC metric
Precision/Recall	~85-90% each	Balanced performance
Convergence	3-5 epochs	Small dataset + aux LM loss
Training Time	~8-12 min	Colab T4, MAX_LEN=64, 100% train

GPT-1: Task 3 — Sentence Similarity (MRPC)¶

GPT-1: Task 3 — Paraphrase Detection (MRPC)¶

GPT-1 Architecture Reference (With Fine-Tuning on Different Tasks)¶

Task Definition: Paraphrase Detection¶

Input Format for MRPC¶

Fine-tuning Architecture for MRPC¶

Training Objective: Task Loss + Auxiliary LM Loss¶

Key Differences: SNLI (3-Class NLI) → MRPC (Binary Paraphrase)¶

MRPC Dataset Statistics¶

Expected Results (Paper Reference)¶

REFERENCES¶

0. Hyperparameters¶

1. Data Preparation & Vocabulary¶

1b. Pre-training Data Pipeline (Causal LM)¶

2. GPT-1 Model Architecture¶

Architectural Note: Pre-LN vs Post-LN¶

References¶

3. Phase 1 — Generative Pre-training (Causal LM)¶

4. Phase 2 — Discriminative Fine-tuning on MRPC¶

5. Final Evaluation on MRPC Test Set¶

6. Qualitative Error Analysis¶

7. Inference Examples¶

Summary¶

Key Results: MRPC Fine-tuning¶

Key Architectural Choices for MRPC¶

MRPC vs SNLI: Task Comparison¶

Limitations & Future Directions¶

References¶

📁 Output File Structure¶

Aspect	SNLI (3-Class NLI)	MRPC (Binary Paraphrase)
Input	`[BOS] premise [SEP] hypothesis [CLF]`	`[BOS] sent₁ [SEP] sent₂ [CLF]`
Classes	3 (entail/neutral/contradict)	2 (paraphrase/not)
Dataset Size	~550k train	~3.7k train (much smaller!)
Primary Metric	Accuracy + F1-macro	Accuracy + F1-binary
Key Challenge	Cross-sentence reasoning	Lexical variation, syntactic reordering
Overfitting Risk	Low (large data)	High (small data → aux LM loss critical)