added training et new improved model structure

architecture.txt

@@ -1,7 +1,7 @@
 Transformer:
     embedding:
         input_embedding (same for encoder and decoder)
-        positional_embedding
+        positional_embedding (learned or sinusoidal)
         dropout
     encoder
         encoder_layer *N

configs/dataset/large_dataset.yaml

@@ -0,0 +1,2 @@
+name: large_dataset
+path: ./data/processed/processed_large_dataset

configs/dataset/medium_dataset.yaml

@@ -0,0 +1,2 @@
+name: medium_dataset
+path: ./data/processed/processed_medium_dataset

configs/dataset/small_dataset.yaml

@@ -0,0 +1,2 @@
+name: small_dataset
+path: ./data/processed/processed_small_dataset

configs/generation/generation.yaml

@@ -0,0 +1,9 @@
+max_new_tokens: 64
+temperature: 1
+top_k: 10
+top_p: 0.6
+do_sample: true
+presence_penalty: 0.6
+frequency_penalty: 0.2
+no_repeat_ngram: 4
+min_steps_before_eos: 2

configs/infer_mode.yaml

@@ -1,10 +1,11 @@
 defaults:
   - model: test_model
-  - runtime: runtime
+  - runtime: local
+  - tokenizer: bpe_8k
+  - generation: generation
   - _self_
-input_text: ""
-max_new_tokens: 64
-temperature: 1.0
-top_k: 4
-top_p: 0.7
-do_sample: True
+input_text: "how are you today"
+hydra:
+  run:
+    dir: .
+  output_subdir: null

configs/model/large_model.yaml

@@ -0,0 +1,14 @@
+name: large_model
+best_checkpoint_path: ./outputs/checkpoints/large_model/best.ckpt
+latest_checkpoint_path: ./outputs/checkpoints/large_model/latest.ckpt
+tokenizer: bpe_32k
+d_model: 512
+num_heads: 8
+num_layers: 6
+d_ff: 2048
+dropout_rate: 0.1
+vocab_size: 32768
+max_seq_len: 256
+pad_id: 0
+bos_id: 1
+eos_id: 2

configs/model/medium_model.yaml

@@ -0,0 +1,14 @@
+name: medium_model
+best_checkpoint_path: ./outputs/checkpoints/medium_model/best.ckpt
+latest_checkpoint_path: ./outputs/checkpoints/medium_model/latest.ckpt
+tokenizer: bpe_16k
+d_model: 512
+num_heads: 8
+num_layers: 6
+d_ff: 2048
+dropout_rate: 0.1
+vocab_size: 16384
+max_seq_len: 256
+pad_id: 0
+bos_id: 1
+eos_id: 2

configs/model/small_model.yaml

@@ -0,0 +1,14 @@
+name: small_model
+best_checkpoint_path: ./outputs/checkpoints/small_model/best.ckpt
+latest_checkpoint_path: ./outputs/checkpoints/small_model/latest.ckpt
+tokenizer: bpe_8k
+d_model: 256
+num_heads: 8
+num_layers: 4
+d_ff: 1024
+dropout_rate: 0.1
+vocab_size: 8192
+max_seq_len: 128
+pad_id: 0
+bos_id: 1
+eos_id: 2

configs/model/test_model.yaml

@@ -1,9 +1,13 @@
-d_model: 256
-num_heads: 4
-num_layers: 4
+name: test_model
+best_checkpoint_path: ./outputs/checkpoints/test_model/best.pt
+latest_checkpoint_path: ./outputs/checkpoints/test_model/latest.pt
+tokenizer: bpe_8k
+d_model: 512
+num_heads: 8
+num_layers: 8
 d_ff: 1024
 dropout_rate: 0.1
-vocab_size: 32100
+vocab_size: 8192
 max_seq_len: 128
 pad_id: 0
 bos_id: 1

configs/optim/Adamw.yaml

@@ -0,0 +1,6 @@
+name: AdamW
+lr: 8e-4
+betas: [0.9, 0.98]
+eps: 1e-8
+weight_decay: 0.08
+no_decay_on_bias_norm_embed: true

configs/runtime/local.yaml

@@ -0,0 +1,7 @@
+seed: 42
+device: cuda
+output_dir: ./outputs/
+data_dir: ./data/
+tokenizer_dir: ./tokenizer/
+cache_dir: ./data/raw
+checkpoint_path: ./outputs/checkpoints/

configs/runtime/runtime.yaml

@@ -1,3 +0,0 @@
-tokenizer_name: t5-small
-weights_path: artifacts/checkpoints/initial_model_weights.pth
-artifact_dir: outputs/${now:%Y-%m-%d}/${now:%H-%M-%S}

configs/sched/cosine.yaml

@@ -0,0 +1,5 @@
+name: cosine_warmup
+warmup_ratio: 0.08
+min_lr_ratio: 0.1
+max_lr_ratio: 1
+hold_min: true

configs/tokenizer/bpe_16k.yaml

@@ -0,0 +1,15 @@
+name: bpe_16k
+path: ./tokenizer/bpe_16k.json
+corpus: ./data/corpus/corpus.txt
+dataset: concatenated_dataset
+vocab_size: 16384
+max_seq_len: 128
+pad_token: <pad>
+bos_token: <bos>
+eos_token: <eos>
+unk_token: <unk>
+special_tokens: ["<pad>", "<bos>", "<eos>", "<unk>"]
+pad_id: 0
+bos_id: 1
+eos_id: 2
+unk_id: 3

configs/tokenizer/bpe_32k.yaml

@@ -0,0 +1,15 @@
+name: bpe_32k
+path: ./tokenizer/bpe_32k.json
+corpus: ./data/corpus/corpus.txt
+dataset: concatenated_dataset
+vocab_size: 32768
+max_seq_len: 256
+pad_token: <pad>
+bos_token: <bos>
+eos_token: <eos>
+unk_token: <unk>
+special_tokens: ["<pad>", "<bos>", "<eos>", "<unk>"]
+pad_id: 0
+bos_id: 1
+eos_id: 2
+unk_id: 3

configs/tokenizer/bpe_4k.yaml

@@ -0,0 +1,15 @@
+name: bpe_4k
+path: ./tokenizer/bpe_4k.json
+corpus: ./data/corpus/corpus.txt
+dataset: concatenated_dataset
+vocab_size: 4096
+max_seq_len: 128
+pad_token: <pad>
+bos_token: <bos>
+eos_token: <eos>
+unk_token: <unk>
+special_tokens: ["<pad>", "<bos>", "<eos>", "<unk>"]
+pad_id: 0
+bos_id: 1
+eos_id: 2
+unk_id: 3

configs/tokenizer/bpe_8k.yaml

@@ -0,0 +1,15 @@
+name: bpe_8k
+path: ./tokenizer/bpe_8k.json
+corpus: ./data/corpus/corpus.txt
+dataset: concatenated_dataset
+vocab_size: 8192
+max_seq_len: 128
+pad_token: <pad>
+bos_token: <bos>
+eos_token: <eos>
+unk_token: <unk>
+special_tokens: ["<pad>", "<bos>", "<eos>", "<unk>"]
+pad_id: 0
+bos_id: 1
+eos_id: 2
+unk_id: 3

configs/train_mode.yaml

@@ -1,9 +1,13 @@
-seed: 42
-batch_size: 16
-epochs: 3
-lr: 3e-4
-weight_decay: 0.01
-grad_clip: 1.0
-precision: "fp32"   # one of: fp32, fp16, bf16
-checkpoint_dir: "artifacts/checkpoints"
-log_dir: "artifacts/mlruns"
+defaults:
+  - model: test_model
+  - tokenizer: bpe_16k
+  - dataset: medium_dataset
+  - optim: Adamw
+  - sched: cosine
+  - runtime: local
+  - trainer: default
+  - _self_
+hydra:
+  run:
+    dir: .        # run in current directory
+  output_subdir: null

configs/trainer/default.yaml

@@ -0,0 +1,14 @@
+epochs: 80
+batch_size: 128
+shuffle: true
+precision: "bf16"
+gradient_accumulation: 2
+clip_grad_norm: 0.5
+eval_interval: 1
+save_interval: 1
+debug_interval: 1000
+mixed_precision: true
+gradient_checkpointing: true
+label_smoothing: 0.05
+export_best_to_model_path: false
+resume: start_over

environment.yml

@@ -10,3 +10,8 @@ dependencies:
       - transformers
       - omegaconf
       - hydra-core
+      - jupyterlab
+      - tokenizers
+      - datasets
+      - fastapi
+      - uvicorn[standard]

pyproject.toml

@@ -5,10 +5,13 @@ description = "A transformer implementation"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
-  "torch==2.8",
+  "torch>=2.6",
   "transformers",
   "omegaconf",
   "hydra-core",
+  "jupyterlab",
+  "tokenizers",
+  "datasets",
 ]
 
 [project.optional-dependencies]
@@ -20,6 +23,10 @@ dev = [
   "mypy",
   "pre-commit",
 ]
+server = [
+  "fastapi",
+  "uvicorn[standard]",
+]
 
 [tool.setuptools]
 package-dir = {"" = "src"}

scripts/inference.py

@@ -1,30 +1,42 @@
+"""CLI entry point for running transformer inference with deterministic sampling."""
+
+from __future__ import annotations
+
 import hydra
 import torch
 from hydra.utils import to_absolute_path
 from omegaconf import DictConfig, OmegaConf
-from transformers import AutoTokenizer
+from transformers import PreTrainedTokenizerFast
 
-from transformer import (
-    BasicEncDecCfg,
-    BasicEncoderDecoderTransformer,
-    InferAppCfg,
-)
+from transformer import BasicEncoderDecoderTransformer
+from transformer.configs import InferAppCfg, ModelCfg, TokenizerCfg
+from transformer.utils import check_tokenizer_model_compatibility, set_global_seed
 
 
-@hydra.main(config_path="../configs", config_name="infer_mode", version_base=None)
-def main(cfg: DictConfig):
+def inference(cfg: DictConfig):
     scfg_temp = OmegaConf.merge(OmegaConf.structured(InferAppCfg), cfg)
     scfg: InferAppCfg = OmegaConf.to_object(scfg_temp)
-    model_cfg = BasicEncDecCfg(**vars(scfg.model))
+    model_cfg = ModelCfg(**vars(scfg.model))
+    tokenizer_cfg = TokenizerCfg(**vars(scfg.tokenizer))
+
+    # Ensure reproducible sampling when requested
+    set_global_seed(scfg.runtime.seed)
+
+    # Compatibility checks
+    check_tokenizer_model_compatibility(model_cfg, tokenizer_cfg)
 
     transformer = BasicEncoderDecoderTransformer(model_cfg)
