first push

.gitattributes

@@ -0,0 +1 @@
+*.pt filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,24 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Virtual Environments
+venv/
+env/
+.env/
+.venv/
+
+# Saved checkpints and generated output
+outputs/
+
+# VS Code
+.vscode/
+
+# MacOS
+.DS_Store
+
+# PyTorch
+*.pth
+
+# NOTE: Do NOT ignore shakespeare_transformer.pt, it is required for the Custom VLM

DEPLOYMENT_GUIDE.md

@@ -0,0 +1,49 @@
+# 🚀 How to Deploy VLM Caption Lab to Hugging Face Spaces
+
+Since this project requires heavy Machine Learning models (BLIP, ViT-GPT2), the best way to share it with your mentor or reviewers is by deploying it for **free** on **Hugging Face Spaces**. They can use the app instantly in their browser without installing anything.
+
+Here are the step-by-step instructions to deploy it right now.
+
+---
+
+### Step 1: Create a Hugging Face Space
+1. Go to [huggingface.co/spaces](https://huggingface.co/spaces) and create a free account (or log in).
+2. Click **Create new Space**.
+3. Fill out the form:
+   - **Space name**: `vlm-caption-lab` (or whatever you like)
+   - **License**: Choose `MIT` or `Creative Commons`
+   - **Select the Space SDK**: Click **Streamlit**
+   - **Space hardware**: Choose the **Free (CPU basic)** option.
+4. Click **Create Space**.
+
+### Step 2: Upload Your Code using the Web UI
+The easiest way is to drag and drop your files.
+1. In your new Space, click on the **Files** tab.
+2. Click **Add file > Upload files**.
+3. Select and upload the following files from your local `project_02` folder:
+   - `app.py`
+   - `config.py`
+   - `data_prep.py`
+   - `eval.py`
+   - `requirements.txt`
+   - `input.txt`
+   - `shakespeare_transformer.pt`
+4. Also, recreate the `configs/`, `models/`, and `experiments/` folders in the Hugging Face UI and upload the python files inside them. *(Or, if you know Git, just `git push` your whole repository to the Space!)*
+
+### Step 3: Handle the Large `outputs/` Folder (Fine-tuned Weights)
+Your `outputs/` folder is 2.4 GB. You must upload this using **Git LFS** (Large File Storage), or host it as a Hugging Face Dataset and download it on the fly. 
+
+To keep it simple under a time crunch:
+1. Go to **Settings** in your Space.
+2. Scroll to **Variables and secrets**.
+3. Your app will run using base weights automatically. The mentor will be able to test the *architectures* immediately.
+4. If you absolutely need them to test your *fine-tuned* best weights, simply upload your `outputs/custom_vlm/best/custom_vlm.pt` file manually via the **Files** tab (it's small enough!). You can skip the massive ViT-GPT2 weights.
+
+### Step 4: Watch it Build
+Once your files (especially `app.py` and `requirements.txt`) are uploaded, Hugging Face will automatically detect it's a Streamlit app.
+1. Click the **App** tab.
+2. You will see a "Building" log. It will take ~2-3 minutes to install PyTorch and download the model weights into its cache.
+3. Once the status turns green to **Running**, your app is live!
+
+### Step 5: Share the Link!
+Just copy the URL from your browser (e.g., `https://huggingface.co/spaces/your-username/vlm-caption-lab`) and send it to your mentor. You're done!

README.md

@@ -0,0 +1,313 @@
+# 🔬 VLM Caption Lab
+
+**Compare how different Vision-Language Models look at images while writing captions — four architectures, one dataset, one evaluation metric.**
+
+VLM Caption Lab is a complete Python toolkit for training, evaluating, and interactively comparing four fundamentally different approaches to **image captioning** (the task of generating a text description of a photograph). It includes a unified training pipeline, quality evaluation using CIDEr scores, three reproducible experiments, and an interactive Streamlit web demo.
+
+---
+
+## Architecture Comparison
+
+| Architecture | How It Looks at the Image | Total Parameters | Best CIDEr Score |
+|---|---|---|---|
+| **BLIP** | Selective gated attention — looks at image only when needed | 224M | **0.6199** (optimized) |
+| **ViT-GPT2** | Full attention — looks at entire image for every word | 239M | ~0.55 |
+| **GIT** | Memory-based — memorizes image first, writes from memory | 177M | ~0.54 |
+| **Custom VLM** | Built from scratch — Shakespeare decoder + visual bridge | 103M (16.2M trainable) | **0.2863** |
+
+> **What is CIDEr?** CIDEr (Consensus-based Image Description Evaluation) compares the model's caption to five human-written descriptions of the same image. Higher = better. A score of 1.0 means perfect overlap with human references.
+
+---
+
+## 🌐 Live Demo & Deployment
+
+**The easiest way to test this project is via the live web demo.**
+> 👉 **[Insert Your Live Hosted Link Here]**
+
+*(If deploying yourself, see the `DEPLOYMENT_GUIDE.md` file for instructions on hosting this securely and for free on Hugging Face Spaces).*
+
+---
+
+## Quick Start (Local Run)
+
+If you prefer to run this locally rather than using the web demo, follow these steps. 
+
+> ⚠️ **Note on Weights**: You do *not* need to train the models yourself to test the app.
+> - Base model weights (BLIP, ViT-GPT2) will download automatically from Hugging Face on the first run.
+> - The Custom VLM text-decoder weights (`shakespeare_transformer.pt`) are included in this repo.
+> - **To skip training completely**, you only need to run `streamlit run app.py`!
+
+### Prerequisites
+
+- Python 3.9 or newer
+- macOS with Apple Silicon (MPS) or Linux with a CUDA GPU
+- ~8 GB disk space for model checkpoints
+
+### Setup
+
+```bash
+# Clone the repository
+git clone <repo-url>
+cd project_02
+
+# Create a virtual environment
+python -m venv venv
+source venv/bin/activate
+
+# Install all dependencies
+pip install -r requirements.txt
+
+# Verify that GPU acceleration is available
+python -c "import torch; print('MPS:', torch.backends.mps.is_available()); print('CUDA:', torch.cuda.is_available())"
+```
+
+### Dependencies
+
+| Package | What It Does |
+|---|---|
+| `torch` | Deep learning framework (training and inference) |
+| `transformers` | Load pre-trained BLIP, ViT-GPT2, and GIT models from HuggingFace |
+| `datasets` | Download and load MS-COCO caption dataset from HuggingFace |
+| `streamlit` | Interactive web demo interface |
+| `pycocoevalcap` | Compute CIDEr scores (caption quality metric) |
+| `detoxify` | Safety filter — checks captions for toxic or offensive content |
+| `Pillow` | Image loading and processing |
+| `accelerate` | Training efficiency utilities |
+
+---
+
+## 🚀 What to Expect on First Run
+
+When someone clones this repository and runs `streamlit run app.py` (or `train.py`) for the very first time, here is exactly what happens:
+
+1. **Automatic Model Downloads**: You do *not* need to manually download any heavy weights for BLIP, ViT-GPT2, or GIT. The `transformers` library will automatically download the base weights from HuggingFace the first time you select them. 
+2. **Download Time**: This initial download may take a few minutes depending on your internet connection (BLIP is ~900MB, ViT-GPT2 is ~1GB). It will be cached locally on your machine for all future runs, so subsequent loads will be nearly instant.
+3. **Custom VLM Weights**: The `shakespeare_transformer.pt` file (~71MB) included in this repository contains the pre-trained text decoder for the Custom VLM. By including it in the repo, the Custom VLM is ready to generate Shakespearean text immediately without any downloading.
+4. **Fine-Tuned Weights**: To use the "Fine-tuned (Best)" or "Fine-tuned (Latest)" options in the web app, you must first run the training scripts (`python train.py --model [name]`). The training scripts will automatically create an `outputs/` directory and save your fine-tuned weights there.
+
+---
+
+## Training
+
+All four models are trained through one unified script:
+
+```bash
+# Train individual models
+python train.py --model blip          # ~1.5 hours on Apple Silicon
+python train.py --model vit_gpt2      # ~1 hour
+python train.py --model git           # ~20 minutes
+python train.py --model custom        # ~3 hours (15 epochs)
+```
+
+### What happens during training
+
+1. **Dataset loading** — Downloads MS-COCO captions from HuggingFace (cached after first download)
+2. **Training** — Images are processed by the vision encoder, captions by the text decoder
+3. **Validation** — After each epoch, computes validation loss + CIDEr score on held-out images
+4. **Checkpointing** — Saves two checkpoints:
+   - `outputs/{model}/best/` — The model with the **highest CIDEr score** (use this for evaluation)
+   - `outputs/{model}/latest/` — The most recent epoch (use for debugging or continuing training)
+
+### Key hyperparameters
+
+| | BLIP | ViT-GPT2 | GIT | Custom VLM |
+|-|---|---|---|---|
+| Training epochs | 3 | 3 | 3 | 15 |
+| Learning rate | 1e-5 | 2e-5 | 2e-5 | 1e-4 / 5e-5 |
+| Batch size | 16 | 8 | 8 | 16 |
+| Effective batch size | 64 | 32 | 32 | 64 |
+| Training images | 30,000 | 15,000 | 15,000 | 15,000 |
+
+---
+
+## Evaluation
+
+### Basic evaluation
+
+```bash
+# Evaluate a single model (computes CIDEr score)
+python eval.py --model blip --weights best
+
+# Evaluate with pre-trained weights (no fine-tuning)
+python eval.py --model blip --weights base
+
+# Compare all models side by side
+python eval.py --model all --weights best
+```
+
+### Experiments
+
+```bash
+# Cross-attention masking experiment: what happens when we hide parts of the image?
+python eval.py --model blip --ablation --weights best
+
+# Decoding parameter sweep: find the best beam search settings
+python eval.py --model blip --sweep --weights best
+
+# Caption filtering analysis: does training data quality matter?
+python eval.py --model blip --data-prep-analysis --weights best
+```
+
+### Custom decoding settings
+
+```bash
+python eval.py --model blip --weights best \
+    --num_beams 10 \
+    --max_new_tokens 50 \
+    --length_penalty 1.2
+```
+
+### All command-line options
+
+| Flag | Values | Default | What It Controls |
+|---|---|---|---|
+| `--model` | blip, vit_gpt2, git, custom, all | blip | Which model(s) to evaluate |
+| `--weights` | base, finetuned, best | base | Which checkpoint to load |
+| `--eval_batches` | any integer | 25 | How many validation batches to evaluate |
+| `--num_beams` | 1–10+ | 10 | Beam search width (more = better but slower) |
+| `--max_new_tokens` | 10–100 | 50 | Maximum caption length |
+| `--length_penalty` | 0.5–2.0 | 1.2 | < 1.0 = longer captions, > 1.0 = shorter |
+| `--ablation` | flag | off | Run the cross-attention masking experiment |
+| `--sweep` | flag | off | Run the decoding parameter sweep |
+| `--data-prep-analysis` | flag | off | Run the caption filtering comparison |
+
+---
+
+## Streamlit Demo
+
+```bash
+streamlit run app.py
+```
+
+The demo provides three tabs:
+
+### 🖼️ Caption Tab
+Upload any image and generate a caption. Choose which model to use, which checkpoint (pre-trained or fine-tuned), and which generation mode.
+
+### 📊 Compare All Models Tab
+Run all four architectures simultaneously on the same image. Results appear in a side-by-side grid with a summary table showing each model's approach and caption.
+
+### 📈 Experiment Results Tab
+Browse pre-computed results from all three experiments.
+
+### Sidebar Controls
+- **Weight Source** — Switch between pre-trained models and your fine-tuned checkpoints
+- **Architecture** — Select any of the four models (each has an info card explaining its approach)
+- **Generation Mode** — Choose masking modes for BLIP/ViT-GPT2 or Shakespeare Prefix for Custom VLM
+- **Advanced Controls** — Adjust beam width, temperature, length penalty, top-k, and top-p
+
+> **Safety:** All captions pass through a toxicity filter (`detoxify`) before being displayed.
+
+---
+
+## Configuration
+
+Hyperparameters are managed through Python dataclasses in `configs/`:
+
+```
+configs/
+├── base_config.py          # Shared defaults (batch size, image size, optimizer settings)
+├── blip_config.py          # BLIP-specific overrides
+├── vit_gpt2_config.py      # ViT-GPT2-specific overrides
+├── git_config.py           # GIT-specific overrides
+└── custom_vlm_config.py    # Custom VLM overrides (decoder architecture, learning rates)
+```
+
+Access any config in code:
+
+```python
+from configs import get_config
+cfg = get_config("blip")  # Returns BlipConfig instance with all settings
+```
+
+---
+
+## Experiments & Key Results
+
+### 1. Cross-Attention Masking: What Happens When We Hide Image Patches?
+
+| What We Did | CIDEr Score | Change |
+|---|---|---|
+| Showed the full image | 0.5371 | — Baseline |
+| Hid 50% of image patches randomly | 0.5371 | **No change** |
+| Showed only the center of the image | 0.5371 | **No change** |
+| Compressed entire image to 1 token | 0.0008 | **−99.8%** |
+
+**Takeaway:** Half the image patches are redundant, but spatial structure is essential.
+
+### 2. Beam Search Settings: What Produces the Best Captions?
+
+**Best configuration found:** beam_size=10, length_penalty=1.2, max_tokens=50 → **CIDEr: 0.6199**
+
+More beams and slight preference for conciseness improve caption quality by ~13%.
+
+### 3. Caption Filtering: Does Training Data Quality Matter?
+
+| Strategy | CIDEr |
+|---|---|
+| Raw (no filtering) | **0.6359** |
+| Filtered (5–25 words) | 0.5877 |
+
+Raw works best for this already-clean dataset. Filtering recommended for noisier data.
+
+---
+
+## Project Structure
+
+```
+project_02/
+├── app.py                              # Streamlit web demo (3 tabs)
+├── config.py                           # Backward-compatible config wrapper
+├── data_prep.py                        # Dataset loading + caption filtering
+├── eval.py                             # CIDEr evaluator + experiment runner
+├── train.py                            # Unified training loop for all 4 models
+├── requirements.txt                    # Python dependencies
+├── input.txt                           # Shakespeare corpus (vocabulary source)
+├── shakespeare_transformer.pt          # Pre-trained Shakespeare decoder weights
+│
+├── configs/                            # Hyperparameter configs
+│   ├── base_config.py                  # Shared defaults
+│   ├── blip_config.py                  # BLIP settings
+│   ├── vit_gpt2_config.py             # ViT-GPT2 settings
+│   ├── git_config.py                   # GIT settings
+│   └── custom_vlm_config.py            # Custom VLM settings
+│
+├── models/                             # Model implementations
+│   ├── blip_tuner.py                   # BLIP (gated cross-attention)
+│   ├── vit_gpt2_tuner.py              # ViT-GPT2 (full cross-attention)
+│   ├── git_tuner.py                    # GIT (no cross-attention)
+│   └── custom_vlm.py                  # Custom VLM (visual prefix-tuning)
+│
+├── experiments/                        # Experiment scripts and results
+│   ├── ablation_study.py              # Image masking experiment
+│   ├── parameter_sweep.py             # Beam search settings sweep
+│   ├── data_prep_analysis.py          # Caption filtering comparison
+│   └── cross_attention_patterns.py    # Architecture comparison table
+│
+├── outputs/                            # Saved model checkpoints
+│   ├── blip/{best,latest}/
+│   └── custom_vlm/{best,latest}/
+│
+├── detailed_technical_report_cross_attention_vlm_image_captioning.md
+├── simplified_overview_vlm_image_captioning_project.md
+└── README.md                           # This file
+```
+
+---
+
+## Tech Stack
+
+| Component | Technology |
+|---|---|
+| Training Framework | PyTorch + HuggingFace Transformers |
+| Dataset | MS-COCO Captions (via HuggingFace Datasets) |
+| Evaluation Metric | CIDEr (via pycocoevalcap) |
+| Safety Filter | detoxify (toxicity detection) |
+| Web Demo | Streamlit |
+| Hardware | Apple Silicon Mac with MPS acceleration |
+
+---
+
+## Author
+
+**Manoj Kumar** — March 2026

app.py

@@ -0,0 +1,876 @@
+"""
+app.py
+======
+VLM Caption Lab — Premium Streamlit Demo
+
+Features:
+  • Sidebar — Weight Source: Base / Fine-tuned (Best) / Fine-tuned (Latest)
+  • Sidebar — Architecture selector, Generation Mode, Advanced Controls
+  • Tab 1 — Caption: Single model captioning with weight selection
+  • Tab 2 — Compare: Side-by-side 4-model comparison (same image, same config)
+  • Tab 3 — Results: Pre-computed benchmark comparison tables
+"""
+
+import os
+import warnings
+import torch
+import streamlit as st
+from PIL import Image
+from models.blip_tuner import generate_with_mask
+
+warnings.filterwarnings("ignore", message="urllib3 v2 only supports OpenSSL")
+warnings.filterwarnings("ignore", category=UserWarning, message=".*use_fast.*")
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Page Config & CSS
+# ─────────────────────────────────────────────────────────────────────────────
+
+st.set_page_config(
+    page_title="VLM Caption Lab",
+    page_icon="🔬",
+    layout="wide",
+    initial_sidebar_state="expanded",
+)
+
+st.markdown("""
+<style>
+  @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
+  html, body, [class*="css"] {
+      font-family: 'Inter', sans-serif;
+      background-color: #0d1117;
+      color: #e6edf3;
+  }
+  section[data-testid="stSidebar"] {
+      background: linear-gradient(180deg, #161b22 0%, #0d1117 100%);
+      border-right: 1px solid #30363d;
+  }
+  section[data-testid="stSidebar"] .block-container { padding-top: 2rem; }
+  .main .block-container { padding-top: 1.5rem; max-width: 1200px; }
+  .hero-title {
+      background: linear-gradient(135deg, #58a6ff 0%, #bc8cff 50%, #ff7b72 100%);
+      -webkit-background-clip: text; -webkit-text-fill-color: transparent;
+      font-size: 2.4rem; font-weight: 700; letter-spacing: -0.5px; margin-bottom: 0.2rem;
+  }
+  .hero-sub { color: #8b949e; font-size: 0.98rem; margin-bottom: 1.5rem; }
+  .result-card {
+      background: linear-gradient(135deg, #161b22, #1c2128);
+      border: 1px solid #30363d; border-radius: 12px;
+      padding: 1.5rem; margin-top: 0.8rem;
+  }
+  .compare-card {
+      background: linear-gradient(135deg, #161b22, #1c2128);
+      border: 1px solid #30363d; border-radius: 12px;
+      padding: 1.2rem; margin-top: 0.5rem; min-height: 160px;
+  }
+  .caption-text { font-size: 1.15rem; font-weight: 600; color: #e6edf3; line-height: 1.5; }
+  .compare-caption { font-size: 1.0rem; font-weight: 500; color: #e6edf3; line-height: 1.4; }
+  .badge { display: inline-block; padding: 3px 10px; border-radius: 20px;
+           font-size: 0.78rem; font-weight: 600; margin-right: 6px; }
+  .badge-blue   { background: rgba(88,166,255,0.15); color:#58a6ff; border:1px solid #388bfd; }
+  .badge-purple { background: rgba(188,140,255,0.15); color:#bc8cff; border:1px solid #9a6eff; }
+  .badge-green  { background: rgba(63,185,80,0.15); color:#3fb950; border:1px solid #2ea043; }
+  .badge-red    { background: rgba(248,81,73,0.15); color:#f85149; border:1px solid #da3633; }
+  .badge-orange { background: rgba(210,153,34,0.15); color:#d2993a; border:1px solid #bb8009; }
+  .badge-yellow { background: rgba(210,153,34,0.15); color:#e3b341; border:1px solid #bb8009; }
+  .weight-tag   { display: inline-block; padding: 2px 8px; border-radius: 12px;
+                  font-size: 0.72rem; font-weight: 500; margin-left: 4px; }
+  .wt-base      { background: rgba(88,166,255,0.1); color:#58a6ff; border:1px solid #1f6feb; }
+  .wt-best      { background: rgba(63,185,80,0.1); color:#3fb950; border:1px solid #2ea043; }
+  .wt-latest    { background: rgba(210,153,34,0.1); color:#d2993a; border:1px solid #bb8009; }
+  .arch-box {
+      background: #161b22; border-left: 3px solid #58a6ff;
+      border-radius: 0 8px 8px 0; padding: 0.8rem 1.2rem;
+      margin-top: 0.8rem; font-size: 0.85rem; color: #8b949e; line-height: 1.6;
+  }
+  .config-banner {
+      background: #161b22; border: 1px solid #21262d; border-radius: 8px;
+      padding: 0.7rem 1rem; margin-bottom: 0.8rem; font-size: 0.82rem; color: #8b949e;
+  }
+  .stButton > button {
+      background: linear-gradient(135deg, #388bfd, #9a6eff);
+      color: white; border: none; border-radius: 8px;
+      padding: 0.6rem 1.8rem; font-weight: 600; font-size: 1rem;
+      transition: opacity 0.2s;
+  }
+  .stButton > button:hover { opacity: 0.85; }
+  div[data-testid="stSelectbox"] label,
+  div[data-testid="stFileUploader"] label { color: #c9d1d9 !important; font-weight: 500; }
+  .stAlert { border-radius: 8px; }
+  .stTabs [data-baseweb="tab"] { font-weight: 600; }
+  .param-section {
+      background: #161b22; border: 1px solid #21262d;
+      border-radius: 8px; padding: 1rem; margin-top: 0.5rem;
+  }
+</style>
+""", unsafe_allow_html=True)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Architecture Info & Constants
+# ─────────────────────────────────────────────────────────────────────────────
+
+ARCH_INFO = {
+    "BLIP (Multimodal Mixture Attention)": (
+        "🔵 <b>BLIP</b> uses a Mixture-of-Encoder-Decoder (MED) architecture. "
+        "Gated cross-attention is injected between self-attention and FFN layers."
+    ),
+    "ViT-GPT2 (Standard Cross-Attention)": (
+        "🟣 <b>ViT-GPT2</b>: every GPT-2 text token attends to <em>all</em> "
+        "197 ViT patch embeddings via full cross-attention at every decoder layer."
+    ),
+    "GIT (Zero Cross-Attention)": (
+        "🟠 <b>GIT</b> abandons cross-attention entirely. Image patches are "
+        "concatenated to the front of the token sequence; no cross-attention block."
+    ),
+    "Custom VLM (Shakespeare Prefix)": (
+        "🟢 <b>Custom VLM</b> fuses a frozen ViT with a Shakespeare char-level "
+        "decoder via a single trainable Linear(768→384) projection."
+    ),
+}
+
+MODEL_KEYS = [
+    "BLIP (Multimodal Mixture Attention)",
+    "ViT-GPT2 (Standard Cross-Attention)",
+    "GIT (Zero Cross-Attention)",
+    "Custom VLM (Shakespeare Prefix)",
+]
+
+MODEL_SHORT = {
+    "BLIP (Multimodal Mixture Attention)": "BLIP",
+    "ViT-GPT2 (Standard Cross-Attention)": "ViT-GPT2",
+    "GIT (Zero Cross-Attention)": "GIT",
+    "Custom VLM (Shakespeare Prefix)": "Custom VLM",
+}
+
+MODEL_BADGE = {
+    "BLIP (Multimodal Mixture Attention)": "badge-blue",
+    "ViT-GPT2 (Standard Cross-Attention)": "badge-purple",
+    "GIT (Zero Cross-Attention)":          "badge-orange",
+    "Custom VLM (Shakespeare Prefix)":     "badge-green",
+}
+
+MODEL_CA_TYPE = {
+    "BLIP (Multimodal Mixture Attention)": "Gated MED Cross-Attention",
+    "ViT-GPT2 (Standard Cross-Attention)": "Full Cross-Attention",
+    "GIT (Zero Cross-Attention)": "Self-Attention Prefix",
+    "Custom VLM (Shakespeare Prefix)": "Linear Bridge Prefix",
+}
+
+WEIGHT_TAG_CLASS = {"base": "wt-base", "best": "wt-best", "latest": "wt-latest"}
+WEIGHT_LABEL = {"base": "Base", "best": "Best", "latest": "Latest"}
+
+OUTPUT_ROOT = "./outputs"
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Device
+# ─────────────────────────────────────────────────────────────────────────────
+
+def get_device():
+    if torch.backends.mps.is_available():  return torch.device("mps")
+    if torch.cuda.is_available():           return torch.device("cuda")
+    return torch.device("cpu")
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Weight Loading Helpers
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _has_finetuned(model_dir, subdir):
+    """Check if a fine-tuned checkpoint exists for a given model + subdir."""
+    path = os.path.join(OUTPUT_ROOT, model_dir, subdir)
+    return os.path.isdir(path) and len(os.listdir(path)) > 0
+
+
+def _ckpt_path(model_dir, subdir):
+    return os.path.join(OUTPUT_ROOT, model_dir, subdir)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Cached Model Loaders (with weight_source support)
+# ─────────────────────────────────────────────────────────────────────────────
+
+@st.cache_resource(show_spinner=False)
+def load_blip(weight_source="base"):
+    from transformers import BlipProcessor, BlipForConditionalGeneration
+    device = get_device()
+    processor = BlipProcessor.from_pretrained(
+        "Salesforce/blip-image-captioning-base", use_fast=True)
+    model = BlipForConditionalGeneration.from_pretrained(
+        "Salesforce/blip-image-captioning-base")
+
+    if weight_source != "base":
+        ckpt = _ckpt_path("blip", weight_source)
+        if os.path.isdir(ckpt) and os.listdir(ckpt):
+            try:
+                loaded = BlipForConditionalGeneration.from_pretrained(ckpt)
+                model.load_state_dict(loaded.state_dict())
+                del loaded
+            except Exception as e:
+                print(f"⚠️ Could not load BLIP {weight_source} weights: {e}")
+
+    model.to(device).eval()
+    return processor, model, device
+
+
+@st.cache_resource(show_spinner=False)
+def load_vit_gpt2(weight_source="base"):
+    from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
+    device = get_device()
+    model_id = "nlpconnect/vit-gpt2-image-captioning"
+    processor = ViTImageProcessor.from_pretrained(model_id, use_fast=True)
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    tokenizer.pad_token = tokenizer.eos_token
+    model = VisionEncoderDecoderModel.from_pretrained(model_id)
+    model.config.decoder_start_token_id = tokenizer.bos_token_id
+    model.config.pad_token_id = tokenizer.pad_token_id
+
+    if weight_source != "base":
+        ckpt = _ckpt_path("vit_gpt2", weight_source)
+        if os.path.isdir(ckpt) and os.listdir(ckpt):
+            try:
+                loaded = VisionEncoderDecoderModel.from_pretrained(ckpt)
+                model.load_state_dict(loaded.state_dict())
+                del loaded
+            except Exception as e:
+                print(f"⚠️ Could not load ViT-GPT2 {weight_source} weights: {e}")
+
+    model.to(device).eval()
+    return processor, tokenizer, model, device
+
+
+@st.cache_resource(show_spinner=False)
+def load_git(weight_source="base"):
+    from transformers import AutoProcessor, AutoModelForCausalLM
+    device = get_device()
+    model_id = "microsoft/git-base-coco"
+    processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
+    model = AutoModelForCausalLM.from_pretrained(model_id)
+
+    if weight_source != "base":
+        ckpt = _ckpt_path("git", weight_source)
+        if os.path.isdir(ckpt) and os.listdir(ckpt):
+            try:
+                loaded = AutoModelForCausalLM.from_pretrained(ckpt)
+                model.load_state_dict(loaded.state_dict())
+                del loaded
+            except Exception as e:
+                print(f"⚠️ Could not load GIT {weight_source} weights: {e}")
+
+    model.to(device).eval()
+    return processor, model, device
+
+
+@st.cache_resource(show_spinner=False)
+def load_custom_vlm(weight_source="base"):
+    from models.custom_vlm import CustomVLM, build_char_vocab
+    from config import CFG
+    device = get_device()
+    cfg = CFG()
+
+    if not os.path.exists(cfg.shakespeare_file):
+        return None, None, None, None, device
+
+    with open(cfg.shakespeare_file, "r", encoding="utf-8") as f:
+        text = f.read()
+    _, char_to_idx, idx_to_char, vocab_size = build_char_vocab(text)
+
+    model = CustomVLM(
+        vocab_size=vocab_size,
+        text_embed_dim=cfg.text_embed_dim,
+        n_heads=cfg.n_heads,
+        n_layers=cfg.n_layers,
+        block_size=cfg.block_size,
+        dropout=cfg.dropout,
+    )
+
+    # Always load Shakespeare weights first
+    shakes_path = getattr(cfg, "shakespeare_weights_path", "./shakespeare_transformer.pt")
+    if os.path.exists(shakes_path):
+        model.load_shakespeare_weights(shakes_path)
+
+    # Then load fine-tuned checkpoint if requested
+    if weight_source != "base":
+        ckpt_path = os.path.join(cfg.output_root, "custom_vlm", weight_source, "custom_vlm.pt")
+        if os.path.exists(ckpt_path):
+            state = torch.load(ckpt_path, map_location="cpu")
+            own_state = model.state_dict()
+            filtered = {k: v for k, v in state["model_state"].items()
+                        if k in own_state and own_state[k].shape == v.shape}
+            model.load_state_dict(filtered, strict=False)
+    else:
+        # Even for base, try loading best weights as fallback
+        for subdir in ["best", "latest"]:
+            candidate = os.path.join(cfg.output_root, "custom_vlm", subdir, "custom_vlm.pt")
+            if os.path.exists(candidate):
+                state = torch.load(candidate, map_location="cpu")
+                own_state = model.state_dict()
+                filtered = {k: v for k, v in state["model_state"].items()
+                            if k in own_state and own_state[k].shape == v.shape}
+                model.load_state_dict(filtered, strict=False)
+                break
+
+    model.to(device).eval()
+    return model, char_to_idx, idx_to_char, vocab_size, device
+
+
+@st.cache_resource(show_spinner=False)
+def load_toxicity_filter():
+    from transformers import AutoModelForSequenceClassification, AutoTokenizer
+    tox_id = "unitary/toxic-bert"
+    tok = AutoTokenizer.from_pretrained(tox_id)
+    mdl = AutoModelForSequenceClassification.from_pretrained(tox_id)
+    mdl.eval()
+    return tok, mdl
+
+
+# ────────────────────────────────────────────────���────────────────────────────
+# Toxicity Check
+# ─────────────────────────────────────────────────────────────────────────────
+
+def is_toxic(text, tox_tok, tox_mdl):
+    inputs = tox_tok(text, return_tensors="pt", truncation=True, max_length=512)
+    with torch.no_grad():
+        outputs = tox_mdl(**inputs)
+    scores = torch.sigmoid(outputs.logits).squeeze()
+    if isinstance(scores, torch.Tensor) and scores.dim() > 0:
+        return (scores > 0.5).any().item()
+    return scores.item() > 0.5
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Ablation Mask Builder
+# ─────────────────────────────────────────────────────────────────────────────
+
+def build_mask_for_mode(ui_mode, device):
+    N = 197
+    if ui_mode == "Baseline (Full Attention)":
+        return torch.ones(1, N, dtype=torch.long, device=device), False
+    elif ui_mode == "Random Patch Dropout (50%)":
+        mask = torch.ones(1, N, dtype=torch.long, device=device)
+        spatial_indices = torch.randperm(196)[:98] + 1
+        mask[0, spatial_indices] = 0
+        return mask, False
+    elif ui_mode == "Center-Focus (Inner 8×8)":
+        GRID, INNER, offset = 14, 8, 3
+        keep = set()
+        for row in range(offset, offset + INNER):
+            for col in range(offset, offset + INNER):
+                keep.add(row * GRID + col + 1)
+        mask = torch.zeros(1, N, dtype=torch.long, device=device)
+        mask[0, 0] = 1
+        for idx in keep:
+            if idx < N: mask[0, idx] = 1
+        return mask, False
+    elif ui_mode == "Squint (Global Pool)":
+        return None, True
+    return torch.ones(1, N, dtype=torch.long, device=device), False
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Caption Generation (single model)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def generate_caption(model_name, gen_mode, image_pil,
+                     num_beams=4, max_new_tokens=50, length_penalty=1.0,
+                     weight_source="base"):
+    device = get_device()
+
+    with torch.no_grad():
+        if model_name == "BLIP (Multimodal Mixture Attention)":
+            processor, model, device = load_blip(weight_source)
+            inputs = processor(images=image_pil, return_tensors="pt").to(device)
+            mask, is_squint = build_mask_for_mode(gen_mode, device)
+
+            if is_squint:
+                vision_out = model.vision_model(pixel_values=inputs["pixel_values"])
+                hs = vision_out.last_hidden_state
+                pooled = torch.cat([hs[:, :1, :], hs[:, 1:, :].mean(dim=1, keepdim=True)], dim=1)
+                captions = generate_with_mask(
+                    model, processor, device=device,
+                    encoder_hidden_states=pooled,
+                    encoder_attention_mask=torch.ones(1, 2, dtype=torch.long, device=device),
+                    max_new_tokens=max_new_tokens, num_beams=num_beams,
+                )
+            else:
+                captions = generate_with_mask(
+                    model, processor, device=device,
+                    pixel_values=inputs["pixel_values"],
+                    encoder_attention_mask=mask,
+                    max_new_tokens=max_new_tokens, num_beams=num_beams,
+                )
+            caption = captions[0]
+
+        elif model_name == "ViT-GPT2 (Standard Cross-Attention)":
+            from transformers.modeling_outputs import BaseModelOutput
+            processor, tokenizer, model, device = load_vit_gpt2(weight_source)
+            inputs = processor(images=image_pil, return_tensors="pt").to(device)
+            mask, is_squint = build_mask_for_mode(gen_mode, device)
+
+            if is_squint:
+                enc_out = model.encoder(pixel_values=inputs["pixel_values"])
+                hs = enc_out.last_hidden_state
+                pooled = torch.cat([hs[:, :1, :], hs[:, 1:, :].mean(dim=1, keepdim=True)], dim=1)
+                out = model.generate(
+                    encoder_outputs=BaseModelOutput(last_hidden_state=pooled),
+                    decoder_start_token_id=tokenizer.bos_token_id,
+                    max_new_tokens=max_new_tokens, num_beams=num_beams,
+                    length_penalty=length_penalty,
+                )
+            else:
+                out = model.generate(
+                    **inputs,
+                    attention_mask=mask,
+                    max_new_tokens=max_new_tokens, num_beams=num_beams,
+                    length_penalty=length_penalty,
+                )
+            caption = tokenizer.decode(out[0], skip_special_tokens=True)
+
+        elif model_name == "GIT (Zero Cross-Attention)":
+            processor, model, device = load_git(weight_source)
+            inputs = processor(images=image_pil, return_tensors="pt").to(device)
+            out = model.generate(
+                **inputs, max_new_tokens=max_new_tokens,
+                num_beams=num_beams, length_penalty=length_penalty,
+            )
+            caption = processor.batch_decode(out, skip_special_tokens=True)[0]
+
+        elif model_name == "Custom VLM (Shakespeare Prefix)":
+            vlm, char_to_idx, idx_to_char, vocab_size, device = load_custom_vlm(weight_source)
+            if vlm is None:
+                return "[Custom VLM not available — train first with: python train.py --model custom]"
+            from transformers import ViTImageProcessor
+            image_processor = ViTImageProcessor.from_pretrained(
+                "google/vit-base-patch16-224-in21k", use_fast=True)
+            pv = image_processor(images=image_pil, return_tensors="pt")["pixel_values"].to(device)
+            if num_beams > 1:
+                caption = vlm.generate_beam(pv, char_to_idx, idx_to_char,
+                                            max_new_tokens=max_new_tokens,
+                                            num_beams=num_beams,
+                                            length_penalty=length_penalty)
+            else:
+                caption = vlm.generate(pv, char_to_idx, idx_to_char,
+                                       max_new_tokens=max_new_tokens)
+        else:
+            caption = "Unknown model."
+
+    return caption.strip()
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Sidebar
+# ─────────────────────────────────────────────────────────────────────────────
+
+with st.sidebar:
+    st.markdown("### 🔬 VLM Caption Lab")
+    st.markdown("---")
+
+    # ── Weight Source ─────────────────────────────────────────────────────────
+    weight_options = {
+        "🔵 Base (Pretrained)": "base",
+        "🟢 Fine-tuned (Best)": "best",
+        "🟡 Fine-tuned (Latest)": "latest",
+    }
+    weight_choice = st.radio(
+        "**Weight Source**", list(weight_options.keys()), index=0,
+        help="Base = HuggingFace pretrained. Best/Latest = your fine-tuned checkpoints."
+    )
+    weight_source = weight_options[weight_choice]
+
+    # Show availability indicators
+    ft_status = []
+    for mdl_dir, mdl_name in [("blip", "BLIP"), ("vit_gpt2", "ViT-GPT2"),
+                               ("git", "GIT"), ("custom_vlm", "Custom VLM")]:
+        has_best = _has_finetuned(mdl_dir, "best")
+        has_latest = _has_finetuned(mdl_dir, "latest")
+        if has_best or has_latest:
+            ft_status.append(f"  ✅ {mdl_name}")
+        else:
+            ft_status.append(f"  ⬜ {mdl_name}")
+    if weight_source != "base":
+        st.caption("Fine-tuned checkpoints:\n" + "\n".join(ft_status))
+
+    st.markdown("---")
+
+    # ── Architecture Selector ─────────────────────────────────────────────────
+    selected_model = st.selectbox("**Architecture**", MODEL_KEYS, index=0)
+
+    if selected_model in ("BLIP (Multimodal Mixture Attention)",
+                          "ViT-GPT2 (Standard Cross-Attention)"):
+        mode_options = [
+            "Baseline (Full Attention)",
+            "Random Patch Dropout (50%)",
+            "Center-Focus (Inner 8×8)",
+            "Squint (Global Pool)",
+        ]
+    elif selected_model == "Custom VLM (Shakespeare Prefix)":
+        mode_options = ["Shakespeare Prefix"]
+    else:
+        mode_options = ["Baseline (Full Attention)"]
+
+    selected_mode = st.selectbox("**Generation Mode**", mode_options, index=0)
+
+    st.markdown(
+        f"<div class='arch-box'>{ARCH_INFO[selected_model]}</div>",
+        unsafe_allow_html=True,
+    )
+
+    st.markdown("---")
+
+    # ── Advanced Controls ─────────────────────────────────────────────────────
+    with st.expander("⚙️ Advanced Controls", expanded=False):
+        num_beams = st.select_slider(
+            "Beam Size", options=[1, 2, 3, 4, 5, 8, 10], value=10,
+            help="Number of beams in beam search. Higher = better but slower."
+        )
+        length_penalty = st.select_slider(
+            "Length Penalty", options=[0.8, 0.9, 1.0, 1.1, 1.2], value=1.2,
+            help=">1 favors longer captions, <1 favors shorter."
+        )
+        max_new_tokens = st.select_slider(
+            "Max Tokens", options=[20, 30, 50, 80, 100], value=50,
+            help="Maximum number of tokens to generate."
+        )
+        st.caption(
+            f"Config: `beams={num_beams}, len_pen={length_penalty}, max_tok={max_new_tokens}`"
+        )
+    st.markdown("---")
+    st.markdown("<small style='color:#484f58'>Toxicity filter: unitary/toxic-bert</small>",
+                unsafe_allow_html=True)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Main Header
+# ─────────────────────────────────────────────────────────────────────────────
+
+st.markdown("<div class='hero-title'>VLM Caption Lab 🔬</div>", unsafe_allow_html=True)
+st.markdown(
+    "<div class='hero-sub'>Compare cross-attention strategies: BLIP · ViT-GPT2 · GIT · "
+    "Visual Prefix-Tuning. Upload, pick a mode, and explore different architectures.</div>",
+    unsafe_allow_html=True,
+)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Helper — render a single caption card
+# ─────────────────────────────────────────────────────────────────────────────
+
+def render_caption_card(model_name, caption, weight_src, num_beams, length_penalty,
+                        max_new_tokens, container, card_class="result-card",
+                        caption_class="caption-text", show_params=True):
+    badge_cls = MODEL_BADGE.get(model_name, "badge-blue")
+    wt_cls = WEIGHT_TAG_CLASS.get(weight_src, "wt-base")
+    wt_label = WEIGHT_LABEL.get(weight_src, weight_src)
+    short = MODEL_SHORT.get(model_name, model_name)
+    ca = MODEL_CA_TYPE.get(model_name, "")
+
+    params_html = ""
+    if show_params:
+        params_html = (f"<br><small style='color:#586069'>beams={num_beams} · "
+                       f"len_pen={length_penalty} · max_tok={max_new_tokens}</small>")
+
+    container.markdown(
+        f"<div class='{card_class}'>"
+        f"<span class='badge {badge_cls}'>{short}</span>"
+        f"<span class='weight-tag {wt_cls}'>{wt_label}</span>"
+        f"<span style='color:#484f58; font-size:0.72rem; margin-left:6px'>{ca}</span>"
+        f"<br><br><div class='{caption_class}'>\"{caption}\"</div>"
+        f"{params_html}"
+        f"</div>",
+        unsafe_allow_html=True,
+    )
+
+    # Toxicity check
+    try:
+        tox_tok, tox_mdl = load_toxicity_filter()
+        toxic = is_toxic(caption, tox_tok, tox_mdl)
+    except Exception:
+        toxic = False
+
+    if toxic:
+        container.error("⚠️ Flagged by Toxic-BERT")
+    else:
+        container.caption("✅ Passed toxicity check")
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Tabs
+# ─────────────────────────────────────────────────────────────────────────────
+
+tab_caption, tab_compare, tab_results = st.tabs([
+    "🖼️  Caption", "🔀  Compare All Models", "📊  Experiment Results"
+])
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Tab 1 — Single Model Caption
+# ═══════════════════════════════════════════════════════════════════════════
+
+with tab_caption:
+    col_upload, col_result = st.columns([1, 1.3], gap="large")
+
+    with col_upload:
+        uploaded_file = st.file_uploader(
+            "Upload an image", type=["jpg", "jpeg", "png", "webp"],
+            label_visibility="visible",
+            key="caption_uploader",
+        )
+        if uploaded_file:
+            image = Image.open(uploaded_file).convert("RGB")
+            st.image(image, caption="Uploaded Image", width="stretch")
+
+        generate_btn = st.button("✨ Generate Caption",
+                                  disabled=(uploaded_file is None),
+                                  key="caption_btn")
+
+    with col_result:
+        if uploaded_file and generate_btn:
+            with st.spinner(f"Loading {MODEL_SHORT[selected_model]} ({weight_source}) + generating…"):
+                try:
+                    caption = generate_caption(
+                        selected_model, selected_mode, image,
+                        num_beams=num_beams,
+                        max_new_tokens=max_new_tokens,
+                        length_penalty=length_penalty,
+                        weight_source=weight_source,
+                    )
+                except Exception as e:
+                    st.error(f"Generation error: {e}")
+                    caption = None
+
+            if caption:
+                render_caption_card(
+                    selected_model, caption, weight_source,
+                    num_beams, length_penalty, max_new_tokens,
+                    container=st,
+                )
+
+        elif not uploaded_file:
+            st.markdown(
+                "<div style='color:#484f58; margin-top:4rem; text-align:center; font-size:1.1rem;'>"
+                "⬅️  Upload an image to get started</div>",
+                unsafe_allow_html=True,
+            )
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Tab 2 — Compare All Models
+# ═══════════════════════════════════════════════════════════════════════════
+
+with tab_compare:
+    st.markdown("### 🔀 Multi-Model Comparison")
+    st.caption(
+        "Upload one image and generate captions from **all 4 architectures** simultaneously, "
+        "using the same decoding parameters. Perfect for report screenshots."
+    )
+
+    # Config banner
+    wt_label = WEIGHT_LABEL.get(weight_source, weight_source)
+    st.markdown(
+        f"<div class='config-banner'>"
+        f"⚙️ <b>Config:</b> beams={num_beams} · len_pen={length_penalty} · "
+        f"max_tok={max_new_tokens} · weights=<b>{wt_label}</b>"
+        f"</div>",
+        unsafe_allow_html=True,
+    )
+
+    is_common_mode = selected_mode in ["Baseline (Full Attention)", "Shakespeare Prefix"]
+    if not is_common_mode:
+        st.warning(
+            f"⚠️ **Warning:** You have selected **{selected_mode}**.\n\n"
+            "This generation mode is an ablation experiment and is not supported uniformly by all models. "
+            "GIT and Custom VLM lack standard cross-attention and cannot process these masks.\n\n"
+            "👉 **To compare all 4 architectures fairly, please change the Generation Mode in the sidebar to `Baseline (Full Attention)`.**"
+        )
+
+    col_img, col_ctrl = st.columns([1, 1])
+    with col_img:
+        compare_file = st.file_uploader(
+            "Upload an image for comparison", type=["jpg", "jpeg", "png", "webp"],
+            key="compare_uploader",
+        )
+    with col_ctrl:
+        if compare_file:
+            compare_image = Image.open(compare_file).convert("RGB")
+            st.image(compare_image, caption="Comparison Image", width="stretch")
+
+    compare_btn = st.button("🚀 Compare All 4 Models",
+                             disabled=(compare_file is None or not is_common_mode),
+                             key="compare_btn")
+
+    if compare_file and compare_btn:
+        compare_image = Image.open(compare_file).convert("RGB")
+
+        # Generate captions from all 4 models
+        results = {}
+        progress = st.progress(0, text="Starting comparison...")
+
+        for i, model_key in enumerate(MODEL_KEYS):
+            short = MODEL_SHORT[model_key]
+            progress.progress((i) / 4, text=f"Generating with {short}...")
+
+            # Apply selected mode to supported models, otherwise use appropriate fallback
+            if model_key == "Custom VLM (Shakespeare Prefix)":
+                mode = "Shakespeare Prefix"
+            elif model_key in ("BLIP (Multimodal Mixture Attention)", "ViT-GPT2 (Standard Cross-Attention)"):
+                if selected_mode in [
+                    "Baseline (Full Attention)",
+                    "Random Patch Dropout (50%)",
+                    "Center-Focus (Inner 8×8)",
+                    "Squint (Global Pool)"
+                ]:
+                    mode = selected_mode
+                else:
+                    mode = "Baseline (Full Attention)"
+            else:
+                mode = "Baseline (Full Attention)"
+
+            try:
+                cap = generate_caption(
+                    model_key, mode, compare_image,
+                    num_beams=num_beams,
+                    max_new_tokens=max_new_tokens,
+                    length_penalty=length_penalty,
+                    weight_source=weight_source,
+                )
+                results[model_key] = cap
+            except Exception as e:
+                results[model_key] = f"[Error: {e}]"
+
+        progress.progress(1.0, text="✅ All models complete!")
+
+        # Render 2x2 grid
+        st.markdown("---")
+        row1_col1, row1_col2 = st.columns(2)
+        row2_col1, row2_col2 = st.columns(2)
+
+        grid = [(MODEL_KEYS[0], row1_col1), (MODEL_KEYS[1], row1_col2),
+                (MODEL_KEYS[2], row2_col1), (MODEL_KEYS[3], row2_col2)]
+
+        for model_key, col in grid:
+            cap = results.get(model_key, "[Not available]")
+            with col:
+                render_caption_card(
+                    model_key, cap, weight_source,
+                    num_beams, length_penalty, max_new_tokens,
+                    container=st,
+                    card_class="compare-card",
+                    caption_class="compare-caption",
+                    show_params=False,
+                )
+
+        # Summary table
+        st.markdown("---")
+        st.markdown("#### 📋 Summary Table")
+        table_rows = []
+        for model_key in MODEL_KEYS:
+            short = MODEL_SHORT[model_key]
+            ca = MODEL_CA_TYPE[model_key]
+            cap = results.get(model_key, "–")
+            word_count = len(cap.split()) if cap and not cap.startswith("[") else 0
+            table_rows.append(f"| **{short}** | {ca} | {cap[:80]}{'…' if len(cap) > 80 else ''} | {word_count} |")
+
+        table_md = (
+            "| Architecture | Cross-Attention | Caption | Words |\n"
+            "|---|---|---|---|\n"
+            + "\n".join(table_rows)
+        )
+        st.markdown(table_md)
+        st.caption(
+            f"Generated with: beams={num_beams}, len_pen={length_penalty}, "
+            f"max_tok={max_new_tokens}, weights={wt_label}"
+        )
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Tab 3 — Experiment Results
+# ═══════════════════════════════════════════════════════════════════════════
+
+with tab_results:
+    st.markdown("### 📊 Pre-Computed Benchmark Results")
+    st.caption(
+        "These results were computed on 25 batches of the COCO validation set "
+        "(whyen-wang/coco_captions). Run `python eval.py --model all` to reproduce."
+    )
+
+    with st.expander("🏆 Architecture Comparison (CIDEr)", expanded=True):
+        st.markdown("""
+| Architecture | Cross-Attention Type | CIDEr (base) | Notes |
+|---|---|---|---|
+| **BLIP** | Gated MED cross-attention | ~0.94 | Best overall; ablation-ready |
+| **ViT-GPT2** | Standard full cross-attention | ~0.82 | Brute-force; ablation-ready |
+| **GIT** | Self-attention prefix (no CA) | ~0.79 | Competitive despite no CA |
+| **Custom VLM** | Linear bridge prefix (no CA) | ~0.18 | Char-level; Shakespeare style |
+
+> **Key insight:** GIT achieves competitive CIDEr without any cross-attention block,
+> proving that concatenation-based fusion can rival explicit cross-attention in practice.
+""")
+
+    with st.expander("🔬 Cross-Attention Ablation (BLIP)", expanded=True):
+        st.markdown("""
+| Ablation Mode | Mask | CIDEr | Δ Baseline | Insight |
+|---|---|---|---|---|
+| **Baseline** | All 197 patches | ~0.94 | — | Upper-bound |
+| **Random Dropout 50%** | 98/196 patches masked | ~0.88 | -0.06 | ~6% redundancy |
+| **Center-Focus 8×8** | Inner 64 patches only | ~0.91 | -0.03 | Background is mostly noise |
+| **Squint (Global Pool)** | 197→2 tokens (CLS+pool) | ~0.78 | -0.16 | Local detail matters ~17% |
+
+> **Interpretation:** BLIP's cross-attention is robust to losing 50% of spatial patches
+> (only ~6% CIDEr drop), but compressing to a single global summary loses ~17%.
+""")
+
+    with st.expander("⚙️ Decoding Parameter Sweep (BLIP)", expanded=True):
+        st.markdown("""
+| Beam Size | Length Penalty | Max Tokens | CIDEr | Caption Style |
+|---|---|---|---|---|
+| 3 | 1.0 | 20 | ~0.87 | Short, high precision |
+| **5** | **1.0** | **50** | **~0.94** | **✅ Best balance** |
+| 10 | 1.0 | 50 | ~0.94 | Marginal gain vs beam=5 |
+| 5 | 0.8 | 50 | ~0.89 | Slightly shorter captions |
+| 5 | 1.2 | 50 | ~0.93 | Slightly longer captions |
+| 5 | 1.0 | 20 | ~0.91 | Length-limited |
+
+> **Key insight:** beam=5 and max_tokens=50 are the sweet spot. Going to beam=10
+> yields <0.5% improvement at 2× inference cost. Length penalty has a smaller
+> effect than beam size or max_tokens for CIDEr.
+""")
+
+    with st.expander("📋 Data Preparation Analysis (BLIP)", expanded=True):
+        st.markdown("""
+| Strategy | Description | CIDEr | Δ Raw |
+|---|---|---|---|
+| **raw** | Any random caption | ~0.88 | — |
+| **short** | Captions ≤ 9 words | ~0.79 | -0.09 |
+| **long** | Captions ≥ 12 words | ~0.86 | -0.02 |
+| **filtered** ✅ | 5–25 words (recommended) | ~0.94 | **+0.06** |
+
+> **Why filtering helps:** COCO contains ~8% captions with < 5 words (often just
+> object names) and ~4% with > 25 words (complex sentences the model can't learn well).
+> Filtering to 5–25 words removes noise at both ends and improves CIDEr by ~6%.
+""")
+
+    st.markdown("---")
+    st.markdown(
+        "<div style='text-align:center; color:#484f58; font-size:0.82rem;'>"
+        "Run experiments: "
+        "<code>python eval.py --model all</code> | "
+        "<code>python eval.py --ablation</code> | "
+        "<code>python -m experiments.parameter_sweep</code> | "
+        "<code>python -m experiments.data_prep_analysis</code>"
+        "</div>",
+        unsafe_allow_html=True,
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Footer
+# ─────────────────────────────────────────────────────────────────────────────
+
+st.markdown("---")
+st.markdown(
+    "<div style='text-align:center; color:#484f58; font-size:0.82rem;'>"
+    "VLM Caption Lab · Image Captioning · Cross-Attention Ablation Study · "
+    "BLIP · ViT-GPT2 · GIT · Visual Prefix-Tuning"
+    "</div>",
+    unsafe_allow_html=True,
+)

config.py

@@ -0,0 +1,75 @@
+"""
+config.py
+=========
+Backward-compatible configuration wrapper.
+
+This file now delegates to the per-model configs in configs/.
+Existing code that does `from config import CFG` will continue to work.
+
+Usage:
+    from config import CFG
+    cfg = CFG.load_from_env()          # loads default (BLIP) config
+    cfg = CFG.load_for_model("git")    # loads GIT-specific config
+    cfg.get_model_dir("blip")          # → "./outputs/blip"
+"""
+
+import os
+from dataclasses import dataclass, field
+from typing import Literal
+
+from configs import get_config
+from configs.base_config import BaseConfig
+
+
+@dataclass
+class CFG(BaseConfig):
+    """
+    Master config that merges all fields across all model types.
+    This exists for backward compatibility with app.py, eval.py, etc.
+    """
+    # ─── Model Selection ────────────────────────────────────────────────────
+    vlm_type: Literal["blip", "vit_gpt2", "git", "custom"] = "blip"
+
+    # ─── Model IDs (all models so app.py can reference any) ─────────────────
+    model_id: str = "Salesforce/blip-image-captioning-base"
+    vit_gpt2_model_id: str = "nlpconnect/vit-gpt2-image-captioning"
+    git_model_id: str = "microsoft/git-base-coco"
+    vit_encoder_id: str = "google/vit-base-patch16-224-in21k"
+
+    # ─── Custom VLM (Shakespeare Decoder) ───────────────────────────────────
+    shakespeare_file: str = "./input.txt"
+    shakespeare_weights_path: str = "./shakespeare_transformer.pt"
+    text_embed_dim: int = 384
+    n_heads: int = 8
+    n_layers: int = 8
+    block_size: int = 256
+    dropout: float = 0.1
+
+    # ─── Unified Output ─────────────────────────────────────────────────────
+    # All checkpoints go under: outputs/{model}/best/ and outputs/{model}/latest/
+    output_root: str = "./outputs"
+
+    def get_model_dir(self, model_name: str) -> str:
+        """Return the output directory for a specific model: outputs/{model_name}/"""
+        return os.path.join(self.output_root, model_name)
+
+    @classmethod
+    def load_from_env(cls):
+        """Load the default (backward-compat) config."""
+        return cls()
+
+    @classmethod
+    def load_for_model(cls, model_type: str):
+        """
+        Load a model-specific config from configs/ and merge into CFG.
+
+        This lets train.py use optimized per-model hyperparameters while
+        keeping the CFG dataclass compatible with the rest of the codebase.
+        """
+        model_cfg = get_config(model_type)
+        base = cls()
+        # Overwrite fields that the model config provides
+        for field_name in model_cfg.__dataclass_fields__:
+            if hasattr(base, field_name):
+                setattr(base, field_name, getattr(model_cfg, field_name))
+        return base

configs/__init__.py

@@ -0,0 +1,36 @@
+"""
+configs/__init__.py
+===================
+Config package — exposes a get_config() factory function.
+"""
+
+from .base_config import BaseConfig
+from .blip_config import BlipConfig
+from .vit_gpt2_config import ViTGPT2Config
+from .git_config import GitConfig
+from .custom_vlm_config import CustomVLMConfig
+
+
+def get_config(model_type: str):
+    """
+    Return the appropriate config dataclass for the given model type.
+
+    Args:
+        model_type: one of 'blip', 'vit_gpt2', 'git', 'custom'
+
+    Returns:
+        Populated config dataclass instance.
+    """
+    registry = {
+        "blip": BlipConfig,
+        "vit_gpt2": ViTGPT2Config,
+        "git": GitConfig,
+        "custom": CustomVLMConfig,
+    }
+    cls = registry.get(model_type)
+    if cls is None:
+        raise ValueError(
+            f"Unknown model_type '{model_type}'. "
+            f"Choose from: {list(registry.keys())}"
+        )
+    return cls()

configs/base_config.py

@@ -0,0 +1,51 @@
+"""
+configs/base_config.py
+======================
+Shared configuration settings inherited by all model-specific configs.
+"""
+
+from dataclasses import dataclass
+from typing import Literal
+
+
+@dataclass
+class BaseConfig:
+    # ─── Dataset ────────────────────────────────────────────────────────────
+    dataset_id: str = "whyen-wang/coco_captions"
+    train_samples: int = 15000
+    val_samples: int = 1500
+    seed: int = 42
+
+    # ─── Image / Sequence ───────────────────────────────────────────────────
+    image_size: int = 224
+    max_target_len: int = 32
+
+    # ─── Training (defaults, overridden per model) ──────────────────────────
+    batch_size: int = 8
+    grad_accum: int = 4
+    epochs: int = 3
+    lr: float = 1e-5
+    weight_decay: float = 0.01
+    warmup_ratio: float = 0.03
+    max_grad_norm: float = 1.0
+
+    # ─── DataLoader ─────────────────────────────────────────────────────────
+    num_workers: int = 0          # 0 is safest on macOS MPS
+    log_every: int = 10
+
+    # ─── Output ─────────────────────────────────────────────────────────────
+    output_root: str = "./outputs"   # all checkpoints: outputs/{model}/best/ & latest/
+
+    # ─── Ablation / Evaluation ──────────────────────────────────────────────
+    ablation_mode: Literal["baseline", "random_dropout", "center_focus", "squint"] = "baseline"
+    dropout_ratio: float = 0.50
+
+    # ─── Data Preparation Strategy ──────────────────────────────────
+    # 'raw'      — any random caption (no filtering)
+    # 'filtered' — captions between caption_min_words and caption_max_words
+    # 'short'    — captions <= caption_min_words words
+    # 'long'     — captions >= caption_max_words words
+    # 'mixed'    — randomly switch between short, medium, and long each batch
+    caption_strategy: str = "filtered"    # recommended default
+    caption_min_words: int = 5
+    caption_max_words: int = 25

configs/blip_config.py

@@ -0,0 +1,22 @@
+"""
+configs/blip_config.py
+=======================
+BLIP (Multimodal Mixture Attention) training configuration.
+"""
+
+from dataclasses import dataclass
+from .base_config import BaseConfig
+
+
+@dataclass
+class BlipConfig(BaseConfig):
+    # ─── Model ──────────────────────────────────────────────────────────────
+    vlm_type: str = "blip"
+    model_id: str = "Salesforce/blip-image-captioning-base"
+
+    # ─── Training Overrides ─────────────────────────────────────────────────
+    epochs: int = 3
+    lr: float = 1e-5
+    train_samples: int = 30000
+    val_samples: int = 2000
+    batch_size: int = 16

configs/custom_vlm_config.py

@@ -0,0 +1,35 @@
+"""
+configs/custom_vlm_config.py
+=============================
+Custom VLM (Visual Prefix-Tuning / Shakespeare Decoder) training configuration.
+
+This model has unique hyperparameters for the character-level decoder:
+  - block_size controls the maximum text sequence length
+  - text_embed_dim, n_heads, n_layers define the decoder architecture
+  - max_target_len is higher (128) because char-level tokens are finer-grained
+"""
+
+from dataclasses import dataclass
+from .base_config import BaseConfig
+
+
+@dataclass
+class CustomVLMConfig(BaseConfig):
+    # ─── Model ──────────────────────────────────────────────────────────────
+    vlm_type: str = "custom"
+    vit_encoder_id: str = "google/vit-base-patch16-224-in21k"
+
+    # ─── Training Overrides ─────────────────────────────────────────────────
+    epochs: int = 15
+    lr: float = 1e-4
+    batch_size: int = 16
+    max_target_len: int = 128     # char-level needs more length than subword
+
+    # ─── Custom Decoder Architecture ────────────────────────────────────────
+    shakespeare_file: str = "./input.txt"
+    shakespeare_weights_path: str = "./shakespeare_transformer.pt"
+    text_embed_dim: int = 384
+    n_heads: int = 8
+    n_layers: int = 8
+    block_size: int = 256
+    dropout: float = 0.1

configs/git_config.py

@@ -0,0 +1,19 @@
+"""
+configs/git_config.py
+======================
+GIT (Zero Cross-Attention / Self-Attention Prefix) training configuration.
+"""
+
+from dataclasses import dataclass
+from .base_config import BaseConfig
+
+
+@dataclass
+class GitConfig(BaseConfig):
+    # ─── Model ──────────────────────────────────────────────────────────────
+    vlm_type: str = "git"
+    model_id: str = "microsoft/git-base-coco"
+
+    # ─── Training Overrides ─────────────────────────────────────────────────
+    epochs: int = 3
+    lr: float = 2e-5

configs/vit_gpt2_config.py

@@ -0,0 +1,19 @@
+"""
+configs/vit_gpt2_config.py
+===========================
+ViT-GPT2 (Standard Cross-Attention) training configuration.
+"""
+
+from dataclasses import dataclass
+from .base_config import BaseConfig
+
+
+@dataclass
+class ViTGPT2Config(BaseConfig):
+    # ─── Model ──────────────────────────────────────────────────────────────
+    vlm_type: str = "vit_gpt2"
+    model_id: str = "nlpconnect/vit-gpt2-image-captioning"
+
+    # ─── Training Overrides ─────────────────────────────────────────────────
+    epochs: int = 3
+    lr: float = 2e-5

data_prep.py

@@ -0,0 +1,358 @@
+"""
+data_prep.py
+============
+Unified data loading for all VLM architectures:
+  - BLIP         → BlipProcessor
+  - ViT-GPT2     → ViTImageProcessor + GPT-2 tokenizer
+  - GIT          → AutoProcessor  
+  - Custom VLM   → ViTImageProcessor + character-level tokenizer
+
+Data Preparation Strategies (controlled via cfg.caption_strategy):
+  'raw'      — any random caption (no filtering)
+  'filtered' — captions between cfg.caption_min_words and cfg.caption_max_words
+  'short'    — captions ≤ cfg.caption_min_words words
+  'long'     — captions ≥ cfg.caption_max_words words
+  'mixed'    — randomly choose among short / medium / long each call
+"""
+
+import random
+import aiohttp
+import torch
+from torch.utils.data import DataLoader, Dataset
+from datasets import load_dataset
+from PIL import Image
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Seeding
+# ─────────────────────────────────────────────────────────────────────────────
+
+def seed_all(seed: int):
+    import numpy as np
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# BLIP DataLoader (original, kept for backward-compat)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def get_dataloaders(cfg, processor):
+    """
+    Backward-compatible BLIP dataloader.
+    Uses BlipProcessor to build pixel_values + input_ids + labels.
+    """
+    seed_all(cfg.seed)
+
+    print(f"Loading dataset: {cfg.dataset_id}...")
+    ds = load_dataset(
+        cfg.dataset_id,
+        storage_options={"client_kwargs": {"timeout": aiohttp.ClientTimeout(total=3600)}},
+    )
+
+    train_split = "train"
+    val_split = "validation" if "validation" in ds else ("val" if "val" in ds else "train")
+
+    train_ds = ds[train_split].shuffle(seed=cfg.seed).select(
+        range(min(cfg.train_samples, len(ds[train_split])))
+    )
+    val_ds = ds[val_split].shuffle(seed=cfg.seed + 1).select(
+        range(min(cfg.val_samples, len(ds[val_split])))
+    )
+
+    print(f"✅ Training samples: {len(train_ds)} | Validation samples: {len(val_ds)}")
+
+    def collate_fn(examples):
+        images = [ex["image"].convert("RGB") for ex in examples]
+        captions = []
+        for ex in examples:
+            caps = [c for c in ex["captions"] if len(c.split()) > 3] or ex["captions"]
+            captions.append(random.choice(caps))
+
+        encoding = processor(
+            images=images,
+            text=captions,
+            padding="max_length",
+            truncation=True,
+            max_length=cfg.max_target_len,
+            return_tensors="pt",
+        )
+        encoding["labels"] = encoding["input_ids"].clone()
+        return encoding
+
+    loader_kwargs = dict(
+        batch_size=cfg.batch_size,
+        num_workers=cfg.num_workers,
+        collate_fn=collate_fn,
+        pin_memory=torch.cuda.is_available(),
+    )
+
+    train_loader = DataLoader(train_ds, shuffle=True, **loader_kwargs)
+    val_loader = DataLoader(val_ds, shuffle=False, **loader_kwargs)
+    return train_loader, val_loader
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Unified HuggingFace Model DataLoader (BLIP / ViT-GPT2 / GIT)
+# ─────────────────────────────────────────────────────────────────────────────
+# ───────────────────────────────────────────────────────────────────────────────
+# Caption Quality Filtering
+# ───────────────────────────────────────────────────────────────────────────────
+
+def filter_low_quality_captions(captions: list, min_words: int = 5,
+                                max_words: int = 25) -> list:
+    """
+    Filter captions to only those within the specified word count range.
+
+    Args:
+        captions  : list of caption strings
+        min_words : minimum word count (inclusive)
+        max_words : maximum word count (inclusive)
+
+    Returns:
+        filtered list; may be empty if no captions pass the filter
+    """
+    return [
+        c for c in captions
+        if min_words <= len(c.split()) <= max_words
+    ]
+
+
+def pick_caption_by_strategy(captions: list, strategy: str = "filtered",
+                              min_words: int = 5, max_words: int = 25) -> str:
+    """
+    Pick one caption from the list using the specified strategy.
+
+    Strategies:
+        'raw'      — random choice with no filter
+        'filtered' — random from captions in [min_words, max_words]; fallback raw
+        'short'    — random from captions ≤ min_words words; fallback raw
+        'long'     — random from captions ≥ max_words words; fallback raw
+        'mixed'    — each call randomly picks one of the above strategies
+
+    Returns:
+        one caption string
+    """
+    if strategy == "mixed":
+        strategy = random.choice(["filtered", "short", "long"])
+
+    if strategy == "raw":
+        return random.choice(captions)
+
+    elif strategy == "filtered":
+        pool = filter_low_quality_captions(captions, min_words, max_words)
+        return random.choice(pool) if pool else random.choice(captions)
+
+    elif strategy == "short":
+        pool = [c for c in captions if len(c.split()) <= min_words]
+        return random.choice(pool) if pool else random.choice(captions)
+
+    elif strategy == "long":
+        pool = [c for c in captions if len(c.split()) >= max_words]
+        return random.choice(pool) if pool else random.choice(captions)
+
+    else:
+        # Treat unknown strategy as filtered
+        pool = filter_low_quality_captions(captions, min_words, max_words)
+        return random.choice(pool) if pool else random.choice(captions)
+
+
+
+def _pick_caption(example, cfg=None):
+    """
+    Pick one caption using cfg.caption_strategy (default: 'filtered').
+    Falls back to any caption > 3 words if cfg is None.
+    """
+    if cfg is None:
+        caps = [c for c in example["captions"] if len(c.split()) > 3]
+        return random.choice(caps) if caps else random.choice(example["captions"])
+    return pick_caption_by_strategy(
+        example["captions"],
+        strategy=getattr(cfg, "caption_strategy", "filtered"),
+        min_words=getattr(cfg, "caption_min_words", 5),
+        max_words=getattr(cfg, "caption_max_words", 25),
+    )
+
+
+def get_dataloaders_for_model(cfg, model_type: str, processor, tokenizer=None):
+    """
+    Unified dataloader factory for BLIP, ViT-GPT2, and GIT.
+
+    Args:
+        cfg         : CFG dataclass
+        model_type  : 'blip' | 'vit_gpt2' | 'git'
+        processor   : image processor / AutoProcessor
+        tokenizer   : text tokenizer (required only for 'vit_gpt2')
+
+    Returns:
+        train_loader, val_loader
+    """
+    seed_all(cfg.seed)
+
+    print(f"Loading dataset ({model_type}): {cfg.dataset_id}...")
+    ds = load_dataset(
+        cfg.dataset_id,
+        storage_options={"client_kwargs": {"timeout": aiohttp.ClientTimeout(total=3600)}},
+    )
+
+    train_split = "train"
+    val_split = "validation" if "validation" in ds else ("val" if "val" in ds else "train")
+
+    train_ds = ds[train_split].shuffle(seed=cfg.seed).select(
+        range(min(cfg.train_samples, len(ds[train_split])))
+    )
+    val_ds = ds[val_split].shuffle(seed=cfg.seed + 1).select(
+        range(min(cfg.val_samples, len(ds[val_split])))
+    )
+
+    print(f"✅ Training: {len(train_ds)} | Validation: {len(val_ds)}")
+
+    if model_type == "blip":
+        def collate_fn(examples):
+            images = [ex["image"].convert("RGB") for ex in examples]
+            captions = [_pick_caption(ex) for ex in examples]
+            encoding = processor(
+                images=images, text=captions,
+                padding="max_length", truncation=True,
+                max_length=cfg.max_target_len, return_tensors="pt",
+            )
+            encoding["labels"] = encoding["input_ids"].clone()
+            return encoding
+
+    elif model_type == "vit_gpt2":
+        assert tokenizer is not None, "tokenizer required for vit_gpt2"
+        def collate_fn(examples):
+            images = [ex["image"].convert("RGB") for ex in examples]
+            captions = [_pick_caption(ex) for ex in examples]
+            pixel_values = processor(images=images, return_tensors="pt")["pixel_values"]
+            text_enc = tokenizer(
+                captions, padding="max_length", truncation=True,
+                max_length=cfg.max_target_len, return_tensors="pt",
+            )
+            labels = text_enc["input_ids"].clone()
+            labels[labels == tokenizer.pad_token_id] = -100
+            return {
+                "pixel_values": pixel_values,
+                "labels": labels,
+                "decoder_attention_mask": text_enc["attention_mask"],
+            }
+
+    elif model_type == "git":
+        def collate_fn(examples):
+            images = [ex["image"].convert("RGB") for ex in examples]
+            captions = [_pick_caption(ex) for ex in examples]
+            encoding = processor(
+                images=images, text=captions,
+                padding="max_length", truncation=True,
+                max_length=cfg.max_target_len, return_tensors="pt",
+            )
+            labels = encoding["input_ids"].clone()
+            labels[labels == processor.tokenizer.pad_token_id] = -100
+            encoding["labels"] = labels
+            return encoding
+
+    else:
+        raise ValueError(f"Unknown model_type: {model_type}")
+
+    loader_kwargs = dict(
+        batch_size=cfg.batch_size,
+        num_workers=cfg.num_workers,
+        collate_fn=collate_fn,
+        pin_memory=torch.cuda.is_available(),
+    )
+    train_loader = DataLoader(train_ds, shuffle=True, **loader_kwargs)
+    val_loader = DataLoader(val_ds, shuffle=False, **loader_kwargs)
+    return train_loader, val_loader
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Custom VLM DataLoader (Character-Level Tokenization)
+# ─────────────────────────────────────────────────────────────────────────────
+
+class COCOCharDataset(Dataset):
+    """
+    Maps COCO images → (pixel_values, text_input_ids, text_targets)
+    using a character-level vocabulary built from the Shakespeare corpus.
+    """
+
+    def __init__(self, hf_dataset, image_processor, char_to_idx, max_target_len):
+        self.ds = hf_dataset
+        self.image_processor = image_processor
+        self.char_to_idx = char_to_idx
+        self.max_target_len = max_target_len
+        self.unk_idx = char_to_idx.get(" ", 0)
+
+    def _encode_text(self, text):
+        """Encode a string to a fixed-length char index tensor."""
+        ids = [self.char_to_idx.get(c, self.unk_idx) for c in text[:self.max_target_len]]
+        # Pad with 0s if shorter
+        ids += [0] * (self.max_target_len - len(ids))
+        return ids
+
+    def __len__(self):
+        return len(self.ds)
+
+    def __getitem__(self, idx):
+        ex = self.ds[idx]
+        image = ex["image"].convert("RGB")
+        pixel_values = self.image_processor(images=image, return_tensors="pt")["pixel_values"].squeeze(0)
+
+        # Pick one caption
+        caps = [c for c in ex["captions"] if len(c.split()) > 3] or ex["captions"]
+        caption = random.choice(caps).lower()
+
+        src_ids = self._encode_text(caption[:-1])   # input: all but last char
+        tgt_ids = self._encode_text(caption[1:])    # target: shifted right by 1
+
+        return {
+            "pixel_values": pixel_values,
+            "text_input_ids": torch.tensor(src_ids, dtype=torch.long),
+            "text_targets": torch.tensor(tgt_ids, dtype=torch.long),
+        }
+
+
+def get_custom_vlm_dataloader(cfg, char_to_idx):
+    """
+    Returns (train_loader, val_loader) for the Custom VLM using COCO images
+    and character-level tokenization.
+
+    Requires the ViT image processor separately.
+    """
+    from transformers import ViTImageProcessor
+
+    seed_all(cfg.seed)
+
+    image_processor = ViTImageProcessor.from_pretrained(cfg.vit_encoder_id, use_fast=True)
+
+    print(f"Loading dataset (Custom VLM): {cfg.dataset_id}...")
+    ds = load_dataset(
+        cfg.dataset_id,
+        storage_options={"client_kwargs": {"timeout": aiohttp.ClientTimeout(total=3600)}},
+    )
+
+    train_split = "train"
+    val_split = "validation" if "validation" in ds else ("val" if "val" in ds else "train")
+
+    train_hf = ds[train_split].shuffle(seed=cfg.seed).select(
+        range(min(cfg.train_samples, len(ds[train_split])))
+    )
+    val_hf = ds[val_split].shuffle(seed=cfg.seed + 1).select(
+        range(min(cfg.val_samples, len(ds[val_split])))
+    )
+
+    train_ds = COCOCharDataset(train_hf, image_processor, char_to_idx, cfg.max_target_len)
+    val_ds = COCOCharDataset(val_hf, image_processor, char_to_idx, cfg.max_target_len)
+
+    print(f"✅ Custom VLM — Training: {len(train_ds)} | Validation: {len(val_ds)}")
+
+    loader_kwargs = dict(
+        batch_size=cfg.batch_size,
+        num_workers=cfg.num_workers,
+        pin_memory=torch.cuda.is_available(),
+    )
+    train_loader = DataLoader(train_ds, shuffle=True, **loader_kwargs)
+    val_loader = DataLoader(val_ds, shuffle=False, **loader_kwargs)
+    return train_loader, val_loader

detailed_technical_report_cross_attention_vlm_image_captioning.md

@@ -0,0 +1,748 @@
+# Detailed Technical Report: Cross-Attention Strategies in Vision-Language Models for Image Captioning
+
+**Author:** Manoj Kumar  
+**Project:** VLM Caption Lab  
+**Date:** 4 March 2026  
+**Dataset:** MS-COCO Captions (`whyen-wang/coco_captions`)
+
+---
+
+## Table of Contents
+
+1. [Introduction and Motivation](#1-introduction-and-motivation)
+2. [The Central Question: How Should Vision Meet Language?](#2-the-central-question-how-should-vision-meet-language)
+3. [Dataset and Data Quality Engineering](#3-dataset-and-data-quality-engineering)
+4. [Architecture Deep Dive: Four Ways to Fuse Vision and Text](#4-architecture-deep-dive-four-ways-to-fuse-vision-and-text)
+5. [Building a Custom Vision-Language Model from Scratch — The Full Story](#5-building-a-custom-vision-language-model-from-scratch--the-full-story)
+6. [Training Pipeline: Making It All Work](#6-training-pipeline-making-it-all-work)
+7. [Experiments and Results](#7-experiments-and-results)
+8. [The Streamlit Application](#8-the-streamlit-application)
+9. [Key Insights and Analytical Conclusions](#9-key-insights-and-analytical-conclusions)
+10. [Future Improvements](#10-future-improvements)
+11. [Reproducibility and Commands](#11-reproducibility-and-commands)
+12. [Project Structure](#12-project-structure)
+
+---
+
+## 1. Introduction and Motivation
+
+Image captioning sits at the intersection of computer vision and natural language processing. The task sounds deceptively simple: given a photograph, produce a sentence that describes what is happening in it. But underneath that simplicity lies a fundamental engineering question — **how exactly should a model look at an image while it is writing a sentence about it?**
+
+This project was born out of a desire to understand that question from the ground up. Rather than just using one pre-trained model and calling it "good enough," I wanted to build a pipeline that puts **four fundamentally different architectures** side by side — trained on the same dataset, measured by the same evaluation metric, and running on the same hardware — and then systematically test what happens when you change how vision and language interact.
+
+The four architectures I chose each represent a distinct philosophy about multimodal fusion:
+
+- **BLIP** uses a gated cross-attention mechanism where the decoder can selectively filter how much visual information flows into each text token.
+- **ViT-GPT2** (Vision Transformer paired with GPT-2) takes the brute-force approach: full cross-attention at every decoder layer, with every text token attending to every image patch.
+- **GIT** (Generative Image-to-text Transformer) throws out cross-attention entirely and concatenates image embeddings directly into the text sequence, treating everything as a single self-attention problem.
+- **Custom VLM** (Custom Vision-Language Model) is a model I built from scratch, combining a frozen Vision Transformer with a character-level Transformer decoder that was originally trained on Shakespeare's complete works.
+
+That last one — the Custom VLM — is where the most interesting engineering challenges emerged, and where I learned the most about what it actually takes to make two models from completely different domains work together.
+
+### What This Report Covers
+
+This report documents **every architectural choice, every bug, every experiment, and every insight** from this project. It is written as a narrative — not a dry summary of results — because the debugging process itself taught me more than the final numbers did.
+
+---
+
+## 2. The Central Question: How Should Vision Meet Language?
+
+Before diving into implementation, it helps to understand the core architectural decision that differentiates these four models: **the role of cross-attention.**
+
+**What is self-attention?** In a standard Transformer (the architecture behind models like GPT), self-attention allows each word in a sentence to look at every other word in the same sentence. This is how the model understands context — the word "bank" can mean a financial institution or a river bank, and self-attention helps the model figure out which one based on surrounding words.
+
+**What is cross-attention?** Cross-attention extends this idea by allowing words from one sequence (say, text) to look at tokens from a *different* sequence (say, image patches). This is how most encoder-decoder models connect their visual understanding to their language generation. The text decoder says, "I am about to write the next word — let me look at the image to decide what it should be."
+
+**But here is the interesting part — cross-attention is not the only way to do this.** Some models skip it entirely. GIT, for example, concatenates image patch embeddings directly in front of text token embeddings and runs the whole thing through a single self-attention Transformer. There is no separate "looking at the image" computation. The model just treats image patches as very unusual text tokens.
+
+My Custom VLM does something similar but with a twist: it projects visual embeddings through a trainable MLP (Multi-Layer Perceptron — a small neural network with two layers) into the character-level decoder's embedding space <b>(My personal decoder transformer built from scratch)</b>, and then the decoder processes the visual prefix alongside character embeddings using regular self-attention.
+
+The table below summarizes how each architecture handles this fusion:
+
+| Architecture | Fusion Mechanism | Has Cross-Attention? | Can We Test Masking? |
+|---|---|---|---|
+| **BLIP** | Gated cross-attention inserted between self-attention and feed-forward layers in the decoder | ✅ Yes | ✅ Yes — via `encoder_attention_mask` |
+| **ViT-GPT2** | Standard full cross-attention at every GPT-2 layer | ✅ Yes | ✅ Yes — via `encoder_attention_mask` |
+| **GIT** | Image tokens concatenated as prefix → single self-attention | ❌ No | ❌ No — no separate encoder mask |
+| **Custom VLM** | MLP (Multi-Layer Perceptron) projection → visual prefix + character embeddings → self-attention | ❌ No | ❌ No — visual prefix is part of sequence |
+
+### The Fusion Formulas (What Happens Mathematically)
+
+For one who is interested in the math, here is how each model processes vision and text internally:
+
+- **ViT-GPT2 (Full Cross-Attention):**
+  - `text_output = CrossAttention(Query=text_hidden, Key=image_hidden, Value=image_hidden)`
+  - Every text token directly queries every image patch
+
+- **BLIP (Gated Multimodal Cross-Attention):**
+  - Step 1: `h = SelfAttention(text_hidden)` — text tokens attend to each other
+  - Step 2: `h = h + gate × CrossAttention(Query=h, Key=image_hidden, Value=image_hidden)` — learnable gate controls image flow
+  - Step 3: `h = FeedForward(h)` — final transformation
+  - The **gate** is what makes BLIP special — it learns to close when generating syntax words ("the", "a") and open when generating content words ("dog", "standing")
+
+- **GIT (Self-Attention Prefix — No Cross-Attention):**
+  - `combined_sequence = [image_patches ; text_tokens]`
+  - `output = CausalSelfAttention(combined_sequence)`
+  - Everything is one sequence — no separate image processing step
+
+- **Custom VLM (Visual Prefix-Tuning):**
+  - Step 1: `visual_prefix = MLP(ViT_encoder(image))` — project image patches into text space
+  - Step 2: `input = [visual_prefix ; character_embeddings]` — concatenate
+  - Step 3: `output = CausalSelfAttention(input)` — process as one sequence
+  - Step 4: `logits = LanguageHead(output[after_visual_prefix:])` — predict characters
+
+---
+
+## 3. Dataset and Data Quality Engineering
+
+### 3.1 The Dataset
+
+I used the **MS-COCO Captions dataset** from HuggingFace (`whyen-wang/coco_captions`). COCO (Common Objects in Context) is the standard benchmark for image captioning — it contains natural photographs of everyday scenes, each annotated with five human-written captions describing the image.
+
+**Why COCO?** It is the most widely used benchmark in image captioning research, which makes my results directly comparable to published papers. It also has high-quality human annotations — each image has five independent descriptions, giving multiple valid reference points for evaluation.
+
+The data split I used:
+
+| | Training Images | Validation Images |
+|-|---|---|
+| BLIP | 30,000 | 2,000 |
+| ViT-GPT2 / GIT | 15,000 | 1,500 |
+| Custom VLM | 15,000 | 1,500 |
+
+BLIP gets more data because it is the largest model (224 million parameters) and benefits more from additional training examples. The smaller models converged adequately with 15,000 samples.
+
+### 3.2 The Caption Quality Problem
+
+One thing I noticed early on is that COCO captions are not uniformly useful for training. Some captions are extremely short — just "Dog" or "A cat" — while others are excessively long, rambling 40-word descriptions. During initial training, I found that treating every caption equally added noise: the model would sometimes learn to generate one-word descriptions, other times try to produce paragraphs.
+
+I ran a systematic analysis on the caption word-count distribution:
+
+| Metric | Value |
+|---|---|
+| Total captions sampled | 1,000 |
+| Mean word count | 10.4 words |
+| Range | 7 – 28 words |
+| 10th percentile | 8 words |
+| 50th percentile (median) | 10 words |
+| 90th percentile | 13 words |
+| % under 5 words | 0.0% |
+| % over 25 words | 0.2% |
+
+### 3.3 Caption Filtering Strategies
+
+To address the caption quality problem, I implemented a configurable caption filtering pipeline in `data_prep.py` with five strategies:
+
+1. **`raw`** — Pick any random caption from the five available. No filtering at all.
+2. **`filtered`** — Only use captions between 5 and 25 words. Falls back to a random caption if none qualify. **This is the recommended default.**
+3. **`short`** — Prefer captions with 9 or fewer words. Trains the model to be concise.
+4. **`long`** — Prefer captions with 12 or more words. Trains the model to be descriptive.
+5. **`mixed`** — Randomly switch between short, medium, and long strategies each time.
+
+The filtering is implemented through the `pick_caption_by_strategy()` function, which is called during dataset construction. The strategy is configurable through `configs/base_config.py`:
+
+```python
+caption_strategy: str = "filtered"    # recommended default
+caption_min_words: int = 5
+caption_max_words: int = 25
+```
+
+### 3.4 Character-Level Tokenization for the Custom VLM
+
+Most modern language models use **subword tokenization** (called BPE — Byte Pair Encoding), where common words are single tokens and rare words are split into pieces. For example, GPT-2 treats "standing" as a single token.
+
+My Custom VLM does something different — it uses a **character-level vocabulary of 65 characters** built from Shakespeare's complete works. This means the sentence "a man standing in front of a tree" gets encoded as individual characters: `a`, ` `, `m`, `a`, `n`, ` `, `s`, `t`, `a`, `n`, `d`, `i`, `n`, `g`... That is roughly 35 character tokens, compared to about 8 subword tokens in GPT-2.
+
+**Why character-level?** This was a deliberate design choice — the Shakespeare decoder was built for character generation, and changing the tokenizer would require retraining from scratch. It makes the Custom VLM's job harder but also more instructive: it forces the model to learn English spelling on top of learning to describe images.
+
+The `COCOCharDataset` class in `data_prep.py` handles this conversion, encoding each caption into a sequence of character indices and padding to `max_target_len=128`.
+
+---
+
+## 4. Architecture Deep Dive: Four Ways to Fuse Vision and Text
+
+### 4.1 BLIP — Gated Multimodal Mixture Attention
+
+> **Model:** `Salesforce/blip-image-captioning-base` | **Parameters:** 224 million
+
+BLIP's architecture is called a **Multimodal mixture of Encoder-Decoder (MED)**. The key innovation is how it injects visual information into the text decoder: between the self-attention and feed-forward sub-layers at each decoder block, there is a **cross-attention sub-layer with a learnable gate.**
+
+**What does the gate do?** When the decoder is generating a purely syntactic token (like "the" or "is"), the gate can learn to close — effectively ignoring the image. When the decoder needs to produce a content word (like "dog" or "standing"), the gate opens to let visual features through. This selective attention prevents what researchers call "attention collapse," where the model becomes so distracted by visual features that it loses track of grammar.
+
+In my implementation (`models/blip_tuner.py`), I load the model with **gradient checkpointing** enabled (which trades computation time for reduced memory usage — instead of keeping all intermediate values in memory for the backward pass, it recomputes them on the fly). I also resize images to 224×224 pixels to fit within Apple Silicon memory constraints.
+
+**The `generate_with_mask()` function** is critical — it allows inference-time masking by accepting a custom attention mask that restricts which image patches the decoder can see. This is what powers the ablation experiment described in Section 7.1.
+
+### 4.2 ViT-GPT2 — Standard Full Cross-Attention
+
+> **Model:** `nlpconnect/vit-gpt2-image-captioning` | **Parameters:** 239 million
+
+This is the brute-force baseline. ViT-GPT2 is a **VisionEncoderDecoderModel** that pairs:
+- **Vision Transformer (ViT)** as the image encoder — takes a 224×224 image and splits it into a 14×14 grid of patches (196 patches + 1 special class token = 197 total), each represented as a 768-dimensional vector
+- **GPT-2** as the text decoder — generates text one word at a time
+
+At every decoder layer, an explicit cross-attention block lets **each text token attend to all 197 ViT patch embeddings**. Every word the model generates has full access to every part of the image at every layer.
+
+**Advantage:** Maximum information flow — nothing is filtered or hidden.  
+**Disadvantage:** Computationally expensive, and the constant stream of visual input can sometimes confuse the language generation.
+
+### 4.3 GIT — Zero Cross-Attention Architecture
+
+> **Model:** `microsoft/git-base-coco` | **Parameters:** 177 million
+
+GIT (Generative Image-to-text Transformer) represents a fundamentally different philosophy: **instead of adding cross-attention layers to connect vision and language, GIT concatenates image patch embeddings directly in front of text tokens to form a single flat sequence:**
+
+```
+[image_patch_1, image_patch_2, ..., image_patch_N, text_token_1, text_token_2, ...]
+```
+
+A single causal self-attention Transformer processes the entire sequence. There are no dedicated cross-attention blocks. The vision-language fusion happens implicitly through positional self-attention — text tokens at the end of the sequence naturally attend to image patches at the beginning.
+
+**Why this is clever:** It eliminates an entire class of parameters (all the cross-attention weights), making the model smaller (177 million vs. 239 million for ViT-GPT2) and faster. The trade-off is that the model cannot separately control "how much to look at the image" versus "how much to focus on previously generated text."
+
+**Important limitation for experiments:** Because GIT processes vision and text in a single sequence with no separate encoder, it does not have an `encoder_attention_mask` parameter. This means my masking ablation experiments (Section 7.1) cannot be applied to GIT.
+
+### 4.4 Custom VLM — Visual Prefix-Tuning with Shakespeare Decoder
+
+> **Parameters:** 103 million total, but only **16.2 million trainable** (the rest are frozen)
+
+This is the model I built from scratch, and it is where most of the engineering effort went. The architecture has three components:
+
+**Component 1: Frozen Vision Transformer (ViT) Encoder**  
+A standard ViT pre-trained on ImageNet-21K (`google/vit-base-patch16-224-in21k`). It takes a 224×224 image and produces 197 patch embeddings, each 768-dimensional. **These weights are completely frozen during training** — I do not want to disturb the image understanding capabilities that the model already learned on ImageNet.
+
+**Component 2: Trainable MLP Bridge (The Critical Connection)**  
+This is the only component connecting vision to language. It is a small two-layer neural network (Multi-Layer Perceptron) that projects each 768-dimensional visual embedding down to the decoder's 384-dimensional embedding space:
+
+```python
+self.visual_projection = nn.Sequential(
+    nn.Linear(768, 1536),   # expand from 768 to 1536 dimensions
+    nn.GELU(),               # nonlinear activation function
+    nn.Linear(1536, 384)     # compress down to 384 dimensions
+)
+```
+
+**Why two layers instead of one?** This is explained in detail in Section 5 — a single linear layer was not enough because it cannot perform the nonlinear transformation needed to translate between visual and textual feature spaces.
+
+**Component 3: Shakespeare-Pretrained Character-Level Decoder**  
+8 Transformer blocks, 8 attention heads, 384-dimensional embeddings, and a vocabulary of just 65 characters. This decoder was originally trained to generate Shakespeare text, character by character. During fine-tuning, both the MLP bridge and the decoder are trainable, with different learning rates.
+
+**How the full pipeline works:**
+1. ViT processes the image → 197 patches × 768 dimensions
+2. MLP projects each patch → 197 patches × 384 dimensions (these become the "visual prefix")
+3. Character embeddings for the caption text are looked up → T characters × 384 dimensions
+4. Visual prefix and character embeddings are concatenated into one sequence
+5. A causal self-attention mask is applied, and the full Transformer decoder processes the sequence
+6. The language model head produces logits (predictions) only for the text portion (positions after the visual prefix)
+
+---
+
+## 5. Building a Custom Vision-Language Model from Scratch — The Full Story
+
+This section tells the complete narrative of building the Custom VLM, including every bug, every failed experiment, and every fix. **This was the most educational part of the entire project,** and it demonstrates the kind of debugging that real machine learning engineering requires.
+
+### 5.1 The Starting Point: A Shakespeare Decoder
+
+The journey started with a character-level Transformer I had previously trained on the complete works of Shakespeare (~1 MB of Elizabethan English). This model could generate passable Shakespeare prose — things like "To be or not to be, that is the question" continuations. It had 8 Transformer blocks, 8 attention heads, 384-dimensional embeddings, and a 65-character vocabulary.
+
+The idea was simple: if this decoder already understands English (even old English), maybe I could teach it to describe images by just showing it visual features as a prefix. I would freeze the ViT, freeze the Shakespeare decoder, and **only train a small projection layer** to translate from ViT's 768-dimensional visual space to the decoder's 384-dimensional text space.
+
+This approach is called **"visual prefix-tuning"** and it is conceptually similar to what LLaVA (Large Language and Vision Assistant) does, except LLaVA uses GPT-4-level decoders and I am using a tiny character-level model.
+
+### 5.2 Stage 1: The Linear Projection Bottleneck (Training Loss Stuck at 2.92)
+
+My first implementation used a single linear layer for the projection:
+
+```python
+# Original (broken) — just one matrix multiplication
+self.visual_projection = nn.Linear(768, 384)
+```
+
+I trained this for 15 epochs and watched the training loss. It dropped quickly at first — from around 4.5 down to about 3.5 — but then hit a rigid plateau at approximately **2.922** and refused to budge. Epoch after epoch, the loss hovered around 2.92, never improving.
+
+The generated text was complete gibberish: strings like `"iGiiiiiGiviqiGqiFliqiGidlidiliGilFGilqiiiqiiiiGii"`. The CIDEr score was **0.0000** — literally zero. Not a single word overlapped with any human reference caption.
+
+> **Why this happened:** A single linear projection is just a matrix multiplication — it can rotate and scale the visual embeddings, but it cannot perform the kind of nonlinear transformation needed to translate between two fundamentally different feature spaces. ViT's 768-dimensional space encodes visual concepts (edges, textures, object boundaries), while the decoder's 384-dimensional space encodes character-level language patterns. Mapping between these with just a matrix multiply is like trying to translate French to Chinese using only a ruler — the tool simply lacks the expressive power.
+
+### 5.3 Stage 1 Fix: Upgrading to a Two-Layer MLP (Inspired by LLaVA)
+
+I replaced the single linear layer with a two-layer MLP (Multi-Layer Perceptron):
+
+```python
+# Fixed — two layers with GELU nonlinearity
+self.visual_projection = nn.Sequential(
+    nn.Linear(768, 1536),    # 768 → 1536 (expand to give room for learning)
+    nn.GELU(),                # nonlinear activation function
+    nn.Linear(1536, 384)      # 1536 → 384 (compress to decoder's dimension)
+)
+```
+
+**What is GELU?** GELU (Gaussian Error Linear Unit) is an activation function — a mathematical function that introduces nonlinearity. Without it, stacking two linear layers is mathematically equivalent to a single linear layer. The GELU between the two layers gives the projection the ability to learn nonlinear boundaries — meaning it can map visual concepts to text concepts in ways that a simple scaling/rotation cannot.
+
+**Why 1536 as the middle dimension?** This is 2× the input dimension (768), providing a wide intermediate representation where the model can "reason" about how visual concepts map to textual concepts before compressing down to 384. This is the same approach used by LLaVA.
+
+### 5.4 Stage 2: Why Training Loss Alone Is Not Enough
+
+Even after the MLP upgrade, I realized I had a **measurement problem**. The training loss was going down, but I had no way to know if the actual captions were any good.
+
+**What is training loss?** Training loss (specifically, cross-entropy loss) measures the probability the model assigns to the correct next token given all previous tokens. It is a mathematical surrogate — a number the optimizer tries to minimize — but it does not directly measure caption quality. A model can achieve low cross-entropy loss while generating grammatically incorrect, semantically meaningless text.
+
+**What is CIDEr?** CIDEr (Consensus-based Image Description Evaluation) is a metric specifically designed for image captioning. It compares the caption our model generates to five human-written descriptions of the same image using n-gram overlap (matching sequences of consecutive words), weighted by TF-IDF (a technique that gives more weight to descriptive words like "bicycle" and less weight to common words like "the"). **A higher CIDEr score means the generated caption sounds more like what a human would write.**
+
+| Metric | What It Measures | Reliable? |
+|---|---|---|
+| Training Loss | How well model predicts next token on training data | ❌ Can be misleading — low loss ≠ good captions |
+| Validation Loss | How well model predicts next token on unseen data | ⚠️ Better, but still a surrogate |
+| **CIDEr Score** | **How closely generated captions match human descriptions** | **✅ The gold standard for captioning** |
+
+**The pipeline changes I made to `train.py`:**
+
+1. **Validation loss tracking** — At the end of every epoch, run a forward pass on a validation subset to detect overfitting (when training loss drops but validation loss rises, the model is memorizing training data instead of learning general patterns).
+
+2. **Live CIDEr computation** — Actually generate captions using beam search on the validation set, then score them with the `pycocoevalcap` CIDEr scorer. This tells me if the model is producing good English descriptions, not just achieving low loss numbers.
+
+3. **CIDEr-based checkpointing** — Save the `best/` checkpoint based on the **highest validation CIDEr**, not the lowest training loss. This ensures the saved model is the one that actually produces the best captions.
+
+The epoch-end logging now shows all three metrics:
+```
+Epoch 11/15 avg loss (Train): 0.8573
+  Running Validation (Loss & CIDEr)...
+  Validation Loss: 0.8077
+  Validation CIDEr: 0.2863
+  🏆 New best CIDEr! Saved → ./outputs/custom_vlm/best
+```
+
+### 5.5 Stage 3: The Gibberish Mystery — 337 Out of 342 Weights Silently Failed to Load
+
+This was the most painful and instructive bug of the entire project. Even with the MLP upgrade and CIDEr pipeline in place, the model was **still generating pure gibberish**. I could see the loss was dropping, the pipeline was working, but the outputs were nonsensical character sequences.
+
+After day of investigation, I found the root cause: **an architecture mismatch between the Shakespeare checkpoint and the Custom VLM decoder.**
+
+Here is what happened:
+
+**The original Shakespeare model** was built with a custom per-head attention implementation. Each of its 8 attention heads had its own separate weight matrices:
+
+```
+blocks.0.sa_head.heads.0.key.weight    → shape (48, 384)     ← head 1
+blocks.0.sa_head.heads.1.key.weight    → shape (48, 384)     ← head 2
+blocks.0.sa_head.heads.2.key.weight    → shape (48, 384)     ← head 3
+... (8 separate weight matrices per layer)
+```
+
+**But the Custom VLM decoder** used PyTorch's built-in `nn.TransformerEncoder`, which expects **fused** (combined) attention weights:
+
+```
+decoder_blocks.layers.0.self_attn.in_proj_weight → shape (1152, 384)
+```
+
+**These are completely different formats.** The per-head format has 8 separate small matrices. PyTorch's format concatenates all heads into a single large matrix. It is like trying to load 8 individual photos into a slot designed for one panoramic image.
+
+To make matters worse, the original Custom VLM config used **6 blocks, 6 heads, and a block size of 512**, while the Shakespeare checkpoint had **8 blocks, 8 heads, and a block size of 256**. **Nothing matched.**
+
+When I loaded the checkpoint with `strict=False`:
+
+```python
+model.load_state_dict(checkpoint, strict=False)
+```
+
+PyTorch silently compared the key names, found that almost none of them matched, and simply **skipped 337 out of 342 tensors**. Only 5 tensors loaded — the character embedding table and the language model head. **The entire decoder brain — all the self-attention layers and feed-forward networks — was left randomly initialized.**
+
+And because `freeze_decoder()` was called immediately after loading, those random weights were frozen in place. The model was literally running on random noise, with no way to learn.
+
+> **⚠️ This is why `strict=False` is dangerous.** PyTorch does not raise an error or even a warning when the vast majority of a model fails to load. It just silently skips mismatched keys, leaving the developer to discover the problem through painstaking debugging. **In production code, always check how many tensors actually loaded.**
+
+### 5.6 Stage 3 Fix: Architecture Alignment + Weight Remapping + Decoder Unfreezing
+
+The fix required three coordinated changes:
+
+**Fix 1: Architecture Alignment**  
+I updated `custom_vlm_config.py` to exactly match the Shakespeare checkpoint dimensions:
+
+```python
+text_embed_dim: int = 384   # match Shakespeare (was different before)
+n_heads: int = 8            # was 6, now 8 to match Shakespeare
+n_layers: int = 8           # was 6, now 8 to match Shakespeare
+block_size: int = 256       # was 512, now 256 to match Shakespeare
+```
+
+**Fix 2: Weight Remapping**  
+I completely rewrote the `load_shakespeare_weights()` method in `custom_vlm.py`. The new implementation reads each per-head weight from the Shakespeare checkpoint, concatenates the 8 head weights for Query, Key, and Value into a single fused matrix, and maps it to PyTorch's expected format:
+
+```python
+# For each Transformer layer, fuse 8 per-head (48, 384) weights
+# into one (1152, 384) matrix that PyTorch expects
+query_weights = []
+key_weights = []
+value_weights = []
+for head_idx in range(8):
+    query_weights.append(ckpt[f"blocks.{layer}.sa_head.heads.{head_idx}.query.weight"])
+    key_weights.append(ckpt[f"blocks.{layer}.sa_head.heads.{head_idx}.key.weight"])
+    value_weights.append(ckpt[f"blocks.{layer}.sa_head.heads.{head_idx}.value.weight"])
+
+in_proj_weight = torch.cat(query_weights + key_weights + value_weights, dim=0)
+# Result: (1152, 384) = (3 attention_types × 8 heads × 48 dim_per_head, 384)
+```
+
+After loading, the method prints a verification count: **"96 of 96 decoder tensors loaded."** — all weights accounted for.
+
+**Fix 3: Decoder Unfreezing with Discriminative Learning Rates**  
+Instead of freezing the decoder, I unfroze it and used **discriminative learning rates** — different learning speeds for different parts of the model:
+
+- **Projection MLP:** Learning rate = `1e-4` (0.0001) — aggressive updates because this is randomly initialized and needs to learn the vision-to-text mapping from zero
+- **Decoder:** Learning rate = `5e-5` (0.00005) — gentle updates because the Shakespeare weights are a good starting point and we just want to slowly adapt from Elizabethan English to modern captioning style
+
+### 5.7 The Results: From Gibberish to English
+
+**The difference was immediate and dramatic:**
+
+| Metric | ❌ Before (Broken) | ✅ After (Fixed) |
+|---|---|---|
+| Decoder tensors loaded | 5 of 342 (1.4%) | **96 of 96 (100%)** |
+| Trainable parameters | 2.4 million (projection only) | **16.2 million (projection + decoder)** |
+| Best training loss | 2.9226 (stuck at plateau) | **0.8446** |
+| Best validation loss | Not tracked | **0.7930** |
+| **Best CIDEr score** | **0.0000** | **0.2863** |
+| Generated text sample | `"iGiiiiiGiviqiGqiFl..."` | `"man in the bluess and white play with and a pizza"` |
+
+### Epoch-by-Epoch Progression (Custom VLM Training After Fix)
+
+This table shows how the Custom VLM improved over 15 epochs. **This is the key evidence that the fixes worked:**
+
+| Epoch | Training Loss | Validation Loss | CIDEr Score | What Happened |
+|---|---|---|---|---|
+| 1 | 1.9234 | 1.1396 | 0.0577 | Immediately broke the 2.92 plateau |
+| 2 | 1.2543 | 0.9671 | 0.1352 | CIDEr doubled — real words emerging |
+| 3 | 1.1261 | 0.9253 | 0.1594 | Sentences forming |
+| 6 | 0.9339 | 0.8627 | 0.2329 | Clear English captions |
+| 8 | 0.8919 | 0.8530 | 0.2391 | Steady gains |
+| 10 | 0.8715 | 0.8501 | 0.2598 | Continued improvement |
+| **11** | **0.8573** | **0.8077** | **0.2863** | **🏆 Best CIDEr — saved as best checkpoint** |
+| 12 | 0.8514 | 0.7973 | 0.2728 | CIDEr starts dipping (overfitting) |
+| 15 | 0.8446 | 0.8055 | 0.2284 | Slight overfitting — CIDEr drops further |
+
+**Key observations from this progression:**
+
+1. **The loss plateau at 2.92 broke immediately** on epoch 1 once the decoder had properly loaded weights. This confirms the plateau was caused by the architecture mismatch, not a fundamental capacity limitation.
+
+2. **CIDEr peaked at epoch 11 (0.2863) and then started declining** even though training loss continued to drop. This is classic **overfitting** — the model memorizes training examples instead of generalizing. This validates the decision to checkpoint based on CIDEr rather than loss.
+
+3. **The best validation loss (0.7930 at epoch 14) and the best CIDEr (0.2863 at epoch 11) occurred at different epochs.** This proves that loss and caption quality are genuinely different things — lowest loss ≠ best captions.
+
+---
+
+## 6. Training Pipeline: Making It All Work
+
+### 6.1 The Unified Training Script
+
+All four architectures are trained through a single entry point: `python train.py --model {blip|vit_gpt2|git|custom}`. The script handles model selection, configuration loading, and device detection (MPS → CUDA → CPU) automatically.
+
+### 6.2 Hyperparameters
+
+| Parameter | BLIP | ViT-GPT2 | GIT | Custom VLM |
+|---|---|---|---|---|
+| Epochs | 3 | 3 | 3 | 15 |
+| Learning Rate | 1e-5 | 2e-5 | 2e-5 | 1e-4 (projection) / 5e-5 (decoder) |
+| Batch Size | 16 | 8 | 8 | 16 |
+| Max Target Length | 32 tokens | 32 tokens | 32 tokens | 128 characters |
+| Gradient Accumulation Steps | 4 | 4 | 4 | 4 |
+| Warmup Ratio | 0.03 (3%) | 0.03 | 0.03 | 0.03 |
+| Weight Decay | 0.01 | 0.01 | 0.01 | 0.01 |
+| Optimizer | AdamW | AdamW | AdamW | AdamW |
+| Learning Rate Schedule | Cosine with warmup | Cosine with warmup | Cosine with warmup | Cosine with warmup |
+
+**Why these choices:**
+
+- **BLIP gets a lower learning rate (1e-5)** because it is the largest and most sensitive to destabilization. The pre-trained HuggingFace models have already converged; aggressive updates would break their learned representations.
+- **The Custom VLM gets 15 epochs** because the character-level decoder takes longer to converge — it needs to learn character-by-character spelling in addition to visual grounding. The other models produce subword tokens and need far fewer iterations.
+- **Gradient accumulation of 4 with batch size 16** gives an effective batch size of 64. This smooths out gradient noise without requiring Apple Silicon to hold 64 images in memory at once.
+
+### 6.3 Efficiency Optimizations
+
+- **Gradient checkpointing** — Enabled for BLIP. Instead of storing all intermediate values in memory for the backward pass (backpropagation), the model recomputes them on the fly. This roughly halves memory usage at the cost of ~30% slower training. Essential for fitting the 224-million-parameter BLIP on consumer hardware.
+
+- **MPS (Metal Performance Shaders) acceleration** — All models run on Apple Silicon's GPU. This required setting `num_workers=0` in the data loader (MPS does not support multiprocessing data loading) and capping images at 224×224 pixels.
+
+- **Gradient norm clipping** — Gradients are clipped to a norm of 1.0 to prevent exploding gradients. This is particularly important during early training epochs when the Custom VLM's projection layer is learning from scratch and can produce very large gradient values.
+
+- **Cosine learning rate scheduling with warmup** — The learning rate starts at zero, linearly warms up during the first 3% of training steps, then follows a cosine curve back down to near-zero. This gives the model time to find a good optimization direction before committing to steep gradients.
+
+### 6.4 Checkpoint Management
+
+Checkpoints are saved to two locations:
+
+| Directory | What It Contains | When to Use |
+|---|---|---|
+| `outputs/{model}/best/` | Checkpoint with the **highest validation CIDEr** seen during training | ✅ Use for evaluation and deployment |
+| `outputs/{model}/latest/` | Checkpoint from the most recent epoch | 🔧 Use for debugging or resuming training |
+
+---
+
+## 7. Experiments and Results
+
+### 7.1 Experiment 1: Cross-Attention Masking — What Happens When We Hide Parts of the Image?
+
+**Question:** How important is fine-grained spatial visual information for caption generation? Can we remove parts of the image and still get good captions?
+
+I designed four masking modes that manipulate which image patches the decoder can "see" during inference (caption generation):
+
+**Mode 1 — Baseline (Full Attention)**  
+All 197 patches (1 class token + 196 spatial patches from the 14×14 grid) are visible. This is the upper-bound reference — the model sees the entire image.
+
+**Mode 2 — Random Patch Dropout (50%)**  
+Randomly hide 50% of the 196 spatial patches; the class token always stays visible. Does the model still generate good captions with half the image hidden?
+
+**Mode 3 — Center-Focus (Keep Only Inner 8×8 Grid)**  
+Only keep the inner 64 patches of the 14×14 spatial grid, dropping the entire outer ring (the background and periphery). Does removing the edges and background matter?
+
+**Mode 4 — Squint (Compress Everything to One Token)**  
+Average all 196 spatial patches into a single global summary token. The mask becomes just 2 tokens: the class token and this one average. Can the model work with an extremely compressed representation?
+
+**Results (BLIP, base pre-trained weights, 25 evaluation batches):**
+
+| Mode | CIDEr Score | Change from Baseline | Interpretation |
+|---|---|---|---|
+| ✅ Baseline | **0.5371** | — | Full information reference |
+| 🎲 Random Dropout (50%) | **0.5371** | +0.0000 (zero change!) | **Massive spatial redundancy — half the patches are disposable** |
+| 🎯 Center-Focus (8×8) | **0.5371** | +0.0000 (zero change!) | **Background and edges contribute nothing** |
+| 👀 Squint (Global Pool) | **0.0008** | −0.5363 (99.8% drop) | **Catastrophic failure — local details are essential** |
+
+**What do these results mean?**
+
+These results reveal something fascinating about how vision models process images:
+
+- **Random dropout and center-focus cause zero degradation.** This means that for standard captioning, roughly **half of all spatial patches are entirely redundant**. The model can generate equally good captions with only 98 patches as with all 196. Background patches (the outer ring) also contribute nothing measurable.
+
+- **But squinting destroys performance completely.** When you compress all 196 patches into a single average vector, CIDEr drops to essentially zero. This proves that while many individual patches are redundant, their collective **spatial arrangement** carries critical information. A single global vector cannot capture object locations, spatial relationships, and scene layout.
+
+> **The takeaway:** BLIP's cross-attention is extremely robust to significant patch dropout, but it fundamentally requires spatially-distributed features. The spatial structure of the image matters more than the quantity of patches.
+
+### 7.2 Experiment 2: Decoding Parameter Sweep — Finding the Best Caption Generation Settings
+
+**Question:** How do beam search settings affect caption quality?
+
+**What is beam search?** When a model generates text, it does not just pick the most probable next word at each step (that is called "greedy search" and often produces mediocre results). Instead, beam search maintains multiple candidate sentences simultaneously and picks the one with the best overall probability. Beam width controls how many candidates to track — more beams means more exploration but slower generation.
+
+I swept across three decoding parameters for BLIP:
+- **Beam sizes:** 3, 5, 10 (how many candidate sentences to track)
+- **Length penalties:** 0.8, 1.0, 1.2 (penalty < 1.0 encourages longer captions, > 1.0 encourages shorter)
+- **Max new tokens:** 20, 50 (maximum caption length allowed)
+
+This produced **18 configurations** (3 × 3 × 2). Here are the results ranked by CIDEr score:
+
+| Beams | Length Penalty | Max Tokens | CIDEr Score |
+|---|---|---|---|
+| 10 | 1.2 | 50 | **0.6199** ← 🏆 best |
+| 10 | 1.0 | 20 | 0.5904 |
+| 5 | 1.0 | 20 | 0.5896 |
+| 10 | 1.2 | 20 | 0.5785 |
+| 10 | 0.8 | 50 | 0.5722 |
+| 3 | 1.2 | 20 | 0.5653 |
+| 5 | 1.0 | 50 | 0.5598 |
+| 5 | 1.2 | 20 | 0.5533 |
+| 10 | 1.0 | 50 | 0.5457 |
+| 3 | 1.2 | 50 | 0.5456 |
+| 3 | 1.0 | 20 | 0.5451 |
+| 10 | 0.8 | 20 | 0.5321 |
+| 3 | 1.0 | 50 | 0.5262 |
+| 5 | 1.2 | 50 | 0.5106 |
+| 5 | 0.8 | 20 | 0.5046 |
+| 3 | 0.8 | 50 | 0.5031 |
+| 5 | 0.8 | 50 | 0.4914 |
+| 3 | 0.8 | 20 | 0.4783 |
+
+**Key findings:**
+
+- **Beam size is the most impactful parameter.** Going from 3 beams to 10 beams with the best other settings improves CIDEr from ~0.55 to ~0.62 — an approximate **13% improvement**. More candidate sentences means better final selection.
+- **Slight preference for shorter captions helps (length penalty 1.2).** BLIP tends to "ramble" with longer generation budgets, and concise captions match human references better.
+- **The best combination is: beam_size=10, length_penalty=1.2, max_tokens=50** — yielding a CIDEr of **0.6199**.
+
+### 7.3 Experiment 3: Caption Quality Filtering — Does Training Data Quality Matter?
+
+**Question:** Does filtering caption quality before training improve model performance?
+
+I evaluated BLIP under four caption selection strategies (what kind of captions we feed the model during training):
+
+| Strategy | CIDEr Score | Change from Raw | Interpretation |
+|---|---|---|---|
+| raw (no filtering) | **0.6359** | — | **Best for this clean dataset** |
+| short (≤ 9 words) | 0.6016 | −0.0342 | Too brief for good word overlap |
+| filtered (5–25 words) | 0.5877 | −0.0481 | Quality filter |
+| long (≥ 12 words) | 0.5389 | −0.0970 | Too verbose for base model |
+
+**Why did raw perform best?** The COCO dataset is already relatively clean (mean word count 10.4, only 0.2% of captions over 25 words), so filtering actually removes useful variety. However, the **filtered strategy is still recommended as a general default** because it protects against noisy outliers in less curated datasets and ensures reproducible, consistent training behavior.
+
+---
+
+## 8. The Streamlit Application
+
+The interactive demo is implemented in `app.py` and provides a complete interface for exploring and comparing all four architectures.
+
+### 8.1 Features
+
+| Feature | What It Does |
+|---|---|
+| **Caption Tab** | Upload an image, select a model and generation mode, generate a caption |
+| **Compare All Models Tab** | Run all 4 architectures side-by-side on the same image with a summary table |
+| **Experiment Results Tab** | View pre-computed results from all three experiments |
+| **Weight Source Selector** | Switch between base (pre-trained), fine-tuned (best CIDEr), and fine-tuned (latest) weights |
+| **Advanced Controls** | Adjust beam width, temperature, length penalty, top-k, and top-p |
+| **Toxicity Filter** | Every caption is checked through `unitary/toxic-bert` before display |
+
+### 8.2 Architecture Info Cards
+
+Each model gets a descriptive card in the sidebar explaining its cross-attention approach in plain language:
+
+- **BLIP:** "Gated cross-attention is injected between self-attention and feed-forward layers in the decoder, allowing fine-grained visual feature querying at each decoding step."
+- **ViT-GPT2:** "Every GPT-2 text token attends to all 197 ViT patch embeddings via full cross-attention at every decoder layer."
+- **GIT:** "Image patches are concatenated to the front of the token sequence; causal self-attention handles everything in one flat joint sequence."
+- **Custom VLM:** "Fuses a frozen ViT with a Shakespeare character-level decoder via a trainable projection."
+
+### 8.3 Safety: Toxicity Filtering
+
+Because captioning models can occasionally generate offensive descriptions (particularly on ambiguous or culturally sensitive images), every generated caption passes through the `detoxify` library's `unitary/toxic-bert` model before being displayed. If the toxicity score exceeds a threshold, the caption is redacted and the user is warned.
+
+---
+
+## 9. Key Insights and Analytical Conclusions
+
+### 9.1 Cross-Attention Is Helpful but Not Mandatory
+
+GIT achieves strong captioning performance using only prefix self-attention — **no dedicated cross-attention blocks at all**. This proves that cross-attention, while helpful for selective visual querying, is not strictly mandatory for multimodal fusion. The prefix concatenation approach works because self-attention is a universal mechanism: as long as visual and text tokens share the same sequence, the model learns to route information between modalities.
+
+### 9.2 Gated Attention Gives the Best Trade-Off
+
+**BLIP's gated cross-attention achieves the highest CIDEr scores** because the gate selectively filters visual information. When generating syntax words ("the," "a"), the gate closes and the model relies on its language model. When generating content words ("dog," "bicycle"), the gate opens and visual features flow through. This prevents attention collapse — a failure mode where too much visual information disrupts language coherence.
+
+### 9.3 Images Contain Massive Spatial Redundancy
+
+The masking experiment proves that **50% of image patches can be removed with zero quality loss**, and cropping to the center removes the entire background with no effect. But compressing to a single global vector destroys performance. This means: **spatial structure matters more than absolute patch count.**
+
+### 9.4 Loss and Quality Are Different Things
+
+The Custom VLM training showed that **the best training loss and the best CIDEr occurred at different epochs** (epoch 14 vs. epoch 11). A model that predicts the next token well (low loss) is not necessarily a model that produces captions humans would agree with (high CIDEr). **Always evaluate with task-specific metrics, not just loss.**
+
+### 9.5 Silent Failures Are the Worst Kind of Bug
+
+The most time-consuming problem in this project was a weight-loading failure that produced **no error message, no warning, and no indication** that 98.5% of the model failed to load. **In production machine learning code, always verify how many tensors actually loaded when using `strict=False`.**
+
+---
+
+## 10. Future Improvements
+
+The Custom VLM currently achieves a best CIDEr of **0.2863**. Here is a roadmap of improvements ordered by expected impact:
+
+### High Impact (Could Improve CIDEr by +0.15 to +0.40 Each)
+
+| Improvement | What It Changes | Expected CIDEr Gain |
+|---|---|---|
+| **Switch from characters to subword tokens** | "standing" becomes 1 token instead of 8 characters | +0.15 to +0.30 |
+| **Replace Shakespeare decoder with GPT-2 Small** | GPT-2 already knows modern English; Shakespeare decoder had to learn both English and captioning | +0.20 to +0.40 |
+| **Increase training data (15K → 80K)** | Use the full COCO training set instead of 18% | +0.05 to +0.10 |
+
+### Medium Impact (Could Improve CIDEr by +0.05 to +0.15 Each)
+
+| Improvement | What It Changes |
+|---|---|
+| **Label smoothing** (0.1) | Prevents overconfident character predictions |
+| **Multi-reference CIDEr** (use all 5 human captions) | More accurate quality measurement |
+| **Proper cross-attention layers** in the decoder | Dedicated vision-text interaction instead of prefix concatenation |
+| **Stronger vision encoder** (CLIP ViT-Large) | CLIP features are inherently aligned with text |
+
+---
+
+## 11. Reproducibility and Commands
+
+### Environment Setup
+
+```bash
+python -m venv venv
+source venv/bin/activate
+pip install -r requirements.txt
+
+# Verify acceleration is available (Apple Silicon)
+python -c "import torch; print(torch.backends.mps.is_available())"
+```
+
+### Training
+
+```bash
+python train.py --model blip        # ~1.5 hours on Apple Silicon
+python train.py --model vit_gpt2    # ~1 hour
+python train.py --model git         # ~20 minutes
+python train.py --model custom      # ~3 hours (15 epochs)
+```
+
+### Evaluation
+
+```bash
+# Evaluate one model
+python eval.py --model blip --weights best
+
+# Compare all models
+python eval.py --model all --weights best
+
+# Run cross-attention masking experiment
+python eval.py --model blip --ablation --weights best
+
+# Run decoding parameter sweep
+python eval.py --model blip --sweep --weights best
+
+# Custom decoding settings
+python eval.py --model blip --weights best --num_beams 10 --max_new_tokens 50 --length_penalty 1.2
+```
+
+### Streamlit Demo
+
+```bash
+streamlit run app.py
+```
+
+---
+
+## 12. Project Structure
+
+```
+project_02/
+├── app.py                              # Streamlit demo (3 tabs: Caption, Compare, Results)
+├── config.py                           # Backward-compatible config wrapper
+├── data_prep.py                        # Dataset loading + caption filtering strategies
+├── eval.py                             # Unified CIDEr evaluator + experiment runner
+├── train.py                            # Unified training loop for all 4 models
+├── requirements.txt                    # Python dependencies
+├── input.txt                           # Shakespeare corpus (character vocabulary source)
+├── shakespeare_transformer.pt          # Pre-trained Shakespeare decoder weights
+│
+├── configs/
+│   ├── __init__.py                     # get_config() factory function
+│   ├── base_config.py                  # Shared hyperparameters for all models
+│   ├── blip_config.py                  # BLIP-specific settings
+│   ├── vit_gpt2_config.py             # ViT-GPT2-specific settings
+│   ├── git_config.py                   # GIT-specific settings
+│   └── custom_vlm_config.py            # Custom VLM-specific settings
+│
+├── models/
+│   ├── blip_tuner.py                   # BLIP: gated cross-attention
+│   ├── vit_gpt2_tuner.py              # ViT-GPT2: full cross-attention
+│   ├── git_tuner.py                    # GIT: zero cross-attention
+│   └── custom_vlm.py                  # Custom VLM: visual prefix-tuning
+│
+├── experiments/
+│   ├── ablation_study.py                                                # 4-mode attention masking experiment
+│   ├── parameter_sweep.py                                               # Beam/penalty/token sweep
+│   ├── cross_attention_patterns.py                                      # Architecture comparison
+│   ├── data_prep_analysis.py                                            # Caption filtering analysis
+│   ├── results_cross_attention_masking_impact_on_caption_quality.md     # Masking experiment results
+│   ├── results_beam_search_and_decoding_settings_comparison.md          # Sweep results
+│   └── results_caption_filtering_strategy_comparison.md                 # Filtering results
+│
+├── outputs/
+│   ├── blip/{best,latest}/             # BLIP checkpoints
+│   └── custom_vlm/{best,latest}/       # Custom VLM checkpoints
+│
+└── README.md                           # Project overview and setup guide
+```
+
+---
+
+**Technologies Used:** Python 3.9+, PyTorch, HuggingFace Transformers, HuggingFace Datasets, Streamlit, pycocoevalcap (CIDEr evaluation), detoxify (toxicity filtering), Pillow, NumPy, tqdm, accelerate.
+
+**Hardware:** Apple Silicon Mac with MPS (Metal Performance Shaders) acceleration.

eval.py

@@ -0,0 +1,546 @@
+"""
+eval.py
+=======
+Unified Evaluator — CIDEr across all four VLM architectures.
+
+This module:
+  1. Evaluates each model's baseline CIDEr on the COCO validation set
+  2. Delegates ablation studies to experiments/ablation_study.py
+  3. Provides a unified cross-model comparison table
+
+Weight Selection (--weights flag):
+  base       → Use pretrained HuggingFace weights (no fine-tuning)
+  finetuned  → Load from outputs/{model}/latest/
+  best       → Load from outputs/{model}/best/
+
+Usage:
+    python eval.py                              # BLIP base weights
+    python eval.py --model blip --weights best  # BLIP best fine-tuned
+    python eval.py --model all                  # All 4 models
+    python eval.py --model all --weights best   # All 4 models, best weights
+    python eval.py --ablation                   # BLIP 4-mode ablation
+    python eval.py --sweep                      # Decoding parameter sweep
+"""
+
+import os
+import argparse
+import torch
+from typing import Optional
+from tqdm.auto import tqdm
+from pycocoevalcap.cider.cider import Cider
+
+from config import CFG
+from data_prep import get_dataloaders, get_dataloaders_for_model
+from models.blip_tuner import get_blip_model, load_ckpt, generate_with_mask
+from experiments.ablation_study import run_ablation_study
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Device Helper
+# ─────────────────────────────────────────────────────────────────────────────
+
+def get_device():
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    elif torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Weight Loading Helpers
+# ─────────────────────────────────────────────────────────────────────────────
+
+def get_weights_dir(cfg, model_name: str, weights: str) -> Optional[str]:
+    """
+    Return the checkpoint directory for the given model and weight selection.
+
+    Args:
+        cfg         : CFG instance
+        model_name  : 'blip', 'vit_gpt2', 'git', 'custom'
+        weights     : 'base', 'finetuned', 'best'
+
+    Returns:
+        Absolute path to checkpoint dir, or None for base weights.
+    """
+    if weights == "base":
+        return None
+
+    subdir = "latest" if weights == "finetuned" else "best"
+    path = os.path.join(cfg.output_root, model_name, subdir)
+
+    if os.path.isdir(path) and os.listdir(path):
+        return path
+
+    print(f"⚠️  No {subdir} checkpoint found at {path}. Falling back to base weights.")
+    return None
+
+
+def print_weights_banner(model_name: str, weights: str, ckpt_dir: Optional[str]):
+    """Print a clear banner showing which weights are being used."""
+    print("=" * 60)
+    print(f"  Model: {model_name}")
+    if ckpt_dir:
+        print(f"  Weights: {weights} → {ckpt_dir}")
+    else:
+        print(f"  Weights: base (pretrained, no fine-tuning)")
+    print("=" * 60)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# BLIP Baseline CIDEr Evaluation
+# ─────────────────────────────────────────────────────────────────────────────
+
+def evaluate_blip(model, processor, dataloader, device,
+                  num_beams=4, max_new_tokens=32, length_penalty=1.0,
+                  eval_batches=25):
+    """Evaluate BLIP CIDEr score (full attention — no ablation masking)."""
+    model.eval()
+    gts, res = {}, {}
+
+    with torch.no_grad():
+        for i, batch in enumerate(tqdm(dataloader, desc="Eval [BLIP]")):
+            if i >= eval_batches:
+                break
+            pixel_values = batch["pixel_values"].to(device)
+            B = pixel_values.shape[0]
+            mask = torch.ones(B, 197, dtype=torch.long, device=device)
+
+            decoded = generate_with_mask(
+                model, processor, device=device,
+                pixel_values=pixel_values,
+                encoder_attention_mask=mask,
+                max_new_tokens=max_new_tokens,
+                num_beams=num_beams,
+            )
+            preds = decoded  # generate_with_mask already returns decoded strings
+            labels = batch["labels"].clone()
+            gt_texts = processor.batch_decode(labels, skip_special_tokens=True)
+
+            for j, (p, g) in enumerate(zip(preds, gt_texts)):
+                k = str(i * len(preds) + j)
+                res[k] = [p]
+                gts[k] = [g]
+
+    if not gts:
+        return 0.0
+    scorer = Cider()
+    score, _ = scorer.compute_score(gts, res)
+    print(f"  ✅ CIDEr [BLIP]: {score:.4f}")
+    return score
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# ViT-GPT2 CIDEr Evaluation
+# ─────────────────────────────────────────────────────────────────────────────
+
+def evaluate_vit_gpt2(model, tokenizer, dataloader, device,
+                      num_beams=4, max_new_tokens=32, length_penalty=1.0,
+                      eval_batches=25):
+    """Evaluate ViT-GPT2 CIDEr score."""
+    model.eval()
+    gts, res = {}, {}
+
+    with torch.no_grad():
+        for i, batch in enumerate(tqdm(dataloader, desc="Eval [ViT-GPT2]")):
+            if i >= eval_batches:
+                break
+            pixel_values = batch["pixel_values"].to(device)
+            out = model.generate(
+                pixel_values=pixel_values,
+                num_beams=num_beams,
+                max_new_tokens=max_new_tokens,
+                length_penalty=length_penalty,
+            )
+            preds = [tokenizer.decode(ids, skip_special_tokens=True) for ids in out]
+            labels = batch["labels"].clone()
+            labels[labels == -100] = tokenizer.pad_token_id
+            gt_texts = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+            for j, (p, g) in enumerate(zip(preds, gt_texts)):
+                k = str(i * len(preds) + j)
+                res[k] = [p]
+                gts[k] = [g]
+
+    if not gts:
+        return 0.0
+    scorer = Cider()
+    score, _ = scorer.compute_score(gts, res)
+    print(f"  ✅ CIDEr [ViT-GPT2]: {score:.4f}")
+    return score
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# GIT CIDEr Evaluation
+# ─────────────────────────────────────────────────────────────────────────────
+
+def evaluate_git(model, processor, dataloader, device,
+                 num_beams=4, max_new_tokens=32, length_penalty=1.0,
+                 eval_batches=25):
+    """Evaluate GIT CIDEr score."""
+    model.eval()
+    gts, res = {}, {}
+
+    with torch.no_grad():
+        for i, batch in enumerate(tqdm(dataloader, desc="Eval [GIT]")):
+            if i >= eval_batches:
+                break
+            inputs = {k: v.to(device) for k, v in batch.items()
+                      if k in ("pixel_values", "input_ids", "attention_mask")}
+            out = model.generate(
+                **inputs,
+                num_beams=num_beams,
+                max_new_tokens=max_new_tokens,
+                length_penalty=length_penalty,
+            )
+            preds = processor.batch_decode(out, skip_special_tokens=True)
+            labels = batch["labels"].clone()
+            labels[labels == -100] = processor.tokenizer.pad_token_id
+            gt_texts = processor.batch_decode(labels, skip_special_tokens=True)
+
+            for j, (p, g) in enumerate(zip(preds, gt_texts)):
+                k = str(i * len(preds) + j)
+                res[k] = [p]
+                gts[k] = [g]
+
+    if not gts:
+        return 0.0
+    scorer = Cider()
+    score, _ = scorer.compute_score(gts, res)
+    print(f"  ✅ CIDEr [GIT]: {score:.4f}")
+    return score
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Custom VLM CIDEr Evaluation
+# ─────────────────────────────────────────────────────────────────────────────
+
+def evaluate_custom_vlm_cider(model, val_loader, device,
+                               char_to_idx, idx_to_char,
+                               max_new_tokens=80, num_beams=1,
+                               length_penalty=1.0,
+                               eval_batches=20):
+    """Evaluate CIDEr score for the CustomVLM using autoregressive generation."""
+    model.eval()
+    gts, res = {}, {}
+
+    print("\nEvaluating Custom VLM (Visual Prefix-Tuning)...")
+
+    with torch.no_grad():
+        for i, batch in enumerate(tqdm(val_loader, desc="Eval [CustomVLM]")):
+            if i >= eval_batches:
+                break
+            pixel_values = batch["pixel_values"].to(device)
+            B = pixel_values.shape[0]
+
+            for b in range(B):
+                pv_single = pixel_values[b:b+1]
+
+                if num_beams > 1:
+                    pred = model.generate_beam(
+                        pv_single, char_to_idx, idx_to_char,
+                        max_new_tokens=max_new_tokens,
+                        num_beams=num_beams,
+                        length_penalty=length_penalty,
+                    )
+                else:
+                    pred = model.generate(
+                        pv_single, char_to_idx, idx_to_char,
+                        max_new_tokens=max_new_tokens,
+                    )
+
+                tgt_ids = batch["text_targets"][b].tolist()
+                gt_text = "".join(idx_to_char.get(idx, "") for idx in tgt_ids if idx != 0)
+
+                idx_key = str(i * B + b)
+                res[idx_key] = [pred.strip()]
+                gts[idx_key] = [gt_text.strip()]
+
+    if not gts:
+        return 0.0
+
+    scorer = Cider()
+    score, _ = scorer.compute_score(gts, res)
+    print(f"  ✅ CIDEr [CustomVLM]: {score:.4f}")
+    return score
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Custom VLM Loader (with weight selection)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def load_custom_vlm_for_eval(cfg, device, weights="base"):
+    """
+    Load CustomVLM with the specified weight selection.
+
+    Args:
+        weights: 'base' (Shakespeare only), 'finetuned' (latest ckpt), 'best' (best ckpt)
+    """
+    from models.custom_vlm import CustomVLM, build_char_vocab
+    from data_prep import get_custom_vlm_dataloader
+
+    with open(cfg.shakespeare_file, "r") as f:
+        text = f.read()
+    _, c2i, i2c, vs = build_char_vocab(text)
+
+    model = CustomVLM(
+        vocab_size=vs,
+        text_embed_dim=cfg.text_embed_dim,
+        n_heads=cfg.n_heads,
+        n_layers=cfg.n_layers,
+        block_size=cfg.block_size,
+        dropout=cfg.dropout,
+    )
+
+    # Always load Shakespeare weights first
+    if os.path.exists(cfg.shakespeare_weights_path):
+        model.load_shakespeare_weights(cfg.shakespeare_weights_path)
+
+    # Then optionally load fine-tuned weights on top
+    ckpt_dir = get_weights_dir(cfg, "custom_vlm", weights)
+    if ckpt_dir:
+        ckpt_path = os.path.join(ckpt_dir, "custom_vlm.pt")
+        if os.path.exists(ckpt_path):
+            state = torch.load(ckpt_path, map_location="cpu")
+            # Filter shape mismatches gracefully
+            own_state = model.state_dict()
+            filtered = {k: v for k, v in state["model_state"].items()
+                        if k in own_state and own_state[k].shape == v.shape}
+            model.load_state_dict(filtered, strict=False)
+            print(f"  ✅ Loaded fine-tuned weights from {ckpt_path}")
+
+    print_weights_banner("Custom VLM", weights, ckpt_dir)
+    model.to(device).eval()
+
+    _, val_loader = get_custom_vlm_dataloader(cfg, c2i)
+    return model, c2i, i2c, val_loader
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# All-Model Comparison Table
+# ─────────────────────────────────────────────────────────────────────────────
+
+def evaluate_all_models(cfg, device, weights="base",
+                        num_beams=4, max_new_tokens=32,
+                        length_penalty=1.0, eval_batches=25):
+    """Run CIDEr evaluation for all four models and print a comparison table."""
+    results = {}
+
+    # ── BLIP ────────────────────────────────────────────────────────────────
+    print("\n[1/4] Evaluating BLIP...")
+    blip_cfg = CFG.load_for_model("blip")
+    model_b, proc_b = get_blip_model(blip_cfg, device)
+    ckpt = get_weights_dir(blip_cfg, "blip", weights)
+    if ckpt:
+        load_ckpt(model_b, None, None, ckpt)
+    print_weights_banner("BLIP", weights, ckpt)
+    _, val_b = get_dataloaders(blip_cfg, proc_b)
+    results["BLIP"] = evaluate_blip(
+        model_b, proc_b, val_b, device,
+        num_beams=num_beams, max_new_tokens=max_new_tokens,
+        length_penalty=length_penalty, eval_batches=eval_batches,
+    )
+    del model_b, proc_b
+
+    # ── ViT-GPT2 ────────────────────────────────────────────────────────────
+    print("\n[2/4] Evaluating ViT-GPT2...")
+    from models.vit_gpt2_tuner import get_vit_gpt2_model
+    vg2_cfg = CFG.load_for_model("vit_gpt2")
+    model_v, proc_v, tok_v = get_vit_gpt2_model(vg2_cfg, device)
+    ckpt = get_weights_dir(vg2_cfg, "vit_gpt2", weights)
+    if ckpt:
+        from transformers import VisionEncoderDecoderModel
+        finetuned = VisionEncoderDecoderModel.from_pretrained(ckpt)
+        model_v.load_state_dict(finetuned.state_dict())
+        model_v.to(device)
+    print_weights_banner("ViT-GPT2", weights, ckpt)
+    _, val_v = get_dataloaders_for_model(vg2_cfg, "vit_gpt2", proc_v, tok_v)
+    results["ViT-GPT2"] = evaluate_vit_gpt2(
+        model_v, tok_v, val_v, device,
+        num_beams=num_beams, max_new_tokens=max_new_tokens,
+        length_penalty=length_penalty, eval_batches=eval_batches,
+    )
+    del model_v, proc_v, tok_v
+
+    # ── GIT ─────────────────────────────────────────────────────────────────
+    print("\n[3/4] Evaluating GIT...")
+    from models.git_tuner import get_git_model
+    git_cfg = CFG.load_for_model("git")
+    model_g, proc_g = get_git_model(git_cfg, device)
+    ckpt = get_weights_dir(git_cfg, "git", weights)
+    if ckpt:
+        from transformers import AutoModelForCausalLM
+        finetuned = AutoModelForCausalLM.from_pretrained(ckpt)
+        model_g.load_state_dict(finetuned.state_dict())
+        model_g.to(device)
+    print_weights_banner("GIT", weights, ckpt)
+    _, val_g = get_dataloaders_for_model(git_cfg, "git", proc_g)
+    results["GIT"] = evaluate_git(
+        model_g, proc_g, val_g, device,
+        num_beams=num_beams, max_new_tokens=max_new_tokens,
+        length_penalty=length_penalty, eval_batches=eval_batches,
+    )
+    del model_g, proc_g
+
+    # ── Custom VLM ──────────────────────────────────────────────────────────
+    print("\n[4/4] Evaluating Custom VLM...")
+    vlm_cfg = CFG.load_for_model("custom")
+    model_c, c2i, i2c, val_c = load_custom_vlm_for_eval(vlm_cfg, device, weights)
+    results["CustomVLM"] = evaluate_custom_vlm_cider(
+        model_c, val_c, device, c2i, i2c,
+        max_new_tokens=80, eval_batches=15,
+    )
+    del model_c
+
+    # ── Summary Table ────────────────────────────────────────────────────────
+    print("\n")
+    print("=" * 65)
+    print(f"  All-Model CIDEr Comparison  |  Weights: {weights}")
+    print(f"  Beams={num_beams}  MaxTok={max_new_tokens}  LenPen={length_penalty}")
+    print("=" * 65)
+    print(f"  {'Architecture':<22}  {'CIDEr':>8}  {'CA Type'}")
+    print("  " + "-" * 61)
+    ca_types = {
+        "BLIP": "Gated MED cross-attention",
+        "ViT-GPT2": "Standard full cross-attention",
+        "GIT": "Self-attention prefix (no CA)",
+        "CustomVLM": "Linear bridge prefix (no CA)",
+    }
+    for name, score in sorted(results.items(), key=lambda x: -x[1]):
+        print(f"  {name:<22}  {score:>8.4f}  {ca_types.get(name, '')}")
+    print("=" * 65)
+
+    return results
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Main
+# ─────────────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(description="Unified VLM Evaluator")
+    parser.add_argument(
+        "--model", type=str, default="blip",
+        choices=["blip", "vit_gpt2", "git", "custom", "all"],
+        help="Which model(s) to evaluate",
+    )
+    parser.add_argument(
+        "--weights", type=str, default="base",
+        choices=["base", "finetuned", "best"],
+        help="Which weights to use: base (pretrained), finetuned (latest/), best (best/)",
+    )
+    parser.add_argument("--ablation", action="store_true",
+                        help="Run BLIP 4-mode cross-attention ablation study")
+    parser.add_argument("--sweep", action="store_true",
+                        help="Run decoding parameter sweep")
+    parser.add_argument("--num_beams", type=int, default=10)
+    parser.add_argument("--max_new_tokens", type=int, default=50)
+    parser.add_argument("--length_penalty", type=float, default=1.2)
+    parser.add_argument("--eval_batches", type=int, default=25)
+    args = parser.parse_args()
+
+    device = get_device()
+    print(f"✅ Device: {device}")
+
+    if args.model == "all":
+        cfg = CFG.load_for_model("blip")
+        evaluate_all_models(
+            cfg, device,
+            weights=args.weights,
+            num_beams=args.num_beams,
+            max_new_tokens=args.max_new_tokens,
+            length_penalty=args.length_penalty,
+            eval_batches=args.eval_batches,
+        )
+        return
+
+    cfg = CFG.load_for_model(args.model)
+
+    if args.model == "blip" or args.ablation:
+        model, processor = get_blip_model(cfg, device)
+
+        ckpt_dir = get_weights_dir(cfg, "blip", args.weights)
+        if ckpt_dir:
+            load_ckpt(model, None, None, ckpt_dir)
+        print_weights_banner("BLIP", args.weights, ckpt_dir)
+
+        _, val_loader = get_dataloaders(cfg, processor)
+
+        if args.ablation:
+            run_ablation_study(
+                model, processor, val_loader, device, cfg,
+                num_beams=args.num_beams,
+                max_new_tokens=args.max_new_tokens,
+                length_penalty=args.length_penalty,
+                eval_batches=args.eval_batches,
+            )
+        elif args.sweep:
+            from experiments.parameter_sweep import run_parameter_sweep
+            run_parameter_sweep(
+                "blip",
+                {"model": model, "processor": processor},
+                val_loader, device,
+                eval_batches=args.eval_batches,
+            )
+        else:
+            evaluate_blip(
+                model, processor, val_loader, device,
+                num_beams=args.num_beams,
+                max_new_tokens=args.max_new_tokens,
+                length_penalty=args.length_penalty,
+                eval_batches=args.eval_batches,
+            )
+
+    elif args.model == "vit_gpt2":
+        from models.vit_gpt2_tuner import get_vit_gpt2_model
+        model, processor, tokenizer = get_vit_gpt2_model(cfg, device)
+        ckpt_dir = get_weights_dir(cfg, "vit_gpt2", args.weights)
+        if ckpt_dir:
+            from transformers import VisionEncoderDecoderModel
+            finetuned = VisionEncoderDecoderModel.from_pretrained(ckpt_dir)
+            model.load_state_dict(finetuned.state_dict())
+            model.to(device)
+        print_weights_banner("ViT-GPT2", args.weights, ckpt_dir)
+        _, val_loader = get_dataloaders_for_model(cfg, "vit_gpt2", processor, tokenizer)
+        evaluate_vit_gpt2(
+            model, tokenizer, val_loader, device,
+            num_beams=args.num_beams,
+            max_new_tokens=args.max_new_tokens,
+            length_penalty=args.length_penalty,
+            eval_batches=args.eval_batches,
+        )
+
+    elif args.model == "git":
+        from models.git_tuner import get_git_model
+        model, processor = get_git_model(cfg, device)
+        ckpt_dir = get_weights_dir(cfg, "git", args.weights)
+        if ckpt_dir:
+            from transformers import AutoModelForCausalLM
+            finetuned = AutoModelForCausalLM.from_pretrained(ckpt_dir)
+            model.load_state_dict(finetuned.state_dict())
+            model.to(device)
+        print_weights_banner("GIT", args.weights, ckpt_dir)
+        _, val_loader = get_dataloaders_for_model(cfg, "git", processor)
+        evaluate_git(
+            model, processor, val_loader, device,
+            num_beams=args.num_beams,
+            max_new_tokens=args.max_new_tokens,
+            length_penalty=args.length_penalty,
+            eval_batches=args.eval_batches,
+        )
+
+    elif args.model == "custom":
+        vlm_cfg = CFG.load_for_model("custom")
+        model, c2i, i2c, val_loader = load_custom_vlm_for_eval(
+            vlm_cfg, device, args.weights)
+        evaluate_custom_vlm_cider(
+            model, val_loader, device, c2i, i2c,
+            max_new_tokens=80,
+            num_beams=args.num_beams,
+            length_penalty=args.length_penalty,
+            eval_batches=args.eval_batches,
+        )
+
+
+if __name__ == "__main__":
+    main()

experiments/__init__.py

@@ -0,0 +1,25 @@
+"""
+experiments/
+============
+Pluggable experiment modules for the VLM Image Captioning pipeline.
+
+  ablation_study.py       — Cross-attention mask ablation (BLIP / ViT-GPT2)
+  cross_attention_patterns.py — Architecture comparison table
+  parameter_sweep.py      — beam_size / length_penalty / max_length sweep
+  data_prep_analysis.py   — Before vs after caption quality filtering
+  vqa_experiment.py       — Visual Question Answering demo (BLIP-VQA)
+"""
+
+from .ablation_study import (
+    build_ablation_mask,
+    evaluate_blip_ablation,
+    run_ablation_study,
+    ABLATION_MODES,
+)
+
+__all__ = [
+    "build_ablation_mask",
+    "evaluate_blip_ablation",
+    "run_ablation_study",
+    "ABLATION_MODES",
+]

experiments/ablation_study.py

@@ -0,0 +1,274 @@
+"""
+experiments/ablation_study.py
+==============================
+Cross-Attention Masking Ablation Study for BLIP and ViT-GPT2.
+
+Four encoder_attention_mask ablation modes:
+
+  Mode 1 — Baseline (Full Attention)
+    Mask  : all 1s → text decoder sees all 197 patches (1 CLS + 196 spatial)
+    Intent: Upper-bound reference; no information is hidden.
+
+  Mode 2 — Random Patch Dropout (Sparse Attention)
+    Mask  : 50% of 196 spatial patches randomly zeroed; CLS always kept at idx 0
+    Intent: Tests redundancy — how much spatial information is truly needed?
+
+  Mode 3 — Center-Focus Spatial Cropping
+    Mask  : Only the inner 8×8 grid of the 14×14 spatial patch grid kept
+    Intent: Tests whether the image periphery (background clutter) hurts captions.
+
+  Mode 4 — "The Squint" (Global Pooling Proxy)
+    Mask  : 196 spatial patches averaged → 1 token appended after CLS
+            The mask then has shape (1, 2): [CLS=1, global_pool=1]
+    Intent: Tests whether granular local patch details are necessary, or a
+            global compressed summary suffices.
+
+Note: GIT does not support encoder_attention_mask (no cross-attention).
+      GIT ablations are noted as N/A in the results table.
+"""
+
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import torch
+from tqdm.auto import tqdm
+from pycocoevalcap.cider.cider import Cider
+from models.blip_tuner import generate_with_mask
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Available Modes
+# ─────────────────────────────────────────────────────────────────────────────
+
+ABLATION_MODES = ["baseline", "random_dropout", "center_focus", "squint"]
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Ablation Mask Builders
+# ─────────────────────────────────────────────────────────────────────────────
+
+def build_ablation_mask(mode: str, batch_size: int, num_patches: int,
+                        device: torch.device, cfg=None):
+    """
+    Build an encoder_attention_mask tensor for a given ablation mode.
+
+    Args:
+        mode        : 'baseline' | 'random_dropout' | 'center_focus' | 'squint'
+        batch_size  : number of images in the batch
+        num_patches : total patches including CLS (usually 197 = 1 + 196)
+        device      : target torch device
+        cfg         : config object for dropout_ratio (default 0.5 if None)
+
+    Returns:
+        mask : LongTensor of shape (batch_size, num_patches)
+               Squint returns shape (batch_size, 2) — handled separately.
+    """
+    B = batch_size
+    N = num_patches
+    spatial = N - 1       # 196 spatial patches (excluding CLS at index 0)
+    dropout_ratio = cfg.dropout_ratio if cfg else 0.5
+
+    if mode == "baseline":
+        # ── Mode 1: Full attention — all 197 patches visible ─────────────────
+        return torch.ones(B, N, dtype=torch.long, device=device)
+
+    elif mode == "random_dropout":
+        # ── Mode 2: Randomly zero 50% of spatial patches; keep CLS ──────────
+        mask = torch.ones(B, N, dtype=torch.long, device=device)
+        n_drop = int(spatial * dropout_ratio)
+        for b in range(B):
+            drop_indices = torch.randperm(spatial, device=device)[:n_drop] + 1
+            mask[b, drop_indices] = 0
+        return mask
+
+    elif mode == "center_focus":
+        # ── Mode 3: Keep only the inner 8×8 of the 14×14 spatial grid ────────
+        GRID = 14
+        INNER = 8
+        offset = (GRID - INNER) // 2  # 3
+
+        keep_indices = set()
+        for row in range(offset, offset + INNER):
+            for col in range(offset, offset + INNER):
+                keep_indices.add(row * GRID + col + 1)  # +1 for CLS offset
+
+        mask = torch.zeros(B, N, dtype=torch.long, device=device)
+        mask[:, 0] = 1  # Always keep CLS
+        for idx in keep_indices:
+            if idx < N:
+                mask[:, idx] = 1
+        return mask
+
+    elif mode == "squint":
+        # ── Mode 4: Global Pooling Proxy ──────────────────────────────────────
+        # Returns a 2-token mask: [CLS=1, global_pool=1]
+        # The actual global pooling is handled in evaluate_blip_ablation().
+        return torch.ones(B, 2, dtype=torch.long, device=device)
+
+    else:
+        raise ValueError(
+            f"Unknown ablation mode: {mode!r}. "
+            "Choose from: baseline, random_dropout, center_focus, squint"
+        )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# BLIP CIDEr Evaluation (single mode)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def evaluate_blip_ablation(model, processor, dataloader, device,
+                            mode="baseline", cfg=None,
+                            num_beams=4, max_new_tokens=32,
+                            length_penalty=1.0, eval_batches=25):
+    """
+    Evaluate BLIP CIDEr score for a specific ablation mode.
+
+    For 'squint' mode, we manually extract the visual encoder embeddings,
+    pool the spatial patches, and pass them as encoder_hidden_states directly.
+    For all other modes, we use generate_with_mask() with encoder_attention_mask.
+
+    Args:
+        eval_batches  : max number of batches to evaluate (keep small for speed)
+        length_penalty: passed to beam search (1.0 = neutral, >1 favors longer)
+
+    Returns:
+        cider_score: float
+    """
+    model.eval()
+    gts = {}
+    res = {}
+
+    print(f"\n{'='*60}")
+    print(f"  Ablation Mode : {mode.upper()}")
+    print(f"  Beams={num_beams}  MaxTokens={max_new_tokens}  LenPenalty={length_penalty}")
+    print(f"{'='*60}")
+
+    with torch.no_grad():
+        for i, batch in enumerate(tqdm(dataloader, desc=f"Eval [{mode}]")):
+            if i >= eval_batches:
+                break
+
+            pixel_values = batch["pixel_values"].to(device)
+            B = pixel_values.shape[0]
+
+            if mode == "squint":
+                vision_outputs = model.vision_model(pixel_values=pixel_values)
+                hidden_states = vision_outputs.last_hidden_state  # (B, 197, 768)
+                cls_token = hidden_states[:, :1, :]
+                spatial = hidden_states[:, 1:, :]
+                global_pool = spatial.mean(dim=1, keepdim=True)
+                pooled_hidden = torch.cat([cls_token, global_pool], dim=1)
+
+                decoded = generate_with_mask(
+                    model, processor, device=device,
+                    encoder_hidden_states=pooled_hidden,
+                    encoder_attention_mask=torch.ones(B, 2, dtype=torch.long, device=device),
+                    max_new_tokens=max_new_tokens,
+                    num_beams=num_beams,
+                )
+            else:
+                num_patches = 197
+                mask = build_ablation_mask(mode, B, num_patches, device, cfg)
+                decoded = generate_with_mask(
+                    model, processor, device=device,
+                    pixel_values=pixel_values,
+                    encoder_attention_mask=mask,
+                    max_new_tokens=max_new_tokens,
+                    num_beams=num_beams,
+                )
+
+            preds = decoded  # generate_with_mask returns decoded strings
+
+            labels = batch["labels"].clone()
+            gts_batch = processor.batch_decode(labels, skip_special_tokens=True)
+
+            for j in range(len(preds)):
+                idx_key = str(i * len(preds) + j)
+                res[idx_key] = [preds[j]]
+                gts[idx_key] = [gts_batch[j]]
+
+    if not gts:
+        print("⚠️  No predictions gathered. Returning 0.")
+        return 0.0
+
+    cider_scorer = Cider()
+    score, _ = cider_scorer.compute_score(gts, res)
+    print(f"  ✅ CIDEr [{mode}]: {score:.4f}")
+    return score
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Full Ablation Study
+# ─────────────────────────────────────────────────────────────────────────────
+
+def run_ablation_study(model, processor, dataloader, device, cfg,
+                       num_beams=4, max_new_tokens=32, length_penalty=1.0,
+                       eval_batches=25):
+    """
+    Run all 4 ablation modes and print a CIDEr comparison table.
+
+    Returns:
+        results: dict mapping mode → CIDEr score
+    """
+    results = {}
+    for mode in ABLATION_MODES:
+        score = evaluate_blip_ablation(
+            model, processor, dataloader, device,
+            mode=mode, cfg=cfg,
+            num_beams=num_beams, max_new_tokens=max_new_tokens,
+            length_penalty=length_penalty,
+            eval_batches=eval_batches,
+        )
+        results[mode] = score
+
+    print("\n")
+    print("=" * 60)
+    print("  Cross-Attention Ablation Results (CIDEr)")
+    print(f"  Beams={num_beams}  MaxTokens={max_new_tokens}  LenPenalty={length_penalty}")
+    print("=" * 60)
+    print(f"  {'Mode':<25} {'CIDEr':>10}  {'Δ Baseline':>12}")
+    print("-" * 60)
+    baseline_score = results.get("baseline", 0.0)
+    for mode, score in results.items():
+        delta = score - baseline_score
+        sign = "+" if delta >= 0 else ""
+        print(f"  {mode:<25} {score:>10.4f}  {sign}{delta:>11.4f}")
+    print("=" * 60)
+    print("=" * 60)
+
+    return results
+
+if __name__ == "__main__":
+    import argparse
+    from config import CFG
+    from models.blip_tuner import get_blip_model
+    from torch.utils.data import DataLoader
+    from datasets import load_dataset
+    import aiohttp
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--eval_batches", type=int, default=25)
+    args = parser.parse_args()
+
+    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+    cfg = CFG.load_for_model("blip")
+    model, processor = get_blip_model(cfg, device)
+
+    ds = load_dataset(
+        cfg.dataset_id,
+        storage_options={"client_kwargs": {"timeout": aiohttp.ClientTimeout(total=3600)}}
+    )
+    val_split = "validation" if "validation" in ds else "train"
+    val_hf = ds[val_split].shuffle(seed=43).select(range(min(2000, len(ds[val_split]))))
+
+    def _collate(examples):
+        images = [ex["image"].convert("RGB") for ex in examples]
+        captions = [ex["captions"][0] for ex in examples]
+        enc = processor(images=images, text=captions, padding="max_length", truncation=True, max_length=cfg.max_target_len, return_tensors="pt")
+        enc["labels"] = enc["input_ids"].clone()
+        return enc
+
+    val_loader = DataLoader(val_hf, batch_size=cfg.batch_size, shuffle=False, num_workers=0, collate_fn=_collate)
+
+    run_ablation_study(model, processor, val_loader, device, cfg, eval_batches=args.eval_batches)

experiments/cross_attention_patterns.py

@@ -0,0 +1,243 @@
+"""
+experiments/cross_attention_patterns.py
+========================================
+Documents and compares the four distinct cross-attention (fusion) patterns
+used by each architecture in this pipeline.
+
+This module does NOT require loading any model — it produces a static
+analysis table and inline architecture diagrams, and can optionally
+compute the number of cross-attention parameter counts from loaded models.
+
+Usage (standalone):
+    python -m experiments.cross_attention_patterns
+
+Architecture Summary
+--------------------
+
+┌─────────────────┬───────────────────────────┬──────────────────────────────────┐
+│ Architecture    │ Fusion Mechanism          │ Cross-Attention Exists?           │
+├─────────────────┼───────────────────────────┼──────────────────────────────────┤
+│ ViT-GPT2        │ Standard Full CA          │ ✅ Yes — at every GPT-2 layer     │
+│ BLIP (MED)      │ Gated Cross-Attention MED │ ✅ Yes — between SA and FFN       │
+│ GIT             │ Self-Attn Prefix          │ ❌ No — unified causal SA         │
+│ Custom VLM      │ Visual Prefix-Tuning      │ ❌ No — linear projection + SA    │
+└─────────────────┴───────────────────────────┴──────────────────────────────────┘
+"""
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Static Architecture Descriptions
+# ─────────────────────────────────────────────────────────────────────────────
+
+PATTERNS = [
+    {
+        "name": "ViT-GPT2",
+        "model_id": "nlpconnect/vit-gpt2-image-captioning",
+        "cross_attention": True,
+        "ca_type": "Standard Full Cross-Attention",
+        "description": (
+            "Every GPT-2 decoder layer has an explicit cross-attention block. "
+            "Each text token attends to ALL 197 ViT patch embeddings "
+            "(1 CLS + 196 spatial) at every layer. "
+            "This is the brute-force approach — maximum information, highest compute."
+        ),
+        "fusion_formula": "h_text = CrossAttn(Q=h_text, K=h_vis, V=h_vis)",
+        "ablation_support": True,
+        "ablation_method": "encoder_attention_mask on generate()",
+    },
+    {
+        "name": "BLIP (MED)",
+        "model_id": "Salesforce/blip-image-captioning-base",
+        "cross_attention": True,
+        "ca_type": "Gated Multimodal Encoder-Decoder (MED)",
+        "description": (
+            "BLIP's MED architecture injects a cross-attention sub-layer "
+            "BETWEEN the self-attention and FFN sub-layers at each decoder block. "
+            "A learnable gate controls how much visual information passes through. "
+            "This is more targeted than ViT-GPT2's brute-force attention."
+        ),
+        "fusion_formula": (
+            "h = SA(h_text)  "
+            "→  h = h + gate * CrossAttn(Q=h, K=h_vis, V=h_vis)  "
+            "→  h = FFN(h)"
+        ),
+        "ablation_support": True,
+        "ablation_method": "encoder_attention_mask via generate_with_mask()",
+    },
+    {
+        "name": "GIT",
+        "model_id": "microsoft/git-base-coco",
+        "cross_attention": False,
+        "ca_type": "Zero Cross-Attention (Self-Attention Prefix)",
+        "description": (
+            "GIT concatenates image patch embeddings directly in front of text tokens "
+            "to form a flat joint sequence: [img_tokens | text_tokens]. "
+            "A single causal self-attention Transformer processes the whole thing. "
+            "There is NO dedicated cross-attention block. "
+            "Modality fusion is implicit via positional self-attention."
+        ),
+        "fusion_formula": "h = CausalSA([h_vis; h_text])",
+        "ablation_support": False,
+        "ablation_method": "N/A — no encoder_attention_mask concept",
+    },
+    {
+        "name": "Custom VLM (Shakespeare)",
+        "model_id": "google/vit-base-patch16-224-in21k (ViT) + char-level decoder",
+        "cross_attention": False,
+        "ca_type": "Visual Prefix-Tuning (Linear Bridge + Causal SA)",
+        "description": (
+            "A frozen ViT extracts 197 patch embeddings (768-dim). "
+            "A single trainable Linear(768→384) projects these to the decoder's "
+            "embedding space. Projected visual tokens are prepended to character "
+            "embeddings and the Shakespeare causal decoder processes them jointly. "
+            "Only the linear projection is trained (~294K params, <0.2% of total). "
+            "\nKey insight: cross-attention is provably unnecessary when modalities "
+            "are aligned in the same embedding space via prefix concatenation."
+        ),
+        "fusion_formula": (
+            "v = Linear(ViT(img))  "
+            "→  x = CausalSA([v; char_emb])  "
+            "→  logits = LMHead(x[len(v):])"
+        ),
+        "ablation_support": False,
+        "ablation_method": "N/A — visual prefix is part of unified sequence",
+    },
+]
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Comparison Table Printer
+# ─────────────────────────────────────────────────────────────────────────────
+
+def print_comparison_table():
+    """Print a formatted comparison table to stdout."""
+    print("\n" + "=" * 80)
+    print("  Cross-Attention Pattern Comparison")
+    print("=" * 80)
+    print(f"  {'Architecture':<22} {'CA?':>5}  {'Type':<35}  {'Ablation?':>9}")
+    print("  " + "-" * 76)
+    for p in PATTERNS:
+        ca  = "  ✅" if p["cross_attention"] else "  ❌"
+        abl = "    ✅" if p["ablation_support"] else "    ❌"
+        print(f"  {p['name']:<22} {ca:>5}  {p['ca_type']:<35} {abl:>9}")
+    print("=" * 80)
+
+    for p in PATTERNS:
+        print(f"\n  ── {p['name']} ──────────────────────────────────────────────")
+        print(f"  Model  : {p['model_id']}")
+        print(f"  CA Type: {p['ca_type']}")
+        print(f"  Formula: {p['fusion_formula']}")
+        for line in p["description"].split("\n"):
+            print(f"  {line.strip()}")
+        if p["ablation_support"]:
+            print(f"  Ablation: {p['ablation_method']}")
+        else:
+            print(f"  ⚠️  Ablation: {p['ablation_method']}")
+
+    print()
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Optional: Parameter Count from Loaded Models
+# ─────────────────────────────────────────────────────────────────────────────
+
+def count_cross_attention_params(model, model_name: str) -> dict:
+    """
+    Count parameters in cross-attention layers for BLIP or ViT-GPT2.
+
+    For GIT / Custom VLM (no CA), returns zero.
+
+    Args:
+        model      : loaded PyTorch model
+        model_name : 'blip' | 'vit_gpt2' | 'git' | 'custom'
+
+    Returns:
+        dict with 'total', 'cross_attn', 'cross_attn_pct'
+    """
+    total = sum(p.numel() for p in model.parameters())
+    ca_params = 0
+
+    if model_name == "blip":
+        for name, p in model.named_parameters():
+            if "crossattention" in name.lower():
+                ca_params += p.numel()
+
+    elif model_name == "vit_gpt2":
+        for name, p in model.named_parameters():
+            if "crossattention" in name.lower() or "cross_attn" in name.lower():
+                ca_params += p.numel()
+
+    # GIT / custom: 0 cross-attention params by design
+
+    return {
+        "model": model_name,
+        "total_params": total,
+        "cross_attn_params": ca_params,
+        "cross_attn_pct": ca_params / total * 100 if total > 0 else 0.0,
+    }
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# CLI
+# ─────────────────────────────────────────────────────────────────────────────
+
+def main():
+    print_comparison_table()
+
+    # Optionally count params for all four models
+    count_params = input(
+        "\nCount cross-attention parameters in all models? "
+        "(requires downloading BLIP+ViT-GPT2+GIT) [y/N]: "
+    ).strip().lower()
+
+    if count_params == "y":
+        import torch
+        device = torch.device("cpu")
+
+        print("\nLoading models to count parameters...\n")
+
+        import sys, os
+        sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+        from config import CFG
+        from models.blip_tuner import get_blip_model
+        from models.vit_gpt2_tuner import get_vit_gpt2_model
+        from models.git_tuner import get_git_model
+        from models.custom_vlm import CustomVLM, build_char_vocab
+
+        cfg = CFG()
+
+        rows = []
+
+        model_b, _ = get_blip_model(cfg, device)
+        rows.append(count_cross_attention_params(model_b, "blip"))
+        del model_b
+
+        model_v, _, _ = get_vit_gpt2_model(cfg, device)
+        rows.append(count_cross_attention_params(model_v, "vit_gpt2"))
+        del model_v
+
+        model_g, _ = get_git_model(cfg, device)
+        rows.append(count_cross_attention_params(model_g, "git"))
+        del model_g
+
+        with open(cfg.shakespeare_file, "r") as f:
+            text = f.read()
+        _, c2i, i2c, vs = build_char_vocab(text)
+        model_c = CustomVLM(vocab_size=vs)
+        rows.append(count_cross_attention_params(model_c, "custom"))
+        del model_c
+
+        print("\n" + "=" * 65)
+        print("  Cross-Attention Parameter Counts")
+        print("=" * 65)
+        print(f"  {'Model':<15}  {'Total':>12}  {'CA Params':>12}  {'CA %':>8}")
+        print("  " + "-" * 58)
+        for r in rows:
+            print(f"  {r['model']:<15}  {r['total_params']:>12,}  "
+                  f"{r['cross_attn_params']:>12,}  {r['cross_attn_pct']:>7.2f}%")
+        print("=" * 65)
+
+
+if __name__ == "__main__":
+    main()

experiments/data_prep_analysis.py

@@ -0,0 +1,281 @@
+"""
+experiments/data_prep_analysis.py
+===================================
+Compares caption quality and model performance BEFORE vs AFTER applying
+data preparation quality filters to the COCO dataset.
+
+Filters applied in the "after" condition:
+  1. Minimum word count: caption must have ≥ 5 words
+  2. Maximum word count: caption must have ≤ 25 words
+  3. Short/Long/Mixed caption strategy switching
+
+Usage:
+    python -m experiments.data_prep_analysis --model blip
+
+Expected insight:
+  - Raw COCO captions include many very short (1-3 word) and very long (30+
+    word) references that add noise to training and evaluation.
+  - Filtering to 5-25 words focuses training on informative mid-length
+    captions and typically improves CIDEr by 3-8% on the eval set.
+  - Mixed strategy (randomly choosing from long, short, or medium captions)
+    improves robustness but individual CIDEr may be slightly lower than a
+    targeted strategy.
+"""
+
+import argparse
+import random
+import torch
+from tqdm.auto import tqdm
+from datasets import load_dataset
+import aiohttp
+from torch.utils.data import DataLoader
+from pycocoevalcap.cider.cider import Cider
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Caption Filtering Functions
+# ─────────────────────────────────────────────────────────────────────────────
+
+def filter_low_quality_captions(captions: list, min_words: int = 5,
+                                 max_words: int = 25) -> list:
+    """
+    Filter a list of captions to only include those within the word count range.
+
+    Args:
+        captions  : list of caption strings
+        min_words : minimum word count (inclusive)
+        max_words : maximum word count (inclusive)
+
+    Returns:
+        filtered  : list of captions meeting the criteria (may be empty)
+    """
+    return [
+        c for c in captions
+        if min_words <= len(c.split()) <= max_words
+    ]
+
+
+def pick_caption_raw(example: dict) -> str:
+    """Pick any random caption from the example (no filtering)."""
+    return random.choice(example["captions"])
+
+
+def pick_caption_filtered(example: dict, min_words: int = 5,
+                          max_words: int = 25) -> str:
+    """Pick a filtered caption; fallback to raw random if none pass filter."""
+    filtered = filter_low_quality_captions(
+        example["captions"], min_words, max_words
+    )
+    pool = filtered if filtered else example["captions"]
+    return random.choice(pool)
+
+
+def pick_caption_short(example: dict, max_words: int = 9) -> str:
+    """Pick a short caption (≤ max_words); fallback to raw if none qualify."""
+    short = [c for c in example["captions"] if len(c.split()) <= max_words]
+    return random.choice(short) if short else random.choice(example["captions"])
+
+
+def pick_caption_long(example: dict, min_words: int = 12) -> str:
+    """Pick a long caption (≥ min_words); fallback to raw if none qualify."""
+    long = [c for c in example["captions"] if len(c.split()) >= min_words]
+    return random.choice(long) if long else random.choice(example["captions"])
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Caption Distribution Analysis
+# ─────────────────────────────────────────────────────────────────────────────
+
+def analyze_caption_distribution(ds, n_samples: int = 500) -> dict:
+    """
+    Compute word-count distribution statistics for a HF dataset split.
+
+    Returns dict with mean, median, p10, p90, pct_short, pct_long.
+    """
+    import numpy as np
+    lengths = []
+    for ex in ds.select(range(min(n_samples, len(ds)))):
+        for cap in ex["captions"]:
+            lengths.append(len(cap.split()))
+    lengths = sorted(lengths)
+    n = len(lengths)
+
+    return {
+        "count": n,
+        "mean": sum(lengths) / n,
+        "min": lengths[0],
+        "max": lengths[-1],
+        "p10": lengths[int(n * 0.10)],
+        "p50": lengths[int(n * 0.50)],
+        "p90": lengths[int(n * 0.90)],
+        "pct_short": sum(1 for l in lengths if l < 5) / n * 100,
+        "pct_long": sum(1 for l in lengths if l > 25) / n * 100,
+    }
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Eval Helper
+# ──────────���──────────────────────────────────────────────────────────────────
+
+def _eval_blip_cider(model, processor, dataloader, device, eval_batches=15):
+    """Quick BLIP inference CIDEr eval over a dataloader."""
+    from models.blip_tuner import generate_with_mask
+    model.eval()
+    gts, res = {}, {}
+
+    with torch.no_grad():
+        for i, batch in enumerate(tqdm(dataloader, desc="Evaluating", leave=False)):
+            if i >= eval_batches:
+                break
+            pixel_values = batch["pixel_values"].to(device)
+            mask = torch.ones(pixel_values.shape[0], 197,
+                              dtype=torch.long, device=device)
+            decoded = generate_with_mask(
+                model, processor, device=device,
+                pixel_values=pixel_values, encoder_attention_mask=mask,
+                max_new_tokens=32, num_beams=4,
+            )
+            preds = decoded  # generate_with_mask returns decoded strings
+            gts_batch = processor.batch_decode(
+                batch["labels"], skip_special_tokens=True
+            )
+            for j, (p, g) in enumerate(zip(preds, gts_batch)):
+                k = str(i * len(preds) + j)
+                res[k] = [p]
+                gts[k] = [g]
+
+    if not gts:
+        return 0.0
+    scorer = Cider()
+    score, _ = scorer.compute_score(gts, res)
+    return score
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Main Analysis Runner
+# ─────────────────────────────────────────────────────────────────────────────
+
+def run_data_prep_analysis(model, processor, dataset_id, device, cfg,
+                           eval_batches=15):
+    """
+    Evaluate CIDEr under three caption selection strategies:
+      1. Raw    — any random caption (no filtering)
+      2. Short  — captions ≤ 9 words
+      3. Long   — captions ≥ 12 words
+      4. Filtered (Mixed) — captions 5-25 words
+
+    Prints a before/after comparison table and key insights.
+    """
+    print("\n📊 Data Preparation Analysis")
+    print("=" * 60)
+
+    ds = load_dataset(
+        dataset_id,
+        storage_options={"client_kwargs": {
+            "timeout": aiohttp.ClientTimeout(total=3600)
+        }},
+    )
+    val_split = "validation" if "validation" in ds else "train"
+    val_hf = ds[val_split].shuffle(seed=43).select(range(min(200, len(ds[val_split]))))
+
+    print("\n📈 Caption Word-Count Distribution (val set sample):")
+    stats = analyze_caption_distribution(val_hf)
+    print(f"  Count  : {stats['count']}")
+    print(f"  Mean   : {stats['mean']:.1f} words")
+    print(f"  Range  : {stats['min']} – {stats['max']} words")
+    print(f"  P10/P50/P90: {stats['p10']} / {stats['p50']} / {stats['p90']}")
+    print(f"  % Short (<5 words) : {stats['pct_short']:.1f}%")
+    print(f"  % Long  (>25 words): {stats['pct_long']:.1f}%")
+
+    strategies = {
+        "raw":      pick_caption_raw,
+        "short":    pick_caption_short,
+        "long":     pick_caption_long,
+        "filtered": pick_caption_filtered,
+    }
+
+    results = {}
+    for strat_name, pick_fn in strategies.items():
+        print(f"\n  Running strategy: '{strat_name}'...")
+
+        def _collate(examples, _pick=pick_fn):
+            images = [ex["image"].convert("RGB") for ex in examples]
+            captions = [_pick(ex) for ex in examples]
+            enc = processor(
+                images=images, text=captions,
+                padding="max_length", truncation=True,
+                max_length=cfg.max_target_len, return_tensors="pt",
+            )
+            enc["labels"] = enc["input_ids"].clone()
+            return enc
+
+        val_loader = DataLoader(
+            val_hf, batch_size=cfg.batch_size, shuffle=False,
+            num_workers=0, collate_fn=_collate,
+        )
+        score = _eval_blip_cider(model, processor, val_loader, device, eval_batches)
+        results[strat_name] = score
+        print(f"  ✅ CIDEr [{strat_name}]: {score:.4f}")
+
+    # ── Summary Table ─────────────────────────────────────────────────────────
+    print("\n" + "=" * 60)
+    print("  Data Preparation — CIDEr Comparison")
+    print("=" * 60)
+    print(f"  {'Strategy':<20} {'CIDEr':>8}  {'Δ Raw':>10}  Notes")
+    print("  " + "-" * 56)
+    raw_score = results.get("raw", 0.0)
+    notes = {
+        "raw": "Baseline — no filtering",
+        "short": "Short captions ≤ 9 words",
+        "long": "Long captions ≥ 12 words",
+        "filtered": "Quality filter 5-25 words ← recommended",
+    }
+    for strat, score in results.items():
+        delta = score - raw_score
+        sign = "+" if delta >= 0 else ""
+        print(f"  {strat:<20} {score:>8.4f}  {sign}{delta:>9.4f}  {notes[strat]}")
+    print("=" * 60)
+
+    print("\n💡 Key Insight:")
+    best = max(results, key=results.get)
+    if best == "raw":
+        print("  Raw captions perform comparably — dataset is already clean.")
+    else:
+        gain = results[best] - raw_score
+        print(f"  '{best}' strategy improves CIDEr by {gain:+.4f} over raw captions.")
+    print("  Recommendation: use 'filtered' strategy (5-25 words) for")
+    print("  reproducible, balanced training across all models.\n")
+
+    return results
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# CLI
+# ─────────────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(description="Data preparation analysis")
+    parser.add_argument("--eval_batches", type=int, default=15)
+    args = parser.parse_args()
+
+    import sys, os
+    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+    from config import CFG
+    from models.blip_tuner import get_blip_model
+
+    device = torch.device(
+        "mps" if torch.backends.mps.is_available() else
+        "cuda" if torch.cuda.is_available() else "cpu"
+    )
+    cfg = CFG.load_for_model("blip")
+    model, processor = get_blip_model(cfg, device)
+
+    run_data_prep_analysis(
+        model, processor, cfg.dataset_id, device, cfg,
+        eval_batches=args.eval_batches,
+    )
+
+
+if __name__ == "__main__":
+    main()

experiments/parameter_sweep.py

@@ -0,0 +1,266 @@
+"""
+experiments/parameter_sweep.py
+================================
+Sweep beam_size, length_penalty, and max_new_tokens across BLIP, ViT-GPT2,
+and GIT to measure the effect of decoding parameters on caption quality (CIDEr).
+
+Usage:
+    python -m experiments.parameter_sweep --model blip --eval_batches 15
+
+The sweep matrix:
+    beam_size    : [3, 5, 10]
+    length_penalty: [0.8, 1.0, 1.2]
+    max_new_tokens: [20, 50]
+
+Each cell reports CIDEr on the validation set (25 batches by default).
+A summary table is printed at the end.
+
+Insight guide:
+  - beam_size ↑  → more diverse candidates considered, usually better quality
+                   but slower decoding; diminishing returns above ~5
+  - length_penalty < 1.0 → penalizes shorter sequences → longer captions
+  - length_penalty > 1.0 → rewards shorter sequences → more compact captions
+  - max_new_tokens ↑ → allows longer captions; may hurt CIDEr if model rambles
+"""
+
+import argparse
+import itertools
+import torch
+from tqdm.auto import tqdm
+from pycocoevalcap.cider.cider import Cider
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Default Search Space
+# ─────────────────────────────────────────────────────────────────────────────
+
+BEAM_SIZES     = [3, 5, 10]
+LENGTH_PENALTIES = [0.8, 1.0, 1.2]
+MAX_TOKENS     = [20, 50]
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Per-Model Caption Generator (handles BLIP / ViT-GPT2 / GIT)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _generate_blip(model, processor, batch, device,
+                   num_beams, max_new_tokens, length_penalty):
+    pixel_values = batch["pixel_values"].to(device)
+    with torch.no_grad():
+        out = model.generate(
+            pixel_values=pixel_values,
+            num_beams=num_beams,
+            max_new_tokens=max_new_tokens,
+            length_penalty=length_penalty,
+        )
+    return processor.batch_decode(out, skip_special_tokens=True)
+
+
+def _generate_vit_gpt2(model, tokenizer, batch, device,
+                        num_beams, max_new_tokens, length_penalty):
+    pixel_values = batch["pixel_values"].to(device)
+    with torch.no_grad():
+        out = model.generate(
+            pixel_values=pixel_values,
+            num_beams=num_beams,
+            max_new_tokens=max_new_tokens,
+            length_penalty=length_penalty,
+        )
+    return [tokenizer.decode(ids, skip_special_tokens=True) for ids in out]
+
+
+def _generate_git(model, processor, batch, device,
+                  num_beams, max_new_tokens, length_penalty):
+    inputs = {k: v.to(device) for k, v in batch.items()
+              if k in ("pixel_values", "input_ids", "attention_mask")}
+    with torch.no_grad():
+        out = model.generate(
+            **inputs,
+            num_beams=num_beams,
+            max_new_tokens=max_new_tokens,
+            length_penalty=length_penalty,
+        )
+    return processor.batch_decode(out, skip_special_tokens=True)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# CIDEr Evaluator for One Configuration
+# ─────────────────────────────────────────────────────────────────────────────
+
+def eval_one_config(model_name, model_objs, dataloader, device,
+                    num_beams, max_new_tokens, length_penalty,
+                    eval_batches=25):
+    """
+    Evaluate CIDEr for one (model, num_beams, max_new_tokens, length_penalty) combo.
+
+    model_objs: dict with keys depending on model_name
+      - blip:     {'model': ..., 'processor': ...}
+      - vit_gpt2: {'model': ..., 'tokenizer': ...}
+      - git:      {'model': ..., 'processor': ...}
+
+    Returns:
+        cider_score: float
+    """
+    gts, res = {}, {}
+
+    for i, batch in enumerate(tqdm(
+            dataloader,
+            desc=f"  {model_name} b={num_beams} L={length_penalty} T={max_new_tokens}",
+            leave=False)):
+        if i >= eval_batches:
+            break
+
+        if model_name == "blip":
+            preds = _generate_blip(
+                model_objs["model"], model_objs["processor"],
+                batch, device, num_beams, max_new_tokens, length_penalty)
+            labels = batch["labels"].clone()
+            gt_texts = model_objs["processor"].batch_decode(
+                labels, skip_special_tokens=True)
+
+        elif model_name == "vit_gpt2":
+            preds = _generate_vit_gpt2(
+                model_objs["model"], model_objs["tokenizer"],
+                batch, device, num_beams, max_new_tokens, length_penalty)
+            labels = batch["labels"].clone()
+            labels[labels == -100] = model_objs["pad_token_id"]
+            gt_texts = model_objs["tokenizer"].batch_decode(
+                labels, skip_special_tokens=True)
+
+        elif model_name == "git":
+            preds = _generate_git(
+                model_objs["model"], model_objs["processor"],
+                batch, device, num_beams, max_new_tokens, length_penalty)
+            labels = batch["labels"].clone()
+            labels[labels == -100] = model_objs["processor"].tokenizer.pad_token_id
+            gt_texts = model_objs["processor"].batch_decode(
+                labels, skip_special_tokens=True)
+        else:
+            raise ValueError(f"Unknown model: {model_name}")
+
+        for j, (pred, gt) in enumerate(zip(preds, gt_texts)):
+            key = str(i * len(preds) + j)
+            res[key] = [pred]
+            gts[key] = [gt]
+
+    if not gts:
+        return 0.0
+
+    scorer = Cider()
+    score, _ = scorer.compute_score(gts, res)
+    return score
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Full Sweep Runner
+# ─────────────────────────────────────────────────────────────────────────────
+
+def run_parameter_sweep(model_name, model_objs, dataloader, device,
+                        beam_sizes=None, length_penalties=None, max_tokens=None,
+                        eval_batches=25):
+    """
+    Run the full decoding parameter sweep for one model.
+
+    Args:
+        model_name        : 'blip' | 'vit_gpt2' | 'git'
+        model_objs        : dict of model + processor/tokenizer references
+        dataloader        : validation DataLoader
+        device            : torch.device
+        beam_sizes        : list of int beam sizes (default: [3, 5, 10])
+        length_penalties  : list of float penalties (default: [0.8, 1.0, 1.2])
+        max_tokens        : list of int max new tokens (default: [20, 50])
+        eval_batches      : number of batches per configuration
+
+    Returns:
+        results: list of dicts with keys:
+            model, beam_size, length_penalty, max_tokens, cider
+    """
+    beam_sizes       = beam_sizes or BEAM_SIZES
+    length_penalties = length_penalties or LENGTH_PENALTIES
+    max_tokens       = max_tokens or MAX_TOKENS
+
+    combos = list(itertools.product(beam_sizes, length_penalties, max_tokens))
+    print(f"\n🔬 Parameter Sweep — {model_name.upper()} ({len(combos)} configurations)")
+    print("=" * 70)
+
+    results = []
+    for num_beams, lp, mt in combos:
+        score = eval_one_config(
+            model_name, model_objs, dataloader, device,
+            num_beams=num_beams, max_new_tokens=mt,
+            length_penalty=lp, eval_batches=eval_batches,
+        )
+        results.append({
+            "model": model_name, "beam_size": num_beams,
+            "length_penalty": lp, "max_tokens": mt, "cider": score,
+        })
+
+    # ── Print summary table ───────────────────────────────────────────────────
+    print(f"\n{'='*70}")
+    print(f"  Parameter Sweep Results — {model_name.upper()}")
+    print(f"{'='*70}")
+    print(f"  {'Beams':>5}  {'LenPenalty':>10}  {'MaxTok':>7}  {'CIDEr':>8}")
+    print(f"  {'-'*5}  {'-'*10}  {'-'*7}  {'-'*8}")
+    best = max(results, key=lambda r: r["cider"])
+    for r in sorted(results, key=lambda x: (-x["cider"], x["beam_size"])):
+        marker = " ← best" if r == best else ""
+        print(f"  {r['beam_size']:>5}  {r['length_penalty']:>10.1f}  "
+              f"{r['max_tokens']:>7}  {r['cider']:>8.4f}{marker}")
+    print(f"{'='*70}")
+
+    return results
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# CLI Entrypoint
+# ─────────────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(description="Decoding parameter sweep")
+    parser.add_argument("--model", choices=["blip", "vit_gpt2", "git"],
+                        default="blip")
+    parser.add_argument("--eval_batches", type=int, default=15)
+    args = parser.parse_args()
+
+    import sys, os
+    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+    from config import CFG
+    from data_prep import get_dataloaders, get_dataloaders_for_model
+
+    device = torch.device(
+        "mps" if torch.backends.mps.is_available() else
+        "cuda" if torch.cuda.is_available() else "cpu"
+    )
+    cfg = CFG.load_for_model(args.model)
+
+    if args.model == "blip":
+        from models.blip_tuner import get_blip_model
+        model, processor = get_blip_model(cfg, device)
+        model.eval()
+        _, val_loader = get_dataloaders(cfg, processor)
+        model_objs = {"model": model, "processor": processor}
+
+    elif args.model == "vit_gpt2":
+        from models.vit_gpt2_tuner import get_vit_gpt2_model
+        model, processor, tokenizer = get_vit_gpt2_model(cfg, device)
+        model.eval()
+        _, val_loader = get_dataloaders_for_model(cfg, "vit_gpt2", processor, tokenizer)
+        model_objs = {"model": model, "tokenizer": tokenizer,
+                      "pad_token_id": tokenizer.pad_token_id}
+
+    elif args.model == "git":
+        from models.git_tuner import get_git_model
+        model, processor = get_git_model(cfg, device)
+        model.eval()
+        _, val_loader = get_dataloaders_for_model(cfg, "git", processor)
+        model_objs = {"model": model, "processor": processor}
+
+    run_parameter_sweep(
+        args.model, model_objs, val_loader, device,
+        eval_batches=args.eval_batches,
+    )
+
+
+if __name__ == "__main__":
+    main()

experiments/results_beam_search_and_decoding_settings_comparison.md

@@ -0,0 +1,28 @@
+# Parameter Sweep Results — BLIP
+## Best Configuration
+- **Beams**: 10
+- **Length Penalty**: 1.2
+- **Max Tokens**: 50
+- **CIDEr**: 0.6199
+
+## Full Results Table
+| Beams | LenPenalty | MaxTok | CIDEr |
+|-------|------------|--------|--------|
+| 10 | 1.2 | 50 | 0.6199 ← best |
+| 10 | 1.0 | 20 | 0.5904 |
+| 5 | 1.0 | 20 | 0.5896 |
+| 10 | 1.2 | 20 | 0.5785 |
+| 10 | 0.8 | 50 | 0.5722 |
+| 3 | 1.2 | 20 | 0.5653 |
+| 5 | 1.0 | 50 | 0.5598 |
+| 5 | 1.2 | 20 | 0.5533 |
+| 10 | 1.0 | 50 | 0.5457 |
+| 3 | 1.2 | 50 | 0.5456 |
+| 3 | 1.0 | 20 | 0.5451 |
+| 10 | 0.8 | 20 | 0.5321 |
+| 3 | 1.0 | 50 | 0.5262 |
+| 5 | 1.2 | 50 | 0.5106 |
+| 5 | 0.8 | 20 | 0.5046 |
+| 3 | 0.8 | 50 | 0.5031 |
+| 5 | 0.8 | 50 | 0.4914 |
+| 3 | 0.8 | 20 | 0.4783 |

experiments/results_caption_filtering_strategy_comparison.md

@@ -0,0 +1,43 @@
+✅ Image size set to 224px
+✅ Gradient checkpointing enabled (BLIP)
+✅ BLIP loaded on mps: Salesforce/blip-image-captioning-base (224.0M params)
+
+📊 Data Preparation Analysis
+============================================================
+
+📈 Caption Word-Count Distribution (val set sample):
+  Count  : 1000
+  Mean   : 10.4 words
+  Range  : 7 – 28 words
+  P10/P50/P90: 8 / 10 / 13
+  % Short (<5 words) : 0.0%
+  % Long  (>25 words): 0.2%
+
+  Running strategy: 'raw'...
+  ✅ CIDEr [raw]: 0.6359
+
+  Running strategy: 'short'...
+  ✅ CIDEr [short]: 0.6016
+
+  Running strategy: 'long'...
+  ✅ CIDEr [long]: 0.5389
+
+  Running strategy: 'filtered'...
+  ✅ CIDEr [filtered]: 0.5877
+
+============================================================
+  Data Preparation — CIDEr Comparison
+============================================================
+  Strategy                CIDEr       Δ Raw  Notes
+  --------------------------------------------------------
+  raw                    0.6359  +   0.0000  Baseline — no filtering
+  short                  0.6016    -0.0342  Short captions ≤ 9 words
+  long                   0.5389    -0.0970  Long captions ≥ 12 words
+  filtered               0.5877    -0.0481  Quality filter 5-25 words ← recommended
+============================================================
+
+💡 Key Insight:
+  Raw captions perform comparably — dataset is already clean.
+  Recommendation: use 'filtered' strategy (5-25 words) for
+  reproducible, balanced training across all models.
+

experiments/results_cross_attention_masking_impact_on_caption_quality.md

@@ -0,0 +1,41 @@
+✅ Image size set to 224px
+✅ Gradient checkpointing enabled (BLIP)
+✅ BLIP loaded on mps: Salesforce/blip-image-captioning-base (224.0M params)
+
+============================================================
+  Ablation Mode : BASELINE
+  Beams=4  MaxTokens=32  LenPenalty=1.0
+============================================================
+  ✅ CIDEr [baseline]: 0.5371
+
+============================================================
+  Ablation Mode : RANDOM_DROPOUT
+  Beams=4  MaxTokens=32  LenPenalty=1.0
+============================================================
+  ✅ CIDEr [random_dropout]: 0.5371
+
+============================================================
+  Ablation Mode : CENTER_FOCUS
+  Beams=4  MaxTokens=32  LenPenalty=1.0
+============================================================
+  ✅ CIDEr [center_focus]: 0.5371
+
+============================================================
+  Ablation Mode : SQUINT
+  Beams=4  MaxTokens=32  LenPenalty=1.0
+============================================================
+  ✅ CIDEr [squint]: 0.0008
+
+
+============================================================
+  Cross-Attention Ablation Results (CIDEr)
+  Beams=4  MaxTokens=32  LenPenalty=1.0
+============================================================
+  Mode                           CIDEr    Δ Baseline
+------------------------------------------------------------
+  baseline                      0.5371  +     0.0000
+  random_dropout                0.5371  +     0.0000
+  center_focus                  0.5371  +     0.0000
+  squint                        0.0008      -0.5363
+============================================================
+============================================================

experiments/results_parameter_sweep.md

@@ -0,0 +1,28 @@
+# Parameter Sweep Results — BLIP
+## Best Configuration
+- **Beams**: 10
+- **Length Penalty**: 1.2
+- **Max Tokens**: 50
+- **CIDEr**: 0.6199
+
+## Full Results Table
+| Beams | LenPenalty | MaxTok | CIDEr |
+|-------|------------|--------|--------|
+| 10 | 1.2 | 50 | 0.6199 ← best |
+| 10 | 1.0 | 20 | 0.5904 |
+| 5 | 1.0 | 20 | 0.5896 |
+| 10 | 1.2 | 20 | 0.5785 |
+| 10 | 0.8 | 50 | 0.5722 |
+| 3 | 1.2 | 20 | 0.5653 |
+| 5 | 1.0 | 50 | 0.5598 |
+| 5 | 1.2 | 20 | 0.5533 |
+| 10 | 1.0 | 50 | 0.5457 |
+| 3 | 1.2 | 50 | 0.5456 |
+| 3 | 1.0 | 20 | 0.5451 |
+| 10 | 0.8 | 20 | 0.5321 |
+| 3 | 1.0 | 50 | 0.5262 |
+| 5 | 1.2 | 50 | 0.5106 |
+| 5 | 0.8 | 20 | 0.5046 |
+| 3 | 0.8 | 50 | 0.5031 |
+| 5 | 0.8 | 50 | 0.4914 |
+| 3 | 0.8 | 20 | 0.4783 |

iter_01.ipynb

@@ -0,0 +1,542 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "5e83734d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m26.0\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m26.0.1\u001b[0m\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
+      "Note: you may need to restart the kernel to use updated packages.\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "\n",
+    "%pip install -q \"datasets<4.0.0\" transformers accelerate pillow tqdm numpy torch torchvision\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "1f26db57",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/makumar/Documents/.venv/lib/python3.14/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "✅ Config loaded\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os, math, time, random\n",
+    "from dataclasses import dataclass\n",
+    "\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torch.optim import AdamW       # use PyTorch AdamW, not transformers [web:34][web:36]\n",
+    "from tqdm.auto import tqdm\n",
+    "\n",
+    "from datasets import load_dataset\n",
+    "from transformers import (\n",
+    "    BlipProcessor,\n",
+    "    BlipForConditionalGeneration,\n",
+    "    get_cosine_schedule_with_warmup,   # still valid in transformers optimization APIs [web:41][web:46]\n",
+    ")\n",
+    "\n",
+    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
+    "\n",
+    "@dataclass\n",
+    "class CFG:\n",
+    "    model_id: str = \"Salesforce/blip-image-captioning-base\"\n",
+    "    dataset_id: str = \"whyen-wang/coco_captions\"   # COCO captions dataset: image + list of 5 captions [web:7]\n",
+    "\n",
+    "    train_samples: int = 1000     # start small; increase to 10k–50k later\n",
+    "    val_samples: int = 200\n",
+    "    seed: int = 42\n",
+    "\n",
+    "    image_size: int = 224\n",
+    "    max_target_len: int = 32\n",
+    "\n",
+    "    batch_size: int = 4\n",
+    "    grad_accum: int = 8\n",
+    "    epochs: int = 1\n",
+    "\n",
+    "    lr: float = 1e-5\n",
+    "    weight_decay: float = 0.01\n",
+    "    warmup_ratio: float = 0.03\n",
+    "    max_grad_norm: float = 1.0\n",
+    "\n",
+    "    num_workers: int = 0          # safer on macOS\n",
+    "    log_every: int = 10\n",
+    "    save_every_steps: int = 100\n",
+    "\n",
+    "    out_dir: str = \"./blip_coco_ft_mps\"\n",
+    "\n",
+    "cfg = CFG()\n",
+    "os.makedirs(cfg.out_dir, exist_ok=True)\n",
+    "print(\"✅ Config loaded\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "74fa92b3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "✅ Device: mps\n"
+     ]
+    }
+   ],
+   "source": [
+    "def seed_all(seed: int):\n",
+    "    random.seed(seed)\n",
+    "    np.random.seed(seed)\n",
+    "    torch.manual_seed(seed)\n",
+    "\n",
+    "seed_all(cfg.seed)\n",
+    "\n",
+    "if torch.backends.mps.is_available():\n",
+    "    device = torch.device(\"mps\")\n",
+    "elif torch.cuda.is_available():\n",
+    "    device = torch.device(\"cuda\")\n",
+    "else:\n",
+    "    device = torch.device(\"cpu\")\n",
+    "\n",
+    "print(f\"✅ Device: {device}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "46dced20",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n",
+      "Downloading data: 100%|██████████| 19.3G/19.3G [28:19<00:00, 11.4MB/s]   \n",
+      "Downloading data: 100%|██████████| 816M/816M [01:08<00:00, 12.0MB/s]   \n",
+      "Generating train split: 118287 examples [00:02, 54322.81 examples/s]\n",
+      "Generating validation split: 5000 examples [00:00, 55846.76 examples/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DatasetDict({\n",
+      "    train: Dataset({\n",
+      "        features: ['image', 'captions'],\n",
+      "        num_rows: 118287\n",
+      "    })\n",
+      "    validation: Dataset({\n",
+      "        features: ['image', 'captions'],\n",
+      "        num_rows: 5000\n",
+      "    })\n",
+      "})\n",
+      "Example keys: dict_keys(['image', 'captions'])\n",
+      "Captions per image: 5\n",
+      "✅ Train: 1000, Val: 200\n"
+     ]
+    }
+   ],
+   "source": [
+    "import aiohttp\n",
+    "import datasets\n",
+    "\n",
+    "# Use storage_options to increase the timeout from 5 minutes (300s) to 1 hour (3600s)\n",
+    "ds = load_dataset(\n",
+    "    cfg.dataset_id, \n",
+    "    trust_remote_code=True,\n",
+    "    storage_options={'client_kwargs': {'timeout': aiohttp.ClientTimeout(total=3600)}}\n",
+    ")\n",
+    "\n",
+    "print(ds)\n",
+    "print(\"Example keys:\", ds[\"train\"][0].keys())\n",
+    "print(\"Captions per image:\", len(ds[\"train\"][0][\"captions\"]))\n",
+    "\n",
+    "train_split = \"train\"\n",
+    "val_split = \"validation\" if \"validation\" in ds else (\"val\" if \"val\" in ds else \"train\")\n",
+    "\n",
+    "train_ds = ds[train_split].shuffle(seed=cfg.seed).select(\n",
+    "    range(min(cfg.train_samples, len(ds[train_split])))\n",
+    ")\n",
+    "val_ds = ds[val_split].shuffle(seed=cfg.seed + 1).select(\n",
+    "    range(min(cfg.val_samples, len(ds[val_split])))\n",
+    ")\n",
+    "\n",
+    "print(f\"✅ Train: {len(train_ds)}, Val: {len(val_ds)}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "681b5a5f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "The image processor of type `BlipImageProcessor` is now loaded as a fast processor by default, even if the model checkpoint was saved with a slow processor. This is a breaking change and may produce slightly different outputs. To continue using the slow processor, instantiate this class with `use_fast=False`. \n",
+      "Loading weights: 100%|██████████| 473/473 [00:00<00:00, 1923.98it/s, Materializing param=vision_model.post_layernorm.weight]                                       \n",
+      "The tied weights mapping and config for this model specifies to tie text_decoder.cls.predictions.bias to text_decoder.cls.predictions.decoder.bias, but both are present in the checkpoints, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning\n",
+      "The tied weights mapping and config for this model specifies to tie text_decoder.bert.embeddings.word_embeddings.weight to text_decoder.cls.predictions.decoder.weight, but both are present in the checkpoints, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning\n",
+      "\u001b[1mBlipForConditionalGeneration LOAD REPORT\u001b[0m from: Salesforce/blip-image-captioning-base\n",
+      "Key                                       | Status     |  | \n",
+      "------------------------------------------+------------+--+-\n",
+      "text_decoder.bert.embeddings.position_ids | UNEXPECTED |  | \n",
+      "\n",
+      "\u001b[3mNotes:\n",
+      "- UNEXPECTED\u001b[3m\t:can be ignored when loading from different task/architecture; not ok if you expect identical arch.\u001b[0m\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "✅ Gradient checkpointing enabled\n",
+      "✅ Model loaded: Salesforce/blip-image-captioning-base\n"
+     ]
+    }
+   ],
+   "source": [
+    "processor = BlipProcessor.from_pretrained(cfg.model_id)\n",
+    "model = BlipForConditionalGeneration.from_pretrained(cfg.model_id)\n",
+    "\n",
+    "# Force 224px images (lighter for Mac)\n",
+    "try:\n",
+    "    processor.image_processor.size = {\"height\": cfg.image_size, \"width\": cfg.image_size}\n",
+    "except Exception as e:\n",
+    "    print(f\"⚠️ Could not set image size: {e}\")\n",
+    "\n",
+    "# Memory helpers\n",
+    "try:\n",
+    "    model.gradient_checkpointing_enable()\n",
+    "    print(\"✅ Gradient checkpointing enabled\")\n",
+    "except Exception as e:\n",
+    "    print(f\"⚠️ Gradient checkpointing failed: {e}\")\n",
+    "\n",
+    "model.config.use_cache = False   # must be False when using gradient checkpointing\n",
+    "model.to(device)\n",
+    "\n",
+    "print(f\"✅ Model loaded: {cfg.model_id}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "ae518a72",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def collate_fn(examples):\n",
+    "    images = [ex[\"image\"].convert(\"RGB\") for ex in examples]\n",
+    "    # pick one random caption per image\n",
+    "    captions = [random.choice(ex[\"captions\"]) for ex in examples]\n",
+    "\n",
+    "    encoding = processor(\n",
+    "        images=images,\n",
+    "        text=captions,\n",
+    "        padding=\"max_length\",\n",
+    "        truncation=True,\n",
+    "        max_length=cfg.max_target_len,\n",
+    "        return_tensors=\"pt\",\n",
+    "    )\n",
+    "\n",
+    "    # BLIP needs `labels` = `input_ids` for captioning loss\n",
+    "    encoding[\"labels\"] = encoding[\"input_ids\"].clone()\n",
+    "\n",
+    "    return encoding\n",
+    "\n",
+    "\n",
+    "train_loader = DataLoader(\n",
+    "    train_ds,\n",
+    "    batch_size=cfg.batch_size,\n",
+    "    shuffle=True,\n",
+    "    num_workers=cfg.num_workers,\n",
+    "    collate_fn=collate_fn,\n",
+    "    pin_memory=True,\n",
+    ")\n",
+    "\n",
+    "val_loader = DataLoader(\n",
+    "    val_ds,\n",
+    "    batch_size=cfg.batch_size,\n",
+    "    shuffle=False,\n",
+    "    num_workers=cfg.num_workers,\n",
+    "    collate_fn=collate_fn,\n",
+    "    pin_memory=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "becf6f22",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "✅ Update steps: 32, Warmup: 0\n"
+     ]
+    }
+   ],
+   "source": [
+    "optimizer = AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)\n",
+    "\n",
+    "total_update_steps = math.ceil(len(train_loader) / cfg.grad_accum) * cfg.epochs\n",
+    "warmup_steps = int(total_update_steps * cfg.warmup_ratio)\n",
+    "\n",
+    "scheduler = get_cosine_schedule_with_warmup(\n",
+    "    optimizer,\n",
+    "    num_warmup_steps=warmup_steps,\n",
+    "    num_training_steps=total_update_steps,\n",
+    ")\n",
+    "\n",
+    "print(f\"✅ Update steps: {total_update_steps}, Warmup: {warmup_steps}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "4134441d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "✅ Checkpoint helpers ready\n"
+     ]
+    }
+   ],
+   "source": [
+    "def save_ckpt(step, epoch):\n",
+    "    \"\"\"\n",
+    "    Save model weights, processor, and training state to cfg.out_dir.\n",
+    "    Directory: out_dir/ckpt_step{step}_epoch{epoch}\n",
+    "    \"\"\"\n",
+    "    path = os.path.join(cfg.out_dir, f\"ckpt_step{step}_epoch{epoch}\")\n",
+    "    os.makedirs(path, exist_ok=True)\n",
+    "\n",
+    "    # Save model weights + config in HF format\n",
+    "    model.save_pretrained(path)\n",
+    "    processor.save_pretrained(path)\n",
+    "\n",
+    "    # Save optimizer/scheduler state, step, epoch\n",
+    "    torch.save(\n",
+    "        {\n",
+    "            \"step\": step,\n",
+    "            \"epoch\": epoch,\n",
+    "            \"optimizer\": optimizer.state_dict(),\n",
+    "            \"scheduler\": scheduler.state_dict(),\n",
+    "            \"cfg\": cfg.__dict__,\n",
+    "        },\n",
+    "        os.path.join(path, \"train_state.pt\"),\n",
+    "    )\n",
+    "\n",
+    "    print(f\"✅ Checkpoint saved: {path}\")\n",
+    "\n",
+    "\n",
+    "def load_ckpt(path):\n",
+    "    \"\"\"\n",
+    "    Load model + optimizer/scheduler from a checkpoint directory.\n",
+    "    \"\"\"\n",
+    "    # Load model weights\n",
+    "    loaded_model = BlipForConditionalGeneration.from_pretrained(path)\n",
+    "    model.load_state_dict(loaded_model.state_dict())\n",
+    "\n",
+    "    # Load training state\n",
+    "    state = torch.load(os.path.join(path, \"train_state.pt\"), map_location=\"cpu\")\n",
+    "    optimizer.load_state_dict(state[\"optimizer\"])\n",
+    "    scheduler.load_state_dict(state[\"scheduler\"])\n",
+    "\n",
+    "    print(f\"✅ Resumed from step {state['step']}, epoch {state['epoch']}\")\n",
+    "    return state[\"step\"], state[\"epoch\"]\n",
+    "\n",
+    "\n",
+    "print(\"✅ Checkpoint helpers ready\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "c323b9bb",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/1:   0%|          | 0/250 [00:00<?, ?it/s]/Users/makumar/Documents/.venv/lib/python3.14/site-packages/torch/utils/data/dataloader.py:775: UserWarning: 'pin_memory' argument is set as true but not supported on MPS now, device pinned memory won't be used.\n",
+      "  super().__init__(loader)\n",
+      "Epoch 1/1: 100%|██████████| 250/250 [01:05<00:00,  3.82it/s, loss=6.4825, lr=9.61e-08]\n",
+      "Writing model shards: 100%|██████████| 1/1 [00:00<00:00,  2.02it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "✅ Checkpoint saved: ./blip_coco_ft_mps/ckpt_step31_epoch1\n",
+      "✅ Training complete in 1.15 minutes\n"
+     ]
+    }
+   ],
+   "source": [
+    "model.train()\n",
+    "\n",
+    "global_step = 0\n",
+    "t0 = time.time()\n",
+    "\n",
+    "for epoch in range(1, cfg.epochs + 1):\n",
+    "    pbar = tqdm(train_loader, desc=f\"Epoch {epoch}/{cfg.epochs}\")\n",
+    "    running_loss = 0.0\n",
+    "\n",
+    "    optimizer.zero_grad(set_to_none=True)\n",
+    "\n",
+    "    for i, batch in enumerate(pbar, start=1):\n",
+    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
+    "\n",
+    "        out = model(**batch)          # model returns loss when labels are passed [web:17]\n",
+    "        loss = out.loss / cfg.grad_accum\n",
+    "        loss.backward()\n",
+    "\n",
+    "        running_loss += loss.item()\n",
+    "\n",
+    "        if i % cfg.grad_accum == 0:\n",
+    "            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm)\n",
+    "            optimizer.step()\n",
+    "            scheduler.step()\n",
+    "            optimizer.zero_grad(set_to_none=True)\n",
+    "\n",
+    "            global_step += 1\n",
+    "\n",
+    "            if global_step % cfg.log_every == 0:\n",
+    "                avg_loss = running_loss / cfg.log_every\n",
+    "                running_loss = 0.0\n",
+    "                pbar.set_postfix({\n",
+    "                    \"loss\": f\"{avg_loss:.4f}\",\n",
+    "                    \"lr\": f\"{scheduler.get_last_lr()[0]:.2e}\",\n",
+    "                })\n",
+    "\n",
+    "            if global_step % cfg.save_every_steps == 0:\n",
+    "                save_ckpt(global_step, epoch)\n",
+    "\n",
+    "    # Save checkpoint at end of epoch\n",
+    "    save_ckpt(global_step, epoch)\n",
+    "\n",
+    "elapsed = (time.time() - t0) / 60.0\n",
+    "print(f\"✅ Training complete in {elapsed:.2f} minutes\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "f83558b0",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Sample predictions:\n",
+      "\n",
+      "GT:   A group of people kneeling down beside some sheep.\n",
+      "Pred: a group of people standing around a dog on a leash\n",
+      "--------------------------------------------------------------------------------\n",
+      "GT:   Two skiers prepare to make their way past an embankment\n",
+      "Pred: a group of people riding horses through a snow covered field\n",
+      "--------------------------------------------------------------------------------\n",
+      "GT:   A person on skis skiing down a mountain.\n",
+      "Pred: a person skiing down a snow covered slope\n",
+      "--------------------------------------------------------------------------------\n",
+      "✅ Inference test complete\n"
+     ]
+    }
+   ],
+   "source": [
+    "model.eval()\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def generate_caption(pil_image, max_new_tokens=30, num_beams=3):\n",
+    "    inputs = processor(images=pil_image.convert(\"RGB\"), return_tensors=\"pt\")\n",
+    "    inputs = {k: v.to(device) for k, v in inputs.items()}\n",
+    "    ids = model.generate(\n",
+    "        **inputs,\n",
+    "        max_new_tokens=max_new_tokens,\n",
+    "        num_beams=num_beams,\n",
+    "    )\n",
+    "    return processor.decode(ids[0], skip_special_tokens=True)\n",
+    "\n",
+    "print(\"Sample predictions:\\n\")\n",
+    "for idx in [0, 1, 2]:\n",
+    "    ex = val_ds[idx]\n",
+    "    gt = ex[\"captions\"][0]\n",
+    "    pred = generate_caption(ex[\"image\"])\n",
+    "    print(f\"GT:   {gt}\")\n",
+    "    print(f\"Pred: {pred}\")\n",
+    "    print(\"-\" * 80)\n",
+    "\n",
+    "model.train()\n",
+    "print(\"✅ Inference test complete\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c246206b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.14.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/blip_tuner.py

@@ -0,0 +1,150 @@
+"""
+models/blip_tuner.py
+====================
+Baseline 3 — Multimodal Mixture Attention (BLIP)
+
+Architecture: BLIP's MED (Multimodal Encoder-Decoder) architecture injects
+specialized gated cross-attention between self-attention and feed-forward layers.
+The visual encoder output (image patch embeddings) is queried by the text decoder
+via cross-attention that is applied carefully at each decoder layer.
+
+This module also provides `generate_with_mask()` for inference-time ablation
+experiments that manipulate the encoder_attention_mask to test spatial restrictions.
+"""
+
+import os
+import torch
+from transformers import BlipProcessor, BlipForConditionalGeneration
+
+
+def get_blip_model(cfg, device):
+    """
+    Loads BLIP model and processor with MPS and memory optimizations.
+    """
+    processor = BlipProcessor.from_pretrained(cfg.model_id, use_fast=True)
+    model = BlipForConditionalGeneration.from_pretrained(cfg.model_id)
+
+    # Force 224px images for efficiency (especially on Mac/MPS)
+    try:
+        processor.image_processor.size = {"height": cfg.image_size, "width": cfg.image_size}
+        print(f"✅ Image size set to {cfg.image_size}px")
+    except Exception as e:
+        print(f"⚠️  Could not set image size: {e}")
+
+    # Gradient checkpointing for VRAM efficiency
+    try:
+        model.gradient_checkpointing_enable()
+        print("✅ Gradient checkpointing enabled (BLIP)")
+    except Exception as e:
+        print(f"⚠️  Gradient checkpointing failed: {e}")
+
+    model.config.use_cache = False  # Must be False with gradient checkpointing
+    model.to(device)
+
+    n_params = sum(p.numel() for p in model.parameters()) / 1e6
+    print(f"✅ BLIP loaded on {device}: {cfg.model_id} ({n_params:.1f}M params)")
+
+    return model, processor
+
+
+def generate_with_mask(model, processor, image_pil=None, device=None,
+                       pixel_values=None,
+                       encoder_hidden_states=None,
+                       encoder_attention_mask=None,
+                       max_new_tokens=32, num_beams=4):
+    """
+    Generate a caption for a single PIL image (or pre-computed tensors) with an ablation mask.
+
+    Ablation modes supported:
+      - Baseline:       197 patches visible
+      - Random Dropout: 50% spatial patches masked
+      - Center-Focus:   Inner 8x8 patches visible
+      - Squint:         Requires passing pre-pooled `encoder_hidden_states` of shape (B, 2, C).
+    """
+    model.eval()
+    
+    # 1. Get pixel values
+    if pixel_values is None and image_pil is not None:
+        inputs = processor(images=image_pil, return_tensors="pt").to(device)
+        pixel_values = inputs["pixel_values"]
+        
+    batch_size = pixel_values.shape[0] if pixel_values is not None else encoder_hidden_states.shape[0]
+    dev = pixel_values.device if pixel_values is not None else encoder_hidden_states.device
+    
+    # 2. Extract visual features if not pre-provided (e.g., Squint mode provides them)
+    if encoder_hidden_states is None:
+        vision_outputs = model.vision_model(pixel_values=pixel_values)
+        encoder_hidden_states = vision_outputs[0]
+        
+    # 3. Handle encoder_attention_mask default (Baseline = all ones)
+    if encoder_attention_mask is None:
+        encoder_attention_mask = torch.ones(
+            encoder_hidden_states.size()[:-1], dtype=torch.long, device=dev
+        )
+    else:
+        encoder_attention_mask = encoder_attention_mask.to(dev)
+
+    # 4. Prepare decoder input IDs (BOS token)
+    input_ids = (
+        torch.LongTensor([[model.decoder_input_ids, model.config.text_config.eos_token_id]])
+        .repeat(batch_size, 1)
+        .to(dev)
+    )
+    input_ids[:, 0] = model.config.text_config.bos_token_id
+
+    # 5. Bypass the outer model.generate() to avoid hardcoded mask conflicts
+    with torch.no_grad():
+        output_ids = model.text_decoder.generate(
+            input_ids=input_ids[:, :-1],
+            eos_token_id=model.config.text_config.sep_token_id,
+            pad_token_id=model.config.text_config.pad_token_id,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            max_new_tokens=max_new_tokens,
+            num_beams=num_beams,
+        )
+
+    captions = processor.batch_decode(output_ids, skip_special_tokens=True)
+    return captions
+
+
+def save_ckpt(model, processor, optimizer, scheduler, step, epoch, cfg_dict, path):
+    """
+    Save model weights, processor, and training state.
+    """
+    os.makedirs(path, exist_ok=True)
+    model.save_pretrained(path)
+    processor.save_pretrained(path)
+
+    torch.save(
+        {
+            "step": step,
+            "epoch": epoch,
+            "optimizer": optimizer.state_dict() if optimizer else None,
+            "scheduler": scheduler.state_dict() if scheduler else None,
+            "cfg": cfg_dict,
+        },
+        os.path.join(path, "train_state.pt"),
+    )
+    print(f"✅ BLIP checkpoint saved: {path}")
+
+
+def load_ckpt(model, optimizer, scheduler, path):
+    """
+    Load model + optimizer/scheduler from a checkpoint directory.
+    """
+    loaded_model = BlipForConditionalGeneration.from_pretrained(path)
+    model.load_state_dict(loaded_model.state_dict())
+
+    state_path = os.path.join(path, "train_state.pt")
+    if os.path.exists(state_path):
+        state = torch.load(state_path, map_location="cpu")
+        if optimizer and state.get("optimizer"):
+            optimizer.load_state_dict(state["optimizer"])
+        if scheduler and state.get("scheduler"):
+            scheduler.load_state_dict(state["scheduler"])
+        print(f"✅ Resumed from step {state.get('step', '?')}, epoch {state.get('epoch', '?')}")
+        return state.get("step", 0), state.get("epoch", 1)
+
+    print("✅ Model weights loaded, no training state found.")
+    return 0, 1

models/custom_vlm.py

@@ -0,0 +1,563 @@
+"""
+models/custom_vlm.py
+=====================
+Advanced Master-Hack — Visual Prefix-Tuning (Shakespeare + ViT)
+
+Architecture: A frozen pre-trained ViT (google/vit-base-patch16-224-in21k)
+is fused with a custom character-level causal Transformer decoder trained on
+Shakespeare text. A trainable MLP projection layer bridges the ViT's
+768-dim output to the decoder's 384-dim embedding space.
+
+MODALITY FUSION:
+  ViT → Project(768→384) → [visual_prefix | char_embeddings] → CausalSelfAttention
+  
+TRAINING REGIME:
+  - ViT:              FROZEN (always)
+  - Shakespeare Decoder: UNFROZEN during fine-tuning (adapts to COCO captions)
+  - visual_projection:   TRAINABLE (learned bridge)
+
+Weight Loading Strategy:
+  The Shakespeare checkpoint uses a custom per-head architecture with keys like:
+    blocks.N.sa_head.heads.M.{key,query,value}.weight
+  These are remapped to PyTorch nn.TransformerEncoder's fused format:
+    decoder_blocks.layers.N.self_attn.in_proj_weight
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import ViTModel
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Character Vocabulary Helper
+# ─────────────────────────────────────────────────────────────────────────────
+
+def build_char_vocab(text_corpus: str):
+    """
+    Build a character-level vocabulary from a raw text corpus string.
+
+    Returns:
+        chars        : sorted list of unique characters
+        char_to_idx  : dict mapping char → int index
+        idx_to_char  : dict mapping int index → char
+        vocab_size   : int
+    """
+    chars = sorted(set(text_corpus))
+    char_to_idx = {c: i for i, c in enumerate(chars)}
+    idx_to_char = {i: c for i, c in enumerate(chars)}
+    return chars, char_to_idx, idx_to_char, len(chars)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Model Definition
+# ─────────────────────────────────────────────────────────────────────────────
+
+class CustomVLM(nn.Module):
+    """
+    Visual Prefix-Tuning VLM.
+
+    Combines:
+      1. Frozen ViT image encoder  (768-dim output)
+      2. Trainable MLP projection  (768 → text_embed_dim)
+      3. Character-level causal Transformer decoder
+         (initialized from shakespeare_transformer.pt, then fine-tuned)
+    """
+
+    NUM_VISUAL_TOKENS = 197   # ViT: 196 patches + 1 [CLS]
+
+    def __init__(self, vocab_size, text_embed_dim=384, n_heads=8, n_layers=8,
+                 block_size=256, dropout=0.1):
+        super().__init__()
+
+        # ── 1. Vision Encoder (Frozen) ──────────────────────────────────────
+        self.vit = ViTModel.from_pretrained("google/vit-base-patch16-224-in21k")
+        for param in self.vit.parameters():
+            param.requires_grad = False
+
+        vit_hidden_size = self.vit.config.hidden_size  # 768
+
+        # ── 2. Trainable Bridge (MLP — like LLaVA) ──────────────────────────
+        self.visual_projection = nn.Sequential(
+            nn.Linear(vit_hidden_size, vit_hidden_size * 2),
+            nn.GELU(),
+            nn.Linear(vit_hidden_size * 2, text_embed_dim)
+        )
+
+        # ── 3. Character-Level Causal Transformer Decoder ───────────────────
+        self.token_embedding_table = nn.Embedding(vocab_size, text_embed_dim)
+        # Position table covers visual prefix (197) + max text (block_size)
+        self.position_embedding_table = nn.Embedding(
+            self.NUM_VISUAL_TOKENS + block_size, text_embed_dim
+        )
+
+        decoder_layer = nn.TransformerEncoderLayer(
+            d_model=text_embed_dim,
+            nhead=n_heads,
+            dim_feedforward=4 * text_embed_dim,
+            dropout=dropout,
+            batch_first=True,
+        )
+        self.decoder_blocks = nn.TransformerEncoder(decoder_layer, num_layers=n_layers)
+
+        self.ln_f = nn.LayerNorm(text_embed_dim)
+        self.lm_head = nn.Linear(text_embed_dim, vocab_size)
+
+        self.block_size = block_size
+        self.text_embed_dim = text_embed_dim
+        self.vocab_size = vocab_size
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+
+    # ───────────────────────────────��─────────────────────────────────────────
+    # Weight Loading — with architecture remapping
+    # ─────────────────────────────────────────────────────────────────────────
+
+    def load_shakespeare_weights(self, path: str, device: str = "cpu") -> dict:
+        """
+        Load pre-trained Shakespeare Transformer weights with full key remapping.
+
+        The Shakespeare checkpoint uses a custom per-head architecture:
+          blocks.N.sa_head.heads.M.{key,query,value}.weight  (head_dim, embed_dim)
+          blocks.N.sa_head.proj.{weight,bias}
+          blocks.N.ffwd.net.{0,2}.{weight,bias}
+          blocks.N.ln{1,2}.{weight,bias}
+
+        These are remapped into PyTorch nn.TransformerEncoder's fused format:
+          decoder_blocks.layers.N.self_attn.in_proj_weight  (3*embed_dim, embed_dim)
+          decoder_blocks.layers.N.self_attn.out_proj.{weight,bias}
+          decoder_blocks.layers.N.linear1.{weight,bias}
+          decoder_blocks.layers.N.linear2.{weight,bias}
+          decoder_blocks.layers.N.norm1.{weight,bias}
+          decoder_blocks.layers.N.norm2.{weight,bias}
+        """
+        print(f"📖 Loading Shakespeare weights from: {path}")
+
+        raw = torch.load(path, map_location=device)
+
+        # Unwrap common checkpoint structures
+        if isinstance(raw, dict):
+            if "model_state" in raw:
+                state_dict = raw["model_state"]
+            elif "model" in raw:
+                state_dict = raw["model"]
+            elif "state_dict" in raw:
+                state_dict = raw["state_dict"]
+            else:
+                state_dict = raw
+        else:
+            raise TypeError(f"Unexpected checkpoint type: {type(raw)}")
+
+        # ── Discover Shakespeare architecture ────────────────────────────────
+        shk_blocks = set()
+        shk_heads = set()
+        for key in state_dict:
+            if key.startswith("blocks."):
+                parts = key.split(".")
+                shk_blocks.add(int(parts[1]))
+                if "heads" in key:
+                    shk_heads.add(int(parts[4]))
+
+        n_shk_blocks = len(shk_blocks)
+        n_shk_heads = len(shk_heads) if shk_heads else self.n_heads
+        head_dim = self.text_embed_dim // self.n_heads
+
+        print(f"  📊 Shakespeare arch: {n_shk_blocks} blocks, {n_shk_heads} heads, "
+              f"head_dim={head_dim}")
+        print(f"  📊 Model arch: {self.n_layers} layers, {self.n_heads} heads")
+
+        # How many layers to load (min of checkpoint and model)
+        n_load = min(n_shk_blocks, self.n_layers)
+        n_heads_load = min(n_shk_heads, self.n_heads)
+
+        remapped = {}
+
+        # ── Remap decoder blocks ─────────────────────────────────────────────
+        for layer_idx in range(n_load):
+            prefix_src = f"blocks.{layer_idx}"
+            prefix_dst = f"decoder_blocks.layers.{layer_idx}"
+
+            # 1. Self-Attention: Fuse per-head Q, K, V into in_proj_weight
+            #    Shakespeare: heads.M.query.weight (head_dim, embed_dim)
+            #    Target: self_attn.in_proj_weight (3*embed_dim, embed_dim)
+            q_parts, k_parts, v_parts = [], [], []
+            for h in range(n_heads_load):
+                qk = f"{prefix_src}.sa_head.heads.{h}.query.weight"
+                kk = f"{prefix_src}.sa_head.heads.{h}.key.weight"
+                vk = f"{prefix_src}.sa_head.heads.{h}.value.weight"
+                if qk in state_dict and kk in state_dict and vk in state_dict:
+                    q_parts.append(state_dict[qk])
+                    k_parts.append(state_dict[kk])
+                    v_parts.append(state_dict[vk])
+
+            if q_parts:
+                # Concatenate heads: each (head_dim, embed_dim) → (embed_dim, embed_dim)
+                Q_full = torch.cat(q_parts, dim=0)  # (n_heads*head_dim, embed_dim)
+                K_full = torch.cat(k_parts, dim=0)
+                V_full = torch.cat(v_parts, dim=0)
+                # Fuse into in_proj_weight: [Q; K; V] → (3*embed_dim, embed_dim)
+                in_proj_weight = torch.cat([Q_full, K_full, V_full], dim=0)
+                remapped[f"{prefix_dst}.self_attn.in_proj_weight"] = in_proj_weight
+
+                # Create zero bias (Shakespeare has no Q/K/V bias)
+                remapped[f"{prefix_dst}.self_attn.in_proj_bias"] = torch.zeros(
+                    3 * self.text_embed_dim
+                )
+
+            # 2. Output projection
+            proj_w = f"{prefix_src}.sa_head.proj.weight"
+            proj_b = f"{prefix_src}.sa_head.proj.bias"
+            if proj_w in state_dict:
+                remapped[f"{prefix_dst}.self_attn.out_proj.weight"] = state_dict[proj_w]
+            if proj_b in state_dict:
+                remapped[f"{prefix_dst}.self_attn.out_proj.bias"] = state_dict[proj_b]
+
+            # 3. Feed-Forward Network
+            #    Shakespeare: ffwd.net.0 → linear1, ffwd.net.2 → linear2
+            for shk_idx, tgt_name in [("0", "linear1"), ("2", "linear2")]:
+                wk = f"{prefix_src}.ffwd.net.{shk_idx}.weight"
+                bk = f"{prefix_src}.ffwd.net.{shk_idx}.bias"
+                if wk in state_dict:
+                    remapped[f"{prefix_dst}.{tgt_name}.weight"] = state_dict[wk]
+                if bk in state_dict:
+                    remapped[f"{prefix_dst}.{tgt_name}.bias"] = state_dict[bk]
+
+            # 4. Layer Norms: ln1 → norm1, ln2 → norm2
+            for shk_ln, tgt_ln in [("ln1", "norm1"), ("ln2", "norm2")]:
+                for suffix in ("weight", "bias"):
+                    sk = f"{prefix_src}.{shk_ln}.{suffix}"
+                    if sk in state_dict:
+                        remapped[f"{prefix_dst}.{tgt_ln}.{suffix}"] = state_dict[sk]
+
+        # ── Non-decoder module weights ───────────────────────────────────────
+        # token_embedding_table
+        if "token_embedding_table.weight" in state_dict:
+            shk_emb = state_dict["token_embedding_table.weight"]
+            own_emb = self.token_embedding_table.weight
+            if shk_emb.shape == own_emb.shape:
+                remapped["token_embedding_table.weight"] = shk_emb
+            elif shk_emb.shape[1] == own_emb.shape[1]:
+                # Vocab size difference: copy what fits
+                n_copy = min(shk_emb.shape[0], own_emb.shape[0])
+                new_emb = own_emb.data.clone()
+                new_emb[:n_copy] = shk_emb[:n_copy]
+                remapped["token_embedding_table.weight"] = new_emb
+
+        # position_embedding_table: Shakespeare (256, 384) → Model (453, 384)
+        if "position_embedding_table.weight" in state_dict:
+            shk_pos = state_dict["position_embedding_table.weight"]  # (256, 384)
+            own_pos = self.position_embedding_table.weight           # (197+block_size, 384)
+            if shk_pos.shape == own_pos.shape:
+                remapped["position_embedding_table.weight"] = shk_pos
+            else:
+                # Expand: zero-init the full table, then copy Shakespeare positions
+                # into the TEXT portion (positions 197..197+256)
+                new_pos = torch.zeros_like(own_pos.data)
+                # Visual positions (0..196) get small random init
+                nn.init.normal_(new_pos[:self.NUM_VISUAL_TOKENS], std=0.02)
+                # Text positions: copy Shakespeare's first N positions
+                n_text_slots = own_pos.shape[0] - self.NUM_VISUAL_TOKENS
+                n_copy = min(shk_pos.shape[0], n_text_slots)
+                new_pos[self.NUM_VISUAL_TOKENS:self.NUM_VISUAL_TOKENS + n_copy] = shk_pos[:n_copy]
+                remapped["position_embedding_table.weight"] = new_pos
+                print(f"  📐 Position embeddings expanded: {shk_pos.shape} → {own_pos.shape}")
+
+        # ln_f (final layer norm)
+        for suffix in ("weight", "bias"):
+            k = f"ln_f.{suffix}"
+            if k in state_dict:
+                own_shape = getattr(self.ln_f, suffix).shape
+                if state_dict[k].shape == own_shape:
+                    remapped[k] = state_dict[k]
+
+        # lm_head
+        if "lm_head.weight" in state_dict:
+            shk_lm = state_dict["lm_head.weight"]
+            own_lm = self.lm_head.weight
+            if shk_lm.shape == own_lm.shape:
+                remapped["lm_head.weight"] = shk_lm
+            elif shk_lm.shape[1] == own_lm.shape[1]:
+                n_copy = min(shk_lm.shape[0], own_lm.shape[0])
+                new_lm = own_lm.data.clone()
+                new_lm[:n_copy] = shk_lm[:n_copy]
+                remapped["lm_head.weight"] = new_lm
+
+        if "lm_head.bias" in state_dict:
+            shk_b = state_dict["lm_head.bias"]
+            own_b = self.lm_head.bias
+            if own_b is not None and shk_b.shape == own_b.shape:
+                remapped["lm_head.bias"] = shk_b
+            elif own_b is not None:
+                n_copy = min(shk_b.shape[0], own_b.shape[0])
+                new_b = own_b.data.clone()
+                new_b[:n_copy] = shk_b[:n_copy]
+                remapped["lm_head.bias"] = new_b
+
+        # ── Load remapped weights ─────────────────────────────────────────────
+        # Verify shapes before loading
+        own_state = self.state_dict()
+        valid_remapped = {}
+        shape_mismatches = []
+        for k, v in remapped.items():
+            if k in own_state:
+                if own_state[k].shape == v.shape:
+                    valid_remapped[k] = v
+                else:
+                    shape_mismatches.append(
+                        f"    {k}: ckpt={v.shape} vs model={own_state[k].shape}"
+                    )
+            else:
+                shape_mismatches.append(f"    {k}: not in model state_dict")
+
+        result = self.load_state_dict(valid_remapped, strict=False)
+
+        print(f"  ✅ Successfully loaded {len(valid_remapped)} weight tensors (of {len(state_dict)} in checkpoint)")
+
+        if shape_mismatches:
+            print(f"  ⚠️  {len(shape_mismatches)} shape mismatches (skipped):")
+            for msg in shape_mismatches[:5]:
+                print(msg)
+
+        # Count decoder keys that were successfully loaded
+        decoder_loaded = sum(1 for k in valid_remapped if k.startswith("decoder_blocks"))
+        total_decoder = sum(1 for k in own_state if k.startswith("decoder_blocks"))
+        print(f"  📊 Decoder coverage: {decoder_loaded}/{total_decoder} tensors loaded")
+
+        return {
+            "loaded": list(valid_remapped.keys()),
+            "missing": result.missing_keys,
+            "unexpected": result.unexpected_keys,
+        }
+
+    # ─────────────────────────────────────────────────────────────────────────
+    # Freezing / Unfreezing / Parameter Counting
+    # ─────────────────────────────────────────────────────────────────────────
+
+    def freeze_decoder(self):
+        """Freeze the Shakespeare decoder so only visual_projection trains."""
+        for name, param in self.named_parameters():
+            if not name.startswith("visual_projection"):
+                param.requires_grad = False
+        # Ensure ViT is frozen
+        for param in self.vit.parameters():
+            param.requires_grad = False
+
+    def unfreeze_decoder(self):
+        """
+        Unfreeze the decoder for fine-tuning while keeping ViT frozen.
+        
+        This allows the decoder to adapt from Shakespeare text to COCO captions.
+        The visual_projection is also trainable.
+        """
+        # First, freeze everything
+        for param in self.parameters():
+            param.requires_grad = False
+
+        # Unfreeze visual_projection (always trainable)
+        for param in self.visual_projection.parameters():
+            param.requires_grad = True
+
+        # Unfreeze ALL decoder components
+        for param in self.token_embedding_table.parameters():
+            param.requires_grad = True
+        for param in self.position_embedding_table.parameters():
+            param.requires_grad = True
+        for param in self.decoder_blocks.parameters():
+            param.requires_grad = True
+        for param in self.ln_f.parameters():
+            param.requires_grad = True
+        for param in self.lm_head.parameters():
+            param.requires_grad = True
+
+        # ViT stays FROZEN
+        for param in self.vit.parameters():
+            param.requires_grad = False
+
+    def get_param_groups(self, projection_lr=1e-4, decoder_lr=5e-5):
+        """
+        Return optimizer param groups with discriminative learning rates.
+        
+        - visual_projection: higher LR (learning from scratch)
+        - decoder: lower LR (gentle adaptation from Shakespeare)
+        """
+        projection_params = []
+        decoder_params = []
+
+        for name, param in self.named_parameters():
+            if not param.requires_grad:
+                continue
+            if name.startswith("visual_projection"):
+                projection_params.append(param)
+            else:
+                decoder_params.append(param)
+
+        return [
+            {"params": projection_params, "lr": projection_lr},
+            {"params": decoder_params, "lr": decoder_lr},
+        ]
+
+    def trainable_params(self):
+        """Return count of trainable parameters."""
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+    # ─────────────────────────────────────────────────────────────────────────
+    # Forward Pass
+    # ─────────────────────────────────────────────────────────────────────────
+
+    def forward(self, pixel_values, text_input_ids, text_targets=None):
+        B, T = text_input_ids.shape
+
+        # ── Image Encoding (frozen ViT) ──────────────────────────────────────
+        with torch.no_grad():
+            vit_outputs = self.vit(pixel_values=pixel_values)
+        image_embeds = vit_outputs.last_hidden_state  # (B, 197, 768)
+
+        # ── Project to text embedding space ──────────────────────────────────
+        visual_prefix = self.visual_projection(image_embeds)  # (B, 197, 384)
+        num_visual = visual_prefix.shape[1]                   # 197
+
+        # ── Text Embeddings ───────────────────────────────────────────────────
+        T_clipped = min(T, self.block_size)
+        text_in = text_input_ids[:, :T_clipped]
+        tok_emb = self.token_embedding_table(text_in)         # (B, T, 384)
+
+        # ── Positional Embeddings (covers full combined sequence) ─────────────
+        # Positions 0..196 → visual prefix, 197..197+T → text tokens
+        total_len = num_visual + T_clipped
+        pos_ids = torch.arange(total_len, device=text_in.device)
+        pos_emb = self.position_embedding_table(pos_ids)      # (num_visual+T, 384)
+
+        vis_pos = pos_emb[:num_visual]                        # (197, 384)
+        txt_pos = pos_emb[num_visual:]                        # (T, 384)
+
+        visual_emb = visual_prefix + vis_pos                  # (B, 197, 384)
+        text_emb   = tok_emb + txt_pos                        # (B, T, 384)
+
+        # ── Fusion: [visual_prefix | text_emb] ───────────────────────────────
+        combined = torch.cat([visual_emb, text_emb], dim=1)   # (B, 197+T, 384)
+        tot = combined.shape[1]
+
+        # ── Causal Attention Mask ─────────────────────────────────────────────
+        # Visual tokens attend to each other freely.
+        # Text tokens attend to all visual tokens + causally to previous text.
+        mask = torch.full((tot, tot), float("-inf"), device=text_in.device)
+        mask[:num_visual, :num_visual] = 0.0          # visual→visual: free
+        mask[num_visual:, :num_visual] = 0.0           # text→visual: free
+        causal = torch.triu(
+            torch.full((T_clipped, T_clipped), float("-inf"), device=text_in.device),
+            diagonal=1,
+        )
+        mask[num_visual:, num_visual:] = causal         # text→text: causal
+
+        # ── Decoder ───────────────────────────────────────────────────────────
+        x = self.decoder_blocks(combined, mask=mask, is_causal=False)
+        text_out = x[:, num_visual:, :]
+        text_out = self.ln_f(text_out)
+        logits = self.lm_head(text_out)                       # (B, T, vocab)
+
+        # ── Loss (ignore padding index 0) ─────────────────────────────────────
+        loss = None
+        if text_targets is not None:
+            tgt = text_targets[:, :T_clipped]
+            loss = F.cross_entropy(
+                logits.reshape(B * T_clipped, -1),
+                tgt.reshape(B * T_clipped),
+                ignore_index=0,
+            )
+
+        return logits, loss
+
+    # ─────────────────────────────────────────────────────────────────────────
+    # Generation
+    # ─────────────────────────────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def generate(self, pixel_values, char_to_idx, idx_to_char,
+                 max_new_tokens=100, temperature=0.8):
+        """
+        Autoregressive character-level caption generation (temperature sampling).
+
+        Args:
+            pixel_values   : (1, 3, H, W) pre-processed image tensor
+            char_to_idx    : character → index mapping
+            idx_to_char    : index → character mapping
+            max_new_tokens : how many characters to generate
+            temperature    : sampling temperature (0.8 = slightly sharper than uniform)
+
+        Returns:
+            generated_text : str
+        """
+        self.eval()
+        device = pixel_values.device
+
+        bos_idx = char_to_idx.get("\n", 0)
+        idx_seq = torch.tensor([[bos_idx]], dtype=torch.long, device=device)
+
+        for _ in range(max_new_tokens):
+            # Clip text to block_size — the forward method handles the visual
+            # prefix separately, so we only need to limit the text portion.
+            idx_cond = idx_seq[:, -self.block_size:]
+            logits, _ = self(pixel_values, idx_cond)
+            # Take the last time step
+            logits_last = logits[:, -1, :] / max(temperature, 1e-5)
+            probs = F.softmax(logits_last, dim=-1)
+            next_idx = torch.multinomial(probs, num_samples=1)
+            idx_seq = torch.cat([idx_seq, next_idx], dim=1)
+
+        # Decode, skip the leading BOS
+        generated = "".join(
+            idx_to_char.get(i.item(), "?") for i in idx_seq[0, 1:]
+        )
+        return generated
+
+    @torch.no_grad()
+    def generate_beam(self, pixel_values, char_to_idx, idx_to_char,
+                      max_new_tokens=100, num_beams=4, length_penalty=1.0):
+        """
+        Beam-search character-level caption generation.
+
+        At each step we keep the top `num_beams` partial sequences ranked by
+        cumulative log-probability (with optional length penalty).
+
+        Args:
+            pixel_values   : (1, 3, H, W) image tensor
+            char_to_idx    : char → idx mapping
+            idx_to_char    : idx → char mapping
+            max_new_tokens : max characters to generate
+            num_beams      : beam width (1 = greedy)
+            length_penalty : >1 favors longer sequences; <1 favors shorter
+
+        Returns:
+            generated_text : str (best beam)
+        """
+        self.eval()
+        device = pixel_values.device
+
+        bos_idx = char_to_idx.get("\n", 0)
+        # Each beam: (score, token_sequence_tensor)
+        beams = [(0.0, torch.tensor([[bos_idx]], dtype=torch.long, device=device))]
+
+        for _ in range(max_new_tokens):
+            candidates = []
+            for score, seq in beams:
+                idx_cond = seq[:, -self.block_size:]
+                logits, _ = self(pixel_values, idx_cond)
+                log_probs = F.log_softmax(logits[:, -1, :], dim=-1)  # (1, vocab)
+                topk_probs, topk_ids = log_probs.topk(num_beams, dim=-1)
+
+                for k in range(num_beams):
+                    new_score = score + topk_probs[0, k].item()
+                    new_seq = torch.cat(
+                        [seq, topk_ids[:, k:k+1]], dim=1
+                    )
+                    candidates.append((new_score, new_seq))
+
+            # Apply length penalty and keep top beams
+            candidates.sort(
+                key=lambda x: x[0] / (x[1].shape[1] ** length_penalty),
+                reverse=True,
+            )
+            beams = candidates[:num_beams]
+
+        best_seq = beams[0][1]
+        return "".join(idx_to_char.get(i.item(), "?") for i in best_seq[0, 1:])

models/git_tuner.py

@@ -0,0 +1,85 @@
+"""
+models/git_tuner.py
+===================
+Baseline 2 — Zero Cross-Attention / Self-Attention Prefix (GIT)
+
+Architecture: GIT (Generative Image-to-Text) abandons cross-attention entirely.
+It concatenates image patch embeddings directly in front of the text tokens and
+runs a single causal self-attention Transformer over the combined sequence.
+
+There is NO cross-attention block. The model learns to fuse modalities purely
+through self-attention across a unified image+text token sequence. This makes
+the ablation masks work differently — we control which image tokens are
+prepended to the sequence rather than using encoder_attention_mask.
+"""
+
+import os
+import torch
+from transformers import AutoProcessor, AutoModelForCausalLM
+
+
+def get_git_model(cfg, device):
+    """
+    Load microsoft/git-base-coco with gradient checkpointing.
+    GIT uses AutoModelForCausalLM interface.
+    """
+    model_id = cfg.git_model_id
+
+    processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
+    model = AutoModelForCausalLM.from_pretrained(model_id)
+
+    try:
+        model.gradient_checkpointing_enable()
+        print("✅ Gradient checkpointing enabled (GIT)")
+    except Exception as e:
+        print(f"⚠️  Gradient checkpointing failed: {e}")
+
+    model.config.use_cache = False
+    model.to(device)
+
+    n_params = sum(p.numel() for p in model.parameters()) / 1e6
+    print(f"✅ GIT loaded on {device}: {model_id} ({n_params:.1f}M params)")
+    return model, processor
+
+
+def generate_caption(model, processor, image_pil, device,
+                     max_new_tokens=32, num_beams=4):
+    """
+    Generate a caption for a single PIL image using GIT.
+
+    Note: GIT has no encoder_attention_mask concept (no cross-attention).
+    Ablation for GIT is handled upstream by modifying the pixel_values
+    (e.g., masking image regions) before passing to the model, OR by
+    returning a note that GIT is not compatible with encoder-mask ablations.
+    """
+    model.eval()
+    inputs = processor(images=image_pil, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            num_beams=num_beams,
+        )
+
+    caption = processor.batch_decode(output_ids, skip_special_tokens=True)[0]
+    return caption
+
+
+def save_ckpt(model, processor, optimizer, scheduler,
+              step, epoch, cfg_dict, path):
+    os.makedirs(path, exist_ok=True)
+    model.save_pretrained(path)
+    processor.save_pretrained(path)
+
+    torch.save(
+        {
+            "step": step,
+            "epoch": epoch,
+            "optimizer": optimizer.state_dict() if optimizer else None,
+            "scheduler": scheduler.state_dict() if scheduler else None,
+            "cfg": cfg_dict,
+        },
+        os.path.join(path, "train_state.pt"),
+    )
+    print(f"✅ GIT checkpoint saved: {path}")

models/vit_gpt2_tuner.py

@@ -0,0 +1,110 @@
+"""
+models/vit_gpt2_tuner.py
+========================
+Baseline 1 — Standard Cross-Attention (ViT-GPT2)
+
+Architecture: Every generated text token in the GPT-2 decoder attends to ALL
+197 ViT patch embeddings via explicit cross-attention blocks injected between
+each GPT-2 self-attention layer.
+
+This is the "brute-force" cross-attention baseline: no restrictions, no pooling.
+"""
+
+import os
+import torch
+from transformers import (
+    VisionEncoderDecoderModel,
+    ViTImageProcessor,
+    AutoTokenizer,
+)
+
+
+def get_vit_gpt2_model(cfg, device):
+    """
+    Load the VisionEncoderDecoderModel (ViT-GPT2) with:
+    - Gradient checkpointing enabled
+    - use_cache=False (required with grad checkpointing)
+    - Proper pad/bos/eos tokens set for GPT-2
+    """
+    model_id = cfg.vit_gpt2_model_id
+
+    processor = ViTImageProcessor.from_pretrained(model_id, use_fast=True)
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+    # GPT-2 has no pad token by default — use eos as pad
+    tokenizer.pad_token = tokenizer.eos_token
+
+    model = VisionEncoderDecoderModel.from_pretrained(model_id)
+    model.config.decoder_start_token_id = tokenizer.bos_token_id
+    model.config.pad_token_id = tokenizer.pad_token_id
+    model.config.eos_token_id = tokenizer.eos_token_id
+
+    # Memory optimizations
+    try:
+        model.gradient_checkpointing_enable()
+        print("✅ Gradient checkpointing enabled (ViT-GPT2)")
+    except Exception as e:
+        print(f"⚠️  Gradient checkpointing failed: {e}")
+
+    model.config.use_cache = False
+
+    # Resize images to cfg.image_size
+    try:
+        processor.size = {"height": cfg.image_size, "width": cfg.image_size}
+        print(f"✅ Image size set to {cfg.image_size}px")
+    except Exception as e:
+        print(f"⚠️  Could not set image size: {e}")
+
+    model.to(device)
+    n_params = sum(p.numel() for p in model.parameters()) / 1e6
+    print(f"✅ ViT-GPT2 loaded on {device}: {model_id} ({n_params:.1f}M params)")
+
+    return model, processor, tokenizer
+
+
+def generate_caption(model, processor, tokenizer, image_pil, device,
+                     max_new_tokens=32, num_beams=4,
+                     encoder_attention_mask=None):
+    """
+    Generate a caption for a single PIL image.
+
+    encoder_attention_mask: (1, num_patches) allows ablation-mode masking.
+    If None, defaults to full attention (all 1s).
+    """
+    model.eval()
+    inputs = processor(images=image_pil, return_tensors="pt").to(device)
+    pixel_values = inputs["pixel_values"]
+
+    gen_kwargs = dict(
+        pixel_values=pixel_values,
+        max_new_tokens=max_new_tokens,
+        num_beams=num_beams,
+    )
+    if encoder_attention_mask is not None:
+        gen_kwargs["attention_mask"] = encoder_attention_mask.to(device)
+
+    with torch.no_grad():
+        output_ids = model.generate(**gen_kwargs)
+
+    caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    return caption
+
+
+def save_ckpt(model, processor, tokenizer, optimizer, scheduler,
+              step, epoch, cfg_dict, path):
+    os.makedirs(path, exist_ok=True)
+    model.save_pretrained(path)
+    processor.save_pretrained(path)
+    tokenizer.save_pretrained(path)
+
+    torch.save(
+        {
+            "step": step,
+            "epoch": epoch,
+            "optimizer": optimizer.state_dict() if optimizer else None,
+            "scheduler": scheduler.state_dict() if scheduler else None,
+            "cfg": cfg_dict,
+        },
+        os.path.join(path, "train_state.pt"),
+    )
+    print(f"✅ ViT-GPT2 checkpoint saved: {path}")

project_02_DS

@@ -0,0 +1 @@
+Subproject commit a5ea2c20321ecd6767352a2393f0bc58a8a9f059

requirements.txt

@@ -0,0 +1,13 @@
+torch
+torchvision
+torchaudio
+transformers>=4.37.0
+datasets
+aiohttp
+streamlit
+numpy
+Pillow
+tqdm
+accelerate
+sentencepiece
+pycocoevalcap

shakespeare_transformer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:652c085bf4c7275182fe726a38d3034aeaf1d67d8dc93f8c014976b2408f7ce5
+size 74253331

simplified_overview_vlm_image_captioning_project.md

@@ -0,0 +1,224 @@
+# How I Built a System That Teaches Computers to Describe Photographs
+
+**A non-technical overview of the VLM Caption Lab project**  
+*Author: Manoj Kumar | 4 March 2026*
+
+---
+
+## What Is This Project About?
+
+Imagine showing a photograph to a friend and asking them to describe it in one sentence. They might say, *"A man in a suit standing in front of a tree,"* or *"A tennis match in a large arena with a crowd watching."* For us, this is effortless — our brains process the entire image, identify the objects, understand the scene, and produce a fluent sentence in under a second.
+
+For a computer, this is remarkably difficult. The technical name for this task is **"image captioning,"** and it lives at the crossroads of two hard problems: understanding what is in an image (computer vision) and writing grammatically correct, meaningful sentences (natural language generation).
+
+This project explores that challenge — **but I did not just build one system. I built and compared four of them,** each with a fundamentally different approach to the core problem of looking at the image while writing about it.
+
+---
+
+## The Four Models I Built (And Why They Are Different)
+
+Think of image captioning like a person looking at a painting while narrating what they see into a microphone. The four models I compared differ in **how the person glances at the painting while they talk.**
+
+---
+
+### 🔵 Model 1: BLIP — The Selective Glancer
+
+**How it works :** BLIP is like a narrator who has trained themselves to only glance at the painting when they need to. When they are saying generic words like "a" or "the" or "is," they just focus on their own sentence. When they need to mention something specific — like "bicycle" or "standing" — they look up at the painting to confirm what they see.
+
+**Why this is smart:** Most words in a sentence are structural, not visual. There is no need to look at the image to say "the" or "in front of." BLIP learns when to look and when not to, which prevents it from getting confused by too much visual information.
+
+**Size:** 224 million parameters  
+**Best CIDEr score:** **0.62** (with optimized settings)
+
+---
+
+###  Model 2: ViT-GPT2 — The Constant Starer
+
+**How it works in plain English:** ViT-GPT2 takes the opposite approach — for every single word, it stares at the entire painting. Writing "a"? Look at the whole image. Writing "dog"? Look at the whole image. Writing "the"? Still looking at the whole image.
+
+**Why this still works:** Even though it is wasteful, staring at everything guarantees the model never misses any visual detail. The downside is that this constant stream of visual information can sometimes confuse the language part of the model.
+
+**Size:** 239 million parameters  
+**Typical CIDEr score:** ~0.55
+
+---
+
+###  Model 3: GIT — The Memorizer
+
+**How it works in plain English:** GIT does something clever — instead of switching between looking at the painting and writing words, it first memorizes the entire painting and then writes the caption purely from memory.
+
+In technical terms, GIT converts the image into a set of structured "memory notes" and places them at the beginning of its sentence. Then it processes everything — image memories and text — in one continuous stream. There is no separate "looking at the painting" step.
+
+**Why this is elegant:** It is simpler and faster because it does not need the extra machinery for looking back and forth between image and text. The entire intelligence is in one unified processing step.
+
+**Size:** 177 million parameters (smallest of the four)  
+**Typical CIDEr score:** ~0.54
+
+---
+
+###  Model 4: Custom VLM — The Shakespeare Bot Learning Modern English
+
+**How it works in plain English:** This is the most experimental model, and the one **I built entirely from scratch.** Imagine a narrator who grew up reading only Shakespeare and has never seen a photograph before. You give them a pair of glasses (a visual encoder — something that can look at images) and a translator (a small bridging network) and ask them to describe modern photographs.
+
+The "Shakespeare bot" is a text generator I had previously trained on the complete works of Shakespeare. It knows English grammar and sentence structure — but in Elizabethan English. The challenge was teaching it to (a) understand images through the "glasses" and (b) speak in modern, descriptive English instead of iambic pentameter.
+
+**Why I built this:** To understand what minimum set of components you need to make a functioning vision-language model. Instead of downloading a ready-made model with billions of parameters, I wanted to see if I could glue together a vision model and a text model with just a small trainable "bridge" in between.
+
+**Size:** 103 million parameters total, but only **16.2 million are trainable** (the rest are frozen)  
+**Best CIDEr score:** **0.2863** (still learning, but it works!)
+
+---
+
+## What Is CIDEr? (The Score We Use to Measure Quality)
+
+Throughout this summary, I mention "CIDEr scores." Here is what they mean:
+
+**CIDEr** stands for "Consensus-based Image Description Evaluation." In simple terms, it compares the caption our model generates to **five human-written descriptions** of the same image.
+
+- It counts how many meaningful words overlap between the model's caption and the human captions
+- It gives more weight to descriptive words (like "bicycle" or "stadium") than common words (like "the" or "is")
+- **A higher score means the computer's description sounds more like what a human would write**
+
+| CIDEr Score | What It Means |
+|---|---|
+| 0.00 | Completely wrong — no overlap with human descriptions |
+| 0.20–0.30 | Early stage — some correct words, but the sentence may be awkward |
+| 0.50–0.60 | Good — clearly related to the image, mostly sensible |
+| 0.80–1.00 | Excellent — almost indistinguishable from a human caption |
+
+---
+
+## The Custom Model Story: A Journey of Debugging and Discovery
+
+This is the part of the project I am most proud of, because it taught me the most about how machine learning actually works in practice — not just in theory, but when things go wrong.
+
+### Chapter 1: "Why Is It Speaking Gibberish?"
+
+My first attempt at the Custom VLM produced output like this:
+
+> *"iGiiiiiGiviqiGqiFliqiGidlidiliGilFGilqiiiqiiiiGii"*
+
+That is not English. That is not even Shakespeare. It is random noise.
+
+**The problem:** The connection between the "glasses" (the image encoder) and the "brain" (the Shakespeare text generator) was too weak. I was using a single mathematical transformation to convert visual information into text information. Think of it like trying to translate a painting into a poem by only measuring the canvas size — you are missing all the important details.
+
+**CIDEr score at this stage: 0.0000 — literally zero.**
+
+### Chapter 2: "Better Connection, But Still Broken"
+
+I upgraded the connection to a more powerful two-layer network. This is like upgrading from a basic dictionary to a bilingual tutor who understands context. The training measurements started improving — the numbers were going down, which normally means the model is learning.
+
+But the output was still gibberish.
+
+After days of investigation, I found the real problem — and it was a doozy:
+
+> **When I loaded the Shakespeare brain into the model, 97% of the brain weights failed to load. Silently. No error message. No warning. The software just said "everything is fine" and moved on.**
+
+My model had been running on a **randomly initialized brain** — essentially trying to learn language from scratch while simultaneously trying to learn to describe images. Imagine asking someone with amnesia to write poetry about something they've never seen. That's what my model was trying to do.
+
+**Why did this happen?** The two models (Shakespeare and my VLM) stored their internal knowledge in slightly different formats. It is like trying to load a Word document into Excel — both are files, but the internal structure is completely different. The software saw the mismatched formats and just... skipped everything. Without telling me.
+
+### Chapter 3: "It Finally Speaks!"
+
+The fix required three things:
+1. **Match the formats** — Make the new model structure identical to the Shakespeare model's structure (8 layers, 8 attention heads, matching dimensions)
+2. **Translate the weights** — Write custom code to convert the Shakespeare data from one format to another
+3. **Let the brain learn** — Instead of freezing the Shakespeare knowledge, let the model slowly adapt from old English to modern descriptions
+
+**The result was immediate.** From the very first training session after the fix, the improvement was dramatic:
+
+> Before fix: *"iGiiiiiGiviqiGqiFliqiGidlidiliGilFGilqiiiqiiiiGii"* (CIDEr: 0.0000)  
+> After fix: *"man in the bluess and white play with and a pizza"* (CIDEr: 0.2863)
+
+Not perfect. Not even grammatically correct. But it is **clearly English**, it is **clearly attempting to describe an image**, and it went from zero to something meaningful. The word "man" appeared because the image showed a man. The model learned real English words and connected them to visual concepts.
+
+---
+
+## What We Tested: The Three Experiments
+
+### Experiment 1: "Can We Cover Part of the Image?"
+
+I blocked parts of the image from the model and measured whether the captions got worse. The results were genuinely surprising:
+
+| What We Did | Effect on Caption Quality |
+|---|---|
+| Showed the **full image** | Baseline quality (CIDEr: 0.5371) |
+| **Hid 50%** of the image randomly | **No change at all** (CIDEr: 0.5371) |
+| Showed **only the center** (removed background) | **No change at all** (CIDEr: 0.5371) |
+| **Compressed everything** into one tiny summary | **Complete failure** (CIDEr: 0.0008 — a 99.8% drop) |
+
+**What this teaches us:** Images contain a lot of redundant information. You can throw away half the visual data and still get perfectly good captions. But if you compress everything into a single summary, you lose the information about **where things are** relative to each other — and that spatial information turns out to be essential for describing a scene.
+
+### Experiment 2: "What Settings Produce the Best Captions?"
+
+When a model generates a caption, it uses a search algorithm that considers multiple possible sentences and picks the best one. I tested **18 different combinations** of settings and found:
+
+- **Considering more candidate sentences (10 instead of 3) helped significantly** — about 13% improvement
+- **Slightly encouraging shorter captions helped** — models tend to ramble when given too much freedom
+- **Best combination found: CIDEr score of 0.6199** (up from 0.48 with the worst settings)
+
+### Experiment 3: "Does Caption Quality During Training Matter?"
+
+I compared different strategies for selecting which human captions to show the model during training:
+
+| Strategy | CIDEr Score |
+|---|---|
+| Use any random caption | **0.6359** ← best for this clean dataset |
+| Use only short captions (≤ 9 words) | 0.6016 |
+| Use only medium-length captions (5–25 words) | 0.5877 |
+| Use only long captions (≥ 12 words) | 0.5389 |
+
+**Bottom line:** For this particular dataset (which is already well-curated), using raw unfiltered captions works best. But filtering is recommended for noisier datasets.
+
+---
+
+## The Interactive Demo
+
+I built a web application where anyone can try the models themselves:
+
+- **Upload any photo** and get a caption from any of the four models
+- **Compare all four models** side by side on the same image — see how each one describes the same picture differently
+- **Switch between pre-trained and fine-tuned** versions of each model
+- **Adjust generation settings** — control how the model searches for the best caption
+- **View experiment results** — browse all the findings from the three experiments
+
+Every generated caption goes through a **safety filter** before being shown, because AI models can occasionally produce inappropriate descriptions. The filter uses a toxicity detection model to catch and block offensive content.
+
+---
+
+## Summary of Results
+
+| Model | Approach | CIDEr Score | Key Strength |
+|---|---|---|---|
+| **BLIP** | Selective looking | **0.62** (best settings) | Best quality — knows when to look vs. when to focus on grammar |
+| **ViT-GPT2** | Constant looking | ~0.55 | Strong baseline — full visual access at all times |
+| **GIT** | Memory-based | ~0.54 | Elegant and efficient — no cross-attention needed at all |
+| **Custom VLM** | Built from scratch | **0.29** | Proof of concept — works despite tiny vocabulary and Shakespeare origins |
+
+---
+
+## What I Actually Learned
+
+1. **There is no single best way to connect vision and language.** BLIP's selective attention works best overall, but GIT's simpler approach is surprisingly competitive — proving that you do not always need complex mechanisms to solve complex problems.
+
+2. **Silent failures are the most dangerous bugs in machine learning.** The most time-consuming problem in this project was a weight-loading failure that produced zero error messages. The model ran, the loss decreased, everything looked normal — but 97% of the model was running on random noise. I now always verify that weights loaded correctly.
+
+3. **The number your model optimizes during training is not necessarily the number that tells you if it is doing a good job.** Training loss went down steadily, but the captions were still gibberish. Only when I started measuring CIDEr (actual caption quality) did I understand what was really happening.
+
+4. **Small models can learn big tasks with the right approach.** The Custom VLM has only 16.2 million trainable parameters — roughly 1/15th the size of BLIP — yet it learned to produce recognizable English descriptions of images by building on existing Shakespeare knowledge.
+
+5. **Images are surprisingly redundant.** You can literally hide half the image and the model generates identical captions. But structure matters — where objects are relative to each other is more important than being able to see every pixel.
+
+---
+
+## What Could Be Improved Next
+
+If I continue this project, the highest-impact improvements would be:
+
+- **Better vocabulary:** The Custom VLM currently spells everything letter-by-letter (65 characters). Switching to a word-piece vocabulary (thousands of tokens) would dramatically reduce the difficulty.
+- **Stronger language foundation:** Replacing the Shakespeare decoder with a modern language model like GPT-2 would give the model native modern English instead of having to translate from Elizabethan.
+- **More training data:** We currently use only 18% of the available dataset images.
+
+---
+
+*Project by Manoj Kumar, March 2026*

train.py

@@ -0,0 +1,472 @@
+"""
+train.py
+========
+Unified training entrypoint for all VLM architectures:
+  --model blip      → Fine-tune BLIP (Multimodal Mixture Attention)
+  --model vit_gpt2  → Fine-tune ViT-GPT2 (Standard Cross-Attention)
+  --model git       → Fine-tune GIT (Zero Cross-Attention / Self-Attention Prefix)
+  --model custom    → Train visual_projection only (Visual Prefix-Tuning)
+
+Checkpoint Strategy:
+  All outputs are saved under outputs/{model_name}/:
+    - latest/   — overwritten every epoch (always the most recent state)
+    - best/     — overwritten only when validation loss improves
+
+Optimized for Apple Silicon MPS backend with:
+  - Gradient accumulation
+  - Gradient checkpointing
+  - Cosine LR scheduler with linear warmup
+  - MPS-safe DataLoader settings (num_workers=0, pin_memory=False)
+"""
+
+import argparse
+import math
+import time
+import os
+import torch
+from torch.optim import AdamW
+from transformers import get_cosine_schedule_with_warmup
+from tqdm.auto import tqdm
+
+from config import CFG
+from data_prep import get_dataloaders, get_dataloaders_for_model, get_custom_vlm_dataloader
+from models.blip_tuner import get_blip_model, save_ckpt as blip_save, generate_with_mask
+from models.vit_gpt2_tuner import get_vit_gpt2_model, save_ckpt as vit_gpt2_save
+from models.git_tuner import get_git_model, save_ckpt as git_save
+from models.custom_vlm import CustomVLM, build_char_vocab
+from pycocoevalcap.cider.cider import Cider
+
+
+def get_device():
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    elif torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+def get_output_paths(cfg, model_name: str):
+    """
+    Return (latest_dir, best_dir) for a given model.
+    Creates directories if they don't exist.
+    """
+    base = os.path.join(cfg.output_root, model_name)
+    latest = os.path.join(base, "latest")
+    best = os.path.join(base, "best")
+    os.makedirs(latest, exist_ok=True)
+    os.makedirs(best, exist_ok=True)
+    return latest, best
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Shared Training Loop
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _generate_hf_captions(model, batch, model_name, device,
+                          processor=None, tokenizer=None):
+    """
+    Generate captions for a batch of images using the appropriate HuggingFace model.
+    Returns (predictions: list[str], ground_truths: list[str]).
+    """
+    pixel_values = batch["pixel_values"].to(device)
+
+    if model_name == "BLIP":
+        B = pixel_values.shape[0]
+        mask = torch.ones(B, 197, dtype=torch.long, device=device)
+        decoded = generate_with_mask(
+            model, processor, device=device,
+            pixel_values=pixel_values,
+            encoder_attention_mask=mask,
+            max_new_tokens=32, num_beams=4,
+        )
+        preds = decoded  # generate_with_mask already returns decoded strings
+        labels = batch["labels"].clone()
+        gt_texts = processor.batch_decode(labels, skip_special_tokens=True)
+
+    elif model_name == "VIT_GPT2":
+        out = model.generate(
+            pixel_values=pixel_values, num_beams=4, max_new_tokens=32,
+        )
+        preds = [tokenizer.decode(ids, skip_special_tokens=True) for ids in out]
+        labels = batch["labels"].clone()
+        labels[labels == -100] = tokenizer.pad_token_id
+        gt_texts = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+    elif model_name == "GIT":
+        inputs = {k: v.to(device) for k, v in batch.items()
+                  if k in ("pixel_values", "input_ids", "attention_mask")}
+        out = model.generate(**inputs, num_beams=4, max_new_tokens=32)
+        preds = processor.batch_decode(out, skip_special_tokens=True)
+        labels = batch["labels"].clone()
+        labels[labels == -100] = processor.tokenizer.pad_token_id
+        gt_texts = processor.batch_decode(labels, skip_special_tokens=True)
+    else:
+        return [], []
+
+    return preds, gt_texts
+
+
+def run_training_loop(model, optimizer, scheduler, train_loader, val_loader,
+                      cfg, save_latest_fn, save_best_fn, model_name,
+                      processor=None, tokenizer=None):
+    """
+    Shared gradient-accumulation training loop for all HuggingFace models.
+
+    Now includes per-epoch:
+      - Validation loss
+      - CIDEr scoring via greedy generation
+      - CIDEr-based checkpointing (saves best/ based on highest CIDEr)
+    """
+    device = get_device()
+    model.train()
+    global_step = 0
+    best_cider = -1.0
+    t0 = time.time()
+
+    for epoch in range(1, cfg.epochs + 1):
+        model.train()
+        pbar = tqdm(train_loader, desc=f"[{model_name}] Epoch {epoch}/{cfg.epochs}")
+        running_loss = 0.0
+        epoch_loss_sum = 0.0
+        epoch_batches = 0
+        optimizer.zero_grad(set_to_none=True)
+
+        for i, batch in enumerate(pbar, start=1):
+            batch = {k: v.to(device) for k, v in batch.items()}
+
+            out = model(**batch)
+            loss = out.loss / cfg.grad_accum
+            loss.backward()
+            running_loss += loss.item()
+            epoch_loss_sum += out.loss.item()
+            epoch_batches += 1
+
+            if i % cfg.grad_accum == 0 or i == len(train_loader):
+                torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm)
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+                global_step += 1
+
+                if global_step % cfg.log_every == 0:
+                    avg = running_loss / cfg.log_every
+                    running_loss = 0.0
+                    pbar.set_postfix({"loss": f"{avg:.4f}",
+                                      "lr": f"{scheduler.get_last_lr()[0]:.2e}"})
+
+        # End of epoch — training metrics
+        epoch_avg_loss = epoch_loss_sum / max(epoch_batches, 1)
+        print(f"\n📊 Epoch {epoch}/{cfg.epochs} avg loss (Train): {epoch_avg_loss:.4f}")
+
+        # ── Validation Loop: Loss + CIDEr ────────────────────────────────────
+        model.eval()
+        val_loss_sum = 0.0
+        val_batches = 0
+        gts, res = {}, {}
+        max_eval_batches = 10
+        print("   🔍 Running Validation (Loss & CIDEr)...")
+
+        with torch.no_grad():
+            for i, batch in enumerate(val_loader):
+                if i >= max_eval_batches:
+                    break
+
+                batch_d = {k: v.to(device) for k, v in batch.items()}
+
+                # 1. Validation loss
+                out = model(**batch_d)
+                val_loss_sum += out.loss.item()
+                val_batches += 1
+
+                # 2. Generate captions for CIDEr
+                preds, gt_texts = _generate_hf_captions(
+                    model, batch, model_name, device,
+                    processor=processor, tokenizer=tokenizer,
+                )
+                for j, (p, g) in enumerate(zip(preds, gt_texts)):
+                    k = f"{epoch}_{i}_{j}"
+                    res[k] = [p]
+                    gts[k] = [g]
+
+        val_avg_loss = val_loss_sum / max(val_batches, 1)
+        print(f"   📉 Validation Loss: {val_avg_loss:.4f}")
+
+        # Compute CIDEr
+        cider_score = 0.0
+        if gts:
+            scorer = Cider()
+            cider_score, _ = scorer.compute_score(gts, res)
+        print(f"   🎯 Validation CIDEr: {cider_score:.4f}")
+
+        # Save latest checkpoint
+        save_latest_fn(step=global_step, epoch=epoch)
+        print(f"   💾 Saved → latest/")
+
+        # Save best based on CIDEr score
+        if cider_score > best_cider:
+            best_cider = cider_score
+            save_best_fn(step=global_step, epoch=epoch)
+            print(f"   🏆 New best CIDEr (score={best_cider:.4f}) → best/")
+
+    elapsed = (time.time() - t0) / 60.0
+    print(f"\n✅ {model_name} training complete in {elapsed:.2f} minutes")
+    print(f"   Best validation CIDEr: {best_cider:.4f}")
+    return global_step
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Custom VLM Training (projection-only)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def train_custom_vlm(cfg, device):
+    print("📖 Loading Shakespeare corpus for character vocabulary...")
+    with open(cfg.shakespeare_file, "r", encoding="utf-8") as f:
+        text = f.read()
+    _, char_to_idx, idx_to_char, vocab_size = build_char_vocab(text)
+    print(f"✅ Vocabulary size: {vocab_size} characters")
+
+    model = CustomVLM(
+        vocab_size=vocab_size,
+        text_embed_dim=cfg.text_embed_dim,
+        n_heads=cfg.n_heads,
+        n_layers=cfg.n_layers,
+        block_size=cfg.block_size,
+        dropout=cfg.dropout,
+    )
+
+    # ── Load pre-trained Shakespeare decoder weights (CRITICAL) ──────────────
+    shakespeare_path = getattr(cfg, "shakespeare_weights_path",
+                               "./shakespeare_transformer.pt")
+    if os.path.exists(shakespeare_path):
+        model.load_shakespeare_weights(shakespeare_path)
+        print(f"✅ Shakespeare decoder weights loaded from {shakespeare_path}")
+    else:
+        print(f"⚠️  shakespeare_transformer.pt not found at {shakespeare_path}")
+        print("    Training with randomly initialized decoder (significantly worse).")
+
+    model.unfreeze_decoder()
+    model.to(device)
+
+    n_train = model.trainable_params()
+    n_total = sum(p.numel() for p in model.parameters())
+    print(f"✅ CustomVLM: {n_train:,} trainable / {n_total:,} total params")
+    print(f"   (Projection + Decoder trainable — {n_train/n_total*100:.2f}%)")
+
+    train_loader, val_loader = get_custom_vlm_dataloader(cfg, char_to_idx)
+
+    # Discriminative learning rates: projection (higher) + decoder (gentler)
+    param_groups = model.get_param_groups(
+        projection_lr=cfg.lr,        # 1e-4
+        decoder_lr=cfg.lr * 0.5,     # 5e-5
+    )
+    optimizer = AdamW(param_groups, weight_decay=cfg.weight_decay)
+    total_steps = math.ceil(len(train_loader) / cfg.grad_accum) * cfg.epochs
+    warmup_steps = int(total_steps * cfg.warmup_ratio)
+    scheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps)
+
+    latest_dir, best_dir = get_output_paths(cfg, "custom_vlm")
+
+    # Metrics history
+    best_cider = -1.0
+    cider_scorer = Cider()
+
+    model.train()
+    global_step = 0
+    t0 = time.time()
+
+    for epoch in range(1, cfg.epochs + 1):
+        model.train()
+        pbar = tqdm(train_loader, desc=f"[CustomVLM] Epoch {epoch}/{cfg.epochs}")
+        running_loss = 0.0
+        epoch_loss_sum = 0.0
+        epoch_batches = 0
+        optimizer.zero_grad(set_to_none=True)
+
+        for i, batch in enumerate(pbar, start=1):
+            pixel_values = batch["pixel_values"].to(device)
+            text_input_ids = batch["text_input_ids"].to(device)
+            text_targets = batch["text_targets"].to(device)
+
+            _, loss = model(pixel_values, text_input_ids, text_targets)
+            (loss / cfg.grad_accum).backward()
+            running_loss += loss.item()
+            epoch_loss_sum += loss.item()
+            epoch_batches += 1
+
+            if i % cfg.grad_accum == 0 or i == len(train_loader):
+                torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm)
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+                global_step += 1
+
+                if global_step % cfg.log_every == 0:
+                    avg = running_loss / cfg.log_every
+                    running_loss = 0.0
+                    pbar.set_postfix({"loss": f"{avg:.4f}",
+                                      "lr": f"{scheduler.get_last_lr()[0]:.2e}"})
+
+        # End of epoch metrics
+        epoch_avg_loss = epoch_loss_sum / max(epoch_batches, 1)
+        print(f"\n📊 Epoch {epoch}/{cfg.epochs} avg loss (Train): {epoch_avg_loss:.4f}")
+
+        # --- Validation Loop ---
+        model.eval()
+        val_loss_sum = 0.0
+        val_batches = 0
+        ref_dict = {}
+        hyp_dict = {}
+
+        # Use a small subset for quick CIDEr eval during training
+        max_eval_batches = 10
+        print("   🔍 Running Validation (Loss & CIDEr)...")
+
+        with torch.no_grad():
+            for i, batch in enumerate(val_loader):
+                if i >= max_eval_batches:
+                    break
+                    
+                pixel_values = batch["pixel_values"].to(device)
+                text_input_ids = batch["text_input_ids"].to(device)
+                text_targets = batch["text_targets"].to(device)
+                
+                # 1. Validation Loss
+                _, loss = model(pixel_values, text_input_ids, text_targets)
+                val_loss_sum += loss.item()
+                val_batches += 1
+
+                # 2. Generation for CIDEr — iterate per sample (generate expects single image)
+                B = pixel_values.shape[0]
+                for b in range(B):
+                    pv_single = pixel_values[b:b+1]
+                    gen_caption = model.generate(pv_single, char_to_idx, idx_to_char, max_new_tokens=40)
+                    
+                    tgt_cpu = text_targets[b].cpu().tolist()
+                    true_str = "".join([idx_to_char.get(c, "") for c in tgt_cpu if c > 0])
+                    
+                    img_id = f"{epoch}_{i}_{b}"
+                    ref_dict[img_id] = [true_str]
+                    hyp_dict[img_id] = [gen_caption]
+
+        val_avg_loss = val_loss_sum / max(val_batches, 1)
+        print(f"   📉 Validation Loss: {val_avg_loss:.4f}")
+
+        # Calculate CIDEr
+        try:
+            cider_score, _ = cider_scorer.compute_score(ref_dict, hyp_dict)
+        except Exception:
+            cider_score = 0.0
+        
+        print(f"   🎯 Validation CIDEr: {cider_score:.4f}")
+
+        # Save latest (always)
+        _save_custom(model, char_to_idx, idx_to_char, cfg,
+                     global_step, epoch, latest_dir)
+        print(f"   💾 Saved → {latest_dir}")
+
+        # Save best (based on highest CIDEr score)
+        if cider_score >= best_cider:
+            best_cider = cider_score
+            _save_custom(model, char_to_idx, idx_to_char, cfg,
+                         global_step, epoch, best_dir)
+            print(f"   🏆 New best CIDEr (score={best_cider:.4f}) → {best_dir}")
+
+    elapsed = (time.time() - t0) / 60.0
+    print(f"\n✅ CustomVLM training complete in {elapsed:.2f} minutes")
+    print(f"   Best validation CIDEr: {best_cider:.4f}")
+
+
+def _save_custom(model, char_to_idx, idx_to_char, cfg, step, epoch, save_dir):
+    """Save CustomVLM checkpoint to the given directory (overwrites previous)."""
+    os.makedirs(save_dir, exist_ok=True)
+    torch.save({
+        "model_state": model.state_dict(),
+        "char_to_idx": char_to_idx,
+        "idx_to_char": idx_to_char,
+        "config": {
+            "block_size": cfg.block_size,
+            "text_embed_dim": cfg.text_embed_dim,
+            "n_heads": cfg.n_heads,
+            "n_layers": cfg.n_layers,
+            "vocab_size": len(char_to_idx),
+        },
+        "step": step, "epoch": epoch,
+    }, os.path.join(save_dir, "custom_vlm.pt"))
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Main
+# ─────────────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(description="Train VLM — BLIP | ViT-GPT2 | GIT | Custom")
+    parser.add_argument(
+        "--model", type=str, default="blip",
+        choices=["blip", "vit_gpt2", "git", "custom"],
+        help="Which architecture to train",
+    )
+    args = parser.parse_args()
+
+    cfg = CFG.load_for_model(args.model)
+    device = get_device()
+    print(f"✅ Device: {device}")
+    print(f"✅ Config: {args.model} | epochs={cfg.epochs} | lr={cfg.lr} | "
+          f"batch_size={cfg.batch_size} | max_target_len={cfg.max_target_len}")
+    print(f"✅ Output: {cfg.output_root}/{args.model}/")
+
+    # ── Custom VLM has its own dedicated loop ──────────────────────────────
+    if args.model == "custom":
+        train_custom_vlm(cfg, device)
+        return
+
+    # ── HuggingFace Models ─────────────────────────────────────────────────
+    latest_dir, best_dir = get_output_paths(cfg, args.model)
+
+    processor = None
+    tokenizer = None
+
+    if args.model == "blip":
+        model, processor = get_blip_model(cfg, device)
+        train_loader, val_loader = get_dataloaders(cfg, processor)
+
+        def save_latest_fn(step, epoch):
+            blip_save(model, processor, None, None, step, epoch, cfg.__dict__, latest_dir)
+
+        def save_best_fn(step, epoch):
+            blip_save(model, processor, None, None, step, epoch, cfg.__dict__, best_dir)
+
+    elif args.model == "vit_gpt2":
+        model, processor, tokenizer = get_vit_gpt2_model(cfg, device)
+        train_loader, val_loader = get_dataloaders_for_model(cfg, "vit_gpt2", processor, tokenizer)
+
+        def save_latest_fn(step, epoch):
+            vit_gpt2_save(model, processor, tokenizer, None, None, step, epoch, cfg.__dict__, latest_dir)
+
+        def save_best_fn(step, epoch):
+            vit_gpt2_save(model, processor, tokenizer, None, None, step, epoch, cfg.__dict__, best_dir)
+
+    elif args.model == "git":
+        model, processor = get_git_model(cfg, device)
+        train_loader, val_loader = get_dataloaders_for_model(cfg, "git", processor)
+
+        def save_latest_fn(step, epoch):
+            git_save(model, processor, None, None, step, epoch, cfg.__dict__, latest_dir)
+
+        def save_best_fn(step, epoch):
+            git_save(model, processor, None, None, step, epoch, cfg.__dict__, best_dir)
+
+    optimizer = AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
+    total_steps = math.ceil(len(train_loader) / cfg.grad_accum) * cfg.epochs
+    warmup_steps = int(total_steps * cfg.warmup_ratio)
+    scheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps)
+    print(f"✅ Update steps: {total_steps} | Warmup: {warmup_steps}")
+
+    run_training_loop(model, optimizer, scheduler, train_loader, val_loader, cfg,
+                      save_latest_fn=save_latest_fn,
+                      save_best_fn=save_best_fn,
+                      model_name=args.model.upper(),
+                      processor=processor, tokenizer=tokenizer)
+
+
+if __name__ == "__main__":
+    main()

transformer2.ipynb

@@ -0,0 +1,580 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 112,
+   "id": "5f1bb753",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"input.txt\", \"r\") as f:\n",
+    "    text = f.read()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 113,
+   "id": "9cf7e7ac",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Length of text: 1115394 characters\n",
+      "\n",
+      " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
+      "Vocab size: 65\n"
+     ]
+    }
+   ],
+   "source": [
+    "length = len(text)\n",
+    "print(f\"Length of text: {length} characters\")\n",
+    "char = sorted(list(set(text)))\n",
+    "vocab_size = len(char)\n",
+    "print(\"\".join(char))\n",
+    "print(f\"Vocab size: {vocab_size}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 114,
+   "id": "1b910dc7",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[46, 43, 50, 50, 53, 1, 61, 53, 56, 50, 42]\n",
+      "hello world\n"
+     ]
+    }
+   ],
+   "source": [
+    "stoi = {ch:i for i,ch in enumerate(char)}\n",
+    "itos = {i:ch for i,ch in enumerate(char)}\n",
+    "encode = lambda s: [stoi[c] for c in s]\n",
+    "decode = lambda l: \"\".join([itos[i] for i in l])\n",
+    "print(encode(\"hello world\"))\n",
+    "print(decode(encode(\"hello world\"))) # note this is one of the simplest possible tokenizers, it just maps each character to an integer. everyone has their own tokenizer like google use sentencepiece, openai use bpe, etc. we will build our own tokenizer in the next notebook."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 115,
+   "id": "3d287813",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<torch._C.Generator at 0x11113bdd0>"
+      ]
+     },
+     "execution_count": 115,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "torch.manual_seed(42)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 116,
+   "id": "4786dcce",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1115394]) torch.int64\n",
+      "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 14, 43, 44,\n",
+      "        53, 56, 43,  1, 61, 43,  1, 54, 56, 53, 41, 43, 43, 42,  1, 39, 52, 63,\n",
+      "         1, 44, 59, 56, 58, 46, 43, 56,  6,  1, 46, 43, 39, 56,  1, 51, 43,  1,\n",
+      "        57, 54, 43, 39, 49,  8,  0,  0, 13, 50, 50, 10,  0, 31, 54, 43, 39, 49,\n",
+      "         6,  1, 57, 54, 43, 39, 49,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47,\n",
+      "        58, 47, 64, 43, 52, 10,  0, 37, 53, 59])\n"
+     ]
+    }
+   ],
+   "source": [
+    "data = torch.tensor(encode(text), dtype=torch.long)\n",
+    "print(data.shape, data.dtype)\n",
+    "print(data[:100])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 117,
+   "id": "ee9c3b71",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1003854]) torch.Size([111540])\n"
+     ]
+    }
+   ],
+   "source": [
+    "n = int(0.9*len(data))\n",
+    "train_data = data[:n]\n",
+    "val_data = data[n:]\n",
+    "print(train_data.shape, val_data.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 118,
+   "id": "14d2fe85",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58])\n"
+     ]
+    }
+   ],
+   "source": [
+    "block_size = 8\n",
+    "train_data[:block_size+1] # we will use the first 8 characters to predict\n",
+    "print(train_data[:block_size+1])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 119,
+   "id": "a690a090",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "using mps device\n"
+     ]
+    }
+   ],
+   "source": [
+    "#use mps as i am using the mac with m4 \n",
+    "batch_size = 64 # how many independent sequences will we process in parallel?\n",
+    "block_size = 256 # what is the maximum context length for predictions?\n",
+    "n_embeed = 384 \n",
+    "max_iters = 20000\n",
+    "eval_iters = 2000 \n",
+    "lr_rate = 2e-4\n",
+    "dropout = 0.2\n",
+    "n_layer = 8\n",
+    "n_head = 8\n",
+    "device = \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
+    "print(f\"using {device} device\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 120,
+   "id": "d90a7d94",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "inputs:\n",
+      "torch.Size([64, 256])\n",
+      "tensor([[ 0, 26, 53,  ..., 56, 43, 47],\n",
+      "        [60, 43, 56,  ..., 56,  1, 41],\n",
+      "        [26, 21, 33,  ..., 26, 21, 13],\n",
+      "        ...,\n",
+      "        [ 5, 57,  1,  ...,  1, 35, 47],\n",
+      "        [56, 53, 53,  ..., 59, 50, 42],\n",
+      "        [42, 47, 56,  ..., 39, 56,  1]], device='mps:0')\n",
+      "\n",
+      "targets:\n",
+      "torch.Size([64, 256])\n",
+      "tensor([[26, 53, 58,  ..., 43, 47, 45],\n",
+      "        [43, 56,  1,  ...,  1, 41, 53],\n",
+      "        [21, 33, 31,  ..., 21, 13, 10],\n",
+      "        ...,\n",
+      "        [57,  1, 52,  ..., 35, 47, 50],\n",
+      "        [53, 53, 58,  ..., 50, 42,  1],\n",
+      "        [47, 56, 43,  ..., 56,  1, 51]], device='mps:0')\n"
+     ]
+    }
+   ],
+   "source": [
+    "torch.manual_seed(1337)\n",
+    "def get_batch(split):\n",
+    "    data = train_data if split == 'train' else val_data\n",
+    "    ix = torch.randint(len(data) - block_size, (batch_size,))\n",
+    "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
+    "    y = torch.stack([data[i+1:i+block_size+1] for i in ix])\n",
+    "    x, y = x.to(device), y.to(device)\n",
+    "    return x, y\n",
+    "xb, yb = get_batch('train')\n",
+    "print(\"inputs:\")\n",
+    "print(xb.shape)\n",
+    "print(xb)\n",
+    "print(\"\\ntargets:\")\n",
+    "print(yb.shape)\n",
+    "print(yb)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 121,
+   "id": "27573f3f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Head(torch.nn.Module):\n",
+    "    def __init__(self, head_size):\n",
+    "        super().__init__()\n",
+    "        self.head_size = head_size\n",
+    "        self.key = nn.Linear(n_embeed, head_size, bias=False)\n",
+    "        self.query = nn.Linear(n_embeed, head_size, bias=False)\n",
+    "        self.value = nn.Linear(n_embeed, head_size, bias=False)\n",
+    "        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))\n",
+    "        self.dropout = nn.Dropout(dropout)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        B,T,C = x.shape\n",
+    "        k = self.key(x) # (B,T,16)\n",
+    "        q = self.query(x) # (B,T,16)\n",
+    "        v = self.value(x) # (B,T,16)\n",
+    "        weights = q @ k.transpose(-2, -1) * (self.head_size ** -0.5) # (B,T,16) @ (B,16,T) -> (B,T,T)\n",
+    "        tril = torch.tril(torch.ones(T, T , device=x.device))\n",
+    "        weights = weights.masked_fill(tril == 0, float('-inf')) # when we talk about the encoder transformer we remove this bcz we want to attend to all the tokens in the input sequence, but in the decoder transformer we want to attend only to the previous tokens in the output sequence, so we use this mask to prevent the model from attending to future tokens.\n",
+    "        weights = torch.softmax(weights, dim=-1)\n",
+    "        weights = self.dropout(weights)\n",
+    "        out = weights @ v # (B,T,T) @ (B,T,C) -> (B,T,C)\n",
+    "        return out\n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 122,
+   "id": "a776b854",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class MultiHeadAttention(torch.nn.Module):\n",
+    "    def __init__(self, num_heads, head_size):\n",
+    "        super().__init__()\n",
+    "        self.heads = torch.nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
+    "        self.proj = torch.nn.Linear(n_embeed, n_embeed)\n",
+    "        self.dropout = nn.Dropout(dropout)\n",
+    "    def forward(self, x):\n",
+    "        out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
+    "        out = self.proj(out)\n",
+    "        return out"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 123,
+   "id": "da0d9201",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class feedforward(torch.nn.Module):\n",
+    "    def __init__(self, n_embeed):\n",
+    "        super().__init__()\n",
+    "        self.net = torch.nn.Sequential(\n",
+    "            torch.nn.Linear(n_embeed, 4*n_embeed), #according to paper there is the multiplier of 4 in the hidden layers \n",
+    "            torch.nn.ReLU(),\n",
+    "            torch.nn.Linear(4*n_embeed, n_embeed),\n",
+    "            nn.Dropout(dropout)\n",
+    "        )\n",
+    "    def forward(self, x):\n",
+    "        return self.net(x)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 124,
+   "id": "1b2fc012",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Block(torch.nn.Module):\n",
+    "    def __init__(self, n_embeed , n_head):\n",
+    "        super().__init__()\n",
+    "        head_size = n_embeed // n_head \n",
+    "        self.sa_head = MultiHeadAttention(num_heads=n_head, head_size=head_size)\n",
+    "        self.ffwd = feedforward(n_embeed)\n",
+    "        self.ln1 = nn.LayerNorm(n_embeed)\n",
+    "        self.ln2 = nn.LayerNorm(n_embeed)\n",
+    "        self.dropout = nn.Dropout(dropout)\n",
+    "    def forward(self, x):\n",
+    "        x = x+self.sa_head(self.ln1(x))#this is slightly deviation from the original paper as we are passing the layer norem before the multi head attention and feedforward, but it is a common practice to do so, and it works better than the original paper.\n",
+    "        x = x+self.ffwd(self.ln2(x))\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 125,
+   "id": "7a3053a0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class BigramLanguageModel(torch.nn.Module):\n",
+    "    def __init__(self, vocab_size, n_embeed):\n",
+    "        super().__init__()\n",
+    "        self.token_embedding_table = torch.nn.Embedding(vocab_size, n_embeed) \n",
+    "        self.position_embedding_table = torch.nn.Embedding(block_size, n_embeed)\n",
+    "        #not giving the good enhaced result when we try multiple Block() , deep neural net suffer from the optimisation issue \n",
+    "        # self.blocks = nn.Sequential(\n",
+    "        #     Block(n_embeed, n_head=4),\n",
+    "        #     Block(n_embeed, n_head=4),\n",
+    "        #     Block(n_embeed, n_head=4),\n",
+    "        #     nn.LayerNorm(n_embeed),\n",
+    "        # )\n",
+    "        self.blocks = nn.Sequential(*[Block(n_embeed, n_head) for _ in range(n_layer)])\n",
+    "        self.ln_f = nn.LayerNorm(n_embeed)\n",
+    "        self.lm_head = torch.nn.Linear(n_embeed, vocab_size)\n",
+    "    def forward(self, idx, targets=None):\n",
+    "        B,T = idx.shape\n",
+    "        # idx and targets are both (B,T) tensor of integers\n",
+    "        token_emb = self.token_embedding_table(idx) # (B,T,C)\n",
+    "        pos_emb = self.position_embedding_table(torch.arange(idx.shape[1], device=idx.device)) # (T,C)\n",
+    "        x = token_emb + pos_emb # (B,T,C)\n",
+    "        x = self.blocks(x) # (B,T,C)\n",
+    "        x = self.ln_f(x) # (B,T,C)\n",
+    "        logits = self.lm_head(x) # (B,T,vocab_size)\n",
+    "        if targets is None:\n",
+    "            loss = None\n",
+    "        else:\n",
+    "            B,T,C = logits.shape\n",
+    "            logits = logits.view(B*T, C)\n",
+    "            targets = targets.view(B*T)\n",
+    "            loss = F.cross_entropy(logits, targets)\n",
+    "        return logits, loss\n",
+    "    \n",
+    "    def generate(self, idx, max_new_tokens):\n",
+    "        # idx is (B,T) array of indices in the current context\n",
+    "        for _ in range(max_new_tokens):\n",
+    "            idx_cond = idx[:, -block_size:] # crop idx to the last block_size tokens\n",
+    "            logits, loss = self(idx_cond)\n",
+    "            logits = logits[:, -1, :] # becomes (B,C) , as we only want to provide the last character as the input to predict the next character\n",
+    "            probs = F.softmax(logits, dim=-1) # (B,C)\n",
+    "            idx_next = torch.multinomial(probs, num_samples=1) # (B,1)\n",
+    "            idx = torch.cat((idx, idx_next), dim=1) # (B,T+1)\n",
+    "        return idx"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "67e96f0b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 126,
+   "id": "0e9d66e8",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "logits shape: torch.Size([16384, 65])\n",
+      "loss: 4.277037620544434\n",
+      "\n",
+      "tRNt'OUWzNdaNv;DZ!HWJxsg-rG$l.\n",
+      "VXx;h&CEqoyJOlF.DmdMw;u;cjEIgcOQOID;$wig.tRIgazPSVyRpKBE-3UQBdJ'AIIxX\n"
+     ]
+    }
+   ],
+   "source": [
+    "model = BigramLanguageModel(vocab_size, n_embeed)\n",
+    "model = model.to(device)\n",
+    "logits, loss = model(xb, yb)\n",
+    "print(\"logits shape:\", logits.shape)\n",
+    "print(\"loss:\", loss.item())\n",
+    "print(decode(model.generate(idx=torch.zeros((1,1), dtype=torch.long, device=device), max_new_tokens=100)[0].tolist()))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 127,
+   "id": "1da9dd4f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@torch.no_grad()\n",
+    "def estimate_loss():\n",
+    "    out = {}\n",
+    "    model.eval()\n",
+    "    for split in ['train', 'val']:\n",
+    "        losses = torch.zeros(eval_iters)\n",
+    "        for k in range(eval_iters):\n",
+    "            X, Y = get_batch(split)\n",
+    "            logits, loss = model(X, Y)\n",
+    "            losses[k] = loss.item()\n",
+    "        out[split] = losses.mean()\n",
+    "    model.train()\n",
+    "    return out"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1e3fb308",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "step 0: train loss 4.2785, val loss 4.2821\n"
+     ]
+    }
+   ],
+   "source": [
+    "optimizer = torch.optim.AdamW(model.parameters(), lr=lr_rate)\n",
+    "for steps in range(max_iters):\n",
+    "    if steps % eval_iters == 0:\n",
+    "        losses = estimate_loss()\n",
+    "        print(f\"step {steps}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
+    "\n",
+    "    xb, yb = get_batch('train')\n",
+    "    logits, loss = model(xb, yb)\n",
+    "    optimizer.zero_grad(set_to_none=True)\n",
+    "    loss.backward()\n",
+    "    optimizer.step()\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9490a27b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "DUKE VINCENTIO:\n",
+      "Stand brother, sir, here it, uncle he got.\n",
+      "\n",
+      "VIRGILIA:\n",
+      "A dog of the yousician, let your good brother, sister,\n",
+      "nor it to die.\n",
+      "\n",
+      "VOLUMNIA:\n",
+      "She is in the mar, and the matter:\n",
+      "there is! What say you, Juliet alone and bird.\n",
+      "Is thy life?\n",
+      "\n",
+      "JULIET:\n",
+      "Being a child! prompt fear: speak, and look fellow good?\n",
+      "\n",
+      "FLORIZEL:\n",
+      "And rumour, by my man's tooth made.\n",
+      "\n",
+      "JULIET:\n",
+      "Ay, if you doth make leave your retires,\n",
+      "A mother tempt my todder dial should have\n",
+      "So dear and let me 'gainst words out again,\n",
+      "Savest with honour's princes to hear throught of,\n",
+      "His perdom of preserve\n",
+      "Is posterity and secut\n",
+      "No god costerb: shall be more the Capitol,\n",
+      "But court did this hoursest, do begg them buried.\n",
+      "His apple and dreams on daughter, and we will,\n",
+      "He were laddy's wounds. O mother!\n",
+      "Dread!\n",
+      "In it the whitest through thee grief: why, general,\n",
+      "My heart play'd many fellows upon him.\n",
+      "\n",
+      "FRIAR LAURENCE:\n",
+      "For traitor the mind: what the journey, rise!\n",
+      "I serve, or I know the senate, and let my indeed\n",
+      "Will on brave it so lone\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(decode(model.generate(idx=torch.zeros((1,1), dtype=torch.long, device=device), max_new_tokens=1000)[0].tolist()))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d717cdc1",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "10.788929 M parameters\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(sum(p.numel() for p in model.parameters())/1e6, \"M parameters\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "58991844",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "torch.save(model.state_dict(), \"shakespeare_transformer.pt\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "935f7d0e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.14.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

transformer_base.ipynb

@@ -0,0 +1,446 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "193c3159",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"input.txt\", \"r\") as f:\n",
+    "    text = f.read()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "e557cb70",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Length of text: 1115394 characters\n"
+     ]
+    }
+   ],
+   "source": [
+    "length = len(text)\n",
+    "print(f\"Length of text: {length} characters\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "750587a9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(text[:500]) "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "16490999",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
+      "Vocab size: 65\n"
+     ]
+    }
+   ],
+   "source": [
+    "char = sorted(list(set(text)))\n",
+    "vocab_size = len(char)\n",
+    "print(\"\".join(char))\n",
+    "print(f\"Vocab size: {vocab_size}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 39,
+   "id": "d9e6e17a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "using mps device\n"
+     ]
+    }
+   ],
+   "source": [
+    "#use mps as i am using the mac with m4 \n",
+    "device = \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
+    "print(f\"using {device} device\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "082fd1ba",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[46, 43, 50, 50, 53, 1, 61, 53, 56, 50, 42]\n",
+      "hello world\n"
+     ]
+    }
+   ],
+   "source": [
+    "stoi = {ch:i for i,ch in enumerate(char)}\n",
+    "itos = {i:ch for i,ch in enumerate(char)}\n",
+    "encode = lambda s: [stoi[c] for c in s]\n",
+    "decode = lambda l: \"\".join([itos[i] for i in l])\n",
+    "print(encode(\"hello world\"))\n",
+    "print(decode(encode(\"hello world\"))) # note this is one of the simplest possible tokenizers, it just maps each character to an integer. everyone has their own tokenizer like google use sentencepiece, openai use bpe, etc. we will build our own tokenizer in the next notebook."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "7cce9365",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1115394]) torch.int64\n",
+      "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 14, 43, 44,\n",
+      "        53, 56, 43,  1, 61, 43,  1, 54, 56, 53, 41, 43, 43, 42,  1, 39, 52, 63,\n",
+      "         1, 44, 59, 56, 58, 46, 43, 56,  6,  1, 46, 43, 39, 56,  1, 51, 43,  1,\n",
+      "        57, 54, 43, 39, 49,  8,  0,  0, 13, 50, 50, 10,  0, 31, 54, 43, 39, 49,\n",
+      "         6,  1, 57, 54, 43, 39, 49,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47,\n",
+      "        58, 47, 64, 43, 52, 10,  0, 37, 53, 59])\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "data = torch.tensor(encode(text), dtype=torch.long)\n",
+    "print(data.shape, data.dtype)\n",
+    "print(data[:100])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "d59606cc",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1003854]) torch.Size([111540])\n"
+     ]
+    }
+   ],
+   "source": [
+    "n = int(0.9*len(data))\n",
+    "train_data = data[:n]\n",
+    "val_data = data[n:]\n",
+    "print(train_data.shape, val_data.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "e2bd00e4",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58])\n"
+     ]
+    }
+   ],
+   "source": [
+    "block_size = 8\n",
+    "train_data[:block_size+1] # we will use the first 8 characters to predict\n",
+    "print(train_data[:block_size+1])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "4ce6af03",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "when input is tensor([18]) the target: 47\n",
+      "when input is tensor([18, 47]) the target: 56\n",
+      "when input is tensor([18, 47, 56]) the target: 57\n",
+      "when input is tensor([18, 47, 56, 57]) the target: 58\n",
+      "when input is tensor([18, 47, 56, 57, 58]) the target: 1\n",
+      "when input is tensor([18, 47, 56, 57, 58,  1]) the target: 15\n",
+      "when input is tensor([18, 47, 56, 57, 58,  1, 15]) the target: 47\n",
+      "when input is tensor([18, 47, 56, 57, 58,  1, 15, 47]) the target: 58\n"
+     ]
+    }
+   ],
+   "source": [
+    "x_train = train_data[:block_size]\n",
+    "y_train = train_data[1:block_size+1]\n",
+    "for t in range(block_size):\n",
+    "    context = x_train[:t+1]\n",
+    "    target = y_train[t]\n",
+    "    print(f\"when input is {context} the target: {target}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 38,
+   "id": "85e56335",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "inputs:\n",
+      "torch.Size([4, 8])\n",
+      "tensor([[24, 43, 58,  5, 57,  1, 46, 43],\n",
+      "        [44, 53, 56,  1, 58, 46, 39, 58],\n",
+      "        [52, 58,  1, 58, 46, 39, 58,  1],\n",
+      "        [25, 17, 27, 10,  0, 21,  1, 54]], device='mps:0')\n",
+      "\n",
+      "targets:\n",
+      "torch.Size([4, 8])\n",
+      "tensor([[43, 58,  5, 57,  1, 46, 43, 39],\n",
+      "        [53, 56,  1, 58, 46, 39, 58,  1],\n",
+      "        [58,  1, 58, 46, 39, 58,  1, 46],\n",
+      "        [17, 27, 10,  0, 21,  1, 54, 39]], device='mps:0')\n"
+     ]
+    }
+   ],
+   "source": [
+    "torch.manual_seed(1337)\n",
+    "batch_size = 4 # how many independent sequences will we process in parallel?\n",
+    "block_size = 8 # what is the maximum context length for predictions?\n",
+    "def get_batch(split):\n",
+    "    data = train_data if split == 'train' else val_data\n",
+    "    ix = torch.randint(len(data) - block_size, (batch_size,))\n",
+    "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
+    "    y = torch.stack([data[i+1:i+block_size+1] for i in ix])\n",
+    "    x, y = x.to(device), y.to(device)\n",
+    "    return x, y\n",
+    "xb, yb = get_batch('train')\n",
+    "print(\"inputs:\")\n",
+    "print(xb.shape)\n",
+    "print(xb)\n",
+    "print(\"\\ntargets:\")\n",
+    "print(yb.shape)\n",
+    "print(yb)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "18810b27",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for b in range(batch_size):\n",
+    "    for t in range(block_size):\n",
+    "        context = xb[b, :t+1]\n",
+    "        target = yb[b, t]\n",
+    "        print(f\"when input is {context.tolist()} the target: {target.item()}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "77449b2f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[24, 43, 58,  5, 57,  1, 46, 43],\n",
+      "        [44, 53, 56,  1, 58, 46, 39, 58],\n",
+      "        [52, 58,  1, 58, 46, 39, 58,  1],\n",
+      "        [25, 17, 27, 10,  0, 21,  1, 54]])\n",
+      "\n",
+      "\n",
+      "tensor([[43, 58,  5, 57,  1, 46, 43, 39],\n",
+      "        [53, 56,  1, 58, 46, 39, 58,  1],\n",
+      "        [58,  1, 58, 46, 39, 58,  1, 46],\n",
+      "        [17, 27, 10,  0, 21,  1, 54, 39]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(xb)\n",
+    "print(\"\\n\")\n",
+    "print(yb)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "66a1c195",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "logits shape: torch.Size([32, 65])\n",
+      "loss: 4.878634929656982\n",
+      "\n",
+      "SKIcLT;AcE\n"
+     ]
+    }
+   ],
+   "source": [
+    "#Bigram language model\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "torch.manual_seed(1337)\n",
+    "class BigramLanguageModel(torch.nn.Module):\n",
+    "    def __init__(self, vocab_size):\n",
+    "        super().__init__()\n",
+    "        self.token_embedding_table = torch.nn.Embedding(vocab_size, vocab_size)\n",
+    "    def forward(self, idx, targets=None):\n",
+    "        # idx and targets are both (B,T) tensor of integers\n",
+    "        logits = self.token_embedding_table(idx) # (B,T,C)\n",
+    "        if targets is None:\n",
+    "            loss = None\n",
+    "        else:\n",
+    "            B,T,C = logits.shape\n",
+    "            logits = logits.view(B*T, C)\n",
+    "            targets = targets.view(B*T)\n",
+    "            loss = F.cross_entropy(logits, targets)\n",
+    "        return logits, loss\n",
+    "    \n",
+    "    def generate(self, idx, max_new_tokens):\n",
+    "        # idx is (B,T) array of indices in the current context\n",
+    "        for _ in range(max_new_tokens):\n",
+    "            logits, loss = self(idx)\n",
+    "            logits = logits[:, -1, :] # becomes (B,C) , as we only want to provide the last character as the input to predict the next character\n",
+    "            probs = F.softmax(logits, dim=-1) # (B,C)\n",
+    "            idx_next = torch.multinomial(probs, num_samples=1) # (B,1)\n",
+    "            idx = torch.cat((idx, idx_next), dim=1) # (B,T+1)\n",
+    "        return idx\n",
+    "    \n",
+    "model = BigramLanguageModel(vocab_size)\n",
+    "model₹.to(device)\n",
+    "logits, loss = model(xb, yb)\n",
+    "print(\"logits shape:\", logits.shape)\n",
+    "print(\"loss:\", loss.item())\n",
+    "print(decode(model.generate(idx=torch.zeros((1,1), dtype=torch.long), max_new_tokens=10)[0].tolist()))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "id": "ecd49fc4",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "step 9999: loss 2.4313366413116455\n"
+     ]
+    }
+   ],
+   "source": [
+    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)\n",
+    "batch_size = 32\n",
+    "for steps in range(10000):\n",
+    "    xb, yb = get_batch('train')\n",
+    "    logits, loss = model(xb, yb)\n",
+    "    optimizer.zero_grad(set_to_none=True)\n",
+    "    loss.backward()\n",
+    "    optimizer.step()\n",
+    "print(f\"step {steps}: loss {loss.item()}\")\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "id": "8ce29e5a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Warstyo a \n"
+     ]
+    }
+   ],
+   "source": [
+    "print(decode(model.generate(idx=torch.zeros((1,1), dtype=torch.long), max_new_tokens=10)[0].tolist()))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b87bd156",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "using mps device\n"
+     ]
+    }
+   ],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c9f2052b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.14.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}