-    tokenizer = AutoTokenizer.from_pretrained(
-        scfg.runtime.tokenizer_name,
-        use_fast=True,
+    tokenizer = PreTrainedTokenizerFast(
+        tokenizer_file=to_absolute_path(tokenizer_cfg.path),
+        bos_token=cfg.tokenizer.bos_token,
+        eos_token=cfg.tokenizer.eos_token,
+        unk_token=cfg.tokenizer.unk_token,
+        pad_token=cfg.tokenizer.pad_token,
+        model_max_length=tokenizer_cfg.max_seq_len,
     )
 
-    if scfg.runtime.weights_path:
-        ckpt = torch.load(to_absolute_path(scfg.runtime.weights_path), map_location="cpu")
+    if scfg.model.best_checkpoint_path:
+        ckpt = torch.load(to_absolute_path(scfg.model.best_checkpoint_path), map_location="cpu")
         state = ckpt.get("model_state_dict", ckpt)
         transformer.load_state_dict(state, strict=False)
     transformer.eval()
@@ -35,18 +47,29 @@ def main(cfg: DictConfig):
 
     encoded = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
     src_ids = encoded["input_ids"]
-    src_padd_mask = (encoded["attention_mask"] == 0).unsqueeze(1).unsqueeze(2)
+    src_padd_mask = encoded["attention_mask"] == 0
+
     tgt_ids = transformer.generate(
         src_ids,
         src_padd_mask,
-        max_new_tokens=scfg.max_new_tokens,
-        temperature=scfg.temperature,
-        top_k=scfg.top_k,
-        top_p=scfg.top_p,
-        do_sample=scfg.do_sample,
+        max_new_tokens=scfg.generation.max_new_tokens,
+        temperature=scfg.generation.temperature,
+        top_k=scfg.generation.top_k,
+        top_p=scfg.generation.top_p,
+        do_sample=scfg.generation.do_sample,
+        presence_penalty=scfg.generation.presence_penalty,
+        frequency_penalty=scfg.generation.frequency_penalty,
+        no_repeat_ngram=scfg.generation.no_repeat_ngram,
+        min_steps_before_eos=scfg.generation.min_steps_before_eos,
+        seed=scfg.runtime.seed,
     )
-    output = tokenizer.batch_decode(tgt_ids, skip_special_tokens=True)
-    print(output)
+    output = tokenizer.batch_decode(tgt_ids, skip_special_tokens=False)
+    return output
+
+
+@hydra.main(config_path="../configs", config_name="infer_mode", version_base=None)
+def main(cfg: DictConfig):
+    print(inference(cfg))
 
 
 if __name__ == "__main__":

scripts/load_data.ipynb

@@ -0,0 +1,388 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "e2289881",
+   "metadata": {},
+   "source": [
+    "Downloading datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9ef193d1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "_ = load_dataset(\"ahazeemi/iwslt14-en-fr\", cache_dir=\"./data/raw\")\n",
+    "_ = load_dataset(\"Helsinki-NLP/europarl\", \"en-fr\", cache_dir=\"./data/raw\")\n",
+    "# _ = load_dataset(\"wmt/wmt14\", \"fr-en\", cache_dir=\"./data/raw\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cb48dbf1",
+   "metadata": {},
+   "source": [
+    "Loading datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "2e0838f5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/ala-boussoffara/miniconda3/envs/transformer/lib/python3.12/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",
+      "Generating train split: 179435 examples [00:00, 1366541.45 examples/s]\n",
+      "Generating validation split: 903 examples [00:00, 558869.19 examples/s]\n",
+      "Generating test split: 3666 examples [00:00, 876193.43 examples/s]\n",
+      "Generating train split: 2051014 examples [00:01, 2040864.91 examples/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "small_dataset = load_dataset(\"./data/raw/ahazeemi___iwslt14-en-fr\")\n",
+    "medium_dataset = load_dataset(\"./data/raw/Helsinki-NLP___europarl\")\n",
+    "# large_dataset = load_dataset(\"./data/raw/wmt___wmt14\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cb452f91",
+   "metadata": {},
+   "source": [
+    "Preprocessing\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "ca4d1867",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Map: 100%|██████████| 179435/179435 [00:12<00:00, 14545.45 examples/s]\n",
+      "Map: 100%|██████████| 903/903 [00:00<00:00, 13155.18 examples/s]\n",
+      "Map: 100%|██████████| 3666/3666 [00:00<00:00, 14699.07 examples/s]\n",
+      "Filter: 100%|██████████| 179435/179435 [00:01<00:00, 163853.20 examples/s]\n",
+      "Filter: 100%|██████████| 885/885 [00:00<00:00, 123014.38 examples/s]\n",
+      "Filter: 100%|██████████| 3620/3620 [00:00<00:00, 142652.68 examples/s]\n",
+      "Map: 100%|██████████| 2051014/2051014 [02:27<00:00, 13892.40 examples/s]\n",
+      "Filter: 100%|██████████| 2028907/2028907 [00:12<00:00, 163723.53 examples/s]\n",
+      "Saving the dataset (1/1 shards): 100%|██████████| 179435/179435 [00:00<00:00, 186519.65 examples/s]\n",
+      "Saving the dataset (1/1 shards): 100%|██████████| 885/885 [00:00<00:00, 142866.56 examples/s]\n",
+      "Saving the dataset (1/1 shards): 100%|██████████| 3620/3620 [00:00<00:00, 170674.57 examples/s]\n",
+      "Saving the dataset (3/3 shards): 100%|██████████| 1826016/1826016 [00:11<00:00, 159762.08 examples/s]\n",
+      "Saving the dataset (1/1 shards): 100%|██████████| 101445/101445 [00:00<00:00, 193029.72 examples/s]\n",
+      "Saving the dataset (1/1 shards): 100%|██████████| 101446/101446 [00:00<00:00, 191733.84 examples/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "import re\n",
+    "\n",
+    "import yaml\n",
+    "from datasets import DatasetDict\n",
+    "\n",
+    "from transformer.utils import set_global_seed\n",
+    "\n",
+    "with open(\"configs/runtime/local.yaml\", encoding=\"utf-8\") as cfg_file:\n",
+    "    runtime_cfg = yaml.safe_load(cfg_file) or {}\n",
+    "seed = int(runtime_cfg.get(\"seed\", 0))\n",
+    "set_global_seed(seed)\n",
+    "\n",
+    "\n",
+    "def preprocess_data(example, min_len=2, max_len=256, length_ratio=2.0):\n",
+    "    \"\"\"\n",
+    "    Clean and filter a parallel translation pair (English-French).\n",
+    "\n",
+    "    Args:\n",
+    "        example (dict): {\"translation\": {\"en\": \"...\", \"fr\": \"...\"}}\n",
+    "        min_len (int): Minimum tokenized length allowed.\n",
+    "        max_len (int): Maximum tokenized length allowed.\n",
+    "        length_ratio (float): Max src/tgt length ratio to keep.\n",
+    "\n",
+    "    Returns:\n",
+    "        dict | None: Cleaned translation dict, or None if filtered out.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    src = example[\"translation\"][\"en\"]\n",
+    "    tgt = example[\"translation\"][\"fr\"]\n",
+    "\n",
+    "    # ----------------\n",
+    "    # Basic cleanup\n",
+    "    # ----------------\n",
+    "    def clean_text(text: str) -> str:\n",
+    "        # Remove HTML tags\n",
+    "        text = re.sub(r\"<.*?>\", \" \", text)\n",
+    "\n",
+    "        # Remove URLs and emails\n",
+    "        text = re.sub(r\"http\\S+|www\\.\\S+\", \" \", text)\n",
+    "        text = re.sub(r\"\\S+@\\S+\", \" \", text)\n",
+    "\n",
+    "        # Normalize quotes and punctuation\n",
+    "        text = re.sub(r\"[“”]\", '\"', text)\n",
+    "        text = re.sub(r\"[‘’]\", \"'\", text)\n",
+    "        text = re.sub(r\"([!?.,])\\1+\", r\"\\1\", text)  # collapse repeated punct\n",
+    "\n",
+    "        # Remove non-alphabetic junk (keep accents, digits, basic punct)\n",
+    "        text = re.sub(r\"[^a-zA-ZÀ-ÖØ-öø-ÿ0-9\\s.,!?']\", \" \", text)\n",
+    "\n",
+    "        # Collapse whitespace\n",
+    "        text = re.sub(r\"\\s+\", \" \", text).strip()\n",
+    "        return text.lower()\n",
+    "\n",
+    "    src = clean_text(src)\n",
+    "    tgt = clean_text(tgt)\n",
+    "\n",
+    "    # ----------------\n",
+    "    # Length filtering\n",
+    "    # ----------------\n",
+    "    src_len, tgt_len = len(src.split()), len(tgt.split())\n",
+    "    if src_len < min_len or tgt_len < min_len:\n",
+    "        return None\n",
+    "    if src_len > max_len or tgt_len > max_len:\n",
+    "        return None\n",
+    "    if not (1 / length_ratio <= src_len / tgt_len <= length_ratio):\n",
+    "        return None\n",
+    "\n",
+    "    # ----------------\n",
+    "    # Optional sanity check: drop if empty after cleaning\n",
+    "    # ----------------\n",
+    "    if not src or not tgt:\n",
+    "        return None\n",
+    "\n",
+    "    return {\"en\": src, \"fr\": tgt}\n",
+    "\n",
+    "\n",
+    "processed_small_dataset = small_dataset.map(preprocess_data).filter(lambda x: x is not None)\n",
+    "processed_medium_dataset = medium_dataset.map(preprocess_data).filter(lambda x: x is not None)\n",
+    "# processed_large_dataset = large_dataset.map(preprocess_data).filter(lambda x: x is not None)\n",
+    "\n",
+    "splits = processed_medium_dataset[\"train\"].train_test_split(test_size=0.1, seed=seed)\n",
+    "train_split = splits[\"train\"]\n",
+    "temp_split = splits[\"test\"].train_test_split(test_size=0.5, seed=seed)\n",
+    "processed_medium_dataset = DatasetDict(\n",
+    "    {\n",
+    "        \"train\": train_split,\n",
+    "        \"validation\": temp_split[\"train\"],\n",
+    "        \"test\": temp_split[\"test\"],\n",
+    "    }\n",
+    ")\n",
+    "\n",
+    "processed_small_dataset.save_to_disk(\"./data/processed/processed_small_dataset\")\n",
+    "processed_medium_dataset.save_to_disk(\"./data/processed/processed_medium_dataset\")\n",
+    "# processed_large_dataset.save_to_disk(\"./data/processed/processed_large_dataset\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d782d9db",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from collections import Counter\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "from hydra import compose, initialize\n",
+    "from hydra.core.global_hydra import GlobalHydra\n",
+    "from omegaconf import OmegaConf\n",
+    "from transformers import PreTrainedTokenizerFast\n",
+    "\n",
+    "from transformer.configs import TokenizerCfg\n",
+    "\n",
+    "\n",
+    "def load_tokenizer(name: str) -> PreTrainedTokenizerFast:\n",
+    "    \"\"\"Load a tokenizer defined in configs/tokenizer/<name>.yaml.\"\"\"\n",
+    "    GlobalHydra.instance().clear()\n",
+    "    initialize(config_path=\"./configs/tokenizer\", version_base=None)\n",
+    "    cfg = compose(config_name=\"bpe_4k\")\n",
+    "    scfg_temp = OmegaConf.merge(OmegaConf.structured(TokenizerCfg), cfg)\n",
+    "    tokenizer_cfg: TokenizerCfg = OmegaConf.to_object(scfg_temp)\n",
+    "    tokenizer_path = tokenizer_cfg.path\n",
+    "    tokenizer = PreTrainedTokenizerFast(\n",
+    "        tokenizer_file=str(tokenizer_path),\n",
+    "        bos_token=cfg.get(\"bos_token\"),\n",
+    "        eos_token=cfg.get(\"eos_token\"),\n",
+    "        unk_token=cfg.get(\"unk_token\"),\n",
+    "        pad_token=cfg.get(\"pad_token\"),\n",
+    "    )\n",
+    "    tokenizer.model_max_length = cfg.get(\"max_seq_len\", tokenizer.model_max_length)\n",
+    "    return tokenizer\n",
+    "\n",
+    "\n",
+    "def _resolve_tokenizer(tokenizers, field: str) -> PreTrainedTokenizerFast:\n",
+    "    if isinstance(tokenizers, dict):\n",
+    "        tok = tokenizers.get(field)\n",
+    "        if tok is None:\n",
+    "            raise KeyError(f\"No tokenizer provided for field '{field}'\")\n",
+    "        return tok\n",
+    "    if tokenizers is None:\n",
+    "        raise ValueError(\"Tokenizer or field->tokenizer mapping required\")\n",
+    "    return tokenizers\n",
+    "\n",
+    "\n",
+    "def _tokenize_text(\n",
+    "    tokenizer: PreTrainedTokenizerFast, text: str, include_special_tokens: bool\n",
+    ") -> list[str]:\n",
+    "    encoded = tokenizer(\n",
+    "        text,\n",
+    "        add_special_tokens=include_special_tokens,\n",
+    "        return_attention_mask=False,\n",
+    "        return_token_type_ids=False,\n",
+    "    )\n",
+    "    input_ids = encoded[\"input_ids\"]\n",
+    "    if input_ids and isinstance(input_ids[0], list):\n",
+    "        input_ids = input_ids[0]\n",
+    "    return tokenizer.convert_ids_to_tokens(input_ids)\n",
+    "\n",
+    "\n",
+    "def _subset_dataset(ds, sample_size: int | None):\n",
+    "    if sample_size is None or sample_size >= len(ds):\n",
+    "        return ds\n",
+    "    return ds.select(range(sample_size))\n",
+    "\n",
+    "\n",
+    "def show_top_k_tokens(\n",
+    "    dataset_dict,\n",
+    "    tokenizers,\n",
+    "    fields=(\"en\", \"fr\"),\n",
+    "    k: int = 20,\n",
+    "    sample_size: int | None = 2000,\n",
+    "    include_special_tokens: bool = False,\n",
+    "):\n",
+    "    \"\"\"Print the top-k most frequent tokens using configured tokenizer(s).\"\"\"\n",
+    "    for split, ds in dataset_dict.items():\n",
+    "        ds_view = _subset_dataset(ds, sample_size)\n",
+    "        counters = {field: Counter() for field in fields}\n",
+    "        for example in ds_view:\n",
+    "            for field in fields:\n",
+    "                tokenizer = _resolve_tokenizer(tokenizers, field)\n",
+    "                tokens = _tokenize_text(tokenizer, example[field], include_special_tokens)\n",
+    "                counters[field].update(tokens)\n",
+    "        print(f\"\\nSplit: {split}\")\n",
+    "        for field in fields:\n",
+    "            print(f\"  Top {k} tokens for '{field}':\")\n",
+    "            for token, freq in counters[field].most_common(k):\n",
+    "                print(f\"    {token!r}: {freq}\")\n",
+    "\n",
+    "\n",
+    "def plot_token_length_barchart(\n",
+    "    dataset_dict,\n",
+    "    tokenizers,\n",
+    "    fields=(\"en\", \"fr\"),\n",
+    "    sample_size: int | None = 2000,\n",
+    "    include_special_tokens: bool = False,\n",
+    "):\n",
+    "    \"\"\"Display bar charts of token-length distributions using configured tokenizer(s).\"\"\"\n",
+    "    for split, ds in dataset_dict.items():\n",
+    "        ds_view = _subset_dataset(ds, sample_size)\n",
+    "        length_counters = {field: Counter() for field in fields}\n",
+    "        for example in ds_view:\n",
+    "            for field in fields:\n",
+    "                tokenizer = _resolve_tokenizer(tokenizers, field)\n",
+    "                count = len(_tokenize_text(tokenizer, example[field], include_special_tokens))\n",
+    "                length_counters[field][count] += 1\n",
+    "        support = sorted({length for counter in length_counters.values() for length in counter})\n",
+    "        if not support:\n",
+    "            continue\n",
+    "        bin_size = 50\n",
+    "        bin_starts = list(range(0, support[-1] + bin_size, bin_size))\n",
+    "        bin_labels = [f\"{start}-{start + bin_size - 1}\" for start in bin_starts]\n",
+    "        bin_counts = {field: [0] * len(bin_starts) for field in fields}\n",
+    "        for field in fields:\n",
+    "            for length, freq in length_counters[field].items():\n",
+    "                idx = min(len(bin_starts) - 1, length // bin_size)\n",
+    "                bin_counts[field][idx] += freq\n",
+    "        x = np.arange(len(bin_starts))\n",
+    "        width = 0.8 / max(1, len(fields))\n",
+    "        fig, ax = plt.subplots(figsize=(8, 4))\n",
+    "        for idx_field, field in enumerate(fields):\n",
+    "            counts = bin_counts[field]\n",
+    "            ax.bar(x + idx_field * width, counts, width=width, label=field)\n",
+    "        ax.set_xticks(x + width * (len(fields) - 1) / 2)\n",
+    "        ax.set_xticklabels(bin_labels)\n",
+    "        ax.set_xlabel(\"Token count per example\")\n",
+    "        ax.set_ylabel(\"Number of examples\")\n",
+    "        ax.set_title(f\"Example token lengths in '{split}' split\")\n",
+    "        ax.legend()\n",
+    "        fig.tight_layout()\n",
+    "        plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b8ed32b9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizers = {\n",
+    "    \"en\": load_tokenizer(\"bpe_8k\"),\n",
+    "    \"fr\": load_tokenizer(\"bpe_8k\"),\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "04f27397",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Small dataset statistics\n",
+    "show_top_k_tokens(processed_small_dataset, tokenizers, k=20, sample_size=2000)\n",
+    "plot_token_length_barchart(processed_small_dataset, tokenizers, sample_size=2000)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e7819046",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Medium dataset statistics\n",
+    "show_top_k_tokens(processed_medium_dataset, tokenizers, k=20, sample_size=2000)\n",
+    "plot_token_length_barchart(processed_medium_dataset, tokenizers, sample_size=2000)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "transformer",
+   "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.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

scripts/server.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from fastapi import FastAPI
+from hydra import compose, initialize
+from omegaconf import DictConfig
+
+from scripts.inference import inference as run_inference
+
+app = FastAPI(title="Transformer Inference API")
+
+
+@app.get("/healthz")
+def healthz():
+    return {"status": "ok"}
+
+
+@app.get("/generate")
+def generate(text: str | None = None):
+    # Compose the default config exactly as inference.py expects
+    with initialize(config_path="../configs", version_base=None):
+        cfg: DictConfig = compose(config_name="infer_mode")
+    # If a text query parameter is provided, override the config's input_text
+    if text is not None and text.strip():
+        cfg.input_text = text
+    outputs = run_inference(cfg)  # returns a list of strings
+    return {"outputs": outputs}

scripts/tokenizer.ipynb

@@ -0,0 +1,217 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "e360283a",
+   "metadata": {},
+   "source": [
+    "imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "47d34b96",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/ala-boussoffara/miniconda3/envs/transformer/lib/python3.12/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"
+     ]
+    }
+   ],
+   "source": [
+    "from datasets import concatenate_datasets, load_from_disk\n",
+    "from hydra import compose, initialize\n",
+    "from hydra.core.global_hydra import GlobalHydra\n",
+    "from omegaconf import OmegaConf\n",
+    "from tokenizers import Tokenizer, models, pre_tokenizers, trainers\n",
+    "from tokenizers.processors import TemplateProcessing\n",
+    "\n",
+    "from transformer.configs import TokenizerCfg"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9a0c380c",
+   "metadata": {},
+   "source": [
+    "Choosing tokenizer config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "8f13dc2b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "GlobalHydra.instance().clear()\n",
+    "initialize(config_path=\"./configs/tokenizer\", version_base=None)\n",
+    "cfg = compose(config_name=\"bpe_16k\")\n",
+    "scfg_temp = OmegaConf.merge(OmegaConf.structured(TokenizerCfg), cfg)\n",
+    "tokenizer_cfg: TokenizerCfg = OmegaConf.to_object(scfg_temp)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d5ee3d16",
+   "metadata": {},
+   "source": [
+    "Loading datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "d880e59e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "small_dataset = load_from_disk(\"./data/processed/processed_small_dataset\")\n",
+    "medium_dataset = load_from_disk(\"./data/processed/processed_medium_dataset\")\n",
+    "\n",
+    "concatenated_dataset = concatenate_datasets(\n",
+    "    [\n",
+    "        small_dataset[\"train\"],\n",
+    "        small_dataset[\"validation\"],\n",
+    "        small_dataset[\"test\"],\n",
+    "        medium_dataset[\"train\"],\n",
+    "    ]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c6b53e40",
+   "metadata": {},
+   "source": [
+    "preparing corpus"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "f60ac4b2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(tokenizer_cfg.corpus, \"w\", encoding=\"utf-8\") as f:\n",
+    "    for example in globals()[tokenizer_cfg.dataset]:\n",
+    "        en = example[\"translation\"][\"en\"]\n",
+    "        fr = example[\"translation\"][\"fr\"]\n",
+    "        if en and fr:\n",
+    "            f.write(en + \"\\n\")\n",
+    "            f.write(fr + \"\\n\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f7a50455",
+   "metadata": {},
+   "source": [
+    "training the tokenizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "29fc0f7e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "tokenizer = Tokenizer(models.BPE(unk_token=tokenizer_cfg.unk_token))\n",
+    "tokenizer.pre_tokenizer = pre_tokenizers.Sequence(\n",
+    "    [pre_tokenizers.Whitespace(), pre_tokenizers.Punctuation()]\n",
+    ")\n",
+    "trainer = trainers.BpeTrainer(\n",
+    "    vocab_size=tokenizer_cfg.vocab_size, special_tokens=tokenizer_cfg.special_tokens\n",
+    ")\n",
+    "tokenizer.train([tokenizer_cfg.corpus], trainer)\n",
+    "tokenizer.post_processor = TemplateProcessing(\n",
+    "    single=f\"{tokenizer_cfg.bos_token} $A {tokenizer_cfg.eos_token}\",\n",
+    "    special_tokens=[\n",
+    "        (tokenizer_cfg.bos_token, tokenizer_cfg.bos_id),\n",
+    "        (tokenizer_cfg.eos_token, tokenizer_cfg.eos_id),\n",
+    "    ],\n",
+    ")\n",
+    "tokenizer.save(tokenizer_cfg.path)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "537b9d9d",
+   "metadata": {},
+   "source": [
+    "Loading and testing"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "443ed7ee",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "['i', 'here', 'by', 'an', 'n', 'ou', 'ce', 'you', 'as', 'the', 'next', 'presi', 'dent']\n",
+      "PAD: 0\n",
+      "BOS: 1\n",
+      "EOS: 2\n",
+      "UNK: 3\n"
+     ]
+    }
+   ],
+   "source": [
+    "from tokenizers import Tokenizer\n",
+    "\n",
+    "tokenizer = Tokenizer.from_file(tokenizer_cfg.path)\n",
+    "print(\n",
+    "    tokenizer.encode(\"i hereby annouce you as the next president\", add_special_tokens=False).tokens\n",
+    ")\n",
+    "vocab = tokenizer.get_vocab()\n",
+    "\n",
+    "print(\"PAD:\", vocab[\"<pad>\"])\n",
+    "print(\"BOS:\", vocab[\"<bos>\"])\n",
+    "print(\"EOS:\", vocab[\"<eos>\"])\n",
+    "print(\"UNK:\", vocab[\"<unk>\"])"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "transformer",
+   "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.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

scripts/train.ipynb

@@ -0,0 +1,987 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "381e1824",
+   "metadata": {},
+   "source": [
+    "Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5a9e6fb4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import contextlib\n",
+    "import json\n",
+    "import math\n",
+    "import os\n",
+    "from datetime import datetime\n",
+    "from pathlib import Path\n",
+    "from typing import Any\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from datasets import load_from_disk\n",
+    "from hydra import compose, initialize\n",
+    "from hydra.core.global_hydra import GlobalHydra\n",
+    "from hydra.utils import to_absolute_path\n",
+    "from omegaconf import OmegaConf\n",
+    "from torch.nn.utils import clip_grad_norm_\n",
+    "from torch.utils.data import DataLoader\n",
+    "from tqdm import tqdm\n",
+    "from transformers import PreTrainedTokenizerFast\n",
+    "\n",
+    "from transformer import BasicEncoderDecoderTransformer\n",
+    "from transformer.configs import ModelCfg, TokenizerCfg, TrainAppCfg\n",
+    "from transformer.utils import (\n",
+    "    check_tokenizer_model_compatibility,\n",
+    "    debug_transformer_forward,\n",
+    "    make_worker_init_fn,\n",
+    "    set_global_seed,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "512d7a41",
+   "metadata": {},
+   "source": [
+    "Modifying working directory"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d1c043be",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "project_root = Path(__file__).resolve().parent if \"__file__\" in globals() else Path.cwd()\n",
+    "os.chdir(project_root)\n",
+    "print(f\"Working directory set to: {Path.cwd()}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9b3f89d1",
+   "metadata": {},
+   "source": [
+    "Initiliazation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2bd6dd82",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "GlobalHydra.instance().clear()\n",
+    "initialize(config_path=\"./configs\", version_base=None)\n",
+    "cfg = compose(\n",
+    "    config_name=\"train_mode\",\n",
+    "    overrides=[\n",
+    "        \"model=test_model\",\n",
+    "        \"tokenizer=bpe_8k\",\n",
+    "        \"dataset=medium_dataset\",\n",
+    "        \"trainer.batch_size=128\",\n",
+    "        \"trainer.gradient_accumulation=2\",\n",
+    "        \"trainer.epochs=80\",\n",
+    "    ],\n",
+    ")\n",
+    "scfg_temp = OmegaConf.merge(OmegaConf.structured(TrainAppCfg), cfg)\n",
+    "scfg: TrainAppCfg = OmegaConf.to_object(scfg_temp)\n",
+    "set_global_seed(scfg.runtime.seed)\n",
+    "model_cfg = ModelCfg(**vars(scfg.model))\n",
+    "tokenizer_cfg = TokenizerCfg(**vars(scfg.tokenizer))\n",
+    "\n",
+    "# --- Compatibility check ---\n",
+    "check_tokenizer_model_compatibility(model_cfg, tokenizer_cfg)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "24caf2de",
+   "metadata": {},
+   "source": [
+    "Preparing data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8d60e085",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset = load_from_disk(cfg.dataset.path)\n",
+    "\n",
+    "train_dataset = dataset[\"train\"].select(range(200000))\n",
+    "validation_dataset = dataset[\"validation\"]\n",
+    "test_dataset = dataset[\"test\"]\n",
+    "\n",
+    "\n",
+    "def tuple_collate_fn(batch):\n",
+    "    # batch is a list of dicts with 'en' and 'fr' keys\n",
+    "    src = [item[\"en\"] for item in batch]\n",
+    "    tgt = [item[\"fr\"] for item in batch]\n",
+    "    return src, tgt\n",
+    "\n",
+    "\n",
+    "dataloader_generator = torch.Generator().manual_seed(scfg.runtime.seed)\n",
+    "worker_init = make_worker_init_fn(scfg.runtime.seed)\n",
+    "\n",
+    "train_dataloader = DataLoader(\n",
+    "    train_dataset,\n",
+    "    batch_size=scfg.trainer.batch_size,\n",
+    "    shuffle=scfg.trainer.shuffle,\n",
+    "    collate_fn=tuple_collate_fn,\n",
+    "    generator=dataloader_generator,\n",
+    "    worker_init_fn=worker_init,\n",
+    ")\n",
+    "validation_dataloader = DataLoader(\n",
+    "    validation_dataset,\n",
+    "    batch_size=scfg.trainer.batch_size,\n",
+    "    shuffle=False,\n",
+    "    collate_fn=tuple_collate_fn,\n",
+    "    generator=dataloader_generator,\n",
+    "    worker_init_fn=worker_init,\n",
+    ")\n",
+    "test_dataloader = DataLoader(\n",
+    "    test_dataset,\n",
+    "    batch_size=scfg.trainer.batch_size,\n",
+    "    shuffle=False,\n",
+    "    collate_fn=tuple_collate_fn,\n",
+    "    generator=dataloader_generator,\n",
+    "    worker_init_fn=worker_init,\n",
+    ")\n",
+    "\n",
+    "print(\"Number of training samples:\", len(train_dataset))\n",
+    "print(\"Number of training batches:\", len(train_dataloader))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "96c84adf",
+   "metadata": {},
+   "source": [
+    "Model and tokenizer loading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f9da0dab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = torch.device(scfg.runtime.device)\n",
+    "model = BasicEncoderDecoderTransformer(model_cfg).to(device)\n",
+    "\n",
+    "\n",
+    "# xavier initialization\n",
+    "def init_weights(m):\n",
+    "    if isinstance(m, nn.Linear | nn.Embedding):\n",
+    "        torch.nn.init.xavier_uniform_(m.weight)\n",
+    "    elif isinstance(m, nn.LayerNorm):\n",
+    "        m.bias.data.fill_(0)\n",
+    "        m.weight.data.fill_(1.0)\n",
+    "    if isinstance(m, nn.Linear) and m.bias is not None:\n",
+    "        m.bias.data.fill_(0)\n",
+    "\n",
+    "\n",
+    "model.apply(init_weights)\n",
+    "\n",
+    "tokenizer = PreTrainedTokenizerFast(\n",
+    "    tokenizer_file=tokenizer_cfg.path,\n",
+    "    bos_token=tokenizer_cfg.bos_token,\n",
+    "    eos_token=tokenizer_cfg.eos_token,\n",
+    "    unk_token=tokenizer_cfg.unk_token,\n",
+    "    pad_token=tokenizer_cfg.pad_token,\n",
+    "    model_max_length=tokenizer_cfg.max_seq_len,\n",
+    ")\n",
+    "\n",
+    "criterion = nn.CrossEntropyLoss(\n",
+    "    ignore_index=tokenizer_cfg.pad_id, label_smoothing=scfg.trainer.label_smoothing\n",
+    ")\n",
+    "\n",
+    "# print total number of parameters\n",
+    "total_params = sum(p.numel() for p in model.parameters())\n",
+    "print(f\"Total number of parameters: {total_params}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "608a205a",
+   "metadata": {},
+   "source": [
+    "Optimizer setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "372dea7e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def setup_optimizer(model: nn.Module, cfg: Any) -> torch.optim.Optimizer:\n",
+    "    \"\"\"Create an AdamW optimizer with decoupled weight decay groups.\"\"\"\n",
+    "\n",
+    "    decay_params: set[str] = set()\n",
+    "    no_decay_params: set[str] = set()\n",
+    "\n",
+    "    blacklist_weight_modules = (nn.LayerNorm,)\n",
+    "\n",
+    "    for module_name, module in model.named_modules():\n",
+    "        for param_name, param in module.named_parameters(recurse=False):\n",
+    "            if not param.requires_grad:\n",
+    "                continue\n",
+    "\n",
+    "            full_name = f\"{module_name}.{param_name}\" if module_name else param_name\n",
+    "\n",
+    "            if param_name.endswith(\"bias\"):\n",
+    "                no_decay_params.add(full_name)\n",
+    "            elif isinstance(module, blacklist_weight_modules):\n",
+    "                no_decay_params.add(full_name)\n",
+    "            elif param.ndim == 1:\n",
+    "                no_decay_params.add(full_name)\n",
+    "            else:\n",
+    "                decay_params.add(full_name)\n",
+    "\n",
+    "    for name, param in model.named_parameters():\n",
+    "        if not param.requires_grad:\n",
+    "            continue\n",
+    "        lowered = name.lower()\n",
+    "        if \"embedding\" in lowered or (\"pos\" in lowered and \"embed\" in lowered):\n",
+    "            decay_params.discard(name)\n",
+    "            no_decay_params.add(name)\n",
+    "\n",
+    "    param_dict = {name: param for name, param in model.named_parameters() if param.requires_grad}\n",
+    "    decay_group = [param_dict[name] for name in sorted(decay_params) if name in param_dict]\n",
+    "    no_decay_group = [param_dict[name] for name in sorted(no_decay_params) if name in param_dict]\n",
+    "\n",
+    "    optim_cfg = cfg.optim\n",
+    "    param_groups = []\n",
+    "    if decay_group:\n",
+    "        param_groups.append({\"params\": decay_group, \"weight_decay\": optim_cfg.weight_decay})\n",
+    "    if no_decay_group:\n",
+    "        param_groups.append({\"params\": no_decay_group, \"weight_decay\": 0.0})\n",
+    "\n",
+    "    if not param_groups:\n",
+    "        raise ValueError(\"No trainable parameters found when setting up the optimizer.\")\n",
+    "\n",
+    "    return torch.optim.AdamW(\n",
+    "        model.parameters(),\n",
+    "        lr=optim_cfg.lr,\n",
+    "        betas=optim_cfg.betas,\n",
+    "        eps=optim_cfg.eps,\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "optimizer = setup_optimizer(model, scfg)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f64a8b38",
+   "metadata": {},
+   "source": [
+    "AMP and scaler setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4ad52841",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 1) Decide whether AMP is on and which dtype to use\n",
+    "use_amp = scfg.trainer.mixed_precision and (scfg.trainer.precision in [\"fp16\", \"bf16\"])\n",
+    "\n",
+    "# 2) Pick the autocast dtype from precision\n",
+    "amp_dtype = torch.float16 if scfg.trainer.precision == \"fp16\" else torch.bfloat16\n",
+    "\n",
+    "# 3) Build the context manager:\n",
+    "#    - If AMP is on → use autocast(dtype=...)\n",
+    "#    - If AMP is off → use a nullcontext (does nothing)\n",
+    "autocast_ctx = (\n",
+    "    torch.amp.autocast(dtype=amp_dtype, device_type=scfg.runtime.device)\n",
+    "    if use_amp\n",
+    "    else contextlib.nullcontext()\n",
+    ")\n",
+    "\n",
+    "# 4) Create the GradScaler:\n",
+    "#    - enabled=True only for fp16 (needs loss scaling)\n",
+    "#    - for bf16/fp32 it becomes a no-op automatically\n",
+    "scaler = torch.amp.GradScaler(\n",
+    "    enabled=(scfg.trainer.mixed_precision and scfg.trainer.precision == \"fp16\")\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "34e29854",
+   "metadata": {},
+   "source": [
+    "Scheduler setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "479390a8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def setup_scheduler(\n",
+    "    optimizer: torch.optim.Optimizer, trainer_cfg: Any, sched_cfg: Any, steps_per_epoch: int\n",
+    ") -> torch.optim.lr_scheduler.LambdaLR:\n",
+    "    \"\"\"Create a cosine LR scheduler driven by warmup/min LR ratios.\"\"\"\n",
+    "    total_steps = max(1, int(trainer_cfg.epochs) * max(1, int(steps_per_epoch)))\n",
+    "    warmup_ratio = max(0.0, float(getattr(sched_cfg, \"warmup_ratio\", 0.0)))\n",
+    "    min_lr_ratio = max(0.0, float(getattr(sched_cfg, \"min_lr_ratio\", 0.0)))\n",
+    "    hold_min = bool(getattr(sched_cfg, \"hold_min\", True))\n",
+    "    max_lr_ratio = max(0.0, float(getattr(sched_cfg, \"max_lr_ratio\", 1.0)))\n",
+    "    warmup_steps = int(round(total_steps * warmup_ratio))\n",
+    "    warmup_steps = max(0, min(warmup_steps, max(0, total_steps - 1)))\n",
+    "    decay_steps = max(1, total_steps - warmup_steps)\n",
+    "    for group in optimizer.param_groups:\n",
+    "        base_lr = float(group.get(\"initial_lr\", group[\"lr\"]))\n",
+    "        group[\"initial_lr\"] = base_lr\n",
+    "        group[\"min_lr\"] = base_lr * min_lr_ratio\n",
+    "        group[\"max_lr\"] = base_lr * max_lr_ratio\n",
+    "\n",
+    "    def lr_lambda(step: int) -> float:\n",
+    "        if warmup_steps > 0 and step < warmup_steps:\n",
+    "            warmup_progress = step / max(1, warmup_steps)\n",
+    "            warmup_progress = max(0.0, min(1.0, warmup_progress))\n",
+    "            return min_lr_ratio + (1.0 - min_lr_ratio) * warmup_progress\n",
+    "        if decay_steps <= 0:\n",
+    "            return 1.0\n",
+    "        progress = (step - warmup_steps) / decay_steps\n",
+    "        progress = max(0.0, min(1.0, progress))\n",
+    "        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))\n",
+    "        factor = min_lr_ratio + (1.0 - min_lr_ratio) * cosine\n",
+    "        if progress >= 1.0 and not hold_min:\n",
+    "            return 0.0\n",
+    "        return factor\n",
+    "\n",
+    "    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n",
+    "\n",
+    "\n",
+    "steps_per_epoch = math.ceil(len(train_dataloader) / max(1, scfg.trainer.gradient_accumulation))\n",
+    "scheduler = setup_scheduler(optimizer, scfg.trainer, scfg.sched, steps_per_epoch=steps_per_epoch)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "db51c02a",
+   "metadata": {},
+   "source": [
+    "Loading checkpoints"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a9843809",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "output_root = Path(to_absolute_path(scfg.runtime.output_dir))\n",
+    "output_root.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "timestamp = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
+    "run_dir = output_root / f\"{model_cfg.name}_{timestamp}\"\n",
+    "suffix = 1\n",
+    "while run_dir.exists():\n",
+    "    run_dir = output_root / f\"{model_cfg.name}_{timestamp}_{suffix}\"\n",
+    "    suffix += 1\n",
+    "run_dir.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "checkpoints_root = run_dir / \"checkpoints\"\n",
+    "latest_dir = checkpoints_root / \"latest\"\n",
+    "best_dir = checkpoints_root / \"best\"\n",
+    "latest_dir.mkdir(parents=True, exist_ok=True)\n",
+    "best_dir.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "resume_latest_source = getattr(model_cfg, \"latest_checkpoint_path\", \"\")\n",
+    "resume_best_source = getattr(model_cfg, \"best_checkpoint_path\", \"\")\n",
+    "best_checkpoint_export_target = (\n",
+    "    Path(to_absolute_path(resume_best_source)) if resume_best_source else None\n",
+    ")\n",
+    "\n",
+    "latest_checkpoint_path = latest_dir / \"checkpoint.pt\"\n",
+    "best_checkpoint_path = best_dir / \"checkpoint.pt\"\n",
+    "\n",
+    "\n",
+    "config_snapshot_path = run_dir / \"config.yaml\"\n",
+    "OmegaConf.save(config=cfg, f=str(config_snapshot_path))\n",
+    "\n",
+    "manifest_path = run_dir / \"manifest.json\"\n",
+    "\n",
+    "resume_mode = getattr(scfg.trainer, \"resume\", \"start_over\")\n",
+    "resume_mode = resume_mode.replace(\" \", \"_\").lower()\n",
+    "if resume_mode not in {\"start_over\", \"best\", \"latest\"}:\n",
+    "    print(f\"Unknown resume mode '{resume_mode}', defaulting to 'start_over'.\")\n",
+    "    resume_mode = \"start_over\"\n",
+    "\n",
+    "resume_checkpoint_path = None\n",
+    "if resume_mode == \"latest\" and resume_latest_source:\n",
+    "    resume_checkpoint_path = Path(to_absolute_path(resume_latest_source))\n",
+    "elif resume_mode == \"best\" and resume_best_source:\n",
+    "    resume_checkpoint_path = Path(to_absolute_path(resume_best_source))\n",
+    "\n",
+    "manifest_data = {\n",
+    "    \"run_dir\": str(run_dir),\n",
+    "    \"model\": model_cfg.name,\n",
+    "    \"resume\": {\n",
+    "        \"mode\": resume_mode,\n",
+    "        \"source\": str(resume_checkpoint_path) if resume_checkpoint_path else None,\n",
+    "        \"status\": \"pending\",\n",
+    "    },\n",
+    "    \"latest\": None,\n",
+    "    \"best\": None,\n",
+    "}\n",
+    "manifest_path.write_text(json.dumps(manifest_data, indent=2))\n",
+    "\n",
+    "start_epoch = 0\n",
+    "completed_optimizer_steps = 0\n",
+    "train_losses: list[float] = []\n",
+    "val_losses: list[float] = []\n",
+    "learning_rates: list[float] = []\n",
+    "resume_batch_idx = 0\n",
+    "resume_running_loss = 0.0\n",
+    "best_val_loss = float(\"inf\")\n",
+    "best_val_epoch = -1\n",
+    "\n",
+    "loaded_checkpoint = None\n",
+    "if resume_checkpoint_path and resume_checkpoint_path.is_file():\n",
+    "    print(f\"Loading checkpoint from {resume_checkpoint_path}\")\n",
+    "    loaded_checkpoint = torch.load(resume_checkpoint_path, map_location=\"cpu\", weights_only=False)\n",
+    "    model_state = loaded_checkpoint.get(\"model_state_dict\")\n",
+    "    if model_state is not None:\n",
+    "        model.load_state_dict(model_state)\n",
+    "    optim_state = loaded_checkpoint.get(\"optimizer_state_dict\")\n",
+    "    if optim_state is not None:\n",
+    "        optimizer.load_state_dict(optim_state)\n",
+    "    sched_state = loaded_checkpoint.get(\"scheduler_state_dict\")\n",
+    "    if sched_state is not None:\n",
+    "        scheduler.load_state_dict(sched_state)\n",
+    "    scaler_state = loaded_checkpoint.get(\"scaler_state_dict\")\n",
+    "    if scaler_state is not None:\n",
+    "        scaler.load_state_dict(scaler_state)\n",
+    "    train_losses = list(loaded_checkpoint.get(\"train_losses\", train_losses))\n",
+    "    val_losses = list(loaded_checkpoint.get(\"val_losses\", val_losses))\n",
+    "    learning_rates = list(loaded_checkpoint.get(\"learning_rates\", learning_rates))\n",
+    "    start_epoch = int(loaded_checkpoint.get(\"epoch\", start_epoch))\n",
+    "    completed_optimizer_steps = int(\n",
+    "        loaded_checkpoint.get(\"optimizer_step\", completed_optimizer_steps)\n",
+    "    )\n",
+    "    resume_batch_idx = int(loaded_checkpoint.get(\"batch_idx\", resume_batch_idx))\n",
+    "    resume_running_loss = float(loaded_checkpoint.get(\"running_loss\", resume_running_loss))\n",
+    "    best_val_loss = float(\n",
+    "        loaded_checkpoint.get(\n",
+    "            \"best_val_loss\",\n",
+    "            loaded_checkpoint.get(\"best_train_loss\", best_val_loss),\n",
+    "        )\n",
+    "    )\n",
+    "    best_val_epoch = int(\n",
+    "        loaded_checkpoint.get(\n",
+    "            \"best_val_epoch\",\n",
+    "            loaded_checkpoint.get(\"best_train_epoch\", best_val_epoch),\n",
+    "        )\n",
+    "    )\n",
+    "    for state in optimizer.state.values():\n",
+    "        for key, value in state.items():\n",
+    "            if isinstance(value, torch.Tensor):\n",
+    "                state[key] = value.to(device)\n",
+    "    model.to(device)\n",
+    "else:\n",
+    "    if resume_checkpoint_path:\n",
+    "        print(f\"Checkpoint '{resume_checkpoint_path}' not found. Starting fresh.\")\n",
+    "    else:\n",
+    "        print(\"Starting fresh (no resume checkpoint provided).\")\n",
+    "\n",
+    "if loaded_checkpoint is not None:\n",
+    "    manifest_data[\"resume\"][\"status\"] = \"loaded\"\n",
+    "elif resume_checkpoint_path:\n",
+    "    manifest_data[\"resume\"][\"status\"] = \"missing\"\n",
+    "else:\n",
+    "    manifest_data[\"resume\"][\"status\"] = \"fresh\"\n",
+    "manifest_path.write_text(json.dumps(manifest_data, indent=2))\n",
+    "\n",
+    "\n",
+    "def _update_manifest(section: str, payload: dict) -> None:\n",
+    "    manifest_data[section] = payload\n",
+    "    manifest_path.write_text(json.dumps(manifest_data, indent=2))\n",
+    "\n",
+    "\n",
+    "if loaded_checkpoint is not None:\n",
+    "    torch.save(loaded_checkpoint, latest_checkpoint_path)\n",
+    "    _update_manifest(\n",
+    "        \"latest\",\n",
+    "        {\n",
+    "            \"path\": str(latest_checkpoint_path),\n",
+    "            \"epoch\": start_epoch,\n",
+    "            \"optimizer_step\": completed_optimizer_steps,\n",
+    "            \"batch_idx\": resume_batch_idx,\n",
+    "            \"running_loss\": resume_running_loss,\n",
+    "            \"timestamp\": datetime.now().isoformat(),\n",
+    "        },\n",
+    "    )\n",
+    "    if resume_mode == \"best\":\n",
+    "        torch.save(loaded_checkpoint, best_checkpoint_path)\n",
+    "        _update_manifest(\n",
+    "            \"best\",\n",
+    "            {\n",
+    "                \"path\": str(best_checkpoint_path),\n",
+    "                \"epoch\": start_epoch,\n",
+    "                \"optimizer_step\": completed_optimizer_steps,\n",
+    "                \"val_loss\": loaded_checkpoint.get(\n",
+    "                    \"best_val_loss\", loaded_checkpoint.get(\"best_train_loss\")\n",
+    "                ),\n",
+    "                \"train_loss\": loaded_checkpoint.get(\"metadata\", {}).get(\"train_loss\"),\n",
+    "                \"timestamp\": datetime.now().isoformat(),\n",
+    "            },\n",
+    "        )\n",
+    "\n",
+    "\n",
+    "def save_checkpoint(\n",
+    "    epoch_to_resume: int,\n",
+    "    optimizer_step: int,\n",
+    "    *,\n",
+    "    batch_idx: int,\n",
+    "    running_loss_value: float,\n",
+    "    metadata: dict[str, Any] | None = None,\n",
+    ") -> dict:\n",
+    "    checkpoint = {\n",
+    "        \"epoch\": epoch_to_resume,\n",
+    "        \"optimizer_step\": optimizer_step,\n",
+    "        \"batch_idx\": batch_idx,\n",
+    "        \"running_loss\": running_loss_value,\n",
+    "        \"model_state_dict\": model.state_dict(),\n",
+    "        \"optimizer_state_dict\": optimizer.state_dict(),\n",
+    "        \"scheduler_state_dict\": scheduler.state_dict(),\n",
+    "        \"scaler_state_dict\": scaler.state_dict(),\n",
+    "        \"train_losses\": train_losses,\n",
+    "        \"val_losses\": val_losses,\n",
+    "        \"learning_rates\": learning_rates,\n",
+    "        \"best_val_loss\": best_val_loss,\n",
+    "        \"best_val_epoch\": best_val_epoch,\n",
+    "        \"metadata\": metadata or {},\n",
+    "    }\n",
+    "    torch.save(checkpoint, latest_checkpoint_path)\n",
+    "    _update_manifest(\n",
+    "        \"latest\",\n",
+    "        {\n",
+    "            \"path\": str(latest_checkpoint_path),\n",
+    "            \"epoch\": epoch_to_resume,\n",
+    "            \"optimizer_step\": optimizer_step,\n",
+    "            \"batch_idx\": batch_idx,\n",
+    "            \"running_loss\": running_loss_value,\n",
+    "            \"metadata\": metadata or {},\n",
+    "            \"timestamp\": datetime.now().isoformat(),\n",
+    "        },\n",
+    "    )\n",
+    "    return checkpoint"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ab23bf3c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# src,tgt=next(iter(train_dataloader))\n",
+    "src = [\"it can be a very complicated thing, the ocean\"]\n",
+    "tgt = [\"ca peut être très compliqué, l'océan\"]\n",
+    "encoded_src = tokenizer(\n",
+    "    src, padding=True, truncation=True, return_tensors=\"pt\", add_special_tokens=False\n",
+    ")\n",
+    "encoded_tgt = tokenizer(tgt, padding=True, truncation=True, return_tensors=\"pt\")\n",
+    "src_ids = encoded_src[\"input_ids\"].to(device)\n",
+    "tgt_ids = encoded_tgt[\"input_ids\"].to(device)\n",
+    "src_padd_mask = (src_ids == scfg.tokenizer.pad_id).to(device, dtype=torch.bool)\n",
+    "tgt_padd_mask = (tgt_ids == scfg.tokenizer.pad_id).to(device, dtype=torch.bool)\n",
+    "# initilaize src_ids and tgt_ids randomly with big lengths\n",
+    "\"\"\"src_ids = torch.randint(0, scfg.tokenizer.vocab_size, (2, 50)).to(device, dtype=torch.long)\n",
+    "tgt_ids = torch.randint(0, scfg.tokenizer.vocab_size, (2, 50)).to(device, dtype=torch.long)\n",
+    "src_padd_mask = (src_ids == scfg.tokenizer.pad_id).to(device, dtype=torch.bool)\n",
+    "tgt_padd_mask = (tgt_ids == scfg.tokenizer.pad_id).to(device, dtype=torch.bool)\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "689fe371",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Optional transformer debugging / attention capture\n",
+    "debug_forward_options = {\n",
+    "    \"enabled\": True,  # flip to True to run debug_transformer_forward\n",
+    "    \"show_attention\": False,  # show attention plots interactively\n",
+    "    \"save_attention\": False,  # persist attention artefacts\n",
+    "    \"attention_layers\": None,  # e.g. [0, 1] to restrict encoder/decoder layers\n",
+    "    \"attention_heads\": None,  # e.g. [0, 3] to filter heads\n",
+    "    \"attention_types\": (\"enc_self\", \"dec_self\", \"dec_cross\"),\n",
+    "    \"average_heads\": False,  # average selected heads into a single map\n",
+    "    \"sample_index\": 0,  # which row within the batch to inspect\n",
+    "    \"attention_figsize\": (2.4, 2.4),\n",
+    "    \"skip_special_tokens\": False,\n",
+    "    \"save_dir\": None,  # override directory; defaults to run_dir/debug/attention\n",
+    "    \"return_maps\": False,  # set True to receive raw attention tensors\n",
+    "}\n",
+    "\n",
+    "if debug_forward_options[\"save_dir\"]:\n",
+    "    debug_default_attention_dir = Path(debug_forward_options[\"save_dir\"])\n",
+    "    if not debug_default_attention_dir.is_absolute():\n",
+    "        debug_default_attention_dir = run_dir / debug_default_attention_dir\n",
+    "else:\n",
+    "    debug_default_attention_dir = run_dir / \"debug\" / \"attention\"\n",
+    "debug_default_attention_dir.mkdir(parents=True, exist_ok=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "af2e9c6a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if debug_forward_options.get(\"enabled\", False):\n",
+    "    sample_index = int(debug_forward_options.get(\"sample_index\", 0))\n",
+    "    sample_index = max(0, min(sample_index, src_ids.size(0) - 1))\n",
+    "    debug_result = debug_transformer_forward(\n",
+    "        model=model,\n",
+    "        tokenizer=tokenizer,\n",
+    "        src_ids=src_ids,\n",
+    "        tgt_ids=tgt_ids,\n",
+    "        pad_id=scfg.tokenizer.pad_id,\n",
+    "        device=device,\n",
+    "        batch_index=0,\n",
+    "        sample_index=sample_index,\n",
+    "        logits=None,\n",
+    "        show_attention=debug_forward_options.get(\"show_attention\", False),\n",
+    "        attention_layers=debug_forward_options.get(\"attention_layers\"),\n",
+    "        attention_heads=debug_forward_options.get(\"attention_heads\"),\n",
+    "        attention_types=debug_forward_options.get(\n",
+    "            \"attention_types\", (\"enc_self\", \"dec_self\", \"dec_cross\")\n",
+    "        ),\n",
+    "        average_heads=debug_forward_options.get(\"average_heads\", False),\n",
+    "        save_attention=debug_forward_options.get(\"save_attention\", False),\n",
+    "        save_dir=debug_default_attention_dir,\n",
+    "        run_dir=run_dir,\n",
+    "        attention_figsize=debug_forward_options.get(\"attention_figsize\", (2.4, 2.4)),\n",
+    "        skip_special_tokens=debug_forward_options.get(\"skip_special_tokens\", False),\n",
+    "        return_maps=debug_forward_options.get(\"return_maps\", False),\n",
+    "    )\n",
+    "    print(\"[debug] transformer forward summary saved to\", debug_result.get(\"attention\", {}))\n",
+    "else:\n",
+    "    debug_result = None"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "439d4094",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.to(device)\n",
+    "model.enable_gradient_checkpointing(cfg.trainer.gradient_checkpointing)\n",
+    "\n",
+    "optimizer.zero_grad(set_to_none=True)\n",
+    "\n",
+    "optimizer_step = completed_optimizer_steps\n",
+    "current_resume_batch = resume_batch_idx\n",
+    "current_resume_running_loss = resume_running_loss\n",
+    "\n",
+    "for epoch in range(start_epoch, scfg.trainer.epochs):\n",
+    "    model.train()\n",
+    "    running_loss = current_resume_running_loss if epoch == start_epoch else 0.0\n",
+    "    batches_processed = 0\n",
+    "    last_batch_idx = -1\n",
+    "    for batch_idx, (src, tgt) in enumerate(\n",
+    "        tqdm(train_dataloader, desc=f\"Epoch {epoch + 1} [train]\")\n",
+    "    ):\n",
+    "        if epoch == start_epoch and batch_idx < current_resume_batch:\n",
+    "            continue\n",
+    "        if epoch == start_epoch and current_resume_batch and batch_idx == current_resume_batch:\n",
+    "            current_resume_batch = 0\n",
+    "            current_resume_running_loss = 0.0\n",
+    "\n",
+    "        batches_processed += 1\n",
+    "        last_batch_idx = batch_idx\n",
+    "\n",
+    "        encoded_src = tokenizer(\n",
+    "            src,\n",
+    "            padding=True,\n",
+    "            truncation=True,\n",
+    "            return_tensors=\"pt\",\n",
+    "            add_special_tokens=False,\n",
+    "        )\n",
+    "        encoded_tgt = tokenizer(\n",
+    "            tgt,\n",
+    "            padding=True,\n",
+    "            truncation=True,\n",
+    "            return_tensors=\"pt\",\n",
+    "        )\n",
+    "        tgt_ids = encoded_tgt[\"input_ids\"].to(device)\n",
+    "        decoder_in = tgt_ids[:, :-1]\n",
+    "        labels = tgt_ids[:, 1:]\n",
+    "\n",
+    "        src_ids = encoded_src[\"input_ids\"].to(device)\n",
+    "        src_padd_mask = (encoded_src[\"attention_mask\"] == 0).to(device)\n",
+    "        tgt_padd_mask = decoder_in.eq(scfg.tokenizer.pad_id).to(device)\n",
+    "\n",
+    "        with autocast_ctx:\n",
+    "            logits = model(src_ids, decoder_in, src_padd_mask, tgt_padd_mask)\n",
+    "            loss = criterion(logits.reshape(-1, logits.size(-1)), labels.reshape(-1))\n",
+    "\n",
+    "        running_loss += loss.item()\n",
+    "        loss = loss / scfg.trainer.gradient_accumulation\n",
+    "        scaled_loss = scaler.scale(loss)\n",
+    "        scaled_loss.backward()\n",
+    "\n",
+    "        if (batch_idx + 1) % scfg.trainer.gradient_accumulation == 0:\n",
+    "            scaler.unscale_(optimizer)\n",
+    "            clip_grad_norm_(model.parameters(), cfg.trainer.clip_grad_norm)\n",
+    "            scaler.step(optimizer)\n",
+    "            scaler.update()\n",
+    "            optimizer.zero_grad(set_to_none=True)\n",
+    "            scheduler.step()\n",
+    "            optimizer_step += 1\n",
+    "            save_checkpoint(\n",
+    "                epoch,\n",
+    "                optimizer_step,\n",
+    "                batch_idx=batch_idx + 1,\n",
+    "                running_loss_value=running_loss,\n",
+    "            )\n",
+    "\n",
+    "    if batches_processed and (batches_processed % scfg.trainer.gradient_accumulation != 0):\n",
+    "        scaler.unscale_(optimizer)\n",
+    "        clip_grad_norm_(model.parameters(), cfg.trainer.clip_grad_norm)\n",
+    "        scaler.step(optimizer)\n",
+    "        scaler.update()\n",
+    "        optimizer.zero_grad(set_to_none=True)\n",
+    "        scheduler.step()\n",
+    "        optimizer_step += 1\n",
+    "        save_checkpoint(\n",
+    "            epoch,\n",
+    "            optimizer_step,\n",
+    "            batch_idx=last_batch_idx + 1,\n",
+    "            running_loss_value=running_loss,\n",
+    "        )\n",
+    "\n",
+    "    avg_train_loss = running_loss / len(train_dataloader)\n",
+    "    print(f\"Epoch [{epoch + 1}/{scfg.trainer.epochs}] Train Loss: {avg_train_loss:.4f}\")\n",
+    "\n",
+    "    train_losses.append(avg_train_loss)\n",
+    "    last_lr = scheduler.get_last_lr()\n",
+    "    learning_rates.append(float(last_lr[0]) if last_lr else optimizer.param_groups[0][\"lr\"])\n",
+    "    latest_val_loss: float | None = None\n",
+    "    if epoch % scfg.trainer.eval_interval == 0:\n",
+    "        model.eval()\n",
+    "        val_loss = 0.0\n",
+    "        with torch.no_grad(), autocast_ctx:\n",
+    "            for src, tgt in tqdm(validation_dataloader, desc=f\"Epoch {epoch + 1} [val]\"):\n",
+    "                encoded_src = tokenizer(\n",
+    "                    src,\n",
+    "                    padding=True,\n",
+    "                    truncation=True,\n",
+    "                    return_tensors=\"pt\",\n",
+    "                    add_special_tokens=False,\n",
+    "                )\n",
+    "                encoded_tgt = tokenizer(\n",
+    "                    tgt,\n",
+    "                    padding=True,\n",
+    "                    truncation=True,\n",
+    "                    return_tensors=\"pt\",\n",
+    "                )\n",
+    "\n",
+    "                src_ids = encoded_src[\"input_ids\"].to(device)\n",
+    "                src_padd_mask = (encoded_src[\"attention_mask\"] == 0).to(device)\n",
+    "                tgt_ids = encoded_tgt[\"input_ids\"].to(device)\n",
+    "                decoder_in = tgt_ids[:, :-1]\n",
+    "                labels = tgt_ids[:, 1:]\n",
+    "                tgt_padd_mask = decoder_in.eq(scfg.tokenizer.pad_id)\n",
+    "\n",
+    "                logits = model(src_ids, decoder_in, src_padd_mask, tgt_padd_mask)\n",
+    "                loss = criterion(logits.reshape(-1, logits.size(-1)), labels.reshape(-1))\n",
+    "                val_loss += loss.item()\n",
+    "\n",
+    "        avg_val_loss = val_loss / max(1, len(validation_dataloader))\n",
+    "        print(f\"Epoch [{epoch + 1}/{scfg.trainer.epochs}] Val Loss: {avg_val_loss:.4f}\")\n",
+    "        val_losses.append(avg_val_loss)\n",
+    "        latest_val_loss = avg_val_loss\n",
+    "\n",
+    "        model.train()\n",
+    "\n",
+    "    is_new_best = latest_val_loss is not None and latest_val_loss < best_val_loss\n",
+    "    if is_new_best:\n",
+    "        best_val_loss = latest_val_loss\n",
+    "        best_val_epoch = epoch + 1\n",
+    "\n",
+    "    epoch_metadata = {\"train_loss\": avg_train_loss}\n",
+    "    if latest_val_loss is not None:\n",
+    "        epoch_metadata[\"val_loss\"] = latest_val_loss\n",
+    "\n",
+    "    latest_payload = save_checkpoint(\n",
+    "        epoch + 1,\n",
+    "        optimizer_step,\n",
+    "        batch_idx=0,\n",
+    "        running_loss_value=0.0,\n",
+    "        metadata=epoch_metadata,\n",
+    "    )\n",
+    "\n",
+    "    if is_new_best and best_checkpoint_path is not None:\n",
+    "        torch.save(latest_payload, best_checkpoint_path)\n",
+    "        _update_manifest(\n",
+    "            \"best\",\n",
+    "            {\n",
+    "                \"path\": str(best_checkpoint_path),\n",
+    "                \"epoch\": best_val_epoch,\n",
+    "                \"optimizer_step\": optimizer_step,\n",
+    "                \"train_loss\": avg_train_loss,\n",
+    "                \"val_loss\": best_val_loss,\n",
+    "                \"metadata\": latest_payload.get(\"metadata\", {}),\n",
+    "                \"timestamp\": datetime.now().isoformat(),\n",
+    "            },\n",
+    "        )\n",
+    "        print(\n",
+    "            f\"New best checkpoint saved at {best_checkpoint_path} \"\n",
+    "            f\"(epoch {best_val_epoch}, val {best_val_loss:.4f}, train {avg_train_loss:.4f})\"\n",
+    "        )\n",
+    "\n",
+    "    current_resume_batch = 0\n",
+    "    current_resume_running_loss = 0.0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d55a574f",
+   "metadata": {},
+   "source": [
+    "Exporting"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad75105a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"if scfg.trainer.export_best_to_model_path and best_checkpoint_export_target:\n",
+    "    if best_checkpoint_path.exists():\n",
+    "        best_checkpoint_export_target.parent.mkdir(parents=True, exist_ok=True)\n",
+    "        shutil.copy2(best_checkpoint_path, best_checkpoint_export_target)\n",
+    "        print(f\"[export] Best checkpoint copied to {best_checkpoint_export_target}\")\n",
+    "    else:\n",
+    "        print(f\"[warn] No best checkpoint found at {best_checkpoint_path} to export.\")\"\"\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4b21a8a3",
+   "metadata": {},
+   "source": [
+    "Plotting"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6875ddbe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plots_dir = run_dir / \"plots\"\n",
+    "plots_dir.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "if train_losses:\n",
+    "    epochs = list(range(1, len(train_losses) + 1))\n",
+    "    fig, ax = plt.subplots(figsize=(8, 5))\n",
+    "    ax.plot(epochs, train_losses, label=\"Training Loss\")\n",
+    "    if val_losses:\n",
+    "        val_epochs = list(range(1, len(val_losses) + 1))\n",
+    "        ax.plot(val_epochs, val_losses, label=\"Validation Loss\")\n",
+    "    ax.set_xlabel(\"Epochs\")\n",
+    "    ax.set_ylabel(\"Loss\")\n",
+    "    ax.set_title(\"Training, Validation Loss & Learning Rate\")\n",
+    "    lines = ax.get_lines()\n",
+    "    labels = [line.get_label() for line in lines]\n",
+    "    if learning_rates:\n",
+    "        lr_epochs = epochs[: len(learning_rates)]\n",
+    "        ax2 = ax.twinx()\n",
+    "        ax2.plot(\n",
+    "            lr_epochs, learning_rates, label=\"Learning Rate\", color=\"tab:green\", linestyle=\"--\"\n",
+    "        )\n",
+    "        ax2.set_ylabel(\"Learning Rate\")\n",
+    "        lr_lines = ax2.get_lines()\n",
+    "        lines = list(lines) + list(lr_lines)\n",
+    "        labels += [line.get_label() for line in lr_lines]\n",
+    "    ax.legend(lines, labels, loc=\"upper right\")\n",
+    "    plot_path = plots_dir / \"loss_curves.png\"\n",
+    "    fig.savefig(plot_path, dpi=150, bbox_inches=\"tight\")\n",
+    "    print(f\"[saved] Loss curves -> {plot_path}\")\n",
+    "    plt.show()\n",
+    "else:\n",
+    "    print(\"No training history available to plot.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f8652b1c",
+   "metadata": {},
+   "source": [
+    "Inference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "45a55eaf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"from scripts.inference import inference\n",
+    "GlobalHydra.instance().clear()\n",
+    "initialize(config_path=\"./configs\", version_base=None)\n",
+    "cfg = compose(config_name=\"infer_mode\")\n",
+    "print(inference(cfg))\"\"\""
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "transformer",
+   "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.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/transformer/__init__.py

@@ -1,4 +1,3 @@
-from .configs import BasicEncDecCfg, InferAppCfg, TrainAppCfg
 from .transformer import BasicEncoderDecoderTransformer
 
-__all__ = ["BasicEncoderDecoderTransformer", "BasicEncDecCfg", "TrainAppCfg", "InferAppCfg"]
+__all__ = ["BasicEncoderDecoderTransformer"]

src/transformer/configs.py

@@ -1,78 +1,142 @@
+"""Typed dataclass containers mirroring the YAML configuration structure."""
+
+from __future__ import annotations
+
 from dataclasses import dataclass
 
+__all__ = [
+    "ModelCfg",
+    "DatasetCfg",
+    "OptimCfg",
+    "SchedCfg",
+    "TokenizerCfg",
+    "TrainerCfg",
+    "RuntimeCfg",
+    "GenerationCfg",
+    "TrainAppCfg",
+    "InferAppCfg",
+]
+
 
 @dataclass
-class BasicEncDecCfg:
-    vocab_size: int
+class ModelCfg:
+    name: str
+    best_checkpoint_path: str
+    latest_checkpoint_path: str
+    tokenizer: str
     d_model: int
-    max_seq_len: int
+    num_layers: int
     num_heads: int
     d_ff: int
-    num_layers: int
-    dropout_rate: float = 0.1
-    pad_id: int = 0
-    bos_id: int = 2
-    eos_id: int = 1
+    dropout_rate: float
+    max_seq_len: int
+    vocab_size: int
+    pad_id: int
+    bos_id: int
+    eos_id: int
 
 
 @dataclass
-class DataCfg:
-    train_path: str
-    val_path: str
-    num_workers: int = 4
-    batch_size: int = 32
-    max_length: int = 512
+class DatasetCfg:
+    name: str
+    path: str
 
 
 @dataclass
 class OptimCfg:
-    lr: float = 3e-4
-    betas: tuple[float, float] = (0.9, 0.95)
-    weight_decay: float = 0.1
-    eps: float = 1e-8
+    name: str
+    lr: float
+    betas: tuple[float, float]
+    weight_decay: float
+    eps: float
+    no_decay_on_bias_norm_embed: bool
 
 
 @dataclass
 class SchedCfg:
-    type: str = "cosine"
-    warmup_steps: int = 2000
-    total_steps: int = 200000
-    min_lr: float = 1e-5
+    name: str
+    warmup_ratio: float
+    min_lr_ratio: float
+    max_lr_ratio: float
+    hold_min: bool
+
+
+@dataclass
+class TokenizerCfg:
+    name: str
+    path: str
+    corpus: str
+    dataset: str
+    vocab_size: int
+    max_seq_len: int
+    pad_token: str
+    bos_token: str
+    eos_token: str
+    unk_token: str
+    special_tokens: list[str]
+    pad_id: int
+    bos_id: int
+    eos_id: int
+    unk_id: int
 
 
 @dataclass
 class TrainerCfg:
-    epochs: int = 10
-    gradient_accumulation: int = 1
-    eval_interval: int = 1000
-    save_dir: str = "outputs/run"
+    epochs: int
+    batch_size: int
+    shuffle: bool
+    precision: str
+    gradient_accumulation: int
+    clip_grad_norm: float
+    eval_interval: int
+    save_interval: int
+    debug_interval: int
+    mixed_precision: bool
+    gradient_checkpointing: bool
+    label_smoothing: float
+    export_best_to_model_path: bool
+    resume: str
 
 
 @dataclass
 class RuntimeCfg:
-    tokenizer_name: str = "t5-small"
-    weights_path: str = ""
-    artifact_dir: str = "outputs/run"
+    seed: int
+    device: str
+    output_dir: str
+    data_dir: str
+    cache_dir: str
+    tokenizer_dir: str
+    checkpoint_path: str
+
+
+@dataclass
+class GenerationCfg:
+    max_new_tokens: int
+    temperature: float
+    top_k: int | None  # None disables top-k
+    top_p: float | None
+    do_sample: bool
+    presence_penalty: float
+    frequency_penalty: float
+    no_repeat_ngram: int | None
+    min_steps_before_eos: int
 
 
 @dataclass
 class TrainAppCfg:
-    seed: int
-    model: BasicEncDecCfg
-    data: DataCfg
+    model: ModelCfg
+    dataset: DatasetCfg
     optim: OptimCfg
-    scheduler: SchedCfg
+    sched: SchedCfg
     trainer: TrainerCfg
     runtime: RuntimeCfg
+    tokenizer: TokenizerCfg
 
 
 @dataclass
 class InferAppCfg:
-    model: BasicEncDecCfg
+    model: ModelCfg
     runtime: RuntimeCfg
-    input_text: str = ""
-    max_new_tokens: int = 64
-    temperature: float = 1.0
-    top_k: int = 0
-    top_p: float = 1.0
-    do_sample: bool = False
+    tokenizer: TokenizerCfg
+    generation: GenerationCfg
+    input_text: str

src/transformer/modules/attention.py

@@ -1,105 +1,132 @@
-import torch
-import torch.nn as nn
-from torch import Tensor
-
-from transformer.utils import calculate_attention, join_heads, split_heads
-
-
-class MultiHeadAttention(nn.Module):
-    """
-    Multi-head attention: linear projections -> split heads -> scaled dot-product
-    attention (via utils) -> merge heads -> output projection (+ dropout).
-
-    Args:
-        d_model (int): Model dimension (>0). Must be divisible by num_heads.
-        num_heads (int): Number of attention heads (>0).
-        dropout_rate (float): Dropout probability in (0,1).
-
-    Inputs:
-        query, key, value: (B, S, D) with D == d_model
-        mask (optional): Tensor broadcastable to (B, H, S_q, S_k), either boolean
-                         (True = masked) or additive float mask.
-
-    Returns:
-        Tensor: (B, S_q, D)
-    """
-
-    def __init__(self, d_model: int, num_heads: int, dropout_rate: float):
-        super().__init__()
-        # ---- type checks
-        if not isinstance(num_heads, int):
-            raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
-        if not isinstance(d_model, int):
-            raise TypeError(f"d_model must be an int, got {type(d_model)}")
-        if not isinstance(dropout_rate, float):
-            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
-
-        # ---- value checks
-        if num_heads <= 0:
-            raise ValueError(f"num_heads must be strictly greater than 0, got {num_heads}")
-        if d_model <= 0:
-            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
-        if d_model % num_heads != 0:
-            raise ValueError("d_model must be divisible by num_heads")
-        if not (0 <= dropout_rate <= 1):
-            raise ValueError(f"dropout_rate must be between 0 and 1 included, got {dropout_rate}")
-
-        self.d_model = d_model
-        self.num_heads = num_heads
-        self.d_head = d_model // num_heads
-        self.dropout_rate = dropout_rate
-
-        self.query_linear = nn.Linear(d_model, d_model, bias=False)
-        self.key_linear = nn.Linear(d_model, d_model, bias=False)
-        self.value_linear = nn.Linear(d_model, d_model, bias=False)
-        self.output_linear = nn.Linear(d_model, d_model, bias=True)
-        self.dropout = nn.Dropout(dropout_rate)
-
-    def forward(
-        self,
-        query: Tensor,
-        key: Tensor,
-        value: Tensor,
-        mask: Tensor | None,
-    ) -> Tensor:
-        # ---- basic type checks
-        if not isinstance(query, torch.Tensor):
-            raise TypeError(f"query must be a torch.Tensor, got {type(query)}")
-        if not isinstance(key, torch.Tensor):
-            raise TypeError(f"key must be a torch.Tensor, got {type(key)}")
-        if not isinstance(value, torch.Tensor):
-            raise TypeError(f"value must be a torch.Tensor, got {type(value)}")
-        if mask is not None and not isinstance(mask, torch.Tensor):
-            raise TypeError(f"mask must be a torch.Tensor or None, got {type(mask)}")
-
-        # ---- shape checks for q/k/v (3D [B,S,D] and D == d_model)
-        if query.dim() != 3 or key.dim() != 3 or value.dim() != 3:
-            raise ValueError(
-                "query/key/value must be 3D tensors of shape (B, S, D); "
-                f"got q={tuple(query.shape)}, k={tuple(key.shape)}, v={tuple(value.shape)}"
-            )
-        Bq, Sq, Dq = query.shape
-        Bk, Sk, Dk = key.shape
-        Bv, Sv, Dv = value.shape
-
-        if not (Dq == Dk == Dv == self.d_model):
-            raise ValueError(
-                f"Last dimension must equal d_model={self.d_model}; got Dq={Dq}, Dk={Dk}, Dv={Dv}"
-            )
-        if not (Bq == Bk == Bv):
-            raise ValueError(f"Batch size mismatch: q={Bq}, k={Bk}, v={Bv}")
-        if Sk != Sv:
-            raise ValueError(f"Key/Value seq length mismatch: Sk={Sk} vs Sv={Sv}")
-
-        # ---- project and split into heads -> (B, H, S, Dh)
-        q = split_heads(self.query_linear(query), self.num_heads)
-        k = split_heads(self.key_linear(key), self.num_heads)
-        v = split_heads(self.value_linear(value), self.num_heads)
-
-        # ---- attention (utils handles mask broadcasting/device and numeric stability)
-        p = self.dropout_rate if self.training else 0.0
-        attn = calculate_attention(q, k, v, mask, attn_dropout_p=p)  # (B, H, Sq, Dh)
-
-        # ---- merge heads, project out, dropout
-        out = join_heads(attn)  # (B, Sq, D)
-        return self.dropout(self.output_linear(out))
+"""Multi-head attention building block with explicit validation helpers."""
+
+from __future__ import annotations
+
+from typing import cast
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from transformer.utils import (
+    calculate_attention,
+    combine_masks,
+    create_qk_padding_mask,
+    join_heads,
+    split_heads,
+)
+
+__all__ = ["MultiHeadAttention"]
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-head attention: linear projections -> split heads -> scaled dot-product
+    attention (via utils) -> merge heads -> output projection (+ dropout).
+
+    Args:
+        d_model (int): Model dimension (>0). Must be divisible by num_heads.
+        num_heads (int): Number of attention heads (>0).
+        dropout_rate (float): Dropout probability in (0,1).
+
+    Inputs:
+        query, key, value: (B, S, D) with D == d_model
+        mask (optional): Boolean tensor broadcastable to (B, H, S_q, S_k) where True entries are masked.
+
+    Returns:
+        Tensor: (B, S_q, D)
+    """
+
+    def __init__(self, d_model: int, num_heads: int, dropout_rate: float):
+        super().__init__()
+        # ---- type checks
+        if not isinstance(num_heads, int):
+            raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be an int, got {type(d_model)}")
+        if not isinstance(dropout_rate, float):
+            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
+
+        # ---- value checks
+        if num_heads <= 0:
+            raise ValueError(f"num_heads must be strictly greater than 0, got {num_heads}")
+        if d_model <= 0:
+            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
+        if d_model % num_heads != 0:
+            raise ValueError("d_model must be divisible by num_heads")
+        if not (0 <= dropout_rate < 1):
+            raise ValueError(f"dropout_rate must be between 0 and 1 excluded, got {dropout_rate}")
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.d_head = d_model // num_heads
+        self.dropout_rate = dropout_rate
+
+        self.query_linear = nn.Linear(d_model, d_model, bias=False)
+        self.key_linear = nn.Linear(d_model, d_model, bias=False)
+        self.value_linear = nn.Linear(d_model, d_model, bias=False)
+        self.output_linear = nn.Linear(d_model, d_model, bias=True)
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        q_mask: Tensor,
+        k_mask: Tensor,
+        causal_mask: Tensor | None = None,
+    ) -> Tensor:
+        """Run multi-head attention with boolean padding and optional causal masks."""
+
+        # ---- basic type checks
+        if not isinstance(query, torch.Tensor):
+            raise TypeError(f"query must be a torch.Tensor, got {type(query)}")
+        if not isinstance(key, torch.Tensor):
+            raise TypeError(f"key must be a torch.Tensor, got {type(key)}")
+        if not isinstance(value, torch.Tensor):
+            raise TypeError(f"value must be a torch.Tensor, got {type(value)}")
+        if not isinstance(q_mask, torch.Tensor):
+            raise TypeError(f"q_mask must be a torch.Tensor, got {type(q_mask)}")
+        if k_mask is None or not isinstance(k_mask, torch.Tensor):
+            raise TypeError(f"k_mask must be a torch.Tensor, got {type(k_mask)}")
+
+        if query.dim() != 3 or key.dim() != 3 or value.dim() != 3:
+            raise ValueError(
+                "query/key/value must be 3D tensors of shape (B, S, D); "
+                f"got q={tuple(query.shape)}, k={tuple(key.shape)}, v={tuple(value.shape)}"
+            )
+        Bq, Sq, Dq = query.shape
+        Bk, Sk, Dk = key.shape
+        Bv, Sv, Dv = value.shape
+
+        if not (Dq == Dk == Dv == self.d_model):
+            raise ValueError(
+                f"Last dimension must equal d_model={self.d_model}; got Dq={Dq}, Dk={Dk}, Dv={Dv}"
+            )
+        if not (Bq == Bk == Bv):
+            raise ValueError(f"Batch size mismatch: q={Bq}, k={Bk}, v={Bv}")
+        if Sk != Sv:
+            raise ValueError(f"Key/Value seq length mismatch: Sk={Sk} vs Sv={Sv}")
+
+        # ---- padding mask validation ----
+        if q_mask.dtype != torch.bool:
+            raise TypeError(f"q_mask must be boolean, got {q_mask.dtype}")
+        if k_mask.dtype != torch.bool:
+            raise TypeError(f"k_mask must be boolean, got {k_mask.dtype}")
+        if q_mask.dim() != 4 or k_mask.dim() != 4:
+            raise ValueError(
+                "q_mask and k_mask must be 4D tensors shaped (B, H, 1, S); "
+                f"got {tuple(q_mask.shape)} and {tuple(k_mask.shape)}"
+            )
+
+        # ---- project and split into heads -> (B, H, S, Dh)
+        q = split_heads(self.query_linear(query), self.num_heads)
+        k = split_heads(self.key_linear(key), self.num_heads)
+        v = split_heads(self.value_linear(value), self.num_heads)
+        pad_mask = create_qk_padding_mask(q_mask, k_mask)
+        combined_mask = combine_masks(pad_mask, causal_mask)
+
+        p = self.dropout_rate if self.training else 0.0
+        attn = cast(Tensor, calculate_attention(q, k, v, combined_mask, dropout_p=p))
+        out = join_heads(attn)  # (B, Sq, D)
+        return self.output_linear(out)

src/transformer/modules/decoder.py

@@ -1,99 +1,171 @@
-import torch
-import torch.nn as nn
-from torch import Tensor
-
-from transformer.modules.attention import MultiHeadAttention
-from transformer.modules.feedforward import FeedForwardLayer
-from transformer.utils import combine_masks
-
-
-class DecoderLayer(nn.Module):
-    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout_rate: float):
-        super().__init__()
-
-        if not isinstance(d_model, int):
-            raise TypeError(f"d_model must be an int, got {type(d_model)}")
-        if not isinstance(num_heads, int):
-            raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
-        if not isinstance(d_ff, int):
-            raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
-        if not isinstance(dropout_rate, float):
-            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
-
-        if not d_model > 0:
-            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
-        if not num_heads > 0:
-            raise ValueError(f"num_heads must be strictly greater than 0, got {num_heads}")
-        if not d_ff > 0:
-            raise ValueError(f"d_ff must be strictly greater than 0, got {d_ff}")
-        if not (0.0 <= dropout_rate < 1.0):
-            raise ValueError(f"dropout_rate must be between 0 and 1 excluded, got {dropout_rate}")
-
-        self.self_attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
-        self.feed_forward = FeedForwardLayer(d_model, d_ff)
-        self.cross_attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
-        self.norm1 = nn.LayerNorm(d_model)
-        self.dropout1 = nn.Dropout(dropout_rate)
-        self.norm2 = nn.LayerNorm(d_model)
-        self.dropout2 = nn.Dropout(dropout_rate)
-        self.norm3 = nn.LayerNorm(d_model)
-        self.dropout3 = nn.Dropout(dropout_rate)
-
-    def forward(
-        self,
-        x: Tensor,
-        y: Tensor,
-        src_padding_mask: Tensor | None,
-        tgt_padding_mask: Tensor | None,
-        tgt_causal_mask: Tensor | None,
-    ) -> Tensor:
-        if not isinstance(x, torch.Tensor):
-            raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
-        if not isinstance(y, torch.Tensor):
-            raise TypeError(f"y must be a torch.Tensor, got {type(y)}")
-        if not (x.dim() == 3):
-            raise ValueError(
-                f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
-            )
-        if not (y.dim() == 3):
-            raise ValueError(
-                f"y must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(y.shape)}"
-            )
-        if not (x.shape[0] == y.shape[0] and x.shape[-1] == y.shape[-1]):
-            raise ValueError(
-                "Batch size or d_model mismatch between encoder memory and decoder input"
-            )
-
-        tgt_mask = combine_masks(tgt_padding_mask, tgt_causal_mask)
-        y = self.norm1(y + self.dropout1(self.self_attention_layer(y, y, y, tgt_mask)))
-        y = self.norm2(y + self.dropout2(self.cross_attention_layer(y, x, x, src_padding_mask)))
-        y = self.norm3(y + self.dropout3(self.feed_forward(y)))
-        return y
-
-
-class TransformerDecoder(nn.Module):
-    def __init__(
-        self, d_model: int, num_heads: int, d_ff: int, num_layers: int, dropout_rate: float
-    ):
-        super().__init__()
-
-        if not isinstance(num_layers, int):
-            raise TypeError(f"num_layers must be an int, got {type(num_layers)}")
-        if not num_layers > 0:
-            raise ValueError(f"num_layers must be strictly greater than 0, got {num_layers}")
-
-        self.layers = nn.ModuleList(
-            [DecoderLayer(d_model, num_heads, d_ff, dropout_rate) for _ in range(num_layers)]
-        )
-
-    def forward(
-        self,
-        x: Tensor,
-        y: Tensor,
-        src_padding_mask: Tensor | None,
-        tgt_padding_mask: Tensor | None,
-        tgt_causal_mask: Tensor | None,
-    ) -> Tensor:
-        for layer in self.layers:
-            y = layer(x, y, src_padding_mask, tgt_padding_mask, tgt_causal_mask)
-        return y
+"""Decoder stack used by the encoder-decoder transformer."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as ckpt
+from torch import Tensor
+
+from transformer.modules.attention import MultiHeadAttention
+from transformer.modules.feedforward import FeedForwardLayer
+
+__all__ = ["DecoderLayer", "TransformerDecoder"]
+
+
+class DecoderLayer(nn.Module):
+    """Single decoder block with self-attention, cross-attention, and feed-forward."""
+
+    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout_rate: float) -> None:
+        super().__init__()
+
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be an int, got {type(d_model)}")
+        if not isinstance(num_heads, int):
+            raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
+        if not isinstance(d_ff, int):
+            raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
+        if not isinstance(dropout_rate, float):
+            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
+
+        if d_model <= 0:
+            raise ValueError("d_model must be strictly greater than 0")
+        if num_heads <= 0:
+            raise ValueError("num_heads must be strictly greater than 0")
+        if d_ff <= 0:
+            raise ValueError("d_ff must be strictly greater than 0")
+        if not 0.0 <= dropout_rate < 1.0:
+            raise ValueError("dropout_rate must be in [0, 1)")
+
+        self.self_attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
+        self.cross_attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
+        self.feed_forward = FeedForwardLayer(d_model, d_ff, dropout_rate)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout_rate)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout2 = nn.Dropout(dropout_rate)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout3 = nn.Dropout(dropout_rate)
+
+    def forward(
+        self,
+        x: Tensor,
+        y: Tensor,
+        src_padding_mask: Tensor,
+        tgt_padding_mask: Tensor,
+        tgt_causal_mask: Tensor | None,
+    ) -> Tensor:
+        """Run one decoder layer using post-layer-normalisation formulation."""
+
+        if not isinstance(x, torch.Tensor):
+            raise TypeError("x must be a torch.Tensor")
+        if x.dim() != 3:
+            raise ValueError(
+                f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
+            )
+        if not isinstance(y, torch.Tensor):
+            raise TypeError("y must be a torch.Tensor")
+        if y.dim() != 3:
+            raise ValueError(
+                f"y must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(y.shape)}"
+            )
+        if x.shape[0] != y.shape[0] or x.shape[-1] != y.shape[-1]:
+            raise ValueError("Encoder memory and decoder input must match in batch and d_model")
+
+        for mask_name, mask in (
+            ("src_padding_mask", src_padding_mask),
+            ("tgt_padding_mask", tgt_padding_mask),
+        ):
+            if not isinstance(mask, torch.Tensor):
+                raise TypeError(f"{mask_name} must be a torch.Tensor")
+            if mask.dtype != torch.bool or mask.dim() != 4:
+                raise TypeError(
+                    f"{mask_name} must be boolean with shape (B, H, 1, S);"
+                    f" got dtype {mask.dtype} and shape {tuple(mask.shape)}"
+                )
+
+        # Self-attention
+        self_attn_out = self.self_attention_layer(
+            y,
+            y,
+            y,
+            tgt_padding_mask,
+            tgt_padding_mask,
+            tgt_causal_mask,
+        )
+        y = self.norm1(y + self.dropout1(self_attn_out))
+
+        # Cross-attention (encoder memory as keys/values)
+        cross_attn_out = self.cross_attention_layer(
+            y,
+            x,
+            x,
+            tgt_padding_mask,
+            src_padding_mask,
+        )
+        y = self.norm2(y + self.dropout2(cross_attn_out))
+
+        # Feed-forward block
+        ff_out = self.feed_forward(y)
+        return self.norm3(y + self.dropout3(ff_out))
+
+
+class TransformerDecoder(nn.Module):
+    """Stack of decoder layers with optional activation checkpointing."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        d_ff: int,
+        num_layers: int,
+        dropout_rate: float,
+    ) -> None:
+        super().__init__()
+
+        if not isinstance(num_layers, int):
+            raise TypeError(f"num_layers must be an int, got {type(num_layers)}")
+        if num_layers <= 0:
+            raise ValueError("num_layers must be strictly greater than 0")
+
+        self.layers = nn.ModuleList(
+            [DecoderLayer(d_model, num_heads, d_ff, dropout_rate) for _ in range(num_layers)]
+        )
+        self.use_ckpt = False
+
+    def forward(
+        self,
+        x: Tensor,
+        y: Tensor,
+        src_padding_mask: Tensor,
+        tgt_padding_mask: Tensor,
+        tgt_causal_mask: Tensor | None,
+    ) -> Tensor:
+        """Run the decoder stack for all time steps."""
+
+        if not isinstance(x, torch.Tensor):
+            raise TypeError("x must be a torch.Tensor")
+        if x.dim() != 3:
+            raise ValueError(
+                f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
+            )
+        if not isinstance(y, torch.Tensor):
+            raise TypeError("y must be a torch.Tensor")
+        if y.dim() != 3:
+            raise ValueError(
+                f"y must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(y.shape)}"
+            )
+        if x.shape[0] != y.shape[0] or x.shape[-1] != y.shape[-1]:
+            raise ValueError("Encoder memory and decoder input must match in batch and d_model")
+
+        for layer in self.layers:
+            if self.use_ckpt:
+
+                def _fn(y_, *, _layer=layer):
+                    return _layer(x, y_, src_padding_mask, tgt_padding_mask, tgt_causal_mask)
+
+                y = ckpt.checkpoint(_fn, y, use_reentrant=False)
+            else:
+                y = layer(x, y, src_padding_mask, tgt_padding_mask, tgt_causal_mask)
+        return y

src/transformer/modules/embedding.py

@@ -1,121 +1,146 @@
-import torch
-import torch.nn as nn
-from torch import Tensor
-
-from transformer.utils import sinusoidal_positional_encoding
-
-
-class InputEmbedding(nn.Module):
-    """
-    Token + positional embedding module.
-
-    Args:
-        vocab_size (int): Size of vocabulary (>0).
-        d_model (int): Embedding dimension (>0).
-        sequence_length (int): Maximum sequence length (≥0).
-        pad_id (int): Padding token ID (0 <= pad_id < vocab_size).
-
-    Input:
-        x (LongTensor): shape (batch_size, seq_len), token IDs.
-
-    Output:
-        Tensor: shape (batch_size, seq_len, d_model), embeddings.
-
-    Notes:
-        - Zero-length sequences are allowed (returns [B, 0, D]).
-        - Positional encodings are sinusoidal and added to token embeddings.
-    """
-
-    def __init__(
-        self,
-        vocab_size: int,
-        d_model: int,
-        sequence_length: int,
-        pad_id: int,
-        dropout_rate: float = 0.0,
-    ):
-        super().__init__()
-        # --- type checks
-        if not isinstance(vocab_size, int):
-            raise TypeError(f"vocab_size must be int, got {type(vocab_size)}")
-        if not isinstance(d_model, int):
-            raise TypeError(f"d_model must be int, got {type(d_model)}")
-        if not isinstance(sequence_length, int):
-            raise TypeError(f"sequence_length must be int, got {type(sequence_length)}")
-        if not isinstance(pad_id, int):
-            raise TypeError(f"pad_id must be int, got {type(pad_id)}")
-        if not isinstance(dropout_rate, float):
-            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
-
-        # --- value checks
-        if vocab_size <= 0:
-            raise ValueError(f"vocab_size must be > 0, got {vocab_size}")
-        if d_model <= 0:
-            raise ValueError(f"d_model must be > 0, got {d_model}")
-        if sequence_length < 0:
-            raise ValueError(f"sequence_length must be >= 0, got {sequence_length}")
-        if not (0 <= pad_id < vocab_size):
-            raise ValueError(f"pad_id must be in [0, {vocab_size - 1}], got {pad_id}")
-        if not (0 <= dropout_rate <= 1):
-            raise ValueError(f"dropout_rate must be in [0,1], got {dropout_rate}")
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.sequence_length = sequence_length
-        self.pad_id = pad_id
-        self.dropout_rate = dropout_rate
-
-        self.token_embed = nn.Embedding(vocab_size, d_model, padding_idx=pad_id)
-        self.pos_embed = PositionalEmbedding(sequence_length, d_model)
-        self.dropout = nn.Dropout(dropout_rate)
-
-    def forward(self, x: Tensor) -> Tensor:
-        if not isinstance(x, torch.Tensor):
-            raise TypeError(f"x must be Tensor, got {type(x)}")
-        if x.dtype != torch.long:
-            raise TypeError(f"x must be torch.long, got {x.dtype}")
-        if x.dim() != 2:
-            raise ValueError(f"x must be 2D (B, S), got shape {tuple(x.shape)}")
-
-        tok = self.token_embed(x)  # (B, S, D) — S can be 0
-        return self.dropout(self.pos_embed(tok))  # (B, S, D)
-
-
-class PositionalEmbedding(nn.Module):
-    """
-    Adds sinusoidal positional encodings.
-
-    Args:
-        sequence_length (int): Maximum sequence length (≥0).
-        d_model (int): Embedding dimension (>0).
-
-    Notes:
-        - Zero-length tables are supported (shape [1, 0, D]).
-    """
-
-    def __init__(self, sequence_length: int, d_model: int):
-        super().__init__()
-
-        if sequence_length < 0:
-            raise ValueError(f"sequence_length must be >= 0, got {sequence_length}")
-        if d_model <= 0:
-            raise ValueError(f"d_model must be > 0, got {d_model}")
-
-        self.sequence_length = sequence_length
-        self.d_model = d_model
-
-        pe = sinusoidal_positional_encoding(sequence_length, d_model)  # [S, D]
-        self.register_buffer("pe", pe.unsqueeze(0))  # [1, S, D]
-
-    def forward(self, x: Tensor) -> Tensor:
-        if not isinstance(x, torch.Tensor):
-            raise TypeError(f"x must be Tensor, got {type(x)}")
-        if x.dim() != 3:
-            raise ValueError(f"x must be 3D (B, S, D), got shape {tuple(x.shape)}")
-
-        _, S, D = x.shape
-        if D != self.d_model:
-            raise ValueError(f"d_model mismatch: got {D}, expected {self.d_model}")
-        if S > self.sequence_length:
-            raise ValueError(f"seq_len {S} exceeds max_seq_len {self.sequence_length}")
-
-        return x + self.pe[:, :S].to(dtype=x.dtype, device=x.device)
+"""Embedding modules used by the transformer architecture."""
+
+from __future__ import annotations
+
+from typing import cast
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from transformer.utils import sinusoidal_positional_encoding
+
+__all__ = ["InputEmbedding", "PositionalEmbedding", "LearnedPositionalEmbedding"]
+
+
+class InputEmbedding(nn.Module):
+    """Token and positional embedding with optional dropout."""
+
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int,
+        max_seq_len: int,
+        pad_id: int,
+        dropout_rate: float = 0.0,
+    ) -> None:
+        super().__init__()
+
+        if not isinstance(vocab_size, int):
+            raise TypeError(f"vocab_size must be int, got {type(vocab_size)}")
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be int, got {type(d_model)}")
+        if not isinstance(max_seq_len, int):
+            raise TypeError(f"max_seq_len must be int, got {type(max_seq_len)}")
+        if not isinstance(pad_id, int):
+            raise TypeError(f"pad_id must be int, got {type(pad_id)}")
+        if not isinstance(dropout_rate, float):
+            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
+
+        if vocab_size <= 0:
+            raise ValueError("vocab_size must be > 0")
+        if d_model <= 0:
+            raise ValueError("d_model must be > 0")
+        if max_seq_len <= 0:
+            raise ValueError("max_seq_len must be > 0")
+        if not (0 <= pad_id < vocab_size):
+            raise ValueError(f"pad_id must be in [0, {vocab_size - 1}]")
+        if not 0.0 <= dropout_rate < 1.0:
+            raise ValueError("dropout_rate must be in [0, 1)")
+
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.max_seq_len = max_seq_len
+        self.pad_id = pad_id
+
+        self.token_embed = nn.Embedding(vocab_size, d_model, padding_idx=pad_id)
+        self.pos_embed = PositionalEmbedding(max_seq_len, d_model)
+        self.dropout = nn.Dropout(dropout_rate)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(f"x must be Tensor, got {type(x)}")
+        if x.dtype != torch.long:
+            raise TypeError(f"x must be torch.long, got {x.dtype}")
+        if x.dim() != 2:
+            raise ValueError(f"x must be 2D (B, S), got shape {tuple(x.shape)}")
+        if x.size(1) > self.max_seq_len:
+            raise ValueError(f"Sequence length {x.size(1)} exceeds max_seq_len {self.max_seq_len}")
+
+        tokens = self.token_embed(x) * (self.d_model**0.5)
+        tokens = self.pos_embed(tokens)
+        return self.dropout(tokens)
+
+
+class PositionalEmbedding(nn.Module):
+    """Fixed sinusoidal positional encodings."""
+
+    def __init__(self, max_seq_len: int, d_model: int) -> None:
+        super().__init__()
+
+        if not isinstance(max_seq_len, int):
+            raise TypeError(f"max_seq_len must be int, got {type(max_seq_len)}")
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be int, got {type(d_model)}")
+        if max_seq_len < 0:
+            raise ValueError("max_seq_len must be >= 0")
+        if d_model <= 0:
+            raise ValueError("d_model must be > 0")
+
+        self.max_seq_len = max_seq_len
+        self.d_model = d_model
+
+        pe = sinusoidal_positional_encoding(max_seq_len, d_model)
+        self.register_buffer("pe", pe.unsqueeze(0))
+
+    def forward(self, x: Tensor) -> Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(f"x must be Tensor, got {type(x)}")
+        if x.dim() != 3:
+            raise ValueError(f"x must be 3D (B, S, D), got shape {tuple(x.shape)}")
+
+        _, seq_len, dim = x.shape
+        if dim != self.d_model:
+            raise ValueError(f"d_model mismatch: got {dim}, expected {self.d_model}")
+        if seq_len > self.max_seq_len:
+            raise ValueError(f"seq_len {seq_len} exceeds max_seq_len {self.max_seq_len}")
+
+        pe_buffer = cast(Tensor, self.pe)
+        return x + pe_buffer[:, :seq_len].to(dtype=x.dtype, device=x.device)
+
+
+class LearnedPositionalEmbedding(nn.Module):
+    """Learned positional embeddings compatible with :class:`InputEmbedding`."""
+
+    def __init__(self, max_seq_len: int, d_model: int) -> None:
+        super().__init__()
+
+        if not isinstance(max_seq_len, int):
+            raise TypeError(f"max_seq_len must be int, got {type(max_seq_len)}")
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be int, got {type(d_model)}")
+        if max_seq_len < 0:
+            raise ValueError("max_seq_len must be >= 0")
+        if d_model <= 0:
+            raise ValueError("d_model must be > 0")
+
+        self.max_seq_len = max_seq_len
+        self.d_model = d_model
+        self.pos_table = nn.Embedding(max_seq_len, d_model)
+        nn.init.normal_(self.pos_table.weight, mean=0.0, std=0.02)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(f"x must be Tensor, got {type(x)}")
+        if x.dim() != 3:
+            raise ValueError(f"x must be 3D (B, S, D), got {tuple(x.shape)}")
+
+        batch, seq_len, dim = x.shape
+        if dim != self.d_model:
+            raise ValueError(f"d_model mismatch: got {dim}, expected {self.d_model}")
+        if seq_len > self.max_seq_len:
+            raise ValueError(f"seq_len {seq_len} exceeds max_seq_len {self.max_seq_len}")
+
+        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch, seq_len)
+        pos_emb = self.pos_table(positions)
+        return x + pos_emb

src/transformer/modules/encoder.py

@@ -1,71 +1,108 @@
-import torch
-import torch.nn as nn
-from torch import Tensor
-
-from transformer.modules.attention import MultiHeadAttention
-from transformer.modules.feedforward import FeedForwardLayer
-
-
-class EncoderLayer(nn.Module):
-    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout_rate: float):
-        super().__init__()
-
-        # --- type checks (match project convention)
-        if not isinstance(d_model, int):
-            raise TypeError(f"d_model must be an int, got {type(d_model)}")
-        if not isinstance(num_heads, int):
-            raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
-        if not isinstance(d_ff, int):
-            raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
-        if not isinstance(dropout_rate, float):
-            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
-
-        # --- value checks (match wording)
-        if not d_model > 0:
-            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
-        if not num_heads > 0:
-            raise ValueError(f"num_heads must be strictly greater than 0, got {num_heads}")
-        if not d_ff > 0:
-            raise ValueError(f"d_ff must be strictly greater than 0, got {d_ff}")
-        if not (0.0 <= dropout_rate < 1.0):
-            raise ValueError(f"dropout_rate must be between 0 and 1 excluded, got {dropout_rate}")
-
-        self.attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
-        self.feed_forward = FeedForwardLayer(d_model, d_ff)
-        self.norm1 = nn.LayerNorm(d_model)
-        self.dropout1 = nn.Dropout(dropout_rate)
-        self.norm2 = nn.LayerNorm(d_model)
-        self.dropout2 = nn.Dropout(dropout_rate)
-
-    def forward(self, x: Tensor, src_padding_mask: Tensor | None) -> Tensor:
-        if not isinstance(x, torch.Tensor):
-            raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
-        if not (x.dim() == 3):
-            raise ValueError(
-                f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
-            )
-
-        x = self.norm1(x + self.dropout1(self.attention_layer(x, x, x, src_padding_mask)))
-        x = self.norm2(x + self.dropout2(self.feed_forward(x)))
-        return x
-
-
-class TransformerEncoder(nn.Module):
-    def __init__(
-        self, d_model: int, num_heads: int, d_ff: int, num_layers: int, dropout_rate: float
-    ):
-        super().__init__()
-
-        if not isinstance(num_layers, int):
-            raise TypeError(f"num_layers must be an int, got {type(num_layers)}")
-        if not num_layers > 0:
-            raise ValueError(f"num_layers must be strictly greater than 0, got {num_layers}")
-
-        self.layers = nn.ModuleList(
-            [EncoderLayer(d_model, num_heads, d_ff, dropout_rate) for _ in range(num_layers)]
-        )
-
-    def forward(self, x: Tensor, src_padding_mask: Tensor | None) -> Tensor:
-        for layer in self.layers:
-            x = layer(x, src_padding_mask)
-        return x
+"""Encoder stack used by the transformer architecture."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as ckpt
+from torch import Tensor
+
+from transformer.modules.attention import MultiHeadAttention
+from transformer.modules.feedforward import FeedForwardLayer
+
+__all__ = ["EncoderLayer", "TransformerEncoder"]
+
+
+class EncoderLayer(nn.Module):
+    """Self-attention + feed-forward block with post-layer-normalisation."""
+
+    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout_rate: float) -> None:
+        super().__init__()
+
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be an int, got {type(d_model)}")
+        if not isinstance(num_heads, int):
+            raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
+        if not isinstance(d_ff, int):
+            raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
+        if not isinstance(dropout_rate, float):
+            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
+
+        if d_model <= 0:
+            raise ValueError("d_model must be strictly greater than 0")
+        if num_heads <= 0:
+            raise ValueError("num_heads must be strictly greater than 0")
+        if d_ff <= 0:
+            raise ValueError("d_ff must be strictly greater than 0")
+        if not 0.0 <= dropout_rate < 1.0:
+            raise ValueError("dropout_rate must be in [0, 1)")
+
+        self.attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
+        self.feed_forward = FeedForwardLayer(d_model, d_ff, dropout_rate)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout_rate)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout2 = nn.Dropout(dropout_rate)
+
+    def forward(self, x: Tensor, src_padding_mask: Tensor) -> Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError("x must be a torch.Tensor")
+        if x.dim() != 3:
+            raise ValueError(
+                f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
+            )
+        if not isinstance(src_padding_mask, torch.Tensor):
+            raise TypeError("src_padding_mask must be a torch.Tensor")
+        if src_padding_mask.dtype != torch.bool or src_padding_mask.dim() != 4:
+            raise TypeError(
+                f"src_padding_mask must be boolean with shape (B, H, 1, S);"
+                f" got dtype {src_padding_mask.dtype} and shape {tuple(src_padding_mask.shape)}"
+            )
+
+        attn_out = self.attention_layer(x, x, x, src_padding_mask, src_padding_mask)
+        x = self.norm1(x + self.dropout1(attn_out))
+        ff_out = self.feed_forward(x)
+        return self.norm2(x + self.dropout2(ff_out))
+
+
+class TransformerEncoder(nn.Module):
+    """Stack of encoder layers with optional activation checkpointing."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        d_ff: int,
+        num_layers: int,
+        dropout_rate: float,
+    ) -> None:
+        super().__init__()
+
+        if not isinstance(num_layers, int):
+            raise TypeError(f"num_layers must be an int, got {type(num_layers)}")
+        if num_layers <= 0:
+            raise ValueError("num_layers must be strictly greater than 0")
+
+        self.layers = nn.ModuleList(
+            [EncoderLayer(d_model, num_heads, d_ff, dropout_rate) for _ in range(num_layers)]
+        )
+        self.use_ckpt = False
+
+    def forward(self, x: Tensor, src_padding_mask: Tensor) -> Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError("x must be a torch.Tensor")
+        if x.dim() != 3:
+            raise ValueError(
+                f"x must be a 3D tensor of shape (B, S, D); got shape {tuple(x.shape)}"
+            )
+
+        for layer in self.layers:
+            if self.use_ckpt:
+
+                def _fn(x_, *, _layer=layer):
+                    return _layer(x_, src_padding_mask)
+
+                x = ckpt.checkpoint(_fn, x, use_reentrant=False)
+            else:
+                x = layer(x, src_padding_mask)
+        return x

src/transformer/modules/feedforward.py

@@ -1,59 +1,66 @@
-import torch
-import torch.nn as nn
-
-
-class FeedForwardLayer(nn.Module):
-    """
-    Position-wise feed-forward layer used in Transformer blocks.
-
-    Architecture:
-        fc1: Linear(d_model -> d_ff)
-        activation: ReLU
-        dropout: nn.Dropout(dropout_rate)
-        fc2: Linear(d_ff -> d_model)
-
-    Args:
-        d_model (int): Dimensionality of model embeddings.
-        d_ff (int): Hidden dimensionality of feed-forward layer.
-        dropout_rate (float): Dropout probability between 0 and 1 (exclusive).
-
-    Shape:
-        Input:  (B, S, D) where D == d_model
-        Output: (B, S, D)
-    """
-
-    def __init__(self, d_model: int, d_ff: int, dropout_rate: float = 0.1):
-        super().__init__()
-
-        if not isinstance(d_ff, int):
-            raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
-        if not isinstance(d_model, int):
-            raise TypeError(f"d_model must be an int, got {type(d_model)}")
-        if not isinstance(dropout_rate, float):
-            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
-
-        if d_ff <= 0:
-            raise ValueError(f"d_ff must be strictly greater than 0, got {d_ff}")
-        if d_model <= 0:
-            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
-        if not (0.0 <= dropout_rate < 1.0):
-            raise ValueError(f"dropout_rate must be in [0,1), got {dropout_rate}")
-
-        self.d_model = d_model
-        self.d_ff = d_ff
-        self.dropout_rate = dropout_rate
-
-        self.fc1 = nn.Linear(d_model, d_ff)
-        self.relu = nn.ReLU()
-        self.dropout = nn.Dropout(dropout_rate)
-        self.fc2 = nn.Linear(d_ff, d_model)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if not isinstance(x, torch.Tensor):
-            raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
-        if x.ndim != 3:
-            raise ValueError(f"x must be 3D of shape (B,S,D); got shape {tuple(x.shape)}")
-        if x.shape[-1] != self.d_model:
-            raise ValueError(f"Last dim {x.shape[-1]} must match d_model {self.d_model}")
-
-        return self.fc2(self.dropout(self.relu(self.fc1(x))))
+"""Position-wise feed-forward sub-layer used in transformer blocks."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+__all__ = ["FeedForwardLayer"]
+
+
+class FeedForwardLayer(nn.Module):
+    """
+    Position-wise feed-forward layer used in Transformer blocks.
+
+    Architecture:
+        fc1: Linear(d_model -> d_ff)
+        activation: ReLU
+        dropout: nn.Dropout(dropout_rate)
+        fc2: Linear(d_ff -> d_model)
+
+    Args:
+        d_model (int): Dimensionality of model embeddings.
+        d_ff (int): Hidden dimensionality of feed-forward layer.
+        dropout_rate (float): Dropout probability between 0 and 1 (exclusive).
+
+    Shape:
+        Input:  (B, S, D) where D == d_model
+        Output: (B, S, D)
+    """
+
+    def __init__(self, d_model: int, d_ff: int, dropout_rate: float = 0.1):
+        super().__init__()
+
+        if not isinstance(d_ff, int):
+            raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be an int, got {type(d_model)}")
+        if not isinstance(dropout_rate, float):
+            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
+
+        if d_ff <= 0:
+            raise ValueError(f"d_ff must be strictly greater than 0, got {d_ff}")
+        if d_model <= 0:
+            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
+        if not (0.0 <= dropout_rate < 1.0):
+            raise ValueError(f"dropout_rate must be in [0,1), got {dropout_rate}")
+
+        self.d_model = d_model
+        self.d_ff = d_ff
+        self.dropout_rate = dropout_rate
+
+        self.fc1 = nn.Linear(d_model, d_ff)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(dropout_rate)
+        self.fc2 = nn.Linear(d_ff, d_model)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
+        if x.ndim != 3:
+            raise ValueError(f"x must be 3D of shape (B,S,D); got shape {tuple(x.shape)}")
+        if x.shape[-1] != self.d_model:
+            raise ValueError(f"Last dim {x.shape[-1]} must match d_model {self.d_model}")
+
+        return self.fc2(self.dropout(self.relu(self.fc1(x))))

src/transformer/modules/lm_head.py

@@ -1,47 +1,54 @@
-import torch
-import torch.nn as nn
-
-
-class LMHead(nn.Module):
-    """
-    Language modeling head: projects hidden states to vocabulary logits.
-
-    Args:
-        d_model (int): Model hidden dimension (>0).
-        vocab_size (int): Vocabulary size (>0).
-
-    Input:
-        x (Tensor): shape (B, S, D) with D == d_model. S may be 0.
-
-    Output:
-        logits (Tensor): shape (B, S, V) with V == vocab_size.
-    """
-
-    def __init__(self, d_model: int, vocab_size: int):
-        super().__init__()
-        if not isinstance(vocab_size, int):
-            raise TypeError(f"vocab_size must be an int, got {type(vocab_size)}")
-        if not isinstance(d_model, int):
-            raise TypeError(f"d_model must be an int, got {type(d_model)}")
-        if vocab_size <= 0:
-            raise ValueError(f"vocab_size must be strictly greater than 0, got {vocab_size}")
-        if d_model <= 0:
-            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
-
-        self.d_model = d_model
-        self.vocab_size = vocab_size
-        self.fc = nn.Linear(d_model, vocab_size, bias=False)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if not isinstance(x, torch.Tensor):
-            raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
-        if x.dim() != 3:
-            raise ValueError(
-                f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
-            )
-
-        B, S, D = x.shape
-        if D != self.d_model:
-            raise ValueError(f"Last dim {D} must match d_model {self.d_model}")
-
-        return self.fc(x)
+"""Lightweight language-model head projecting hidden states to vocab logits."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+__all__ = ["LMHead"]
+
+
+class LMHead(nn.Module):
+    """
+    Language modeling head: projects hidden states to vocabulary logits.
+
+    Args:
+        d_model (int): Model hidden dimension (>0).
+        vocab_size (int): Vocabulary size (>0).
+
+    Input:
+        x (Tensor): shape (B, S, D) with D == d_model.
+
+    Output:
+        logits (Tensor): shape (B, S, V) with V == vocab_size.
+    """
+
+    def __init__(self, d_model: int, vocab_size: int):
+        super().__init__()
+        if not isinstance(vocab_size, int):
+            raise TypeError(f"vocab_size must be an int, got {type(vocab_size)}")
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be an int, got {type(d_model)}")
+        if vocab_size <= 0:
+            raise ValueError(f"vocab_size must be strictly greater than 0, got {vocab_size}")
+        if d_model <= 0:
+            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
+
+        self.d_model = d_model
+        self.vocab_size = vocab_size
+        self.fc = nn.Linear(d_model, vocab_size, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
+        if x.dim() != 3:
+            raise ValueError(
+                f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
+            )
+
+        B, S, D = x.shape
+        if D != self.d_model:
+            raise ValueError(f"Last dim {D} must match d_model {self.d_model}")
+
+        return self.fc(x)  # (B, S, V)

src/transformer/transformer.py

@@ -1,239 +1,435 @@
-import torch
-import torch.nn as nn
-from torch import Tensor
-
-from transformer import modules
-from transformer.configs import BasicEncDecCfg
-from transformer.utils import create_causal_mask, sample_from_logits, shift_right
-
-
-class BasicEncoderDecoderTransformer(nn.Module):
-    def __init__(self, cfg: BasicEncDecCfg):
-        super().__init__()
-
-        # Basic type/value guarding for critical cfg fields (mirror project style)
-        if not isinstance(cfg, BasicEncDecCfg):
-            raise TypeError(f"cfg must be a BasicEncDecCfg, got {type(cfg)}")
-        if not isinstance(cfg.vocab_size, int):
-            raise TypeError(f"vocab_size must be an int, got {type(cfg.vocab_size)}")
-        if not cfg.vocab_size > 0:
-            raise ValueError(f"vocab_size must be strictly greater than 0, got {cfg.vocab_size}")
-        if not isinstance(cfg.d_model, int):
-            raise TypeError(f"d_model must be an int, got {type(cfg.d_model)}")
-        if not cfg.d_model > 0:
-            raise ValueError(f"d_model must be strictly greater than 0, got {cfg.d_model}")
-        if not isinstance(cfg.max_seq_len, int):
-            raise TypeError(f"max_seq_len must be an int, got {type(cfg.max_seq_len)}")
-        if not cfg.max_seq_len >= 0:
-            raise ValueError(
-                f"max_seq_len must be greater than or equal to 0, got {cfg.max_seq_len}"
-            )
-        if not isinstance(cfg.pad_id, int):
-            raise TypeError(f"pad_id must be an int, got {type(cfg.pad_id)}")
-        if not (0 <= cfg.pad_id < cfg.vocab_size):
-            raise ValueError(f"pad_id must be in [0, {cfg.vocab_size - 1}], got {cfg.pad_id}")
-        if not isinstance(cfg.bos_id, int):
-            raise TypeError(f"bos_id must be an int, got {type(cfg.bos_id)}")
-        if not (0 <= cfg.bos_id < cfg.vocab_size):
-            raise ValueError(f"bos_id must be in [0, {cfg.vocab_size - 1}], got {cfg.bos_id}")
-        if not isinstance(cfg.eos_id, int):
-            raise TypeError(f"eos_id must be an int, got {type(cfg.eos_id)}")
-        if not (0 <= cfg.eos_id < cfg.vocab_size):
-            raise ValueError(f"eos_id must be in [0, {cfg.vocab_size - 1}], got {cfg.eos_id}")
-
-        self.cfg = cfg
-
-        self.embed = modules.InputEmbedding(
-            cfg.vocab_size, cfg.d_model, cfg.max_seq_len, cfg.pad_id, cfg.dropout_rate
-        )
-        self.encoder = modules.TransformerEncoder(
-            cfg.d_model, cfg.num_heads, cfg.d_ff, cfg.num_layers, cfg.dropout_rate
-        )
-        self.decoder = modules.TransformerDecoder(
-            cfg.d_model, cfg.num_heads, cfg.d_ff, cfg.num_layers, cfg.dropout_rate
-        )
-        self.lm_head = modules.LMHead(cfg.d_model, cfg.vocab_size)
-
-        # tie weights
-        self.embed.token_embed.weight = self.lm_head.fc.weight
-
-        # expose a few attrs for convenience
-        self.max_seq_len = cfg.max_seq_len
-        self.pad_id = cfg.pad_id
-        self.bos_id = cfg.bos_id
-        self.eos_id = cfg.eos_id
-
-    def forward(
-        self,
-        src_ids: Tensor,
-        tgt_ids: Tensor,
-        src_padding_mask: Tensor | None,
-        tgt_padding_mask: Tensor | None,
-    ) -> Tensor:
-        # Validate inputs (match project style)
-        if not isinstance(src_ids, torch.Tensor):
-            raise TypeError(f"src_ids must be a torch.Tensor, got {type(src_ids)}")
-        if not isinstance(tgt_ids, torch.Tensor):
-            raise TypeError(f"tgt_ids must be a torch.Tensor, got {type(tgt_ids)}")
-        if src_ids.dim() != 2:
-            raise ValueError(
-                f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
-            )
-        if tgt_ids.dim() != 2:
-            raise ValueError(
-                f"tgt_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(tgt_ids.shape)}"
-            )
-        if src_ids.dtype != torch.long or tgt_ids.dtype != torch.long:
-            raise TypeError("src_ids and tgt_ids must be torch.long (int64)")
-        if src_padding_mask is not None and not isinstance(src_padding_mask, torch.Tensor):
-            raise TypeError(
-                f"src_padding_mask must be a torch.Tensor or None, got {type(src_padding_mask)}"
-            )
-        if tgt_padding_mask is not None and not isinstance(tgt_padding_mask, torch.Tensor):
-            raise TypeError(
-                f"tgt_padding_mask must be a torch.Tensor or None, got {type(tgt_padding_mask)}"
-            )
-
-        hidden_states = self.encode(src_ids, src_padding_mask)
-        # teacher forcing: shift right with BOS, keep PAD for padding
-        shifted_tgt = shift_right(tgt_ids, self.bos_id, self.pad_id)
-        shifted_tgt_mask = None
-        if tgt_padding_mask is not None:
-            m = tgt_padding_mask.to(dtype=torch.bool, device=shifted_tgt.device)
-
-            # Handle common mask shapes by shifting along the last (time) dimension
-            if m.dim() == 4:  # [B, 1, 1, T]
-                pad_col = torch.zeros_like(m[..., :1])  # [B,1,1,1] False
-                shifted_tgt_mask = torch.cat([pad_col, m[..., :-1]], dim=-1)
-            elif m.dim() == 3:  # [B, 1, T] or [B, H, T]
-                pad_col = torch.zeros_like(m[..., :1])  # [B,*,1] False
-                shifted_tgt_mask = torch.cat([pad_col, m[..., :-1]], dim=-1)
-            elif m.dim() == 2:  # [B, T]
-                pad_col = torch.zeros_like(m[:, :1])  # [B,1] False
-                shifted_tgt_mask = torch.cat([pad_col, m[:, :-1]], dim=-1)
-            else:
-                raise ValueError(f"tgt_padding_mask must have 2, 3, or 4 dims; got {m.dim()}")
-
-        dec_hidden_states = self.decode(
-            hidden_states, shifted_tgt, src_padding_mask, shifted_tgt_mask
-        )
-        logits = self.lm_head(dec_hidden_states)
-        return logits  # (B, T_tgt, V)
-
-    def encode(self, src_ids: Tensor, src_padding_mask: Tensor | None) -> Tensor:
-        if not isinstance(src_ids, torch.Tensor):
-            raise TypeError(f"src_ids must be a torch.Tensor, got {type(src_ids)}")
-        if src_ids.dim() != 2:
-            raise ValueError(
-                f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
-            )
-        if src_ids.dtype != torch.long:
-            raise TypeError("src_ids must be torch.long (int64)")
-        if src_padding_mask is not None and not isinstance(src_padding_mask, torch.Tensor):
-            raise TypeError(
-                f"src_padding_mask must be a torch.Tensor or None, got {type(src_padding_mask)}"
-            )
-
-        x = self.embed(src_ids)  # (B, Sx, D)
-        hidden_states = self.encoder(x, src_padding_mask)
-        return hidden_states  # (B, Sx, D)
-
-    def decode(
-        self,
-        hidden_states: Tensor,
-        tgt_ids: Tensor,
-        src_padding_mask: Tensor | None,
-        tgt_padding_mask: Tensor | None = None,
-    ) -> Tensor:
-        if not isinstance(hidden_states, torch.Tensor):
-            raise TypeError(f"hidden_states must be a torch.Tensor, got {type(hidden_states)}")
-        if hidden_states.dim() != 3:
-            raise ValueError(
-                f"hidden_states must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(hidden_states.shape)}"
-            )
-        if not isinstance(tgt_ids, torch.Tensor):
-            raise TypeError(f"tgt_ids must be a torch.Tensor, got {type(tgt_ids)}")
-        if tgt_ids.dim() != 2:
-            raise ValueError(
-                f"tgt_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(tgt_ids.shape)}"
-            )
-        if tgt_ids.dtype != torch.long:
-            raise TypeError("tgt_ids must be torch.long (int64)")
-        if src_padding_mask is not None and not isinstance(src_padding_mask, torch.Tensor):
-            raise TypeError(
-                f"src_padding_mask must be a torch.Tensor or None, got {type(src_padding_mask)}"
-            )
-        if tgt_padding_mask is not None and not isinstance(tgt_padding_mask, torch.Tensor):
-            raise TypeError(
-                f"tgt_padding_mask must be a torch.Tensor or None, got {type(tgt_padding_mask)}"
-            )
-
-        # causal mask per sequence length (utils returns broadcastable boolean mask)
-        tgt_causal_mask = create_causal_mask(tgt_ids.size(1)).to(tgt_ids.device)
-        y = self.embed(tgt_ids)  # (B, Sy, D)
-        out = self.decoder(hidden_states, y, src_padding_mask, tgt_padding_mask, tgt_causal_mask)
-        return out  # (B, Sy, D)
-
-    def generate(
-        self,
-        src_ids: Tensor,
-        src_padding_mask: Tensor | None,
-        max_new_tokens: int = 20,
-        temperature: float = 1.0,
-        top_k: int | None = None,
-        top_p: float | None = None,
-        do_sample: bool = False,
-    ) -> Tensor:
-        # validate args
-        if not isinstance(src_ids, torch.Tensor):
-            raise TypeError(f"src_ids must be a torch.Tensor, got {type(src_ids)}")
-        if src_ids.dim() != 2:
-            raise ValueError(
-                f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
-            )
-        if src_ids.dtype != torch.long:
-            raise TypeError("src_ids must be torch.long (int64)")
-        if src_padding_mask is not None and not isinstance(src_padding_mask, torch.Tensor):
-            raise TypeError(
-                f"src_padding_mask must be a torch.Tensor or None, got {type(src_padding_mask)}"
-            )
-        if not isinstance(max_new_tokens, int):
-            raise TypeError(f"max_new_tokens must be an int, got {type(max_new_tokens)}")
-        if max_new_tokens < 0:
-            raise ValueError(
-                f"max_new_tokens must be greater than or equal to 0, got {max_new_tokens}"
-            )
-        if not isinstance(temperature, float):
-            raise TypeError(f"temperature must be a float, got {type(temperature)}")
-        if not (temperature > 0.0):
-            raise ValueError(f"temperature must be strictly greater than 0, got {temperature}")
-        if top_k is not None and (not isinstance(top_k, int)):
-            raise TypeError(f"top_k must be an int or None, got {type(top_k)}")
-        if top_k is not None and top_k < 0:
-            raise ValueError(f"top_k must be strictly greater than 0, got {top_k}")
-        if top_p is not None and (not isinstance(top_p, float)):
-            raise TypeError(f"top_p must be a float or None, got {type(top_p)}")
-        if top_p is not None and not (0.0 < top_p <= 1.0):
-            raise ValueError(f"top_p must be in (0, 1], got {top_p}")
-        if not isinstance(do_sample, bool):
-            raise TypeError(f"do_sample must be a bool, got {type(do_sample)}")
-        self.eval()
-        with torch.no_grad():
-            batch_size, _ = src_ids.shape
-            device = src_ids.device
-            tgt_ids = torch.full((batch_size, 1), self.bos_id, device=device, dtype=torch.long)
-            finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
-            hidden_states = self.encode(src_ids, src_padding_mask)
-            for _ in range(max_new_tokens):
-                # decode on current tgt_ids; take last-step logits
-                step_hidden = self.decode(hidden_states, tgt_ids, src_padding_mask)  # (B, T, D)
-                logits = self.lm_head(step_hidden)[:, -1, :]  # (B, V)
-                next_ids = sample_from_logits(
-                    logits, temperature=temperature, top_k=top_k, top_p=top_p, do_sample=do_sample
-                )  # (B,)
-                next_ids = torch.where(finished, torch.full_like(next_ids, self.eos_id), next_ids)
-                tgt_ids = torch.cat([tgt_ids, next_ids.unsqueeze(1)], dim=-1)
-                finished |= next_ids == self.eos_id
-                if finished.all():
-                    return tgt_ids
-            return tgt_ids
+"""Implementation of the encoder-decoder transformer used across the project."""
+
+from __future__ import annotations
+
+import math
+from typing import cast
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from transformer import modules
+from transformer.configs import ModelCfg
+from transformer.utils import broadcast_padding_mask, create_causal_mask, sample_from_logits
+
+__all__ = ["BasicEncoderDecoderTransformer"]
+
+
+class BasicEncoderDecoderTransformer(nn.Module):
+    """Full encoder-decoder model with tied input/output embeddings."""
+
+    def __init__(self, cfg: ModelCfg) -> None:
+        super().__init__()
+
+        if not isinstance(cfg, ModelCfg):
+            raise TypeError(f"cfg must be a ModelCfg, got {type(cfg)}")
+
+        if not isinstance(cfg.vocab_size, int):
+            raise TypeError(f"vocab_size must be an int, got {type(cfg.vocab_size)}")
+        if not cfg.vocab_size > 0:
+            raise ValueError(f"vocab_size must be strictly greater than 0, got {cfg.vocab_size}")
+
+        if not isinstance(cfg.d_model, int):
+            raise TypeError(f"d_model must be an int, got {type(cfg.d_model)}")
+        if not cfg.d_model > 0:
+            raise ValueError(f"d_model must be strictly greater than 0, got {cfg.d_model}")
+
+        if not isinstance(cfg.max_seq_len, int):
+            raise TypeError(f"max_seq_len must be an int, got {type(cfg.max_seq_len)}")
+        if not cfg.max_seq_len > 0:
+            raise ValueError(f"max_seq_len must be greater than 0, got {cfg.max_seq_len}")
+        if not 0 <= cfg.dropout_rate < 1:
+            raise ValueError(f"dropout_rate must be in [0,1[, got {cfg.dropout_rate}")
+
+        if not isinstance(cfg.pad_id, int):
+            raise TypeError(f"pad_id must be an int, got {type(cfg.pad_id)}")
+        if not (0 <= cfg.pad_id < cfg.vocab_size):
+            raise ValueError(f"pad_id must be in [0, {cfg.vocab_size - 1}], got {cfg.pad_id}")
+        if not isinstance(cfg.bos_id, int):
+            raise TypeError(f"bos_id must be an int, got {type(cfg.bos_id)}")
+        if not (0 <= cfg.bos_id < cfg.vocab_size):
+            raise ValueError(f"bos_id must be in [0, {cfg.vocab_size - 1}], got {cfg.bos_id}")
+        if not isinstance(cfg.eos_id, int):
+            raise TypeError(f"eos_id must be an int, got {type(cfg.eos_id)}")
+        if not (0 <= cfg.eos_id < cfg.vocab_size):
+            raise ValueError(f"eos_id must be in [0, {cfg.vocab_size - 1}], got {cfg.eos_id}")
+
+        self.cfg = cfg
+
+        self.embed = modules.InputEmbedding(
+            cfg.vocab_size, cfg.d_model, cfg.max_seq_len, cfg.pad_id, cfg.dropout_rate
+        )
+        self.encoder = modules.TransformerEncoder(
+            cfg.d_model, cfg.num_heads, cfg.d_ff, cfg.num_layers, cfg.dropout_rate
+        )
+        self.decoder = modules.TransformerDecoder(
+            cfg.d_model, cfg.num_heads, cfg.d_ff, cfg.num_layers, cfg.dropout_rate
+        )
+
+        self.lm_head = modules.LMHead(cfg.d_model, cfg.vocab_size)
+
+        # tie weights
+        self.lm_head.fc.weight = self.embed.token_embed.weight
+
+        self.max_seq_len = cfg.max_seq_len
+        self.pad_id = cfg.pad_id
+        self.bos_id = cfg.bos_id
+        self.eos_id = cfg.eos_id
+
+        self.gradient_checkpointing = False
+
+    def enable_gradient_checkpointing(self, enabled: bool = True) -> None:
+        self.gradient_checkpointing = enabled
+        self.encoder.use_ckpt = enabled
+        self.decoder.use_ckpt = enabled
+
+    def forward(
+        self,
+        src_ids: Tensor,
+        tgt_ids: Tensor,
+        src_padding_mask: Tensor,
+        tgt_padding_mask: Tensor,
+    ) -> Tensor:
+        """Compute logits given input/output token ids and boolean padding masks."""
+
+        if not isinstance(src_ids, Tensor):
+            raise TypeError(f"src_ids must be a torch.Tensor, got {type(src_ids)}")
+        if src_ids.dim() != 2:
+            raise ValueError(
+                f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
+            )
+        if src_ids.dtype != torch.long:
+            raise TypeError(f"src_idsmust be torch.long (int64), got {src_ids.dtype}")
+
+        if not isinstance(tgt_ids, Tensor):
+            raise TypeError(f"tgt_ids must be a torch.Tensor, got {type(tgt_ids)}")
+        if tgt_ids.dim() != 2:
+            raise ValueError(
+                f"tgt_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(tgt_ids.shape)}"
+            )
+        if tgt_ids.dtype != torch.long:
+            raise TypeError(f"tgt_ids must be torch.long (int64), got {tgt_ids.dtype}")
+
+        if (
+            not isinstance(src_padding_mask, Tensor)
+            or src_padding_mask.dtype != torch.bool
+            or src_padding_mask.dim() != 2
+        ):
+            raise TypeError("src_padding_mask must be a boolean tensor shaped (B, S)")
+
+        if (
+            not isinstance(tgt_padding_mask, Tensor)
+            or tgt_padding_mask.dtype != torch.bool
+            or tgt_padding_mask.dim() != 2
+        ):
+            raise TypeError("tgt_padding_mask must be a boolean tensor shaped (B, S)")
+        src_padding_mask = broadcast_padding_mask(src_padding_mask, self.cfg.num_heads)
+        tgt_padding_mask = broadcast_padding_mask(tgt_padding_mask, self.cfg.num_heads)
+        hidden_states = self.encode(src_ids, src_padding_mask)
+        dec_hidden_states = self.decode(hidden_states, tgt_ids, src_padding_mask, tgt_padding_mask)
+        logits = self.lm_head(dec_hidden_states)
+        return logits  # (B, T_tgt, V)
+
+    def encode(self, src_ids: Tensor, src_padding_mask: Tensor) -> Tensor:
+        """Encode source tokens into memory representations."""
+
+        if not isinstance(src_ids, torch.Tensor):
+            raise TypeError(f"src_ids must be a torch.Tensor, got {type(src_ids)}")
+        if src_ids.dim() != 2:
+            raise ValueError(
+                f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
+            )
+        if src_ids.dtype != torch.long:
+            raise TypeError("src_ids must be torch.long (int64)")
+
+        if (
+            not isinstance(src_padding_mask, Tensor)
+            or src_padding_mask.dtype != torch.bool
+            or src_padding_mask.dim() != 4
+        ):
+            raise TypeError("src_padding_mask must be a boolean tensor shaped (B, H, 1, S)")
+
+        x = self.embed(src_ids)  # (B, Sx, D)
+        hidden_states = self.encoder(x, src_padding_mask)
+        return hidden_states  # (B, Sx, D)
+
+    def decode(
+        self,
+        hidden_states: Tensor,
+        tgt_ids: Tensor,
+        src_padding_mask: Tensor,
+        tgt_padding_mask: Tensor,
+    ) -> Tensor:
+        if not isinstance(hidden_states, torch.Tensor):
+            raise TypeError(f"hidden_states must be a torch.Tensor, got {type(hidden_states)}")
+        if hidden_states.dim() != 3:
+            raise ValueError(
+                f"hidden_states must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(hidden_states.shape)}"
+            )
+
+        if not isinstance(tgt_ids, torch.Tensor):
+            raise TypeError(f"tgt_ids must be a torch.Tensor, got {type(tgt_ids)}")
+        if tgt_ids.dim() != 2:
+            raise ValueError(
+                f"tgt_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(tgt_ids.shape)}"
+            )
+        if tgt_ids.dtype != torch.long:
+            raise TypeError("tgt_ids must be torch.long (int64)")
+
+        if (
+            not isinstance(src_padding_mask, Tensor)
+            or src_padding_mask.dtype != torch.bool
+            or src_padding_mask.dim() != 4
+        ):
+            raise TypeError("src_padding_mask must be a boolean tensor shaped (B, H, 1, S)")
+
+        if (
+            not isinstance(tgt_padding_mask, Tensor)
+            or tgt_padding_mask.dtype != torch.bool
+            or tgt_padding_mask.dim() != 4
+        ):
+            raise TypeError("tgt_padding_mask must be a boolean tensor shaped (B, H, 1, S)")
+
+        tgt_causal_mask = create_causal_mask(tgt_ids, self.cfg.num_heads)
+
+        y = self.embed(tgt_ids)  # (B, Sy, D)
+        out = self.decoder(hidden_states, y, src_padding_mask, tgt_padding_mask, tgt_causal_mask)
+        return out  # (B, Sy, D)
+
+    def generate(
+        self,
+        src_ids: Tensor,
+        src_padding_mask: Tensor,
+        max_new_tokens: int = 20,
+        temperature: float = 1.0,
+        top_k: int | None = None,
+        top_p: float | None = None,
+        do_sample: bool = False,
+        presence_penalty: float = 0.0,
+        frequency_penalty: float = 0.0,
+        no_repeat_ngram: int | None = None,
+        min_steps_before_eos: int = 0,
+        *,
+        seed: int | None = None,
+        generator: torch.Generator | None = None,
+    ) -> Tensor:
+        if not isinstance(src_ids, torch.Tensor):
+            raise TypeError(f"src_ids must be a torch.Tensor, got {type(src_ids)}")
+        if src_ids.dim() != 2:
+            raise ValueError(
+                f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
+            )
+        if src_ids.dtype != torch.long:
+            raise TypeError("src_ids must be torch.long (int64)")
+
+        if not isinstance(src_padding_mask, torch.Tensor):
+            raise TypeError(
+                f"src_padding_mask must be a torch.Tensor, got {type(src_padding_mask)}"
+            )
+
+        if not isinstance(max_new_tokens, int):
+            raise TypeError(f"max_new_tokens must be an int, got {type(max_new_tokens)}")
+        if max_new_tokens < 0:
+            raise ValueError(
+                f"max_new_tokens must be greater than or equal to 0, got {max_new_tokens}"
+            )
+        if not isinstance(temperature, float):
+            raise TypeError(f"temperature must be a float, got {type(temperature)}")
+        if not (temperature > 0.0):
+            raise ValueError(f"temperature must be strictly greater than 0, got {temperature}")
+        if top_k is not None and not isinstance(top_k, int):
+            raise TypeError(f"top_k must be an int or None, got {type(top_k)}")
+        if top_k is not None and top_k < 0:
+            raise ValueError(f"top_k must be >= 0, got {top_k}")
+        if top_p is not None and not isinstance(top_p, int | float):
+            raise TypeError(f"top_p must be a number or None, got {type(top_p)}")
+        if top_p is not None:
+            top_p = float(top_p)
+            if not (0.0 < top_p <= 1.0):
+                raise ValueError(f"top_p must be in (0, 1], got {top_p}")
+        if not isinstance(do_sample, bool):
+            raise TypeError(f"do_sample must be a bool, got {type(do_sample)}")
+        if not isinstance(presence_penalty, int | float) or not math.isfinite(
+            float(presence_penalty)
+        ):
+            raise ValueError(
+                f"presence_penalty must be a finite number >= 0, got {presence_penalty!r}."
+            )
+        presence_penalty = float(presence_penalty)
+        if presence_penalty < 0.0:
+            raise ValueError(f"presence_penalty must be >= 0, got {presence_penalty}.")
+        if not isinstance(frequency_penalty, int | float) or not math.isfinite(
+            float(frequency_penalty)
+        ):
+            raise ValueError(
+                f"frequency_penalty must be a finite number >= 0, got {frequency_penalty!r}."
+            )
+        frequency_penalty = float(frequency_penalty)
+        if frequency_penalty < 0.0:
+            raise ValueError(f"frequency_penalty must be >= 0, got {frequency_penalty}.")
+        if no_repeat_ngram is not None:
+            if not isinstance(no_repeat_ngram, int):
+                raise ValueError(f"no_repeat_ngram must be int or None, got {no_repeat_ngram!r}.")
+            if no_repeat_ngram < 2:
+                no_repeat_ngram = None  # size <2 is meaningless; treat as disabled
+        if not isinstance(min_steps_before_eos, int) or min_steps_before_eos < 0:
+            raise ValueError(
+                f"min_steps_before_eos must be int >= 0, got {min_steps_before_eos!r}."
+            )
+        if seed is not None and not isinstance(seed, int):
+            raise TypeError(f"seed must be int or None, got {type(seed)}")
+        if generator is not None and not isinstance(generator, torch.Generator):
+            raise TypeError(f"generator must be a torch.Generator or None, got {type(generator)}")
+
+        self.eval()
+        with torch.no_grad():
+            batch_size, _ = src_ids.shape
+            device = src_ids.device
+            tgt_ids = torch.full((batch_size, 1), self.bos_id, device=device, dtype=torch.long)
+            allowed_new = max(0, self.max_seq_len - tgt_ids.size(1))
+            max_new_tokens = min(max_new_tokens, allowed_new)
+            finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
+            src_padding_mask = broadcast_padding_mask(src_padding_mask, self.cfg.num_heads)
+            hidden_states = self.encode(src_ids, src_padding_mask)
+
+            def _make_generator() -> torch.Generator:
+                if device.type == "cpu":
+                    return torch.Generator()
+                return torch.Generator(device=device)
+
+            rng = generator
+            if seed is not None:
+                rng = rng if rng is not None else _make_generator()
+                rng.manual_seed(seed)
+            elif rng is None and do_sample:
+                rng = _make_generator()
+            for step in range(max_new_tokens):
+                # decode on current tgt_ids; take last-step logits
+                tgt_padding_mask = tgt_ids != self.pad_id
+                tgt_padding_mask = broadcast_padding_mask(tgt_padding_mask, self.cfg.num_heads)
+                step_hidden = self.decode(
+                    hidden_states, tgt_ids, src_padding_mask, tgt_padding_mask
+                )  # (B, T, D)
+                logits = self.lm_head(step_hidden)[:, -1, :]  # (B, V)
+                disallowed = [self.pad_id] + ([self.eos_id] if step < min_steps_before_eos else [])
+                next_ids = cast(
+                    Tensor,
+                    sample_from_logits(
+                        logits,
+                        temperature=temperature,
+                        top_k=top_k,
+                        top_p=top_p,
+                        do_sample=do_sample,
+                        disallowed_tokens=disallowed,
+                        repetition_ctx=tgt_ids,
+                        presence_penalty=presence_penalty,
+                        frequency_penalty=frequency_penalty,
+                        no_repeat_ngram_size=no_repeat_ngram,
+                        rng=rng,
+                    ),
+                )  # (B,)
+                next_ids = torch.where(finished, torch.full_like(next_ids, self.eos_id), next_ids)
+                tgt_ids = torch.cat([tgt_ids, next_ids.unsqueeze(1)], dim=-1)
+                finished |= next_ids == self.eos_id
+                if finished.all():
+                    return tgt_ids
+            return tgt_ids
+
+    def debug_generate(
+        self,
+        src_ids: Tensor,
+        src_padding_mask: Tensor | None,
+        tokenizer,
+        max_new_tokens: int = 20,
+        temperature: float = 1.0,
+        top_k: int | None = None,
+        top_p: float | None = None,
+        do_sample: bool = False,
+        *,
+        seed: int | None = None,
+        generator: torch.Generator | None = None,
+    ) -> Tensor:
+        """
+        Like generate(), but prints debug info at each step: input tokens, ids, output ids, tokens, logits, etc.
+        """
+        self.eval()
+        with torch.no_grad():
+            batch_size, _ = src_ids.shape
+            device = src_ids.device
+            tgt_ids = torch.full((batch_size, 1), self.bos_id, device=device, dtype=torch.long)
+            allowed_new = max(0, self.max_seq_len - tgt_ids.size(1))
+            max_new_tokens = min(max_new_tokens, allowed_new)
+            finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
+            if src_padding_mask is not None:
+                if (
+                    src_padding_mask.dtype != torch.bool
+                    or src_padding_mask.dim() != 2
+                    or src_padding_mask.size(0) != batch_size
+                ):
+                    raise TypeError(
+                        "src_padding_mask must be boolean tensor shaped (B, S) when provided"
+                    )
+                src_mask_2d = src_padding_mask
+            else:
+                src_mask_2d = torch.ones_like(src_ids, dtype=torch.bool, device=device)
+            src_mask_4d = broadcast_padding_mask(src_mask_2d, self.cfg.num_heads)
+            hidden_states = self.encode(src_ids, src_mask_4d)
+
+            def _make_generator() -> torch.Generator:
+                if device.type == "cpu":
+                    return torch.Generator()
+                return torch.Generator(device=device)
+
+            rng = generator
+            if seed is not None:
+                rng = rng if rng is not None else _make_generator()
+                rng.manual_seed(seed)
+            elif rng is None and do_sample:
+                rng = _make_generator()
+            # print("[DEBUG] Input ids:", src_ids)
+            print(
+                "[DEBUG] Input tokens:", tokenizer.batch_decode(src_ids, skip_special_tokens=False)
+            )
+            for step in range(max_new_tokens):
+                print(f"[DEBUG] Step {step + 1}")
+                # print("  Output ids so far:", tgt_ids)
+                print(
+                    "  Output tokens so far:",
+                    tokenizer.batch_decode(tgt_ids, skip_special_tokens=False),
+                )
+                tgt_padding_mask_2d = tgt_ids != self.pad_id
+                tgt_padding_mask_4d = broadcast_padding_mask(
+                    tgt_padding_mask_2d, self.cfg.num_heads
+                )
+                step_hidden = self.decode(
+                    hidden_states, tgt_ids, src_mask_4d, tgt_padding_mask_4d
+                )  # (B, T, D)
+                logits = self.lm_head(step_hidden)[:, -1, :]  # (B, V)
+                next_ids = cast(
+                    Tensor,
+                    sample_from_logits(
+                        logits,
+                        temperature=temperature,
+                        top_k=top_k,
+                        top_p=top_p,
+                        do_sample=do_sample,
+                        disallowed_tokens=[self.pad_id],
+                        rng=rng,
+                    ),
+                )  # (B,)
+                next_ids = torch.where(finished, torch.full_like(next_ids, self.eos_id), next_ids)
+                tgt_ids = torch.cat([tgt_ids, next_ids.unsqueeze(1)], dim=-1)
+                finished |= next_ids == self.eos_id
+
+                # print("  Next token ids:", next_ids)
+                print(
+                    "  Next tokens:",
+                    tokenizer.batch_decode(next_ids.unsqueeze(1), skip_special_tokens=False),
+                )
+                # print("  Logits (first 5):", logits[0, :5].cpu().numpy() if logits.shape[0] > 0 else None)
+                if finished.all():
+                    print("[DEBUG] All sequences finished.")
+                    return tgt_ids
+            print("[DEBUG] Max steps reached.")
+            return tgt_ids

src/transformer/utils.py

@@ -1,612 +1,1042 @@
-from __future__ import annotations
-
-import math
-from collections.abc import Iterable
-
-import torch
-import torch.nn.functional as F
-
-
-def split_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor:
-    """
-    Split the last (model) dimension of a 3D torch.Tensor into (num_heads, d_head) and permute to
-    (batch_size, num_heads, seq_length, d_head).
-
-    Args:
-        x (torch.Tensor): Input torch.Tensor of shape (batch_size, seq_length, d_model).
-        num_heads (int): Number of attention heads to split into. Must be a positive integer
-            that divides d_model exactly.
-
-    Returns:
-        torch.Tensor: torch.Tensor of shape (batch_size, num_heads, seq_length, d_head), where
-            d_head = d_model // num_heads. The output torch.Tensor retains the same dtype and
-            device as the input.
-
-    Invariants:
-        - Output device == x.device
-        - Output dtype == x.dtype
-        - batch_size and seq_length are preserved from the input shape.
-
-    Notes:
-        - Zero-length sequences (seq_length == 0) are supported and will return a torch.Tensor
-          with shape (batch_size, num_heads, 0, d_head).
-
-    Raises:
-        TypeError: If x is not a torch.Tensor or num_heads is not an int.
-        ValueError: If x is not 3D, if num_heads <= 0, or if d_model is not divisible by num_heads.
-    """
-    # Type checks
-    if not isinstance(x, torch.Tensor):
-        raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
-    if not isinstance(num_heads, int):
-        raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
-
-    # Shape checks
-    if x.ndim != 3:
-        raise ValueError(
-            f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
-        )
-    if num_heads <= 0:
-        raise ValueError(f"num_heads must be > 0; got {num_heads}")
-
-    batch_size, seq_length, d_model = x.shape
-
-    # Divisibility check
-    if d_model % num_heads != 0:
-        raise ValueError(
-            f"d_model ({d_model}) must be divisible by num_heads ({num_heads}); "
-            f"got remainder {d_model % num_heads}"
-        )
-
-    d_head = d_model // num_heads
-
-    # Reshape and permute: (B, S, D) -> (B, S, H, Dh) -> (B, H, S, Dh)
-    x = x.reshape(batch_size, seq_length, num_heads, d_head)
-    x = x.permute(0, 2, 1, 3)
-
-    return x
-
-
-def calculate_attention(
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    mask: torch.Tensor | None,
-    *,
-    deterministic: bool = False,
-    return_probs: bool = False,
-    attn_dropout_p: float = 0.0,
-):
-    """
-    Scaled dot-product attention with:
-      • device-safe masking (mask auto-moved to query.device),
-      • explicit broadcastability checks for mask,
-      • numerically stable softmax (row-wise max subtraction),
-      • fully-masked rows -> zero probabilities and zero outputs,
-      • fp16/bf16-safe compute via fp32 upcast, output downcast.
-
-    Args:
-        query: (B, H, S_q, D)
-        key:   (B, H, S_kv, D)
-        value: (B, H, S_kv, D)
-        mask: broadcastable to (B, H, S_q, S_kv)
-              - bool: True = masked/ignored position
-              - float: additive bias (e.g., large negative for masked, ALiBi, etc.)
-        deterministic: best-effort reproducibility toggle
-        return_probs: if False, only return attention output
-
-    Returns:
-        attention: (B, H, S_q, D)
-        probs (optional): (B, H, S_q, S_kv)
-    """
-    # ---- Basic shape checks
-    if query.dim() != 4 or key.dim() != 4 or value.dim() != 4:
-        raise ValueError(
-            f"q/k/v must be 4D (B,H,S,D). Got "
-            f"q={tuple(query.shape)}, k={tuple(key.shape)}, v={tuple(value.shape)}"
-        )
-
-    Bq, Hq, Sq, Dq = query.shape
-    Bk, Hk, Sk, Dk = key.shape
-    Bv, Hv, Sv, Dv = value.shape
-
-    if not (Bq == Bk == Bv):
-        raise ValueError(f"Batch mismatch: q={Bq}, k={Bk}, v={Bv}")
-    if not (Hq == Hk == Hv):
-        raise ValueError(f"Heads mismatch: q={Hq}, k={Hk}, v={Hv}")
-    if Sk != Sv:
-        raise ValueError(f"Key/Value seq length mismatch: Sk={Sk} vs Sv={Sv}")
-    if not (Dq == Dk == Dv):
-        raise ValueError(f"Head dimension mismatch: Dq={Dq}, Dk={Dk}, Dv={Dv}")
-
-    target_shape = (Bq, Hq, Sq, Sk)
-
-    # ---- Device checks: q/k/v must be on the same device (fail fast); mask will be auto-moved
-    if not (query.device == key.device == value.device):
-        raise RuntimeError(
-            f"q/k/v must be on the same device, got "
-            f"query={query.device}, key={key.device}, value={value.device}"
-        )
-
-    # ---- Early exit for empty sequences
-    if Sq == 0:
-        empty_attn = query.new_zeros(Bq, Hq, 0, Dq)
-        empty_probs = query.new_zeros(Bq, Hq, 0, Sk)
-        return (empty_attn, empty_probs) if return_probs else empty_attn
-    if Sk == 0:
-        zero_attn = query.new_zeros(Bq, Hq, Sq, Dq)
-        zero_probs = query.new_zeros(Bq, Hq, Sq, 0)
-        return (zero_attn, zero_probs) if return_probs else zero_attn
-
-    # ---- Compute dtype policy: upcast to fp32 for stability if needed
-    input_dtype = query.dtype
-    compute_dtype = torch.float32 if input_dtype in (torch.float16, torch.bfloat16) else input_dtype
-    q = query.to(compute_dtype)
-    k = key.to(compute_dtype)
-    v = value.to(compute_dtype)
-
-    # ---- Determinism (best effort)
-    if deterministic:
-        try:
-            torch.use_deterministic_algorithms(True)
-        except Exception:
-            pass
-
-    # ---- Scores
-    d_head = Dq
-    scores = torch.matmul(q, k.transpose(-1, -2)) / (d_head**0.5)  # (B,H,Sq,Sk)
-
-    # ---- Mask handling (device + broadcastability + semantics)
-    full_row_mask = None  # (B,H,Sq,1) True where the entire key row is invalid
-    if mask is not None:
-        # Move mask to the same device
-        if mask.device != q.device:
-            mask = mask.to(q.device)
-
-        # Check broadcastability to target shape, then expand (without copy)
-        def _expand_or_error(x: torch.Tensor, name: str) -> torch.Tensor:
-            # Try expanding to target_shape; torch.Tensor.expand raises if not broadcastable
-            try:
-                return x.expand(target_shape)
-            except RuntimeError as e:
-                raise ValueError(
-                    f"{name} with shape {tuple(x.shape)} is not broadcastable to {target_shape}. "
-                    f"Broadcasting rules require dimensions to be equal or 1 in the source."
-                ) from e
-
-        if mask.dtype == torch.bool:
-            mask = _expand_or_error(mask, "Boolean mask")
-            # Use a large finite negative to avoid (-inf) - (-inf) during stabilization
-            neg_large = torch.finfo(compute_dtype).min / 4
-            scores = scores.masked_fill(mask, neg_large)
-            full_row_mask = mask.all(dim=-1, keepdim=True)  # (B,H,Sq,1)
-        else:
-            # Additive bias mask
-            mask = _expand_or_error(mask.to(compute_dtype), "Additive mask")
-            scores = scores + mask
-            # If caller used -inf as additive, detect fully-masked rows after addition
-            full_row_mask = torch.isneginf(scores).all(dim=-1, keepdim=True)
-
-    # ---- Numerically stable softmax: subtract row-wise max
-    row_max = torch.amax(scores, dim=-1, keepdim=True)
-    if full_row_mask is None:
-        # If any rows are -inf (rare with finite negatives), neutralize their max to 0
-        row_max = torch.where(torch.isneginf(row_max), torch.zeros_like(row_max), row_max)
-    else:
-        # For rows marked fully masked, set max to 0 to avoid -inf - (-inf)
-        row_max = torch.where(full_row_mask, torch.zeros_like(row_max), row_max)
-
-    stable_scores = scores - row_max
-    probs = F.softmax(stable_scores, dim=-1)
-
-    # ---- Zero-out fully masked rows (probabilities and thus outputs)
-    if full_row_mask is not None:
-        probs = probs * (~full_row_mask).to(probs.dtype)
-    if attn_dropout_p and attn_dropout_p > 0.0:
-        probs = F.dropout(probs, p=attn_dropout_p, training=True)
-    # ---- Weighted sum with values; cast back to input dtype
-    attention = torch.matmul(probs, v).to(input_dtype)
-
-    if return_probs:
-        out_probs_dtype = (
-            input_dtype if input_dtype in (torch.float16, torch.bfloat16) else probs.dtype
-        )
-        return attention, probs.to(out_probs_dtype)
-    else:
-        return attention
-
-
-def join_heads(x: torch.Tensor) -> torch.Tensor:
-    """
-    Merge multi-head attention outputs into model dimension.
-
-    Args:
-        x (torch.Tensor): shape (batch_size, num_heads, seq_length, d_head)
-
-    Returns:
-        torch.Tensor: shape (batch_size, seq_length, d_model) where d_model = num_heads * d_head
-    """
-    # --- Type checks ---
-    if not isinstance(x, torch.Tensor):
-        raise TypeError(f"Expected torch.Tensor, got {type(x)}")
-
-    if x.ndim != 4:
-        raise ValueError(
-            f"Expected 4D torch.Tensor (batch, num_heads, seq_len, d_head), "
-            f"got shape {tuple(x.shape)} with ndim={x.ndim}"
-        )
-
-    batch_size, num_heads, seq_length, d_head = x.shape
-
-    if not all(
-        isinstance(dim, int) and dim >= 0 for dim in (batch_size, num_heads, seq_length, d_head)
-    ):
-        raise ValueError(f"Invalid shape values: {x.shape}")
-
-    d_model = num_heads * d_head
-
-    # --- Safe reshape ---
-    x = x.permute(0, 2, 1, 3).contiguous().view(batch_size, seq_length, d_model)
-
-    return x
-
-
-def sinusoidal_positional_encoding(
-    seq_len: int,
-    dim: int,
-    *,
-    base: float = 10_000.0,
-    dtype: torch.dtype | None = None,
-    device: torch.device | None = None,
-    offset: int = 0,
-    return_positions: bool = False,
-) -> torch.Tensor:
-    """
-    Build a [seq_len, dim] sinusoidal (trig) positional encoding table:
-      PE[pos, 2i]   = sin(pos / (base^(2i/dim)))
-      PE[pos, 2i+1] = cos(pos / (base^(2i/dim)))
-
-    Design goals:
-    - Numerically stable: uses exp with -log(base) to avoid huge powers.
-    - Dtype/device safe: computes frequencies in float32 and casts at the end;
-      supports fp16/bf16 by upcasting, then downcasting.
-    - Odd dimensions handled: the last column is zero-padded when dim is odd.
-    - Supports an integer `offset` so you can continue encodings for cached KV.
-    - No autograd needed: returns a tensor with requires_grad=False.
-
-    Args:
-      seq_len: length of sequence (>=0).
-      dim: embedding dimension (>=1).
-      base: frequency base (default 10k like Vaswani et al.).
-      dtype: desired dtype of the returned table (defaults to torch.get_default_dtype()).
-      device: desired device of the returned table.
-      offset: position offset to start from (useful for incremental decoding).
-      return_positions: if True, also returns the [seq_len] position indices.
-
-    Returns:
-      pe: [seq_len, dim] tensor on `device` with `dtype`.
-      (pos): optional [seq_len] tensor of positions (if return_positions=True).
-
-    Raises:
-      ValueError if seq_len<0 or dim<=0.
-    """
-    if seq_len < 0:
-        raise ValueError(f"seq_len must be >= 0, got {seq_len}")
-    if dim <= 0:
-        raise ValueError(f"dim must be > 0, got {dim}")
-    if dtype is None:
-        dtype = torch.get_default_dtype()
-
-    # Work in float32 for stability regardless of target dtype.
-    work_dtype = torch.float32
-
-    # Positions [offset .. offset+seq_len-1]
-    # Note: torch.arange is dtype-agnostic; compute in float32 later.
-    pos = torch.arange(offset, offset + seq_len, device=device)
-
-    # Frequencies for even indices: exp(-(log(base) * (2i)/dim))
-    # Avoid base**(2i/dim) directly to reduce overflow/underflow risk.
-    half_dim = dim // 2  # number of sin/cos pairs
-    if half_dim > 0:
-        # [half_dim]
-        exponent = torch.arange(0, half_dim, device=device, dtype=work_dtype)
-        inv_freq = torch.exp(-math.log(base) * (2.0 * exponent) / float(dim))  # [half_dim]
-
-        # Outer product: [seq_len, half_dim]
-        # Compute phase = pos[:, None] * inv_freq[None, :]
-        phase = pos.to(work_dtype).unsqueeze(1) * inv_freq.unsqueeze(0)
-
-        sin_part = torch.sin(phase)
-        cos_part = torch.cos(phase)
-
-        # Interleave sin and cos along the last dim
-        pe_even_odd = torch.stack((sin_part, cos_part), dim=-1).reshape(seq_len, 2 * half_dim)
-        if dim % 2 == 1:
-            # Pad the last column with zeros if dim is odd
-            pad = torch.zeros(seq_len, 1, device=device, dtype=work_dtype)
-            pe_work = torch.cat([pe_even_odd, pad], dim=1)
-        else:
-            pe_work = pe_even_odd
-    else:
-        # dim == 1 case → just a zero column (consistent with odd-dim padding above)
-        pe_work = torch.zeros(seq_len, 1, device=device, dtype=work_dtype)
-
-        # If somehow dim==0 was allowed earlier, we'd have raised.
-
-    # Cast once at the end to the requested dtype.
-    pe = pe_work.to(dtype=dtype)
-
-    # Ensure no gradient tracking (users usually register this as a buffer).
-    pe.requires_grad_(False)
-
-    return (pe, pos.to(device)) if return_positions else pe
-
-
-def combine_masks(mask1, mask2):
-    """
-    Combine two boolean masks with logical OR.
-
-    Args:
-        mask1 (torch.Tensor or None): first mask
-        mask2 (torch.Tensor or None): second mask
-
-    Returns:
-        torch.Tensor or None: combined mask or None if both are None
-
-    Notes:
-        - Both masks must be broadcastable to the same shape.
-        - Returns a boolean tensor if inputs are boolean.
-        - If one mask is None, returns the other unchanged.
-    """
-    if mask1 is not None and mask2 is not None:
-        m1 = mask1.to(dtype=torch.bool)
-        m2 = mask2.to(dtype=torch.bool, device=m1.device)
-        if m1.shape != m2.shape:
-            try:
-                return m1 | m2  # allow broadcasting, but safe
-            except RuntimeError as e:
-                raise ValueError(
-                    f"combine_masks: masks not broadcastable: {m1.shape} vs {m2.shape}"
-                ) from e
-        return m1 | m2
-    elif mask1 is not None:
-        return mask1
-    elif mask2 is not None:
-        return mask2
-    else:
-        return None
-
-
-def shift_right(labels: torch.Tensor, bos_id: int = 1, pad_id: int = 0) -> torch.Tensor:
-    """
-    Build decoder input IDs by shifting target labels to the right and
-    prepending a BOS token. Any labels marked -100 (ignore_index) are
-    converted to PAD *after* shifting so they never appear as inputs.
-
-    Args:
-        labels: Long tensor of shape [B, T] or [T].
-        bos_id: Token id to place at position 0.
-        pad_id: Token id used where labels == -100 after shifting.
-
-    Returns:
-        decoder_input_ids: Long tensor of shape [B, T] (same device as labels).
-
-    """
-    if labels.dim() == 1:
-        labels = labels.unsqueeze(0)  # [1, T]
-    if labels.dim() != 2:
-        raise ValueError(f"shift_right expects 1D or 2D tensor, got shape {tuple(labels.shape)}")
-
-    B, T = labels.shape
-    if T == 0:
-        raise ValueError("shift_right: sequence length must be > 0")
-
-    # Work in long dtype on the same device; do NOT modify `labels` in-place.
-    labels = labels.to(torch.long)
-
-    # Start with PADs everywhere, then copy shifted labels, and set BOS.
-    decoder_input_ids = labels.new_full((B, T), pad_id)
-    decoder_input_ids[:, 1:] = labels[:, :-1]
-    decoder_input_ids[:, 0] = bos_id
-
-    # Any positions that inherited -100 from labels should be PAD instead.
-    # (labels contains -100 for "ignore_index" targets; these should never be inputs)
-    mask_ignore = decoder_input_ids.eq(-100)
-    if mask_ignore.any():
-        decoder_input_ids.masked_fill_(mask_ignore, pad_id)
-
-    return decoder_input_ids
-
-
-def _ensure_2d_logits(logits: torch.Tensor) -> torch.Tensor:
-    """
-    Accepts [B, V] or [..., V]. If rank > 2, flattens leading dims into batch.
-    Returns 2D [B', V] along with a flag and original shape for potential future use.
-    """
-    if logits.dim() == 2:
-        return logits
-    if logits.dim() >= 3:
-        V = logits.size(-1)
-        return logits.reshape(-1, V)
-    raise ValueError(
-        f"sample_from_logits: logits must be at least 2D, got shape {tuple(logits.shape)}"
-    )
-
-
-def _apply_allow_deny_mask(
-    logits: torch.Tensor,
-    *,
-    allowed_tokens: Iterable[int] | None,
-    disallowed_tokens: Iterable[int] | None,
-    filter_value: float,
-) -> torch.Tensor:
-    if allowed_tokens is not None:
-        mask = torch.zeros_like(logits, dtype=torch.bool)
-        idx = torch.tensor(list(allowed_tokens), device=logits.device)
-        idx = idx[(idx >= 0) & (idx < logits.size(-1))]
-        if idx.numel() > 0:
-            mask.index_fill_(-1, idx, True)
-        logits = torch.where(mask, logits, torch.full_like(logits, filter_value))
-    if disallowed_tokens is not None:
-        idx = torch.tensor(list(disallowed_tokens), device=logits.device)
-        idx = idx[(idx >= 0) & (idx < logits.size(-1))]
-        if idx.numel() > 0:
-            logits.index_fill_(-1, idx, filter_value)
-    return logits
-
-
-def _top_k_filtering(
-    logits: torch.Tensor,
-    top_k: int | None,
-    *,
-    min_tokens_to_keep: int,
-    filter_value: float,
-) -> torch.Tensor:
-    if top_k is None or top_k <= 0:
-        return logits
-    k = min(max(top_k, min_tokens_to_keep), logits.size(-1))
-    # threshold per row
-    values, _ = torch.topk(logits, k, dim=-1)
-    thresh = values[..., -1, None]
-    return torch.where(logits < thresh, torch.full_like(logits, filter_value), logits)
-
-
-def _top_p_filtering(
-    logits_scaled: torch.Tensor,
-    probs: torch.Tensor,
-    top_p: float,
-    *,
-    min_tokens_to_keep: int,
-    filter_value: float,
-) -> torch.Tensor:
-    if top_p is None or not (0.0 < top_p < 1.0):
-        return logits_scaled
-    # sort by probability
-    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)
-    cum = torch.cumsum(sorted_probs, dim=-1)
-    # tokens to remove: everything after first point where cum > top_p
-    to_remove = cum > top_p
-    # always keep at least min_tokens_to_keep highest-prob tokens
-    to_remove[..., :min_tokens_to_keep] = False
-    # scatter back to original indices
-    scatter_mask = torch.zeros_like(to_remove, dtype=torch.bool).scatter(-1, sorted_idx, to_remove)
-    return torch.where(scatter_mask, torch.full_like(logits_scaled, filter_value), logits_scaled)
-
-
-def sample_from_logits(
-    logits: torch.Tensor,
-    *,
-    do_sample: bool = False,
-    temperature: float = 1.0,
-    top_k: int | None = None,
-    top_p: float | None = None,
-    min_tokens_to_keep: int = 1,
-    allowed_tokens: Iterable[int] | None = None,
-    disallowed_tokens: Iterable[int] | None = None,
-    filter_value: float = -float("inf"),
-    rng: torch.Generator | None = None,
-    return_probs: bool = False,
-) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-    """
-    Select token IDs from logits with optional sampling & filtering.
-
-    Inputs:
-      logits: [..., V] or [B, V]; will be treated as 2D [B', V] internally.
-      do_sample: False -> greedy argmax; True -> multinomial sampling.
-      temperature: >0; values <1 sharpen, >1 flatten.
-      top_k: keep only top-k logits (after temperature); None/<=0 disables.
-      top_p: nucleus (0<p<1): keep smallest set with prob mass >= p (after temperature).
-      min_tokens_to_keep: safety floor so we never filter out *everything*.
-      allowed_tokens: iterable of token IDs to allow (others masked out).
-      disallowed_tokens: iterable of token IDs to ban (masked out).
-      filter_value: value used for filtered logits (-inf by default).
-      rng: optional torch.Generator for reproducible sampling.
-      return_probs: if True, also returns the final probabilities used.
-
-    Returns:
-      token_ids: [B'] int64 (and optionally probs [B', V] if return_probs=True).
-    """
-    if temperature <= 0:
-        raise ValueError("temperature must be > 0")
-
-    # Work on a stable float32 copy; keep device and shape
-    logits2d = _ensure_2d_logits(logits).to(dtype=torch.float32)
-
-    # Apply allow/deny lists first (hard mask)
-    logits2d = _apply_allow_deny_mask(
-        logits2d,
-        allowed_tokens=allowed_tokens,
-        disallowed_tokens=disallowed_tokens,
-        filter_value=filter_value,
-    )
-
-    # Scale by temperature before top-k/top-p
-    logits_scaled = logits2d / float(temperature)
-
-    # Top-k first (logit space)
-    logits_scaled = _top_k_filtering(
-        logits_scaled, top_k=top_k, min_tokens_to_keep=min_tokens_to_keep, filter_value=filter_value
-    )
-
-    # Compute probs (needed for top-p and sampling)
-    probs = F.softmax(logits_scaled, dim=-1)
-
-    # Top-p (probability space) — re-mask logits_scaled accordingly
-    if top_p is not None and 0.0 < top_p < 1.0:
-        logits_scaled = _top_p_filtering(
-            logits_scaled,
-            probs,
-            top_p=top_p,
-            min_tokens_to_keep=min_tokens_to_keep,
-            filter_value=filter_value,
-        )
-        probs = F.softmax(logits_scaled, dim=-1)  # recompute after top-p masking
-
-    if do_sample:
-        # multinomial sampling with optional RNG for reproducibility
-        next_ids = torch.multinomial(probs, num_samples=1, replacement=True, generator=rng).squeeze(
-            -1
-        )
-    else:
-        # greedy
-        next_ids = torch.argmax(probs, dim=-1)
-
-    next_ids = next_ids.to(dtype=torch.long, device=logits.device)  # back to original device
-    if return_probs:
-        return next_ids, probs.to(device=logits.device, dtype=probs.dtype)
-    return next_ids
-
-
-def create_causal_mask(
-    max_seq_len: int,
-    *,
-    device: torch.device | None = None,
-    dtype: torch.dtype = torch.bool,
-) -> torch.Tensor:
-    """
-    Build a standard causal mask to prevent attending to future positions.
-
-    Shape: [1, 1, max_seq_len, max_seq_len]
-
-    mask[i, j] = True means position j should be masked when predicting i.
-
-    Args:
-        max_seq_len: length of the target sequence (must be > 0).
-        device: torch device to place the mask on (defaults to cpu).
-        dtype: dtype of mask (default: torch.bool).
-               Usually boolean, but some fused kernels want float with -inf/0.
-
-    Returns:
-        Tensor of shape [1, 1, max_seq_len, max_seq_len].
-    """
-    if max_seq_len <= 0:
-        raise ValueError(f"max_seq_len must be > 0, got {max_seq_len}")
-
-    # upper-triangular (excluding main diagonal)
-    mask = torch.triu(
-        torch.ones((1, 1, max_seq_len, max_seq_len), device=device, dtype=torch.bool),
-        diagonal=1,
-    )
-    if dtype != torch.bool:
-        mask = mask.to(dtype)
-    return mask
+from __future__ import annotations
+
+import math
+import random
+from collections.abc import Callable, Iterable, Sequence
+from pathlib import Path
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+# -----------------------------------------------------------------------------
+# Global seeding utilities
+# -----------------------------------------------------------------------------
+
+
+def set_global_seed(seed: int, *, deterministic: bool = True) -> None:
+    """Seed Python, NumPy, and PyTorch (CPU and CUDA) for reproducible runs."""
+
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    if deterministic:
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+
+def make_worker_init_fn(seed: int) -> Callable[[int], None]:
+    """Return a DataLoader worker init function that derives unique seeds."""
+
+    def _init_fn(worker_id: int) -> None:
+        worker_seed = seed + worker_id
+        random.seed(worker_seed)
+        np.random.seed(worker_seed)
+        torch.manual_seed(worker_seed)
+
+    return _init_fn
+
+
+# -----------------------------------------------------------------------------
+# Attention building blocks
+# -----------------------------------------------------------------------------
+
+
+def split_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor:
+    """Split the last dimension into ``(num_heads, d_head)`` and permute to (B, H, S, d_head)."""
+
+    if not isinstance(x, torch.Tensor):
+        raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
+    if not isinstance(num_heads, int):
+        raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
+    if x.ndim != 3:
+        raise ValueError(
+            f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
+        )
+    if num_heads <= 0:
+        raise ValueError(f"num_heads must be > 0; got {num_heads}")
+
+    batch_size, seq_length, d_model = x.shape
+    if d_model % num_heads != 0:
+        raise ValueError(
+            f"d_model ({d_model}) must be divisible by num_heads ({num_heads});"
+            f" got remainder {d_model % num_heads}"
+        )
+
+    d_head = d_model // num_heads
+    return x.reshape(batch_size, seq_length, num_heads, d_head).permute(0, 2, 1, 3)
+
+
+def join_heads(x: torch.Tensor) -> torch.Tensor:
+    """Merge ``(num_heads, d_head)`` back into the model dimension."""
+
+    if not isinstance(x, torch.Tensor):
+        raise TypeError(f"Expected torch.Tensor, got {type(x)}")
+    if x.ndim != 4:
+        raise ValueError(
+            f"Expected 4D torch.Tensor (batch, num_heads, seq_len, d_head), got {tuple(x.shape)}"
+        )
+
+    batch_size, num_heads, seq_length, d_head = x.shape
+    return x.permute(0, 2, 1, 3).contiguous().view(batch_size, seq_length, num_heads * d_head)
+
+
+def calculate_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    mask: torch.Tensor | None,
+    *,
+    dropout_p: float = 0.0,
+    return_probs: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """Scaled dot-product attention with optional dropout and probability return.
+
+    Args:
+        query, key, value: ``[B, H, S, Dh]`` tensors on the same device/dtype.
+        mask: Optional boolean tensor broadcastable to ``[B, H, Sq, Sk]`` where
+            ``True`` entries are masked. ``None`` disables masking.
+        dropout_p: Dropout probability applied to the attention weights. Callers
+            should pass ``0.0`` when not training.
+        return_probs: When ``True`` returns ``(context, probs)``; otherwise only
+            the context tensor is returned.
+    """
+
+    if not (query.dim() == key.dim() == value.dim() == 4):
+        raise ValueError(
+            "query, key, value must be 4D tensors shaped (B, H, S, Dh);"
+            f" got q={tuple(query.shape)}, k={tuple(key.shape)}, v={tuple(value.shape)}"
+        )
+    if query.device != key.device or query.device != value.device:
+        raise RuntimeError("query, key, value must be on the same device")
+
+    if mask is not None:
+        if mask.dtype != torch.bool:
+            raise TypeError("mask must be boolean when provided")
+        if mask.device != query.device:
+            mask = mask.to(query.device)
+        target_shape = (query.size(0), query.size(1), query.size(2), key.size(2))
+        if mask.shape != target_shape:
+            try:
+                mask = mask.expand(target_shape)
+            except RuntimeError as exc:
+                raise ValueError(
+                    f"mask with shape {tuple(mask.shape)} not broadcastable to {target_shape}"
+                ) from exc
+
+    p = float(dropout_p)
+    if p < 0 or p >= 1:
+        raise ValueError(f"dropout_p must be in [0, 1), got {dropout_p}")
+
+    return _attention_with_probs(query, key, value, mask, p, return_probs)
+
+
+def _attention_with_probs(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    mask: torch.Tensor | None,
+    dropout_p: float,
+    return_probs: bool,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """Manual attention path supporting optional probability return."""
+
+    head_dim = query.size(-1)
+    work_dtype = torch.float32 if query.dtype in (torch.float16, torch.bfloat16) else query.dtype
+
+    q = query.to(work_dtype)
+    k = key.to(work_dtype)
+    scores = torch.matmul(q, k.transpose(-2, -1))
+    scores.mul_(1.0 / math.sqrt(head_dim))
+
+    full_mask_rows = None
+    if mask is not None:
+        mask = mask.to(scores.device)
+        scores = scores.masked_fill(mask, torch.finfo(work_dtype).min)
+        full_mask_rows = mask.all(dim=-1, keepdim=True)
+        if full_mask_rows.any():
+            scores = scores.masked_fill(full_mask_rows, 0.0)
+
+    row_max = scores.max(dim=-1, keepdim=True).values
+    row_max = torch.where(torch.isfinite(row_max), row_max, torch.zeros_like(row_max))
+    logits = scores - row_max
+    probs = torch.softmax(logits, dim=-1)
+    if full_mask_rows is not None:
+        probs = torch.where(full_mask_rows, torch.zeros_like(probs), probs)
+
+    if dropout_p > 0.0:
+        probs = F.dropout(probs, p=dropout_p, training=True)
+
+    context = torch.matmul(probs.to(value.dtype), value)
+    if return_probs:
+        probs_out = probs.to(work_dtype)
+        return context, probs_out
+    return context
+
+
+# -----------------------------------------------------------------------------
+# Positional encodings
+# -----------------------------------------------------------------------------
+
+
+def sinusoidal_positional_encoding(S: int, D: int) -> torch.Tensor:
+    """Return the classic sinusoidal positional encoding table (shape ``[S, D]``)."""
+
+    if S <= 0 or D <= 0:
+        raise ValueError(f"S and D must be > 0, got S={S}, D={D}")
+
+    positions = torch.arange(S, dtype=torch.float32).unsqueeze(1)
+    div_terms = torch.exp(torch.arange(0, D, 2, dtype=torch.float32) * -(math.log(10000.0) / D))
+
+    pe = torch.zeros(S, D, dtype=torch.float32)
+    pe[:, 0::2] = torch.sin(positions * div_terms)
+    pe[:, 1::2] = torch.cos(positions * div_terms[: D // 2])
+    return pe
+
+
+# -----------------------------------------------------------------------------
+# Sampling utilities
+# -----------------------------------------------------------------------------
+
+
+def _ensure_2d_logits(logits: torch.Tensor) -> torch.Tensor:
+    """Reshape logits so that the last dimension is vocabulary sized and the rest flatten."""
+
+    if logits.dim() == 2:
+        return logits
+    if logits.dim() >= 3:
+        V = logits.size(-1)
+        return logits.reshape(-1, V)
+    raise ValueError(f"sample_from_logits: logits must be at least 2D, got {tuple(logits.shape)}")
+
+
+def _apply_allow_deny_mask(
+    logits: torch.Tensor,
+    *,
+    allowed_tokens: Iterable[int] | None,
+    disallowed_tokens: Iterable[int] | None,
+    filter_value: float,
+) -> torch.Tensor:
+    if allowed_tokens is not None:
+        mask = torch.zeros_like(logits, dtype=torch.bool)
+        idx = torch.tensor(list(allowed_tokens), device=logits.device)
+        idx = idx[(idx >= 0) & (idx < logits.size(-1))]
+        if idx.numel() > 0:
+            mask.index_fill_(-1, idx, True)
+        logits = torch.where(mask, logits, torch.full_like(logits, filter_value))
+    if disallowed_tokens is not None:
+        idx = torch.tensor(list(disallowed_tokens), device=logits.device)
+        idx = idx[(idx >= 0) & (idx < logits.size(-1))]
+        if idx.numel() > 0:
+            logits.index_fill_(-1, idx, filter_value)
+    return logits
+
+
+def _top_k_filtering(
+    logits: torch.Tensor,
+    top_k: int | None,
+    *,
+    min_tokens_to_keep: int,
+    filter_value: float,
+) -> torch.Tensor:
+    if top_k is None or top_k <= 0:
+        return logits
+    k = min(max(top_k, min_tokens_to_keep), logits.size(-1))
+    values, _ = torch.topk(logits, k, dim=-1)
+    threshold = values[..., -1, None]
+    return torch.where(logits < threshold, torch.full_like(logits, filter_value), logits)
+
+
+def _top_p_filtering(
+    logits_scaled: torch.Tensor,
+    probs: torch.Tensor,
+    top_p: float,
+    *,
+    min_tokens_to_keep: int,
+    filter_value: float,
+) -> torch.Tensor:
+    if top_p is None or not (0.0 < top_p < 1.0):
+        return logits_scaled
+
+    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)
+    cumulative = torch.cumsum(sorted_probs, dim=-1)
+    to_remove = cumulative > top_p
+    to_remove[..., :min_tokens_to_keep] = False
+    scatter_mask = torch.zeros_like(to_remove, dtype=torch.bool).scatter(-1, sorted_idx, to_remove)
+    return torch.where(scatter_mask, torch.full_like(logits_scaled, filter_value), logits_scaled)
+
+
+def sample_from_logits(
+    logits: torch.Tensor,
+    *,
+    do_sample: bool = False,
+    temperature: float = 1.0,
+    top_k: int | None = None,
+    top_p: float | None = None,
+    min_tokens_to_keep: int = 1,
+    allowed_tokens: Iterable[int] | None = None,
+    disallowed_tokens: Iterable[int] | None = None,
+    repetition_ctx: torch.Tensor | None = None,
+    presence_penalty: float = 0.0,
+    frequency_penalty: float = 0.0,
+    no_repeat_ngram_size: int | None = None,
+    filter_value: float = -float("inf"),
+    rng: torch.Generator | None = None,
+    return_probs: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """Warp logits (temperature, top-k, nucleus, penalties) and produce next-token ids."""
+
+    if temperature <= 0:
+        raise ValueError("temperature must be > 0")
+
+    logits2d = _ensure_2d_logits(logits).to(dtype=torch.float32)
+    batch, vocab = logits2d.shape
+
+    logits2d = _apply_allow_deny_mask(
+        logits2d,
+        allowed_tokens=allowed_tokens,
+        disallowed_tokens=disallowed_tokens,
+        filter_value=filter_value,
+    )
+
+    if repetition_ctx is not None and (presence_penalty > 0.0 or frequency_penalty > 0.0):
+        if repetition_ctx.dim() != 2 or repetition_ctx.size(0) != batch:
+            raise ValueError(
+                f"repetition_ctx must be [B, T]; got {tuple(repetition_ctx.shape)} with B={batch}"
+            )
+        counts = torch.zeros((batch, vocab), device=logits2d.device, dtype=torch.float32)
+        ctx = repetition_ctx
+        valid = (ctx >= 0) & (ctx < vocab)
+        if valid.any():
+            ids = ctx.masked_select(valid)
+            bidx = (
+                torch.arange(batch, device=logits2d.device)
+                .unsqueeze(1)
+                .expand_as(ctx)
+                .masked_select(valid)
+            )
+            counts.index_put_(
+                (bidx, ids), torch.ones_like(ids, dtype=torch.float32), accumulate=True
+            )
+        if presence_penalty > 0.0:
+            logits2d = logits2d - presence_penalty * (counts > 0).to(logits2d.dtype)
+        if frequency_penalty > 0.0:
+            logits2d = logits2d - frequency_penalty * counts
+
+    if (
+        no_repeat_ngram_size is not None
+        and no_repeat_ngram_size >= 2
+        and repetition_ctx is not None
+    ):
+        ctx_cpu = repetition_ctx.detach().to("cpu")
+        for b, seq in enumerate(ctx_cpu.tolist()):
+            tokens = [t for t in seq if 0 <= t < vocab]
+            if len(tokens) < no_repeat_ngram_size:
+                continue
+            history: dict[tuple[int, ...], set[int]] = {}
+            n = int(no_repeat_ngram_size)
+            for i in range(len(tokens) - (n - 1)):
+                key = tuple(tokens[i : i + n - 1])
+                nxt = tokens[i + n - 1]
+                history.setdefault(key, set()).add(nxt)
+            key = tuple(tokens[-(n - 1) :])
+            banned = history.get(key)
+            if banned:
+                idx = torch.tensor(list(banned), device=logits2d.device, dtype=torch.long)
+                logits2d[b].index_fill_(0, idx, filter_value)
+
+    logits_scaled = logits2d / float(temperature)
+    logits_scaled = _top_k_filtering(
+        logits_scaled, top_k=top_k, min_tokens_to_keep=min_tokens_to_keep, filter_value=filter_value
+    )
+    probs = F.softmax(logits_scaled, dim=-1)
+    if top_p is not None and 0.0 < top_p < 1.0:
+        logits_scaled = _top_p_filtering(
+            logits_scaled,
+            probs,
+            top_p=top_p,
+            min_tokens_to_keep=min_tokens_to_keep,
+            filter_value=filter_value,
+        )
+        probs = F.softmax(logits_scaled, dim=-1)
+
+    if do_sample:
+        next_ids = torch.multinomial(probs, num_samples=1, replacement=True, generator=rng).squeeze(
+            -1
+        )
+    else:
+        next_ids = torch.argmax(probs, dim=-1)
+
+    next_ids = next_ids.to(dtype=torch.long, device=logits.device)
+    if return_probs:
+        return next_ids, probs.to(device=logits.device, dtype=probs.dtype)
+    return next_ids
+
+
+# -----------------------------------------------------------------------------
+# Attention mask helpers
+# -----------------------------------------------------------------------------
+
+
+def create_causal_mask(x: torch.Tensor, num_heads: int) -> torch.Tensor:
+    """Return a boolean causal mask ``[B, num_heads, S, S]`` for decoder self-attention."""
+
+    if x.ndim != 2:
+        raise ValueError(f"Expected input of shape (B, S), got {x.shape}")
+    if num_heads <= 0:
+        raise ValueError(f"num_heads must be > 0, got {num_heads}")
+
+    batch, seq_len = x.shape
+    if seq_len <= 0:
+        raise ValueError(f"Sequence length must be > 0, got {seq_len}")
+
+    base = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=x.device), diagonal=1)
+    return base.view(1, 1, seq_len, seq_len).expand(batch, num_heads, seq_len, seq_len).contiguous()
+
+
+def create_qk_padding_mask(
+    query_attention_mask: torch.Tensor, key_attention_mask: torch.Tensor
+) -> torch.Tensor:
+    """Combine query and key padding masks (boolean) into a ``[B, H, Sq, Sk]`` mask."""
+
+    Bq, Hq, _, Sq = query_attention_mask.shape
+    Bk, Hk, _, Sk = key_attention_mask.shape
+    if Bq != Bk or Hq != Hk:
+        raise ValueError(
+            f"Padding mask batch/head mismatch: query {query_attention_mask.shape} vs key {key_attention_mask.shape}"
+        )
+
+    q_mask = query_attention_mask.to(torch.bool).view(Bq, Hq, Sq, 1)
+    k_mask = key_attention_mask.to(torch.bool).view(Bk, Hk, 1, Sk)
+    return (q_mask | k_mask).expand(Bq, Hq, Sq, Sk)
+
+
+def broadcast_padding_mask(mask: torch.Tensor, num_heads: int) -> torch.Tensor:
+    """Expand a ``[B, S]`` padding mask to ``[B, H, 1, S]``."""
+
+    if not isinstance(mask, torch.Tensor):
+        raise TypeError(f"mask must be a torch.Tensor, got {type(mask)}")
+    if mask.dim() != 2:
+        raise ValueError(f"mask must be [B, S], got shape {tuple(mask.shape)}")
+    if not isinstance(num_heads, int) or num_heads <= 0:
+        raise ValueError(f"num_heads must be a positive int, got {num_heads}")
+
+    batch, seq_len = mask.shape
+    return mask.unsqueeze(1).unsqueeze(2).expand(batch, num_heads, 1, seq_len)
+
+
+def combine_masks(
+    causal_mask: torch.Tensor | None, padding_mask: torch.Tensor | None
+) -> torch.Tensor | None:
+    """Combine causal and padding masks (boolean OR) while handling ``None`` values."""
+
+    if padding_mask is None:
+        return causal_mask.to(torch.bool) if causal_mask is not None else None
+    if causal_mask is None:
+        return padding_mask.to(torch.bool)
+
+    m1 = causal_mask.to(torch.bool)
+    m2 = padding_mask.to(torch.bool).to(device=m1.device)
+    try:
+        return m1 | m2
+    except RuntimeError as exc:
+        raise ValueError(
+            f"Masks not broadcastable: causal {tuple(m1.shape)} vs padding {tuple(m2.shape)}"
+        ) from exc
+
+
+# -----------------------------------------------------------------------------
+# Attention introspection / plotting
+# -----------------------------------------------------------------------------
+
+
+def extract_all_attention_maps(
+    model,
+    src_ids: torch.Tensor,
+    tgt_ids: torch.Tensor,
+    src_padding_2d: torch.Tensor,
+    tgt_padding_2d: torch.Tensor,
+):
+    """Collect attention probability tensors for encoder only/decoder self/cross attention."""
+
+    was_training = model.training
+    model.eval()
+    device = src_ids.device
+    num_heads = model.cfg.num_heads
+
+    src_pad_b = broadcast_padding_mask(src_padding_2d.to(device), num_heads)
+    tgt_pad_b = broadcast_padding_mask(tgt_padding_2d.to(device), num_heads)
+    tgt_causal = create_causal_mask(tgt_ids.to(device), num_heads)
+
+    x = model.embed(src_ids.to(device))
+    y = model.embed(tgt_ids.to(device))
+
+    enc_layers = model.encoder.layers
+    dec_layers = model.decoder.layers
+
+    enc_self_maps: list[Tensor] = []
+    dec_self_maps: list[Tensor] = []
+    dec_cross_maps: list[Tensor] = []
+
+    cur = x
+    for layer in enc_layers:
+        mha = layer.attention_layer
+        q = split_heads(mha.query_linear(cur), mha.num_heads)
+        k = split_heads(mha.key_linear(cur), mha.num_heads)
+        v = split_heads(mha.value_linear(cur), mha.num_heads)
+        mask = create_qk_padding_mask(src_pad_b, src_pad_b)
+        _, probs = calculate_attention(q, k, v, mask, dropout_p=0.0, return_probs=True)
+        enc_self_maps.append(probs)
+        cur = layer(cur, src_pad_b)
+
+    mem_full = cur
+
+    y_in = y
+    for layer in dec_layers:
+        self_mha = layer.self_attention_layer
+        q = split_heads(self_mha.query_linear(y_in), self_mha.num_heads)
+        k = split_heads(self_mha.key_linear(y_in), self_mha.num_heads)
+        v = split_heads(self_mha.value_linear(y_in), self_mha.num_heads)
+        pad = create_qk_padding_mask(tgt_pad_b, tgt_pad_b)
+        combined_mask = combine_masks(pad, tgt_causal)
+        self_ctx, self_probs = calculate_attention(
+            q, k, v, combined_mask, dropout_p=0.0, return_probs=True
+        )
+        dec_self_maps.append(self_probs)
+
+        self_out = self_mha(y_in, y_in, y_in, tgt_pad_b, tgt_pad_b, tgt_causal)
+        y_mid = y_in + layer.dropout1(layer.norm1(self_out))
+
+        cross_mha = layer.cross_attention_layer
+        cq = split_heads(cross_mha.query_linear(y_mid), cross_mha.num_heads)
+        ck = split_heads(cross_mha.key_linear(mem_full), cross_mha.num_heads)
+        cv = split_heads(cross_mha.value_linear(mem_full), cross_mha.num_heads)
+        cmask = create_qk_padding_mask(tgt_pad_b, src_pad_b)
+        _, cross_probs = calculate_attention(cq, ck, cv, cmask, dropout_p=0.0, return_probs=True)
+        dec_cross_maps.append(cross_probs)
+
+        y_in = layer(mem_full, y_in, src_pad_b, tgt_pad_b, tgt_causal)
+
+    maps = {"enc_self": enc_self_maps, "dec_self": dec_self_maps, "dec_cross": dec_cross_maps}
+    model.train(was_training)
+    return maps
+
+
+def attach_tokens(maps_dict, tokenizer, src_ids, tgt_ids):
+    """Attach decoded token strings to a maps dictionary produced by ``extract_all_attention_maps``."""
+
+    src_tok = [
+        tokenizer.convert_ids_to_tokens(row, skip_special_tokens=False) for row in src_ids.tolist()
+    ]
+    tgt_tok = [
+        tokenizer.convert_ids_to_tokens(row, skip_special_tokens=False) for row in tgt_ids.tolist()
+    ]
+    maps_dict["src_tokens"] = src_tok
+    maps_dict["tgt_tokens"] = tgt_tok
+    return maps_dict
+
+
+def _imshow(ax, array2d, title="", vmin=0.0, vmax=1.0):
+    image = ax.imshow(array2d, aspect="auto", vmin=vmin, vmax=vmax)
+    ax.set_title(title, fontsize=9)
+    ax.set_xticks([])
+    ax.set_yticks([])
+    return image
+
+
+def plot_layer_heads_grid(
+    maps_dict: dict,
+    layer_idx: int,
+    batch: int = 0,
+    heads: list[int] | None = None,
+    figsize_per_cell: tuple[float, float] = (2.2, 2.2),
+    show_colorbar: bool = True,
+    vmin: float = 0.0,
+    vmax: float = 1.0,
+    *,
+    show: bool = True,
+    save_path: str | Path | None = None,
+):
+    """Plot a single layer as a 3-row (enc self / dec self / dec cross) grid."""
+
+    enc_layers = maps_dict["enc_self"]
+    dec_self_layers = maps_dict["dec_self"]
+    dec_cross_layers = maps_dict["dec_cross"]
+
+    enc_idx = min(layer_idx, len(enc_layers) - 1)
+    dec_idx = min(layer_idx, len(dec_self_layers) - 1)
+
+    enc = enc_layers[enc_idx][batch]
+    dec_self = dec_self_layers[dec_idx][batch]
+    dec_cross = dec_cross_layers[dec_idx][batch]
+
+    num_heads = enc.size(0)
+    heads = list(range(num_heads)) if heads is None else heads
+
+    rows = 3
+    cols = len(heads)
+    fig_width = max(1, int(round(figsize_per_cell[0] * cols)))
+    fig_height = max(1, int(round(figsize_per_cell[1] * rows)))
+    fig, axes = plt.subplots(
+        rows, cols, figsize=(fig_width, fig_height), squeeze=False, constrained_layout=True
+    )
+
+    images = []
+    for column, head in enumerate(heads):
+        images.append(
+            _imshow(axes[0, column], enc[head].cpu().float().numpy(), f"Enc h{head}", vmin, vmax)
+        )
+        _imshow(
+            axes[1, column], dec_self[head].cpu().float().numpy(), f"Dec self h{head}", vmin, vmax
+        )
+        _imshow(
+            axes[2, column], dec_cross[head].cpu().float().numpy(), f"Dec cross h{head}", vmin, vmax
+        )
+
+    if show_colorbar and images:
+        fig.colorbar(images[0], ax=axes, fraction=0.02, pad=0.01)
+    fig.suptitle(f"Layer {layer_idx}", fontsize=12)
+
+    if save_path is not None:
+        save_path = Path(save_path)
+        save_path.parent.mkdir(parents=True, exist_ok=True)
+        fig.savefig(save_path, dpi=150, bbox_inches="tight")
+
+    if show:
+        plt.show()
+    else:
+        plt.close(fig)
+    return fig
+
+
+def plot_all_layers_all_heads(
+    maps_dict: dict,
+    batch: int = 0,
+    max_layers: int | None = None,
+    heads: list[int] | None = None,
+    figsize_per_cell: tuple[float, float] = (2.0, 2.0),
+    vmin: float = 0.0,
+    vmax: float = 1.0,
+    save_pdf_path: str | None = None,
+    *,
+    show: bool = True,
+):
+    """Render attention grids for every layer (optionally saving a multi-page PDF)."""
+
+    from matplotlib.backends.backend_pdf import PdfPages
+
+    enc_layers = len(maps_dict["enc_self"])
+    dec_layers = len(maps_dict["dec_self"])
+    total_layers = max(enc_layers, dec_layers)
+    if max_layers is not None:
+        total_layers = min(total_layers, max_layers)
+
+    def _render_layers(record_page):
+        for layer in range(total_layers):
+            fig = plot_layer_heads_grid(
+                maps_dict,
+                layer_idx=layer,
+                batch=batch,
+                heads=heads,
+                figsize_per_cell=figsize_per_cell,
+                show_colorbar=True,
+                vmin=vmin,
+                vmax=vmax,
+                show=show,
+            )
+            if record_page is not None:
+                record_page(fig)
+                plt.close(fig)
+
+    if save_pdf_path:
+        with PdfPages(save_pdf_path) as pdf:
+            _render_layers(pdf.savefig)
+        print(f"[saved] multi-page attention viewer -> {save_pdf_path}")
+    else:
+        _render_layers(None)
+
+
+def _format_token_sequence(tokens: list[str]) -> str:
+    return " | ".join(tokens)
+
+
+def _resolve_attention_layers(available: int, requested: Sequence[int] | None) -> list[int]:
+    if requested is None:
+        return list(range(available))
+    return sorted({int(idx) for idx in requested if 0 <= int(idx) < available})
+
+
+def _resolve_attention_heads(available: int, requested: Sequence[int] | None) -> list[int]:
+    if requested is None:
+        return list(range(available))
+    return sorted({int(idx) for idx in requested if 0 <= int(idx) < available})
+
+
+# -----------------------------------------------------------------------------
+# Debug helper
+# -----------------------------------------------------------------------------
+
+
+def debug_transformer_forward(
+    model,
+    tokenizer,
+    src_ids: torch.Tensor,
+    tgt_ids: torch.Tensor,
+    *,
+    logits: torch.Tensor | None = None,
+    pad_id: int,
+    device: str | torch.device = "cpu",
+    batch_index: int = 0,
+    sample_index: int = 0,
+    show_attention: bool = False,
+    save_attention: bool = False,
+    average_heads: bool = False,
+    attention_types: Sequence[str] = ("enc_self", "dec_self", "dec_cross"),
+    attention_layers: Sequence[int] | None = None,
+    attention_heads: Sequence[int] | None = None,
+    attention_figsize: tuple[float, float] = (4.0, 4.0),
+    skip_special_tokens: bool = False,
+    log_fn: Callable[[str], None] | None = print,
+    return_maps: bool = False,
+    save_dir: str | Path | None = None,
+    run_dir: str | Path | None = None,
+) -> dict[str, Any]:
+    """Run a forward pass, compute diagnostics, and optionally render attention heatmaps."""
+
+    if log_fn is None:
+
+        def log_fn(_: str) -> None:  # type: ignore[redefinition]
+            return
+
+    device = torch.device(device)
+
+    if src_ids.dim() != 2 or tgt_ids.dim() != 2:
+        raise ValueError("src_ids and tgt_ids must be rank-2 tensors")
+
+    batch_size = src_ids.size(0)
+    if not (0 <= sample_index < batch_size):
+        raise IndexError(f"sample_index {sample_index} out of range for batch size {batch_size}")
+
+    src_ids = src_ids.to(device)
+    tgt_ids = tgt_ids.to(device)
+
+    if tgt_ids.size(1) < 2:
+        raise ValueError("tgt_ids must contain at least BOS and one target token")
+
+    decoder_in = tgt_ids[:, :-1]
+    labels = tgt_ids[:, 1:]
+
+    src_pad_mask = src_ids.eq(pad_id)
+    tgt_pad_mask = decoder_in.eq(pad_id)
+    tgt_pad_full = tgt_ids.eq(pad_id)
+
+    was_training = model.training
+    attention_requested = show_attention or save_attention or return_maps
+
+    if logits is None:
+        if not callable(model):
+            raise TypeError("model must be callable")
+
+        model.eval()
+        with torch.no_grad():
+            logits = model(src_ids, decoder_in, src_pad_mask, tgt_pad_mask)
+        model.train(was_training)
+    else:
+        logits = logits.to(device)
+
+    if logits.dim() != 3 or logits.size(0) != batch_size:
+        raise ValueError("logits must be of shape [batch, seq_len, vocab]")
+
+    vocab_size = logits.size(-1)
+    mask = labels.ne(pad_id)
+
+    with torch.no_grad():
+        log_probs = F.log_softmax(logits, dim=-1)
+        nll = F.nll_loss(
+            log_probs.reshape(-1, vocab_size),
+            labels.reshape(-1),
+            reduction="sum",
+            ignore_index=pad_id,
+        )
+
+    tokens_total = int(mask.sum().item())
+    avg_nll = (nll / max(tokens_total, 1)).item()
+    perplexity = math.exp(avg_nll) if tokens_total > 0 else float("nan")
+
+    pred_ids = logits.argmax(dim=-1)
+    correct_mask = pred_ids.eq(labels) & mask
+    correct_tokens = int(correct_mask.sum().item())
+    token_accuracy = correct_tokens / max(tokens_total, 1)
+
+    sample_labels = labels[sample_index]
+    sample_preds = pred_ids[sample_index]
+    sample_mask = mask[sample_index]
+
+    sample_tokens_total = int(sample_mask.sum().item())
+    sample_correct_tokens = int((sample_preds.eq(sample_labels) & sample_mask).sum().item())
+    sample_token_accuracy = sample_correct_tokens / max(sample_tokens_total, 1)
+    sample_exact_match = bool(torch.equal(sample_preds[sample_mask], sample_labels[sample_mask]))
+
+    sample_decoder_in = decoder_in[sample_index]
+    sample_pred_sequence = torch.cat([sample_decoder_in[:1], sample_preds], dim=0)
+
+    with torch.no_grad():
+        sample_log_probs = F.log_softmax(logits[sample_index], dim=-1)
+        sample_nll = F.nll_loss(
+            sample_log_probs,
+            sample_labels,
+            reduction="sum",
+            ignore_index=pad_id,
+        )
+
+    sample_avg_nll = (sample_nll / max(sample_tokens_total, 1)).item()
+    sample_perplexity = math.exp(sample_avg_nll) if sample_tokens_total > 0 else float("nan")
+
+    sample_src_ids = src_ids[sample_index]
+    sample_tgt_ids = tgt_ids[sample_index]
+
+    sample_src_tokens = tokenizer.convert_ids_to_tokens(
+        sample_src_ids.tolist(),
+        skip_special_tokens=skip_special_tokens,
+    )
+    sample_tgt_tokens = tokenizer.convert_ids_to_tokens(
+        sample_tgt_ids.tolist(),
+        skip_special_tokens=skip_special_tokens,
+    )
+    sample_pred_tokens = tokenizer.convert_ids_to_tokens(
+        sample_pred_sequence.tolist(),
+        skip_special_tokens=skip_special_tokens,
+    )
+
+    sample_src_text = tokenizer.decode(sample_src_ids.tolist(), skip_special_tokens=True)
+    sample_tgt_text = tokenizer.decode(sample_tgt_ids.tolist(), skip_special_tokens=True)
+    sample_pred_text = tokenizer.decode(sample_pred_sequence.tolist(), skip_special_tokens=True)
+
+    log_fn("\n--- DEBUG TRANSFORMER FORWARD ---")
+    log_fn(f"Batch index : {batch_index}")
+    log_fn(f"Sample index: {sample_index}")
+    log_fn(f"Batch token accuracy: {token_accuracy * 100:.2f}% ({correct_tokens}/{tokens_total})")
+    log_fn(
+        f"Batch NLL / ppl : {avg_nll:.4f} / {perplexity:.4f}"
+        if tokens_total > 0
+        else "Batch NLL / ppl : n/a"
+    )
+    log_fn(
+        f"Sample token accuracy: {sample_token_accuracy * 100:.2f}% ({sample_correct_tokens}/{sample_tokens_total})"
+    )
+    log_fn(
+        f"Sample NLL / ppl : {sample_avg_nll:.4f} / {sample_perplexity:.4f}"
+        if sample_tokens_total > 0
+        else "Sample NLL / ppl : n/a"
+    )
+    log_fn(f"Sample exact match : {sample_exact_match}")
+
+    log_fn("-- Source sequence --")
+    log_fn(f"IDs   : {sample_src_ids.tolist()}")
+    log_fn(f"Tokens: {_format_token_sequence(sample_src_tokens)}")
+    log_fn(f"Text  : {sample_src_text}")
+
+    log_fn("-- Target sequence --")
+    log_fn(f"IDs   : {sample_tgt_ids.tolist()}")
+    log_fn(f"Tokens: {_format_token_sequence(sample_tgt_tokens)}")
+    log_fn(f"Text  : {sample_tgt_text}")
+
+    log_fn("-- Predicted sequence --")
+    log_fn(f"IDs   : {sample_pred_sequence.tolist()}")
+    log_fn(f"Tokens: {_format_token_sequence(sample_pred_tokens)}")
+    log_fn(f"Text  : {sample_pred_text}")
+
+    result: dict[str, Any] = {
+        "batch_index": batch_index,
+        "token_accuracy": token_accuracy,
+        "avg_negative_log_likelihood": avg_nll,
+        "perplexity": perplexity,
+        "tokens_total": tokens_total,
+        "correct_tokens": correct_tokens,
+        "sample": {
+            "index": sample_index,
+            "token_accuracy": sample_token_accuracy,
+            "exact_match": sample_exact_match,
+            "avg_negative_log_likelihood": sample_avg_nll,
+            "perplexity": sample_perplexity,
+            "source_ids": sample_src_ids.tolist(),
+            "target_ids": sample_tgt_ids.tolist(),
+            "predicted_ids": sample_pred_sequence.tolist(),
+            "source_tokens": sample_src_tokens,
+            "target_tokens": sample_tgt_tokens,
+            "predicted_tokens": sample_pred_tokens,
+            "source_text": sample_src_text,
+            "target_text": sample_tgt_text,
+            "predicted_text": sample_pred_text,
+        },
+    }
+
+    if attention_requested:
+        if not hasattr(model, "embed"):
+            raise AttributeError("model must expose an embed() method for debugging")
+        with torch.no_grad():
+            attention_maps = extract_all_attention_maps(
+                model,
+                src_ids,
+                tgt_ids,
+                src_pad_mask,
+                tgt_pad_full,
+            )
+
+        target_types = [t for t in attention_types if t in attention_maps]
+        if not target_types:
+            target_types = [k for k in ("enc_self", "dec_self", "dec_cross") if k in attention_maps]
+
+        target_dir: Path | None = None
+        if save_attention:
+            if save_dir is not None:
+                target_dir = Path(save_dir)
+            elif run_dir is not None:
+                target_dir = Path(run_dir) / "debug" / "attention"
+            else:
+                log_fn(
+                    "[warn] save_attention requested but no save_dir/run_dir provided; skipping save."
+                )
+        if target_dir is not None:
+            target_dir.mkdir(parents=True, exist_ok=True)
+
+        sample_label = f"batch{batch_index}_sample{sample_index}"
+        figure_records: list[dict[str, Any]] = []
+
+        for att_type in target_types:
+            layer_maps = attention_maps.get(att_type, [])
+            if not layer_maps:
+                continue
+
+            layers_to_plot = _resolve_attention_layers(len(layer_maps), attention_layers)
+            for layer_idx in layers_to_plot:
+                layer_tensor = layer_maps[layer_idx][sample_index]
+                head_candidates = _resolve_attention_heads(layer_tensor.size(0), attention_heads)
+                if not head_candidates:
+                    continue
+
+                if average_heads:
+                    averaged = layer_tensor[head_candidates].mean(dim=0).cpu().float().numpy()
+                    fig, ax = plt.subplots(figsize=attention_figsize)
+                    im = ax.imshow(averaged, aspect="auto", vmin=0.0, vmax=1.0)
+                    ax.set_title(f"{att_type} L{layer_idx} (avg heads)")
+                    ax.set_xlabel("Key index")
+                    ax.set_ylabel("Query index")
+                    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+
+                    save_path = None
+                    if target_dir is not None:
+                        save_path = target_dir / f"{sample_label}_{att_type}_L{layer_idx}_avg.png"
+                        fig.savefig(save_path, dpi=150, bbox_inches="tight")
+                    if show_attention:
+                        plt.show()
+                    else:
+                        plt.close(fig)
+
+                    figure_records.append(
+                        {
+                            "type": att_type,
+                            "layer": layer_idx,
+                            "heads": "average",
+                            "path": str(save_path) if save_path else None,
+                        }
+                    )
+                else:
+                    cols = len(head_candidates)
+                    fig, axes = plt.subplots(
+                        1,
+                        cols,
+                        figsize=(attention_figsize[0] * cols, attention_figsize[1]),
+                        squeeze=False,
+                        constrained_layout=True,
+                    )
+                    for col, head_idx in enumerate(head_candidates):
+                        ax = axes[0, col]
+                        head_map = layer_tensor[head_idx].cpu().float().numpy()
+                        im = ax.imshow(head_map, aspect="auto", vmin=0.0, vmax=1.0)
+                        ax.set_title(f"{att_type} L{layer_idx} H{head_idx}")
+                        ax.set_xlabel("Key index")
+                        if col == 0:
+                            ax.set_ylabel("Query index")
+                    fig.colorbar(im, ax=axes.ravel().tolist(), fraction=0.046, pad=0.04)
+
+                    save_path = None
+                    if target_dir is not None:
+                        head_tag = "-".join(str(h) for h in head_candidates)
+                        save_path = (
+                            target_dir / f"{sample_label}_{att_type}_L{layer_idx}_H{head_tag}.png"
+                        )
+                        fig.savefig(save_path, dpi=150, bbox_inches="tight")
+                    if show_attention:
+                        plt.show()
+                    else:
+                        plt.close(fig)
+
+                    figure_records.append(
+                        {
+                            "type": att_type,
+                            "layer": layer_idx,
+                            "heads": head_candidates,
+                            "path": str(save_path) if save_path else None,
+                        }
+                    )
+
+        if target_dir is not None:
+            raw_path = target_dir / f"{sample_label}_attention.pt"
+            torch.save(
+                {k: [v.cpu() for v in tensors] for k, tensors in attention_maps.items()}, raw_path
+            )
+            result.setdefault("attention", {})["raw_path"] = str(raw_path)
+
+            summary_path = target_dir / f"{sample_label}_attention_all_layers.pdf"
+            plot_all_layers_all_heads(
+                attention_maps,
+                batch=sample_index,
+                figsize_per_cell=attention_figsize,
+                save_pdf_path=str(summary_path),
+                show=False,
+            )
+            result.setdefault("attention", {})["summary_path"] = str(summary_path)
+
+        if figure_records:
+            result.setdefault("attention", {})["figures"] = figure_records
+        if return_maps:
+            result["attention_maps"] = attention_maps
+
+    log_fn("--- END DEBUG ---\n")
+    return result
+
+
+# -----------------------------------------------------------------------------
+# Config compatibility checker
+# -----------------------------------------------------------------------------
+
+
+def check_tokenizer_model_compatibility(model_cfg, tokenizer_cfg):
+    """Ensure tokenizer and model configuration agree on core vocabulary settings."""
+
+    if model_cfg.vocab_size != tokenizer_cfg.vocab_size:
+        raise ValueError(
+            f"Vocab size mismatch: model={model_cfg.vocab_size}, tokenizer={tokenizer_cfg.vocab_size}"
+        )
+    if model_cfg.max_seq_len != tokenizer_cfg.max_seq_len:
+        raise ValueError(
+            f"Max sequence length mismatch: model={model_cfg.max_seq_len}, tokenizer={tokenizer_cfg.max_seq_len}"
+        )
+    if model_cfg.pad_id != tokenizer_cfg.pad_id:
+        raise ValueError(
+            f"pad_id mismatch: model={model_cfg.pad_id}, tokenizer={tokenizer_cfg.pad_id}"
+        )
+    if model_cfg.bos_id != tokenizer_cfg.bos_id:
+        raise ValueError(
+            f"bos_id mismatch: model={model_cfg.bos_id}, tokenizer={tokenizer_cfg.bos_id}"
+        )
+    if model_cfg.eos_id != tokenizer_cfg.eos_id:
+        raise ValueError(
+            f"eos_id mismatch: model={model_cfg.eos_id}, tokenizer={tokenizer_cfg.eos_id}"
+        )

tests/units/modules/test_attention.py

@@ -11,6 +11,10 @@ def _rand(B=2, S=5, D=12, device="cpu", dtype=torch.float32):
     return torch.randn(B, S, D, device=device, dtype=dtype)
 
 
+def _pad_mask(batch: int, heads: int, seq_len: int, device) -> torch.Tensor:
+    return torch.zeros(batch, heads, 1, seq_len, dtype=torch.bool, device=device)
+
+
 # =======================
 # Constructor checks
 # =======================
@@ -47,15 +51,18 @@ def test_ctor_happy_path():
 def test_forward_requires_tensors_and_mask_type():
     mha = MultiHeadAttention(32, 4, 0.1)
     q, k, v = _rand(2, 5, 32), _rand(2, 5, 32), _rand(2, 5, 32)
+    mask = _pad_mask(2, mha.num_heads, 5, q.device)
 
     with pytest.raises(TypeError):
-        mha("q", k, v, None)
+        mha("q", k, v, mask, mask)
+    with pytest.raises(TypeError):
+        mha(q, "k", v, mask, mask)
     with pytest.raises(TypeError):
-        mha(q, "k", v, None)
+        mha(q, k, "v", mask, mask)
     with pytest.raises(TypeError):
-        mha(q, k, "v", None)
+        mha(q, k, v, "bad", mask)
     with pytest.raises(TypeError):
-        mha(q, k, v, mask="bad")
+        mha(q, k, v, mask, "bad")
 
 
 def test_forward_rank_and_lastdim_checks():
@@ -63,17 +70,20 @@ def test_forward_rank_and_lastdim_checks():
     q = torch.randn(2, 5, 32)
     k = torch.randn(2, 5, 32)
     v = torch.randn(2, 5, 32)
+    mask = _pad_mask(2, mha.num_heads, 5, q.device)
 
     with pytest.raises(ValueError):
-        mha(q.unsqueeze(0), k, v, None)  # rank 4 for q
+        bad_q = q.unsqueeze(0)
+        bad_q_mask = _pad_mask(bad_q.shape[0], mha.num_heads, bad_q.shape[1], bad_q.device)
+        mha(bad_q, k, v, bad_q_mask, mask)
     with pytest.raises(ValueError):
-        mha(q, k.view(10, 32), v, None)  # rank 2 for k
+        mha(q, k.view(10, 32), v, mask, mask)
     with pytest.raises(ValueError):
-        mha(q[..., :16], k, v, None)  # wrong Dq
+        mha(q[..., :16], k, v, mask, mask)
     with pytest.raises(ValueError):
-        mha(q, k[..., :16], v, None)  # wrong Dk
+        mha(q, k[..., :16], v, mask, mask)
     with pytest.raises(ValueError):
-        mha(q, k, v[..., :16], None)  # wrong Dv
+        mha(q, k, v[..., :16], mask, mask)
 
 
 def test_forward_batch_and_seq_mismatch_checks():
@@ -81,11 +91,15 @@ def test_forward_batch_and_seq_mismatch_checks():
     q = torch.randn(2, 5, 24)
     k = torch.randn(3, 5, 24)  # batch mismatch
     v = torch.randn(2, 5, 24)  # seq mismatch with k
+    q_mask = _pad_mask(q.shape[0], mha.num_heads, q.shape[1], q.device)
+    k_mask = _pad_mask(k.shape[0], mha.num_heads, k.shape[1], k.device)
 
     with pytest.raises(ValueError):
-        mha(q, k, torch.randn(3, 5, 24), None)  # batch mismatch q vs k/v
+        mha(q, k, torch.randn(3, 5, 24), q_mask, k_mask)
     with pytest.raises(ValueError):
-        mha(q, torch.randn(2, 6, 24), v, None)  # Sk != Sv
+        bad_k = torch.randn(2, 6, 24)
+        bad_k_mask = _pad_mask(bad_k.shape[0], mha.num_heads, bad_k.shape[1], bad_k.device)
+        mha(q, bad_k, v, q_mask, bad_k_mask)
 
 
 # =======================
@@ -93,7 +107,7 @@ def test_forward_batch_and_seq_mismatch_checks():
 # =======================
 
 
-@pytest.mark.parametrize("B,S,D,H", [(1, 1, 16, 4), (2, 5, 32, 4), (3, 0, 24, 3)])
+@pytest.mark.parametrize("B,S,D,H", [(1, 1, 16, 4), (2, 5, 32, 4)])
 def test_forward_shapes_and_device_dtype(B, S, D, H):
     device = (
         torch.device(f"cuda:{torch.cuda.current_device()}")
@@ -104,8 +118,9 @@ def test_forward_shapes_and_device_dtype(B, S, D, H):
     q = _rand(B, S, D, device=device)
     k = _rand(B, S, D, device=device)
     v = _rand(B, S, D, device=device)
+    mask = _pad_mask(B, H, S, device)
 
-    out = mha(q, k, v, mask=None)
+    out = mha(q, k, v, mask, mask)
     assert out.shape == (B, S, D)
     assert out.device == device
     assert out.dtype == q.dtype
@@ -118,22 +133,25 @@ def test_boolean_mask_blocks_positions():
     k = _rand(B, S, D)
     v = _rand(B, S, D)
 
-    # mask last two keys for all heads/queries
-    mask = torch.zeros(B, 1, 1, S, dtype=torch.bool)
-    mask[..., -2:] = True
+    q_mask = _pad_mask(B, H, S, q.device)
+    k_mask = _pad_mask(B, H, S, q.device)
+    k_mask[..., -2:] = True
 
-    out1 = mha(q, k, v, mask=None)
-    out2 = mha(q, k, v, mask=mask)
+    out1 = mha(q, k, v, q_mask, q_mask)
+    out2 = mha(q, k, v, q_mask, k_mask)
     assert not torch.allclose(out1, out2)
 
 
-def test_additive_mask_supported():
+def test_causal_mask_blocks_future_positions():
     B, S, D, H = 2, 5, 32, 4
     mha = MultiHeadAttention(D, H, 0.0)
     q, k, v = _rand(B, S, D), _rand(B, S, D), _rand(B, S, D)
-    add = torch.zeros(1, 1, 1, S)
-    out = mha(q, k, v, add)
-    assert out.shape == (B, S, D)
+    base_mask = _pad_mask(B, H, S, q.device)
+    causal = torch.ones(B, H, S, S, dtype=torch.bool, device=q.device).triu(1)
+
+    out_free = mha(q, k, v, base_mask, base_mask, None)
+    out_causal = mha(q, k, v, base_mask, base_mask, causal)
+    assert not torch.allclose(out_free, out_causal)
 
 
 # =======================
@@ -167,8 +185,9 @@ def test_gradients_flow():
     q.requires_grad_(True)
     k.requires_grad_(True)
     v.requires_grad_(True)
+    mask = _pad_mask(2, 4, 5, q.device)
 
-    out = mha(q, k, v, mask=None)
+    out = mha(q, k, v, mask, mask)
     loss = out.pow(2).mean()
     loss.backward()
 

tests/units/modules/test_decoder.py

@@ -45,17 +45,19 @@ def test_transformer_decoder_layers_count_checks():
 # -------- forward path --------
 
 
-@pytest.mark.parametrize("B,Sx,Sy,D,H,FF,L", [(2, 5, 6, 24, 3, 48, 2), (1, 0, 0, 16, 1, 32, 1)])
-def test_decoder_forward_happy_path_and_zero_len(B, Sx, Sy, D, H, FF, L):
+@pytest.mark.parametrize("B,Sx,Sy,D,H,FF,L", [(2, 5, 6, 24, 3, 48, 2)])
+def test_decoder_forward_happy_path(B, Sx, Sy, D, H, FF, L):
     device = _dev()
     dec = TransformerDecoder(D, H, FF, L, 0.1).to(device)
 
     x = torch.randn(B, Sx, D, device=device)
     y = torch.randn(B, Sy, D, device=device)
-
-    src_pad = torch.zeros(B, 1, 1, Sx, dtype=torch.bool, device=device) if Sx > 0 else None
-    tgt_pad = torch.zeros(B, 1, 1, Sy, dtype=torch.bool, device=device) if Sy > 0 else None
-    causal = torch.ones(1, 1, Sy, Sy, dtype=torch.bool, device=device).triu(1) if Sy > 0 else None
+    heads = dec.layers[0].self_attention_layer.num_heads
+    src_pad = torch.zeros(B, heads, 1, Sx, dtype=torch.bool, device=device)
+    tgt_pad = torch.zeros(B, heads, 1, Sy, dtype=torch.bool, device=device)
+    causal = (
+        torch.ones(B, heads, Sy, Sy, dtype=torch.bool, device=device).triu(1) if Sy > 0 else None
+    )
 
     out = dec(x, y, src_pad, tgt_pad, causal)
     assert out.shape == (B, Sy, D)
@@ -64,31 +66,45 @@ def test_decoder_forward_happy_path_and_zero_len(B, Sx, Sy, D, H, FF, L):
 
 def test_decoder_forward_input_checks_and_message_format():
     layer = DecoderLayer(24, 3, 48, 0.1)
+    heads = layer.self_attention_layer.num_heads
+    base_src_mask = torch.zeros(2, heads, 1, 3, dtype=torch.bool)
+    base_tgt_mask = torch.zeros(2, heads, 1, 3, dtype=torch.bool)
 
     with pytest.raises(TypeError):
-        layer("not a tensor", torch.randn(2, 3, 24), None, None, None)
+        layer("not a tensor", torch.randn(2, 3, 24), base_src_mask, base_tgt_mask, None)
     with pytest.raises(TypeError):
-        layer(torch.randn(2, 3, 24), "not a tensor", None, None, None)
+        layer(torch.randn(2, 3, 24), "not a tensor", base_src_mask, base_tgt_mask, None)
 
     with pytest.raises(ValueError) as e1:
-        layer(torch.randn(2, 3, 24, 5), torch.randn(2, 3, 24), None, None, None)  # x rank 4
+        layer(
+            torch.randn(2, 3, 24, 5),
+            torch.randn(2, 3, 24),
+            base_src_mask,
+            base_tgt_mask,
+            None,
+        )  # x rank 4
     assert "x must be a 3D torch.Tensor of shape (B, S, D)" in str(e1.value)
 
     with pytest.raises(ValueError) as e2:
-        layer(torch.randn(2, 3, 24), torch.randn(2, 3, 24, 5), None, None, None)  # y rank 4
+        layer(
+            torch.randn(2, 3, 24),
+            torch.randn(2, 3, 24, 5),
+            base_src_mask,
+            base_tgt_mask,
+            None,
+        )  # y rank 4
     assert "y must be a 3D torch.Tensor of shape (B, S, D)" in str(e2.value)
 
     x = torch.randn(2, 5, 24)
     y = torch.randn(3, 6, 24)  # batch mismatch
+    src_mask_x = torch.zeros(2, heads, 1, 5, dtype=torch.bool)
+    tgt_mask_y = torch.zeros(3, heads, 1, 6, dtype=torch.bool)
     with pytest.raises(ValueError) as e3:
-        layer(x, y, None, None, None)
-    assert "Batch size or d_model mismatch between encoder memory and decoder input" in str(
-        e3.value
-    )
+        layer(x, y, src_mask_x, tgt_mask_y, None)
+    assert "Encoder memory and decoder input must match in batch and d_model" in str(e3.value)
 
     y2 = torch.randn(2, 6, 16)  # d_model mismatch
+    tgt_mask_y2 = torch.zeros(2, heads, 1, 6, dtype=torch.bool)
     with pytest.raises(ValueError) as e4:
-        layer(x, y2, None, None, None)
-    assert "Batch size or d_model mismatch between encoder memory and decoder input" in str(
-        e4.value
-    )
+        layer(x, y2, src_mask_x, tgt_mask_y2, None)
+    assert "Encoder memory and decoder input must match in batch and d_model" in str(e4.value)

tests/units/modules/test_embedding.py

@@ -42,9 +42,9 @@ def test_ctor_value_checks_basic():
         InputEmbedding(10, 8, 16, -1)
 
 
-def test_ctor_allows_zero_length_and_sets_attrs():
-    m = InputEmbedding(11, 8, 0, 0)
-    assert m.sequence_length == 0
+def test_ctor_rejects_zero_length_max_seq():
+    with pytest.raises(ValueError):
+        InputEmbedding(11, 8, 0, 0)
 
 
 def test_padding_row_zero_and_stays_zero_after_step():
@@ -114,7 +114,7 @@ def test_forward_happy_path_shape_dtype_device_and_zero_len():
 def test_forward_adds_positions_not_just_tokens():
     m = InputEmbedding(20, 12, 32, 0)
     x = _rand_ids(B=2, S=5, vocab=20)
-    tok_only = m.token_embed(x)
+    tok_only = m.token_embed(x) * (m.d_model**0.5)
     out = m(x)
     assert not torch.allclose(out, tok_only)
 
@@ -126,7 +126,7 @@ def test_forward_respects_sequence_length_limit():
     x_bad = _rand_ids(B=2, S=6, vocab=16)
     with pytest.raises(ValueError) as ei:
         _ = m(x_bad)
-    assert "exceeds max_seq_len" in str(ei.value)
+    assert "Sequence length" in str(ei.value) and "exceeds max_seq_len" in str(ei.value)
 
 
 # =======================
@@ -141,11 +141,11 @@ def test_positional_ctor_value_checks():
         PositionalEmbedding(16, 0)  # d_model <= 0
 
 
-def test_positional_buffer_registered_and_constant_zero_ok():
-    pe = PositionalEmbedding(0, 10)  # zero supported
+def test_positional_buffer_registered_and_constant_shape():
+    pe = PositionalEmbedding(4, 10)
     assert hasattr(pe, "pe")
     assert isinstance(pe.pe, torch.Tensor)
-    assert pe.pe.shape == (1, 0, 10)
+    assert pe.pe.shape == (1, 4, 10)
     assert pe.pe.requires_grad is False
 
 

tests/units/modules/test_encoder.py

@@ -45,13 +45,14 @@ def test_transformer_encoder_layers_count_checks():
 # -------- forward path --------
 
 
-@pytest.mark.parametrize("B,S,D,H,FF,L", [(2, 5, 24, 3, 48, 2), (1, 0, 16, 1, 32, 1)])
-def test_encoder_forward_happy_path_and_zero_len(B, S, D, H, FF, L):
+@pytest.mark.parametrize("B,S,D,H,FF,L", [(2, 5, 24, 3, 48, 2)])
+def test_encoder_forward_happy_path(B, S, D, H, FF, L):
     device = _dev()
     enc = TransformerEncoder(D, H, FF, L, 0.1).to(device)
 
     x = torch.randn(B, S, D, device=device)
-    src_pad = torch.zeros(B, 1, 1, S, dtype=torch.bool, device=device) if S > 0 else None
+    heads = enc.layers[0].attention_layer.num_heads
+    src_pad = torch.zeros(B, heads, 1, S, dtype=torch.bool, device=device)
 
     out = enc(x, src_pad)
     assert out.shape == (B, S, D)

tests/units/test_transformer.py

@@ -1,8 +1,9 @@
 import pytest
 import torch
 
-from transformer.configs import BasicEncDecCfg
+from transformer.configs import ModelCfg
 from transformer.transformer import BasicEncoderDecoderTransformer
+from transformer.utils import broadcast_padding_mask
 
 
 def _dev():
@@ -16,6 +17,10 @@ def _dev():
 def _tiny_cfg(**over):
     # Small, fast config for tests; zero-length safe via max_seq_len >= 0
     base = dict(
+        name="test",
+        best_checkpoint_path="/tmp/best.ckpt",
+        latest_checkpoint_path="/tmp/latest.ckpt",
+        tokenizer="bpe_8k",
         vocab_size=32,
         d_model=16,
         d_ff=32,
@@ -28,7 +33,15 @@ def _tiny_cfg(**over):
         dropout_rate=0.1,
     )
     base.update(over)
-    return BasicEncDecCfg(**base)
+    return ModelCfg(**base)
+
+
+def _zeros_mask(batch: int, seq_len: int, device) -> torch.Tensor:
+    return torch.zeros(batch, seq_len, dtype=torch.bool, device=device)
+
+
+def _broadcast_mask(mask_2d: torch.Tensor, num_heads: int) -> torch.Tensor:
+    return broadcast_padding_mask(mask_2d, num_heads)
 
 
 # -----------------------
@@ -64,7 +77,7 @@ def test_ctor_cfg_type_and_values():
 # -----------------------
 
 
-@pytest.mark.parametrize("Sx,Sy", [(5, 6), (0, 1)])
+@pytest.mark.parametrize("Sx,Sy", [(5, 6)])
 def test_forward_pipeline_shapes_and_types(Sx, Sy):
     device = _dev()
     cfg = _tiny_cfg(max_seq_len=16)
@@ -74,16 +87,19 @@ def test_forward_pipeline_shapes_and_types(Sx, Sy):
     src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
     tgt = torch.randint(0, cfg.vocab_size, (B, Sy), dtype=torch.long, device=device)
 
-    # padding masks (boolean), broadcastable; allow None when S == 0
-    src_mask = torch.zeros(B, 1, 1, Sx, dtype=torch.bool, device=device) if Sx > 0 else None
-    tgt_mask = torch.zeros(B, 1, 1, Sy, dtype=torch.bool, device=device) if Sy > 0 else None
+    # padding masks (boolean) shaped [B, S]; zero-length handled via empty tensors
+    src_mask_2d = _zeros_mask(B, Sx, device)
+    tgt_mask_2d = _zeros_mask(B, Sy, device)
 
-    logits = model(src, tgt, src_mask, tgt_mask)
+    logits = model(src, tgt, src_mask_2d, tgt_mask_2d)
     assert logits.shape == (B, Sy, cfg.vocab_size)
     assert logits.device == device
     assert logits.dtype == torch.get_default_dtype()
 
     # encode / decode individually
+    src_mask = _broadcast_mask(src_mask_2d, cfg.num_heads)
+    tgt_mask = _broadcast_mask(tgt_mask_2d, cfg.num_heads)
+
     mem = model.encode(src, src_mask)
     assert mem.shape == (B, Sx, cfg.d_model)
     out = model.decode(mem, tgt, src_mask, tgt_mask)
@@ -97,23 +113,28 @@ def test_forward_dtype_checks_and_shape_errors():
     B, Sx, Sy = 2, 4, 5
     src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long)
     tgt = torch.randint(0, cfg.vocab_size, (B, Sy), dtype=torch.long)
+    src_mask = _zeros_mask(B, Sx, src.device)
+    tgt_mask = _zeros_mask(B, Sy, tgt.device)
 
     with pytest.raises(TypeError):
-        model("not a tensor", tgt, None, None)
+        model("not a tensor", tgt, src_mask, tgt_mask)
     with pytest.raises(TypeError):
-        model(src, "not a tensor", None, None)
+        model(src, "not a tensor", src_mask, tgt_mask)
     with pytest.raises(ValueError):
-        model(src.unsqueeze(0), tgt, None, None)  # rank 3 src_ids
+        bad_src = src.unsqueeze(0)
+        bad_src_mask = _zeros_mask(bad_src.shape[0], bad_src.shape[1], bad_src.device)
+        model(bad_src, tgt, bad_src_mask, tgt_mask)
     with pytest.raises(ValueError):
-        model(src, tgt.unsqueeze(-1), None, None)  # rank 3 tgt_ids
+        bad_tgt = tgt.unsqueeze(-1)
+        model(src, bad_tgt, src_mask, tgt_mask)
     with pytest.raises(TypeError):
-        model(src.float(), tgt, None, None)  # wrong dtype
+        model(src.float(), tgt, src_mask, tgt_mask)  # wrong dtype
     with pytest.raises(TypeError):
-        model(src, tgt.int(), None, None)  # wrong dtype
+        model(src, tgt.int(), src_mask, tgt_mask)  # wrong dtype
     with pytest.raises(TypeError):
-        model(src, tgt, src_padding_mask="not a tensor", tgt_padding_mask=None)
+        model(src, tgt, src_padding_mask="not a tensor", tgt_padding_mask=tgt_mask)
     with pytest.raises(TypeError):
-        model(src, tgt, src_padding_mask=None, tgt_padding_mask="not a tensor")
+        model(src, tgt, src_mask, tgt_padding_mask="not a tensor")
 
 
 def test_encode_decode_errors_and_messages():
@@ -121,19 +142,21 @@ def test_encode_decode_errors_and_messages():
     model = BasicEncoderDecoderTransformer(cfg)
 
     src = torch.randint(0, cfg.vocab_size, (2, 4))
+    src_mask = _broadcast_mask(_zeros_mask(src.shape[0], src.shape[1], src.device), cfg.num_heads)
     with pytest.raises(ValueError):
-        model.encode(src.unsqueeze(-1), None)  # rank 3 not allowed
+        model.encode(src.unsqueeze(-1), src_mask)  # rank 3 not allowed
     with pytest.raises(TypeError):
-        model.encode(src.float(), None)  # dtype must be long
+        model.encode(src.float(), src_mask)  # dtype must be long
 
     mem = torch.randn(2, 4, cfg.d_model)
     tgt = torch.randint(0, cfg.vocab_size, (2, 5))
+    tgt_mask = _broadcast_mask(_zeros_mask(tgt.shape[0], tgt.shape[1], tgt.device), cfg.num_heads)
     with pytest.raises(ValueError):
-        model.decode(mem.unsqueeze(0), tgt, None, None)  # mem rank 4
+        model.decode(mem.unsqueeze(0), tgt, src_mask, tgt_mask)  # mem rank 4
     with pytest.raises(ValueError):
-        model.decode(mem, tgt.unsqueeze(-1), None, None)  # tgt rank 3
+        model.decode(mem, tgt.unsqueeze(-1), src_mask, tgt_mask)  # tgt rank 3
     with pytest.raises(TypeError):
-        model.decode(mem, tgt.float(), None, None)  # tgt dtype must be long
+        model.decode(mem, tgt.float(), src_mask, tgt_mask)  # tgt dtype must be long
 
 
 def test_exceeding_max_seq_len_raises():
@@ -141,8 +164,9 @@ def test_exceeding_max_seq_len_raises():
     model = BasicEncoderDecoderTransformer(cfg)
     # src longer than max -> PositionalEmbedding should raise
     src = torch.randint(0, cfg.vocab_size, (2, 5), dtype=torch.long)
+    src_mask = _broadcast_mask(_zeros_mask(src.shape[0], src.shape[1], src.device), cfg.num_heads)
     with pytest.raises(ValueError):
-        model.encode(src, None)
+        model.encode(src, src_mask)
 
 
 # -----------------------
@@ -158,8 +182,8 @@ def test_backward_through_full_forward_and_tied_weights():
     B, Sx, Sy = 2, 6, 5
     src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
     tgt = torch.randint(0, cfg.vocab_size, (B, Sy), dtype=torch.long, device=device)
-    src_mask = torch.zeros(B, 1, 1, Sx, dtype=torch.bool, device=device)
-    tgt_mask = torch.zeros(B, 1, 1, Sy, dtype=torch.bool, device=device)
+    src_mask = _zeros_mask(B, Sx, device)
+    tgt_mask = _zeros_mask(B, Sy, device)
 
     logits = model(src, tgt, src_mask, tgt_mask)  # (B, Sy, V)
     # Dummy labels (language modeling): predict tgt itself; ignore pads via mask
@@ -184,7 +208,7 @@ def test_generate_arg_checks_and_output_shapes():
 
     B, Sx = 2, 4
     src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
-    src_mask = torch.zeros(B, 1, 1, Sx, dtype=torch.bool, device=device)
+    src_mask = _zeros_mask(B, Sx, device)
 
     # type/value errors
     with pytest.raises(TypeError):
@@ -205,14 +229,16 @@ def test_generate_arg_checks_and_output_shapes():
         model.generate(src, src_mask, temperature=0.0)
     with pytest.raises(TypeError):
         model.generate(src, src_mask, top_k="5")
-    with pytest.raises(ValueError):
-        model.generate(src, src_mask, top_k=0)
     with pytest.raises(TypeError):
         model.generate(src, src_mask, top_p="0.9")
     with pytest.raises(ValueError):
         model.generate(src, src_mask, top_p=0.0)
     with pytest.raises(ValueError):
         model.generate(src, src_mask, top_p=1.1)
+    with pytest.raises(TypeError):
+        model.generate(src, src_mask, seed="123")
+    with pytest.raises(TypeError):
+        model.generate(src, src_mask, generator="not a generator")
 
     # happy path: shapes & dtypes; content is stochastic so we don't assert exact ids
     out = model.generate(src, src_mask, max_new_tokens=7, temperature=1.0, top_k=None, top_p=None)
@@ -231,12 +257,83 @@ def test_generate_zero_new_tokens_returns_only_bos():
     model = BasicEncoderDecoderTransformer(cfg).to(device)
 
     src = torch.randint(0, cfg.vocab_size, (2, 3), dtype=torch.long, device=device)
-    out = model.generate(src, None, max_new_tokens=0)
+    src_mask = _zeros_mask(src.shape[0], src.shape[1], device)
+    out = model.generate(src, src_mask, max_new_tokens=0)
     assert out.shape == (2, 1)  # just the BOS token
 
 
+def test_generate_sampling_seed_reproducible():
+    device = _dev()
+    cfg = _tiny_cfg()
+    model = BasicEncoderDecoderTransformer(cfg).to(device).eval()
+
+    B, Sx = 2, 4
+    src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
+    src_mask = _zeros_mask(B, Sx, device)
+
+    out_seed_a = model.generate(
+        src,
+        src_mask,
+        max_new_tokens=5,
+        do_sample=True,
+        temperature=1.0,
+        top_k=None,
+        top_p=None,
+        seed=123,
+    )
+    out_seed_b = model.generate(
+        src,
+        src_mask,
+        max_new_tokens=5,
+        do_sample=True,
+        temperature=1.0,
+        top_k=None,
+        top_p=None,
+        seed=123,
+    )
+    assert torch.equal(out_seed_a, out_seed_b)
+
+    out_seed_c = model.generate(
+        src,
+        src_mask,
+        max_new_tokens=5,
+        do_sample=True,
+        temperature=1.0,
+        top_k=None,
+        top_p=None,
+        seed=124,
+    )
+    assert not torch.equal(out_seed_a, out_seed_c)
+
+    gen1 = torch.Generator(device=device)
+    gen1.manual_seed(999)
+    gen2 = torch.Generator(device=device)
+    gen2.manual_seed(999)
+    out_gen_a = model.generate(
+        src,
+        src_mask,
+        max_new_tokens=5,
+        do_sample=True,
+        temperature=1.0,
+        top_k=None,
+        top_p=None,
+        generator=gen1,
+    )
+    out_gen_b = model.generate(
+        src,
+        src_mask,
+        max_new_tokens=5,
+        do_sample=True,
+        temperature=1.0,
+        top_k=None,
+        top_p=None,
+        generator=gen2,
+    )
+    assert torch.equal(out_gen_a, out_gen_b)
+
+
 # -----------------------
-# Mask broadcast sanity (boolean & additive are handled in utils; we just pass through)
+# Mask broadcast sanity (boolean masks handled in utils; we just pass through)
 # -----------------------
 
 
@@ -249,8 +346,8 @@ def test_forward_accepts_boolean_padding_masks():
     src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
     tgt = torch.randint(0, cfg.vocab_size, (B, Sy), dtype=torch.long, device=device)
 
-    src_pad_bool = torch.zeros(B, 1, 1, Sx, dtype=torch.bool, device=device)
-    tgt_pad_bool = torch.zeros(B, 1, 1, Sy, dtype=torch.bool, device=device)
+    src_pad_bool = _zeros_mask(B, Sx, device)
+    tgt_pad_bool = _zeros_mask(B, Sy, device)
 
     logits = model(src, tgt, src_pad_bool, tgt_pad_bool)
     assert logits.shape == (B, Sy, cfg.vocab_size)
@@ -273,7 +370,9 @@ def test_weight_tying_and_optimizer_step_changes_weights_once():
     B, Sx, Sy = 2, 5, 4
     src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
     tgt = torch.randint(0, cfg.vocab_size, (B, Sy), dtype=torch.long, device=device)
-    logits = model(src, tgt, None, None)  # (B, Sy, V)
+    src_mask = _zeros_mask(B, Sx, device)
+    tgt_mask = _zeros_mask(B, Sy, device)
+    logits = model(src, tgt, src_mask, tgt_mask)  # (B, Sy, V)
 
     # Simple loss to get gradients flowing
     loss = logits.pow(2).mean()
@@ -310,7 +409,9 @@ def test_decode_invariance_to_future_tokens():
 
     B, Sx, Sy = 1, 4, 5
     src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
-    mem = model.encode(src, None)
+    src_mask_2d = _zeros_mask(B, Sx, device)
+    src_mask = _broadcast_mask(src_mask_2d, cfg.num_heads)
+    mem = model.encode(src, src_mask)
 
     # Construct y1 and y2 identical up to t=3, different afterwards
     y1 = torch.randint(0, cfg.vocab_size, (B, Sy), dtype=torch.long, device=device)
@@ -319,8 +420,10 @@ def test_decode_invariance_to_future_tokens():
     if Sy > t + 1:
         y2[:, t + 1 :] = (y1[:, t + 1 :] + 1) % cfg.vocab_size
 
-    out1 = model.decode(mem, y1, None, None)  # (B, Sy, D)
-    out2 = model.decode(mem, y2, None, None)
+    tgt_mask_2d = _zeros_mask(B, Sy, device)
+    tgt_mask = _broadcast_mask(tgt_mask_2d, cfg.num_heads)
+    out1 = model.decode(mem, y1, src_mask, tgt_mask)  # (B, Sy, D)
+    out2 = model.decode(mem, y2, src_mask, tgt_mask)
 
     # Compare hidden states up to and including position t
     assert torch.allclose(out1[:, : t + 1], out2[:, : t + 1], atol=1e-5, rtol=1e-5)

tests/units/test_utils.py

@@ -10,7 +10,6 @@ from transformer.utils import (
     create_causal_mask,
     join_heads,
     sample_from_logits,
-    shift_right,
     sinusoidal_positional_encoding,
     split_heads,
 )
@@ -230,19 +229,6 @@ def test_half_precision_close_to_fp32_or_skip():
 # ----------------------------
 # D. Edge cases
 # ----------------------------
-def test_zero_length_sequences():
-    B, H, D = 2, 3, 8
-    # Sq==0
-    q, k, v = _rand(B, H, 0, 5, D)
-    a, p = calculate_attention(q, k, v, mask=None, return_probs=True)
-    assert a.shape == (B, H, 0, D) and p.shape == (B, H, 0, 5)
-    # Sk==0
-    q, k, v = _rand(B, H, 4, 0, D)
-    a, p = calculate_attention(q, k, v, mask=None, return_probs=True)
-    assert a.shape == (B, H, 4, D) and p.shape == (B, H, 4, 0)
-    assert torch.allclose(a, torch.zeros_like(a))
-
-
 def test_singleton_softmax_is_one():
     q = torch.randn(1, 1, 1, 1)
     k = torch.randn(1, 1, 1, 1)
@@ -266,13 +252,11 @@ def test_gradients_flow_no_inplace_breakage():
         assert t.grad is not None and torch.isfinite(t.grad).all()
 
 
-def test_deterministic_best_effort_same_seed():
+def test_repeated_calls_consistent_without_dropout():
     B, H, Sq, Sk, D = 2, 3, 4, 5, 8
     q, k, v = _rand(B, H, Sq, Sk, D)
-    torch.manual_seed(42)
-    a1 = calculate_attention(q, k, v, mask=None, deterministic=True)
-    torch.manual_seed(42)
-    a2 = calculate_attention(q, k, v, mask=None, deterministic=True)
+    a1 = calculate_attention(q, k, v, mask=None)
+    a2 = calculate_attention(q, k, v, mask=None)
     assert torch.allclose(a1, a2)
 
 
@@ -344,6 +328,7 @@ def test_qkv_device_mismatch_raises_clear_error():
         assert (
             "q/k/v must be on the same device" in msg  # our explicit check
             or "Expected all tensors to be on the same device" in msg  # PyTorch matmul
+            or "query, key, value must be on the same device" in msg  # updated error wording
         )
     finally:
         torch.use_deterministic_algorithms(prev)
@@ -369,23 +354,6 @@ def test_boolean_mask_broadcast_variants_equivalent():
     assert torch.allclose(p2, p_exact, atol=1e-6, rtol=1e-5)
 
 
-def test_additive_mask_broadcast_variants_equivalent():
-    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
-    q, k, v = _rand(B, H, Sq, Sk, D)
-    m_exact = torch.zeros(B, H, Sq, Sk, dtype=torch.float32)
-    m_111Sk = torch.zeros(1, 1, 1, Sk, dtype=torch.float32)
-    m_B1Sq1 = torch.zeros(B, 1, Sq, 1, dtype=torch.float32)
-
-    a_exact, p_exact = calculate_attention(q, k, v, m_exact, return_probs=True)
-    a1, p1 = calculate_attention(q, k, v, m_111Sk, return_probs=True)
-    a2, p2 = calculate_attention(q, k, v, m_B1Sq1, return_probs=True)
-
-    assert torch.allclose(a1, a_exact, atol=1e-6, rtol=1e-5)
-    assert torch.allclose(p1, p_exact, atol=1e-6, rtol=1e-5)
-    assert torch.allclose(a2, a_exact, atol=1e-6, rtol=1e-5)
-    assert torch.allclose(p2, p_exact, atol=1e-6, rtol=1e-5)
-
-
 def test_non_broadcastable_mask_raises_value_error():
     B, H, Sq, Sk, D = 2, 3, 4, 5, 8
     q, k, v = _rand(B, H, Sq, Sk, D)
@@ -432,90 +400,45 @@ def test_join_heads_roundtrip():
 ###___sinusoidal_positional_encoding___###
 
 
-@pytest.mark.parametrize("seq_len,dim", [(0, 1), (1, 1), (1, 2), (7, 4), (17, 33)])
-def test_shapes_and_dtypes(seq_len, dim):
-    # Try a few dtypes; fp16 on CPU can be quirky, so we gate it on CUDA.
-    dtypes = [torch.float32, torch.bfloat16]
-    if torch.cuda.is_available():
-        dtypes.append(torch.float16)
-
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-    for dtype in dtypes:
-        pe = sinusoidal_positional_encoding(seq_len, dim, dtype=dtype, device=device)
-        assert pe.shape == (seq_len, dim)
-        assert pe.dtype == dtype
-        # Returned tensor should be a constant table (no gradients)
-        assert pe.requires_grad is False
+@pytest.mark.parametrize("seq_len,dim", [(1, 1), (1, 2), (7, 4), (17, 33)])
+def test_sinusoidal_shapes_and_dtype(seq_len, dim):
+    pe = sinusoidal_positional_encoding(seq_len, dim)
+    assert pe.shape == (seq_len, dim)
+    assert pe.dtype == torch.float32
+    assert pe.requires_grad is False
 
 
-def test_values_match_reference_small_case():
-    # Small deterministic comparison against the textbook formula.
+def test_sinusoidal_matches_reference_small_case():
     seq_len, dim = 4, 6
-    base = 10_000.0
-    pe = sinusoidal_positional_encoding(seq_len, dim, dtype=torch.float32, device="cpu", base=base)
+    pe = sinusoidal_positional_encoding(seq_len, dim)
 
     ref = torch.zeros_like(pe)
-    # Fill sin/cos pairs explicitly
+    base = 10_000.0
     for pos in range(seq_len):
         for i in range(0, dim // 2):
             denom = base ** (2 * i / dim)
             ref[pos, 2 * i] = math.sin(pos / denom)
             ref[pos, 2 * i + 1] = math.cos(pos / denom)
 
-    # Compare only the complete sin/cos pairs (even columns and their cos twins)
     paired = (dim // 2) * 2
     assert torch.allclose(pe[:, :paired], ref[:, :paired], atol=1e-6, rtol=1e-6)
 
 
-@pytest.mark.parametrize("dim", [1, 3, 8])
-def test_odd_dim_last_column_zero(dim):
-    seq_len = 5
-    pe = sinusoidal_positional_encoding(seq_len, dim)
-    if dim % 2 == 1:
-        assert torch.allclose(
-            pe[:, -1],
-            torch.zeros(seq_len, dtype=pe.dtype, device=pe.device),
-        )
-
-
-def test_offset_windowing_equivalence():
-    # pe[pos] with offset should match a slice of a longer table without offset
-    seq_len, dim, offset = 8, 6, 5
-    base_long = sinusoidal_positional_encoding(seq_len + offset, dim)
-    win = sinusoidal_positional_encoding(seq_len, dim, offset=offset)
-    assert torch.allclose(win, base_long[offset : offset + seq_len])
-
-
-def test_determinism():
-    a = sinusoidal_positional_encoding(32, 64)
-    b = sinusoidal_positional_encoding(32, 64)
+def test_sinusoidal_deterministic():
+    a = sinusoidal_positional_encoding(8, 16)
+    b = sinusoidal_positional_encoding(8, 16)
     assert torch.allclose(a, b)
 
 
-def test_invalid_inputs():
+def test_sinusoidal_invalid_inputs():
+    with pytest.raises(ValueError):
+        sinusoidal_positional_encoding(0, 8)
     with pytest.raises(ValueError):
         sinusoidal_positional_encoding(-1, 8)
     with pytest.raises(ValueError):
         sinusoidal_positional_encoding(1, 0)
 
 
-@pytest.mark.parametrize("base", [1_000.0, 10_000.0, 100_000.0])
-def test_different_bases_change_values(base):
-    seq_len, dim = 6, 8
-    a = sinusoidal_positional_encoding(seq_len, dim, base=base)
-    b = sinusoidal_positional_encoding(seq_len, dim, base=base * 10)
-    # Different bases should not produce identical tables
-    assert not torch.allclose(a, b)
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for strict fp16 checks")
-def test_cuda_half_precision_support():
-    seq_len, dim = 12, 16
-    pe = sinusoidal_positional_encoding(seq_len, dim, dtype=torch.float16, device="cuda")
-    assert pe.device.type == "cuda"
-    assert pe.dtype == torch.float16
-
-
 ###___combine_masks___###
 
 
@@ -560,72 +483,6 @@ def test_dtype_conversion():
     assert torch.equal(out, expected)
 
 
-###___shift_right___###
-
-
-def test_basic_shift_single_batch():
-    labels = torch.tensor([[5, 6, 7]], dtype=torch.long)
-    out = shift_right(labels, bos_id=1, pad_id=0)
-    assert torch.equal(out, torch.tensor([[1, 5, 6]], dtype=torch.long))
-
-
-def test_basic_shift_vector_input():
-    labels = torch.tensor([10, 11, 12], dtype=torch.long)
-    out = shift_right(labels, bos_id=2, pad_id=0)
-    assert out.shape == (1, 3)
-    assert torch.equal(out, torch.tensor([[2, 10, 11]], dtype=torch.long))
-
-
-def test_batch_shift():
-    labels = torch.tensor([[5, 6, 7], [8, 9, 10]], dtype=torch.long)
-    out = shift_right(labels, bos_id=3, pad_id=0)
-    expected = torch.tensor([[3, 5, 6], [3, 8, 9]], dtype=torch.long)
-    assert torch.equal(out, expected)
-
-
-def test_ignores_are_padded_after_shift():
-    # -100 in labels must never appear in inputs; becomes PAD after shifting.
-    labels = torch.tensor([[5, -100, 7, -100]], dtype=torch.long)
-    out = shift_right(labels, bos_id=4, pad_id=99)
-    # positions 2 and 4 in inputs come from -100 -> PAD(99)
-    expected = torch.tensor([[4, 5, 99, 7]], dtype=torch.long)
-    assert torch.equal(out, expected)
-    assert (out == -100).sum() == 0
-
-
-def test_custom_pad_and_bos_ids():
-    labels = torch.tensor([[42, 43, -100, 45]], dtype=torch.long)
-    out = shift_right(labels, bos_id=7, pad_id=3)
-    expected = torch.tensor([[7, 42, 43, 3]], dtype=torch.long)
-    assert torch.equal(out, expected)
-
-
-def test_zero_length_raises():
-    with pytest.raises(ValueError):
-        shift_right(torch.zeros((2, 0), dtype=torch.long))
-
-
-def test_wrong_rank_raises():
-    with pytest.raises(ValueError):
-        shift_right(torch.zeros((2, 3, 4), dtype=torch.long))
-
-
-def test_dtype_device_preserved():
-    device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
-    labels = torch.tensor([[1, 2, 3]], device=device)
-    out = shift_right(labels, bos_id=9, pad_id=0)
-    assert out.device == device
-    assert out.dtype == torch.long
-    assert torch.equal(out, torch.tensor([[9, 1, 2]], device=device, dtype=torch.long))
-
-
-def test_no_inplace_on_labels():
-    labels = torch.tensor([[5, 6, -100]], dtype=torch.long)
-    labels_clone = labels.clone()
-    _ = shift_right(labels, bos_id=1, pad_id=0)
-    assert torch.equal(labels, labels_clone)
-
-
 ###___sample_from_logits___###
 
 
@@ -706,19 +563,18 @@ def test_handles_higher_rank_logits_flattening():
 ###___create_causal_mask___###
 
 
-def test_mask_shape_and_dtype_cpu():
-    L = 5
-    mask = create_causal_mask(L)
-    assert mask.shape == (1, 1, L, L)
+def test_causal_mask_shape_and_dtype_cpu():
+    x = torch.zeros(2, 5, dtype=torch.long)
+    mask = create_causal_mask(x, num_heads=3)
+    assert mask.shape == (2, 3, 5, 5)
     assert mask.dtype == torch.bool
-    assert mask.device.type == "cpu"
+    assert mask.device == x.device
 
 
-def test_mask_upper_triangle():
-    L = 4
-    mask = create_causal_mask(L)
-    # For causal mask: mask[i,j] == True if j > i
-    mat = mask[0, 0].to(torch.int)
+def test_causal_mask_is_upper_triangular():
+    x = torch.zeros(1, 4, dtype=torch.long)
+    mask = create_causal_mask(x, num_heads=2)
+    tri = mask[0, 0].to(torch.int)
     expected = torch.tensor(
         [
             [0, 1, 1, 1],
@@ -727,23 +583,19 @@ def test_mask_upper_triangle():
             [0, 0, 0, 0],
         ]
     )
-    assert torch.equal(mat, expected)
-
-
-@pytest.mark.parametrize("dtype", [torch.bool, torch.float32])
-def test_mask_dtype(dtype):
-    L = 3
-    mask = create_causal_mask(L, dtype=dtype)
-    assert mask.dtype == dtype
+    assert torch.equal(tri, expected)
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_device_cuda():
-    L = 6
-    mask = create_causal_mask(L, device=torch.device("cuda"))
+def test_causal_mask_tracks_device():
+    x = torch.zeros(1, 3, dtype=torch.long, device="cuda")
+    mask = create_causal_mask(x, num_heads=4)
     assert mask.device.type == "cuda"
 
 
-def test_invalid_length_raises():
+def test_causal_mask_rejects_invalid_inputs():
+    with pytest.raises(ValueError):
+        create_causal_mask(torch.zeros(2, 0, dtype=torch.long), num_heads=2)
+
     with pytest.raises(ValueError):
-        create_causal_mask(0)
+        create_causal_mask(torch.zeros(2, 3, dtype=torch.long), num_heads=0)