Q3 v1 notebook (tokens scrubbed)

contrastive_finetune_e5_q3v1.ipynb

@@ -0,0 +1,1305 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Contrastive Fine-Tune of multilingual-e5-large for Swiss Legal Citation Retrieval\n",
+    "\n",
+    "**Goal:** Improve dense retrieval F1 on the Kaggle LLM-Agentic-Legal-IR competition by fine-tuning the encoder on (English query → Swiss German citation) pairs from `train.csv` with hard negatives mined from BM25 over `laws_de.csv`.\n",
+    "\n",
+    "**Why this and not the previous attempts:**\n",
+    "- Past LoRA fine-tunes used generic text MLM losses → no improvement.\n",
+    "- This run uses **MultipleNegativesRankingLoss** with **hard negatives** specifically mined from our own pool (i.e. cits the model currently confuses with gold). This is the failure mode we want to fix, not generic semantic matching.\n",
+    "- Training pairs are **goal-aligned**: same format, same languages, same domain as test time.\n",
+    "\n",
+    "**Hard-negative quality safeguards** (NV-Retriever / GISTEmbedLoss style):\n",
+    "- Cross-query gold filter — exclude any cit that's gold for ANY train query\n",
+    "- Cosine-ratio false-neg drop — if `cos(q,neg) > 0.95 * max cos(q,gold)` it's likely a true gold leaking in\n",
+    "- Unique-qid batch sampler — no two examples in a batch share the same query (MNR loss hygiene)\n",
+    "\n",
+    "**Two-gate eval before re-encoding:**\n",
+    "- Gate A: ΔRecall@40 on val ≥ +0.01\n",
+    "- Gate B: mean gold rank in val retrieval improves (or holds)\n",
+    "\n",
+    "**Smoke tests at every stage** to catch silent failures (corrupted data, wrong prefixes, broken loss, frozen layers, etc.).\n",
+    "\n",
+    "**Expected runtime:** ~30-60 minutes on a single 24GB GPU (RTX 4090 / A6000 / A100).\n",
+    "\n",
+    "**Outputs:**\n",
+    "- `outputs/checkpoint_e5legal_v1/` — fine-tuned model directory\n",
+    "- `outputs/laws_emb_v1.npy` — re-encoded laws corpus (175,933 × 1024)\n",
+    "- `outputs/queries_v1.npz` — re-encoded train/val/test queries\n",
+    "- `outputs/val_baseline_vs_finetuned.csv` — Recall@40 / F1 comparison\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## Cell 1 — Install dependencies (Vast / Colab)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "98756acf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "# All deps in one shot — pinned for compatibility (ST 2.7.0 is last version where model.fit() works without datasets quirks)\n",
+    "!{sys.executable} -m pip install -q -U \"pandas>=2.0\" \"numpy>=1.24,<2.0\" \"scipy>=1.11\" \"scikit-learn>=1.3\" \"sentence-transformers==2.7.0\" \"transformers>=4.34,<4.45\" \"datasets>=2.18\" \"accelerate>=0.30\" \"huggingface_hub>=0.23\" \"bm25s>=0.2\" \"rank-bm25\" \"tqdm>=4.66\" \"kaggle\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 2 — Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys, json, re, random, time, gc\n",
+    "from pathlib import Path\n",
+    "from collections import Counter, defaultdict\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import torch\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from torch.nn import functional as F\n",
+    "\n",
+    "from transformers import AutoTokenizer, AutoModel\n",
+    "from sentence_transformers import SentenceTransformer, InputExample, losses, models\n",
+    "from sentence_transformers.evaluation import InformationRetrievalEvaluator\n",
+    "\n",
+    "from rank_bm25 import BM25Okapi\n",
+    "\n",
+    "random.seed(42); np.random.seed(42); torch.manual_seed(42)\n",
+    "if torch.cuda.is_available(): torch.cuda.manual_seed_all(42)\n",
+    "\n",
+    "print(f'torch {torch.__version__}  cuda {torch.cuda.is_available()}')\n",
+    "print(f'sentence-transformers loaded')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 3 — SMOKE TEST 1: GPU + tiny model forward pass"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert torch.cuda.is_available(), 'No GPU detected — fine-tuning will be too slow on CPU. Spin up Vast/Colab GPU instance.'\n",
+    "device = torch.device('cuda')\n",
+    "props = torch.cuda.get_device_properties(0)\n",
+    "print(f'GPU: {props.name}  mem={props.total_memory/1e9:.1f} GB')\n",
+    "assert props.total_memory > 16e9, f'Need >=16GB VRAM, got {props.total_memory/1e9:.1f} GB'\n",
+    "\n",
+    "# Tiny smoke test: load smallest e5 model and encode 2 sentences\n",
+    "_t = AutoTokenizer.from_pretrained('intfloat/multilingual-e5-small')\n",
+    "_m = AutoModel.from_pretrained('intfloat/multilingual-e5-small').to(device).eval()\n",
+    "_inp = _t(['query: Was ist Vertragsrecht?', 'passage: Vertragsrecht regelt ...'],\n",
+    "          padding=True, truncation=True, return_tensors='pt').to(device)\n",
+    "with torch.no_grad():\n",
+    "    _h = _m(**_inp).last_hidden_state\n",
+    "    # mean pool with mask\n",
+    "    mask = _inp.attention_mask.unsqueeze(-1).float()\n",
+    "    _emb = (_h * mask).sum(1) / mask.sum(1).clamp(min=1e-9)\n",
+    "    _emb = F.normalize(_emb, p=2, dim=1)\n",
+    "    cos = (_emb[0] @ _emb[1].T).item()\n",
+    "print(f'cos(query, passage) on smoke test = {cos:.3f}  (expect ~0.7-0.95 — high since topically aligned)')\n",
+    "del _m, _t, _inp, _h, _emb; gc.collect(); torch.cuda.empty_cache()\n",
+    "print('SMOKE 1 PASS')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 4 — Config (edit these for your run)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9533537c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# === GUARDRAILS APPLIED ===\n",
+    "# (a) Hard-neg false-positive filter via baseline encoder cosine ratio\n",
+    "# (b) Batch sampler with unique qid per batch (MNR hygiene)\n",
+    "# (c) Two-gate eval: ΔRecall@40 ≥ +0.01 AND avg-gold-rank improves\n",
+    "\n",
+    "CFG = {\n",
+    "    # Paths — auto-resolved by env detection cell below; these are placeholders\n",
+    "    'DATA_DIR':   None,    # set by env detection: /workspace/data on Vast, /kaggle/input/... on Kaggle, ./Data locally\n",
+    "    'OUTPUT_DIR': None,    # set by env detection: /workspace/outputs on Vast, /kaggle/working on Kaggle, ./outputs locally\n",
+    "\n",
+    "    # Base encoder. multilingual-e5-large is the default we already use.\n",
+    "    # To start from a swiss-legal-finetuned checkpoint, point this at that path.\n",
+    "    'BASE_MODEL': 'intfloat/multilingual-e5-large',\n",
+    "\n",
+    "    # Training\n",
+    "    'N_HARD_NEGS_PER_POS':  4,    # mined hard negatives per (query, positive) pair\n",
+    "    'BM25_NEG_TOP_K':       50,   # mine hard negs from top-K BM25 hits (excluding gold)\n",
+    "    'EPOCHS':               2,    # 1-3 typical; watch val Recall@40 to pick\n",
+    "    'BATCH_SIZE':           16,   # 16 fits e5-large training on 24GB GPU with AMP. Encoding uses ENCODE_BATCH separately.\n",
+    "    'LEARNING_RATE':        1.5e-5, # gentle — keep top blocks from drifting too far on 1k-pair training set\n",
+    "    'WARMUP_RATIO':         0.1,\n",
+    "    'MAX_LEN_QUERY':        160,\n",
+    "    'MAX_LEN_PASSAGE':      256,\n",
+    "    'GRADIENT_ACCUM':       1,\n",
+    "    # Hard-negative quality filters (NV-Retriever / GISTEmbedLoss style)\n",
+    "    'FALSE_NEG_SIM_RATIO':       0.95,  # drop neg if cos(q,neg) > 0.95 * max cos(q,gold)\n",
+    "    'BATCH_NO_DUP_QID':          True,  # batch sampler ensures unique qid per batch\n",
+    "    'CROSS_QUERY_GOLD_FILTER':   True,  # exclude any cit that is gold for ANY train query\n",
+    "\n",
+    "    # Eval\n",
+    "    'VAL_RECALL_K':         40,   # competition uses ~25-40 cit predictions per query\n",
+    "    'ENCODE_BATCH':         32,   # 32 is safe on 24GB during inference. Increase to 64 if you have headroom.\n",
+    "\n",
+    "    # Reproducibility\n",
+    "    'SEED':                 42,\n",
+    "}\n",
+    "\n",
+    "# Reduce CUDA memory fragmentation. Only effective if set before any CUDA allocation,\n",
+    "# but harmless to set here too.\n",
+    "import os as _os\n",
+    "_os.environ.setdefault('PYTORCH_CUDA_ALLOC_CONF', 'expandable_segments:True')\n",
+    "\n",
+    "# OUTPUT_DIR mkdir happens in Cell 4.5 after env detection resolves the path.\n",
+    "print(json.dumps({k: str(v) if v is not None else None for k, v in CFG.items()}, indent=2))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 4.5 — Environment detection + Kaggle data download\n",
+    "\n",
+    "Auto-detect Kaggle / Vast / local. On Vast.ai or Colab without local data, pull the competition CSVs via the Kaggle API.\n",
+    "\n",
+    "**Set these as environment variables before launching this notebook** (Vast console → Edit → Environment):\n",
+    "- `KAGGLE_USERNAME` — your Kaggle username\n",
+    "- `KAGGLE_KEY` — your Kaggle API key (kaggle.com/settings → API → Create New Token)\n",
+    "- `HF_TOKEN` — your Hugging Face token (huggingface.co/settings/tokens)\n",
+    "\n",
+    "Inline fallbacks below for one-off runs — but rotate any leaked keys after the session."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bbcbf08f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, subprocess\n",
+    "from pathlib import Path\n",
+    "\n",
+    "# ── Environment detection ──\n",
+    "IS_KAGGLE = os.path.exists('/kaggle')\n",
+    "IS_VASTAI = not IS_KAGGLE and os.path.exists('/workspace')\n",
+    "IS_LOCAL  = not IS_KAGGLE and not IS_VASTAI\n",
+    "\n",
+    "print(f'Environment: KAGGLE={IS_KAGGLE}  VASTAI={IS_VASTAI}  LOCAL={IS_LOCAL}')\n",
+    "\n",
+    "# ── API keys ──\n",
+    "# Priority: existing ~/.kaggle/kaggle.json > env vars > inline fallback\n",
+    "_kaggle_json = Path.home() / '.kaggle' / 'kaggle.json'\n",
+    "if _kaggle_json.exists():\n",
+    "    import json as _json\n",
+    "    _creds = _json.loads(_kaggle_json.read_text())\n",
+    "    os.environ['KAGGLE_USERNAME'] = _creds['username']\n",
+    "    os.environ['KAGGLE_KEY']      = _creds['key']\n",
+    "    print(f'Loaded creds from {_kaggle_json} (username={_creds[\"username\"]})')\n",
+    "else:\n",
+    "    # Tokens must come from env vars or kaggle.json — never hardcode in notebook\n",
+    "    if 'KAGGLE_KEY' not in os.environ or 'KAGGLE_USERNAME' not in os.environ:\n",
+    "        raise RuntimeError(\"Set KAGGLE_USERNAME + KAGGLE_KEY env vars or place ~/.kaggle/kaggle.json before running\")\n",
+    "if 'HF_TOKEN' not in os.environ:\n",
+    "    raise RuntimeError(\"Set HF_TOKEN env var before running\")\n",
+    "\n",
+    "# ── Resolve paths per environment ──\n",
+    "if IS_KAGGLE:\n",
+    "    _ca = Path('/kaggle/input/competitions/llm-agentic-legal-information-retrieval')\n",
+    "    _cb = Path('/kaggle/input/llm-agentic-legal-information-retrieval')\n",
+    "    BASE = _ca if _ca.exists() else _cb\n",
+    "    OUT_BASE = Path('/kaggle/working')\n",
+    "elif IS_VASTAI:\n",
+    "    BASE = Path('/workspace/data')\n",
+    "    OUT_BASE = Path('/workspace/outputs')\n",
+    "else:\n",
+    "    BASE = Path('./Data')\n",
+    "    OUT_BASE = Path('./outputs')\n",
+    "\n",
+    "CFG['DATA_DIR']   = str(BASE)\n",
+    "CFG['OUTPUT_DIR'] = str(OUT_BASE)\n",
+    "OUT_BASE.mkdir(parents=True, exist_ok=True)\n",
+    "print(f'BASE       = {BASE}')\n",
+    "print(f'OUTPUT_DIR = {OUT_BASE}')\n",
+    "\n",
+    "# ── Download competition data on Vast if missing ──\n",
+    "if IS_VASTAI and not (BASE / 'laws_de.csv').exists():\n",
+    "    assert os.environ.get('KAGGLE_KEY'), 'KAGGLE_KEY required to download competition data on Vast'\n",
+    "    print('Downloading competition data via Kaggle API...')\n",
+    "    kaggle_dir = Path.home() / '.kaggle'\n",
+    "    kaggle_dir.mkdir(exist_ok=True)\n",
+    "    # Only write kaggle.json if it doesn't exist (don't overwrite user's manual update)\n",
+    "    kj = kaggle_dir / 'kaggle.json'\n",
+    "    if not kj.exists():\n",
+    "        kj.write_text(f'{{\"username\":\"{os.environ[\"KAGGLE_USERNAME\"]}\",\"key\":\"{os.environ[\"KAGGLE_KEY\"]}\"}}')\n",
+    "        os.chmod(kj, 0o600)\n",
+    "    BASE.mkdir(parents=True, exist_ok=True)\n",
+    "    from kaggle.api.kaggle_api_extended import KaggleApi\n",
+    "    api = KaggleApi(); api.authenticate()\n",
+    "    api.competition_download_files('llm-agentic-legal-information-retrieval', path=str(BASE))\n",
+    "    subprocess.check_call(['bash', '-c', f\"cd '{BASE}' && unzip -o '*.zip' && rm -f *.zip\"])\n",
+    "    print('  download complete')\n",
+    "\n",
+    "# ── Verify required files ──\n",
+    "for f in ['laws_de.csv', 'train.csv', 'val.csv', 'test.csv']:\n",
+    "    p = BASE / f\n",
+    "    assert p.exists(), f'Missing {p} — set KAGGLE_KEY or copy data manually'\n",
+    "    print(f'  {f:<20s} {p.stat().st_size/1e6:>6.1f} MB')\n",
+    "\n",
+    "# ── Optional HF login ──\n",
+    "if os.environ.get('HF_TOKEN'):\n",
+    "    from huggingface_hub import login\n",
+    "    login(token=os.environ['HF_TOKEN'], add_to_git_credential=False)\n",
+    "    print('HF logged in.')\n",
+    "\n",
+    "print('\\nSMOKE 1.5 PASS — environment + data ready')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 5 — Load data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "DATA = Path(CFG['DATA_DIR'])\n",
+    "OUT  = Path(CFG['OUTPUT_DIR'])\n",
+    "\n",
+    "laws  = pd.read_csv(DATA / 'laws_de.csv')\n",
+    "train = pd.read_csv(DATA / 'train.csv')\n",
+    "val   = pd.read_csv(DATA / 'val.csv')\n",
+    "test  = pd.read_csv(DATA / 'test.csv')\n",
+    "\n",
+    "# Sanity print\n",
+    "print(f'laws_de: {len(laws):,} rows  cols={list(laws.columns)}')\n",
+    "print(f'train:   {len(train):,} rows  cols={list(train.columns)}')\n",
+    "print(f'val:     {len(val):,} rows')\n",
+    "print(f'test:    {len(test):,} rows')\n",
+    "print('\\n--- laws_de head ---')\n",
+    "print(laws.head(3))\n",
+    "print('\\n--- train head ---')\n",
+    "print(train.head(2))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 6 — SMOKE TEST 2: Schema + gold coverage"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Verify schemas\n",
+    "for df, name, expected in [\n",
+    "    (laws,  'laws_de',  ['citation','text']),\n",
+    "    (train, 'train',    ['query_id','query','gold_citations']),\n",
+    "    (val,   'val',      ['query_id','query','gold_citations']),\n",
+    "    (test,  'test',     ['query_id','query']),\n",
+    "]:\n",
+    "    missing = [c for c in expected if c not in df.columns]\n",
+    "    assert not missing, f'{name}: missing cols {missing}'\n",
+    "\n",
+    "# Gold coverage check: how many train gold cits are present in laws_de?\n",
+    "laws_cit_set = set(laws['citation'].astype(str))\n",
+    "print(f'laws_de unique cits: {len(laws_cit_set):,}')\n",
+    "\n",
+    "all_gold = []\n",
+    "for s in train['gold_citations']:\n",
+    "    all_gold.extend([c.strip() for c in str(s).split(';') if c.strip()])\n",
+    "all_gold_set = set(all_gold)\n",
+    "covered = sum(1 for c in all_gold_set if c in laws_cit_set)\n",
+    "print(f'train gold unique cits: {len(all_gold_set):,}  covered_in_laws: {covered:,} ({100*covered/len(all_gold_set):.1f}%)')\n",
+    "\n",
+    "all_gold_val = []\n",
+    "for s in val['gold_citations']:\n",
+    "    all_gold_val.extend([c.strip() for c in str(s).split(';') if c.strip()])\n",
+    "all_gold_val_set = set(all_gold_val)\n",
+    "covered_val = sum(1 for c in all_gold_val_set if c in laws_cit_set)\n",
+    "print(f'val gold unique cits: {len(all_gold_val_set):,}  covered_in_laws: {covered_val:,} ({100*covered_val/len(all_gold_val_set):.1f}%)')\n",
+    "\n",
+    "# Court cits won't be in laws_de — that's expected.\n",
+    "ART_PAT = re.compile(r'^Art\\.\\s*\\d+[a-z]?')\n",
+    "court_in_train = sum(1 for c in all_gold_set if not ART_PAT.match(c))\n",
+    "print(f'court cits in train gold (ART_PAT no match): {court_in_train}')\n",
+    "print('SMOKE 2 PASS — schema OK, gold coverage measured')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 7 — Build (query, positive_passage) pairs\n",
+    "\n",
+    "For each train query, expand its `gold_citations` into individual positive pairs `(query_text, citation_text)`. We use **only law citations** present in `laws_de` for fine-tuning the law encoder. Court cits are excluded here (they live in a separate corpus and would require a court-specific re-encoding loop)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cit_to_text = dict(zip(laws['citation'].astype(str), laws['text'].astype(str)))\n",
+    "\n",
+    "pairs = []   # list of (query_text, citation, citation_text, query_id)\n",
+    "for _, r in train.iterrows():\n",
+    "    q  = str(r['query']).strip()\n",
+    "    qid = r['query_id']\n",
+    "    for c in str(r['gold_citations']).split(';'):\n",
+    "        c = c.strip()\n",
+    "        if not c: continue\n",
+    "        if c in cit_to_text:\n",
+    "            pairs.append((q, c, cit_to_text[c], qid))\n",
+    "\n",
+    "print(f'(query, positive_law) pairs: {len(pairs):,}')\n",
+    "print('\\nSample pair:')\n",
+    "p = pairs[0]\n",
+    "print(f'  qid: {p[3]}\\n  query: {p[0][:150]}...\\n  cit: {p[1]}\\n  text: {p[2][:200]}...')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 8 — SMOKE TEST 3: Pair sanity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert len(pairs) > 1000, f'Too few pairs ({len(pairs)}) — check gold coverage'\n",
+    "for q, c, t, qid in pairs[:3]:\n",
+    "    assert isinstance(q, str) and len(q) > 5\n",
+    "    assert isinstance(c, str) and ART_PAT.match(c) is not None, f'unexpected cit format: {c}'\n",
+    "    assert isinstance(t, str) and len(t) > 5\n",
+    "# Distribution of golds per query\n",
+    "g_per_q = Counter(p[3] for p in pairs)\n",
+    "print(f'queries with >=1 law gold: {len(g_per_q):,}')\n",
+    "print(f'mean golds per query (in laws): {np.mean(list(g_per_q.values())):.1f}')\n",
+    "print('SMOKE 3 PASS')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 9 — Build BM25 index over laws_de for hard negative mining"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# Q3 v1 — Load pre-mined hard negatives (4-source mining done locally)\n",
+    "# Replaces the original BM25 mining cells. Keeps `pairs` from earlier cell.\n",
+    "# Sources: BM25, diff-code, adjacent, diff-paragraph (with co-citation validity check).\n",
+    "\n",
+    "from huggingface_hub import hf_hub_download\n",
+    "import pandas as pd\n",
+    "import os\n",
+    "\n",
+    "HF_TOKEN = os.environ.get('HF_TOKEN')\n",
+    "parquet_path = hf_hub_download(\n",
+    "    repo_id='Dharun72/swiss-legal-checkpoints',\n",
+    "    filename='q3_train_v1.parquet',\n",
+    "    token=HF_TOKEN,\n",
+    ")\n",
+    "df_pairs = pd.read_parquet(parquet_path)\n",
+    "print(f'Loaded q3_train_v1.parquet: {len(df_pairs):,} rows from {df_pairs[\"qid\"].nunique()} unique queries')\n",
+    "\n",
+    "# qid_to_query_text — needed by smoke + filter cells downstream\n",
+    "qid_to_query_text = dict(zip(train['query_id'], train['query']))\n",
+    "\n",
+    "# qid_to_golds — needed by smoke + filter cells downstream\n",
+    "qid_to_golds = {}\n",
+    "for r in df_pairs.itertuples():\n",
+    "    qid_to_golds.setdefault(r.qid, set()).add(r.gold_cit)\n",
+    "\n",
+    "# qid_to_negidx — list of laws DataFrame row indices for hard negatives\n",
+    "qid_to_negidx = {}\n",
+    "for r in df_pairs.itertuples():\n",
+    "    qid_to_negidx.setdefault(r.qid, []).extend(list(r.hard_neg_idxs))\n",
+    "# Dedupe within each qid\n",
+    "for qid in qid_to_negidx:\n",
+    "    qid_to_negidx[qid] = list(dict.fromkeys(qid_to_negidx[qid]))\n",
+    "print(f'qid_to_negidx: {len(qid_to_negidx)} qids, mean {sum(len(v) for v in qid_to_negidx.values())/len(qid_to_negidx):.1f} negs/qid')\n",
+    "\n",
+    "# laws_cits_arr used downstream by smoke-test cell\n",
+    "laws_cits_arr = laws[\"citation\"].astype(str).values\n",
+    "\n",
+    "# N_NEG and TOP_K used downstream\n",
+    "N_NEG = CFG.get('N_HARD_NEGS_PER_POS', 4)\n",
+    "TOP_K = max(len(v) for v in qid_to_negidx.values())\n",
+    "print(f'N_NEG={N_NEG}  TOP_K (max negs per qid)={TOP_K}')\n",
+    "\n",
+    "# Smoke A: validate that no negatives are gold for that qid\n",
+    "n_violations = 0\n",
+    "for qid, negs in qid_to_negidx.items():\n",
+    "    golds = qid_to_golds.get(qid, set())\n",
+    "    for ni in negs:\n",
+    "        if laws_cits_arr[ni] in golds:\n",
+    "            n_violations += 1\n",
+    "print(f'Smoke A: gold-leak in negatives = {n_violations} (must be 0)')\n",
+    "assert n_violations == 0, \"Hard negative is in gold — bug\"\n",
+    "print('Smoke A PASS — no gold leakage in mined negatives')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 11 — SMOKE TEST 4: Hard negatives are non-gold and informative"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random as _random\n",
+    "_random.seed(7)\n",
+    "sample_qids = _random.sample(list(qid_to_golds.keys()), 3)\n",
+    "for qid in sample_qids:\n",
+    "    golds = qid_to_golds[qid]\n",
+    "    negs = qid_to_negidx[qid][:N_NEG]\n",
+    "    print(f'\\n=== {qid} ===')\n",
+    "    print(f'  query: {train[train.query_id==qid].iloc[0][\"query\"][:200]}...')\n",
+    "    print(f'  golds: {list(golds)[:3]}')\n",
+    "    for ni in negs:\n",
+    "        nc = laws_cits_arr[ni]\n",
+    "        nt = laws.iloc[ni][\"text\"][:120]\n",
+    "        assert nc not in golds, f'NEG {nc} found in gold of {qid} — bug'\n",
+    "        print(f'    neg: {nc}  text: {nt}...')\n",
+    "print('\\nSMOKE 4 PASS — hard negatives valid (non-gold, BM25-similar)')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 11.5 — SMOKE TEST 4.5: False-negative filter via baseline encoder cosine ratio\n",
+    "\n",
+    "Per NV-Retriever / GISTEmbedLoss best practice: a hard negative whose embedding is suspiciously close to gold is probably ITSELF a near-gold (paragraph variant, sibling article, semantic duplicate). Training against it punishes the model for being correct.\n",
+    "\n",
+    "We filter using the BASE encoder's cosine: drop any candidate where `cos(query, neg) > FALSE_NEG_SIM_RATIO * max cos(query, gold)`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1eefd5c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# e5 prefix conventions — used throughout (defined here so earlier cells can reference)\n",
+    "Q_PFX = 'query: '\n",
+    "P_PFX = 'passage: '\n",
+    "\n",
+    "# Helper: aggressively release a GPU model. PyTorch's caching allocator holds\n",
+    "# memory after del; moving to cpu first FORCES freelist release.\n",
+    "def _free_gpu_model(*objs):\n",
+    "    for o in objs:\n",
+    "        try:\n",
+    "            o.to('cpu')\n",
+    "        except Exception:\n",
+    "            pass\n",
+    "    for o in objs:\n",
+    "        del o\n",
+    "    gc.collect(); gc.collect()\n",
+    "    if torch.cuda.is_available():\n",
+    "        torch.cuda.empty_cache()\n",
+    "        torch.cuda.synchronize()\n",
+    "        torch.cuda.empty_cache()\n",
+    "        try: torch.cuda.reset_peak_memory_stats()\n",
+    "        except Exception: pass\n",
+    "\n",
+    "# Need the model loaded before filtering — load it now (will be RELEASED before training)\n",
+    "print(f'Loading base model for filter: {CFG[\"BASE_MODEL\"]}')\n",
+    "word_emb_filter = models.Transformer(CFG['BASE_MODEL'], max_seq_length=CFG['MAX_LEN_PASSAGE'])\n",
+    "pool_filter = models.Pooling(word_emb_filter.get_word_embedding_dimension(), pooling_mode='mean')\n",
+    "norm_filter = models.Normalize()\n",
+    "_filter_model = SentenceTransformer(modules=[word_emb_filter, pool_filter, norm_filter], device='cuda')\n",
+    "print(f'  loaded.')\n",
+    "\n",
+    "# 1) Encode unique queries\n",
+    "unique_qids = list(qid_to_negidx.keys())\n",
+    "qtxts = {qid: qid_to_query_text[qid] for qid in unique_qids}\n",
+    "\n",
+    "print(f'Encoding {len(unique_qids):,} queries...')\n",
+    "q_embs_arr = _filter_model.encode(\n",
+    "    [Q_PFX + str(qtxts[qid]) for qid in unique_qids],\n",
+    "    batch_size=CFG['ENCODE_BATCH'], convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=True,\n",
+    ")\n",
+    "qid_to_qemb = {qid: q_embs_arr[i] for i, qid in enumerate(unique_qids)}\n",
+    "\n",
+    "# 2) Collect cits we need to encode (golds + candidate negs)\n",
+    "needed_cits = set()\n",
+    "for qid in unique_qids:\n",
+    "    for g in qid_to_golds[qid]:\n",
+    "        if g in cit_to_text: needed_cits.add(g)\n",
+    "    for ni in qid_to_negidx[qid]:\n",
+    "        needed_cits.add(laws_cits_arr[ni])\n",
+    "needed_cits = sorted(needed_cits)\n",
+    "print(f'Encoding {len(needed_cits):,} unique cit texts...')\n",
+    "\n",
+    "cit_embs_arr = _filter_model.encode(\n",
+    "    [P_PFX + cit_to_text.get(c, c) for c in needed_cits],\n",
+    "    batch_size=CFG['ENCODE_BATCH'] * 2, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=True,\n",
+    ")\n",
+    "cit_to_emb_idx = {c: i for i, c in enumerate(needed_cits)}\n",
+    "\n",
+    "# 3) Filter — drop any neg with sim ratio > threshold\n",
+    "ratio = CFG['FALSE_NEG_SIM_RATIO']\n",
+    "qid_to_negidx_filt = {}\n",
+    "n_kept_total, n_dropped_total = 0, 0\n",
+    "margin_examples = []\n",
+    "for qid in unique_qids:\n",
+    "    qe = qid_to_qemb[qid]\n",
+    "    golds_present = [g for g in qid_to_golds[qid] if g in cit_to_emb_idx]\n",
+    "    if not golds_present:\n",
+    "        continue\n",
+    "    pos_sims = np.array([float(qe @ cit_embs_arr[cit_to_emb_idx[g]]) for g in golds_present])\n",
+    "    pos_max = float(pos_sims.max())\n",
+    "    threshold = ratio * pos_max\n",
+    "\n",
+    "    keep = []\n",
+    "    for ni in qid_to_negidx[qid]:\n",
+    "        nc = laws_cits_arr[ni]\n",
+    "        ne = cit_embs_arr[cit_to_emb_idx[nc]]\n",
+    "        neg_sim = float(qe @ ne)\n",
+    "        if neg_sim > threshold:\n",
+    "            n_dropped_total += 1\n",
+    "            if len(margin_examples) < 8:\n",
+    "                margin_examples.append((qid, nc, neg_sim, pos_max, threshold))\n",
+    "            continue\n",
+    "        keep.append(ni)\n",
+    "        n_kept_total += 1\n",
+    "        if len(keep) >= N_NEG: break\n",
+    "    if keep:\n",
+    "        qid_to_negidx_filt[qid] = keep\n",
+    "\n",
+    "# Replace\n",
+    "qid_to_negidx = qid_to_negidx_filt\n",
+    "print(f'\\nFilter results:')\n",
+    "print(f'  kept: {n_kept_total:,}  dropped_as_likely_falsenegs: {n_dropped_total:,}  ({100*n_dropped_total/max(1,n_kept_total+n_dropped_total):.1f}%)')\n",
+    "print(f'  queries surviving with >=1 neg: {len(qid_to_negidx):,}/{len(unique_qids):,}')\n",
+    "\n",
+    "print(f'\\nSample dropped negatives (qid, cit, sim_neg, sim_pos_max, threshold):')\n",
+    "for qid, nc, ns, pm, th in margin_examples:\n",
+    "    print(f'  {qid}  {nc:<35s}  neg_sim={ns:.3f}  pos_max={pm:.3f}  threshold={th:.3f}')\n",
+    "\n",
+    "# Aggressive GPU release — move model off GPU then drop refs\n",
+    "_free_gpu_model(_filter_model)\n",
+    "del word_emb_filter, pool_filter, norm_filter, q_embs_arr, cit_embs_arr\n",
+    "gc.collect(); torch.cuda.empty_cache()\n",
+    "\n",
+    "if torch.cuda.is_available():\n",
+    "    free_gb, total_gb = (x/1e9 for x in torch.cuda.mem_get_info())\n",
+    "    print(f'\\nGPU after filter cleanup: free={free_gb:.2f} GB / {total_gb:.2f} GB')\n",
+    "print('SMOKE 4.5 PASS — false-negative filter applied')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 12 — Build training InputExamples"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f7788468",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# sentence-transformers e5 prefixes:\n",
+    "#   query:   queries get \"query: \" prefix\n",
+    "#   passage: passages/citations get \"passage: \" prefix\n",
+    "Q_PFX = 'query: '\n",
+    "P_PFX = 'passage: '\n",
+    "\n",
+    "examples = []\n",
+    "example_qids = []   # parallel array — used by NoDupQidBatchSampler later\n",
+    "for q, c, t, qid in pairs:\n",
+    "    neg_idxs = qid_to_negidx.get(qid, [])\n",
+    "    if not neg_idxs: continue\n",
+    "    sampled = _random.sample(neg_idxs, min(N_NEG, len(neg_idxs)))\n",
+    "    for ni in sampled:\n",
+    "        nt = str(laws.iloc[ni]['text'])\n",
+    "        # MultipleNegativesRankingLoss with explicit hard neg: triplet [anchor, positive, neg]\n",
+    "        examples.append(InputExample(texts=[Q_PFX + q, P_PFX + t, P_PFX + nt]))\n",
+    "        example_qids.append(qid)\n",
+    "\n",
+    "assert len(examples) == len(example_qids)\n",
+    "print(f'Training InputExamples: {len(examples):,}  (unique qids: {len(set(example_qids)):,})')\n",
+    "print(f'  Anchor (sample): {examples[0].texts[0][:80]}...')\n",
+    "print(f'  Positive (sample): {examples[0].texts[1][:80]}...')\n",
+    "print(f'  Hard neg (sample): {examples[0].texts[2][:80]}...')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 13 — Load base encoder as SentenceTransformer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(f'Loading base model: {CFG[\"BASE_MODEL\"]}')\n",
+    "word_emb = models.Transformer(CFG['BASE_MODEL'], max_seq_length=CFG['MAX_LEN_PASSAGE'])\n",
+    "pool = models.Pooling(word_emb.get_word_embedding_dimension(), pooling_mode='mean')\n",
+    "norm = models.Normalize()  # L2 normalize for cosine\n",
+    "model = SentenceTransformer(modules=[word_emb, pool, norm], device='cuda')\n",
+    "print(f'Loaded.  embedding_dim={word_emb.get_word_embedding_dimension()}  device={model.device}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 14 — SMOKE TEST 5: Single forward, check output shape + cosine"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "_q = examples[0].texts[0]\n",
+    "_p = examples[0].texts[1]\n",
+    "_n = examples[0].texts[2]\n",
+    "_emb = model.encode([_q, _p, _n], convert_to_tensor=True, normalize_embeddings=True)\n",
+    "print(f'  emb shape: {_emb.shape}')\n",
+    "cos_qp = (_emb[0] @ _emb[1].T).item()\n",
+    "cos_qn = (_emb[0] @ _emb[2].T).item()\n",
+    "print(f'  baseline cos(Q, Pos)  = {cos_qp:.4f}')\n",
+    "print(f'  baseline cos(Q, Neg)  = {cos_qn:.4f}')\n",
+    "print(f'  margin (pos-neg)      = {cos_qp - cos_qn:+.4f}  (training should INCREASE this)')\n",
+    "assert _emb.shape[1] >= 768\n",
+    "print('SMOKE 5 PASS')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 15 — Build val IR evaluator (Recall@K on val set)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6375fa43",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Manual eval — replaces sentence_transformers.InformationRetrievalEvaluator.\n",
+    "# Why not the ST evaluator?\n",
+    "#   1. ST 2.7's IR evaluator returns 0.0000 for the primary metric (known regression).\n",
+    "#   2. It encodes the full 175k corpus on GPU and holds those tensors after the call,\n",
+    "#      causing OOM on 24GB cards before training even starts.\n",
+    "# This function encodes corpus in chunks, computes ranks on CPU, and releases GPU\n",
+    "# state immediately after each call.\n",
+    "\n",
+    "# Pre-compute val gold sets restricted to laws_de (court cits dropped — they live\n",
+    "# in court_considerations.csv which we don't fine-tune here).\n",
+    "val_qtexts = [Q_PFX + str(q) for q in val['query'].tolist()]\n",
+    "val_gold_idx = []   # list[set[int]] aligned with val rows: indices into laws_cits_arr\n",
+    "for _, r in val.iterrows():\n",
+    "    s = set()\n",
+    "    for c in str(r['gold_citations']).split(';'):\n",
+    "        c = c.strip()\n",
+    "        if c in cit_to_idx:\n",
+    "            s.add(cit_to_idx[c])\n",
+    "    val_gold_idx.append(s)\n",
+    "\n",
+    "n_val_gold = sum(len(s) for s in val_gold_idx)\n",
+    "print(f'Val: {len(val_qtexts)} queries, corpus {len(laws_cits_arr):,} laws, {n_val_gold} law-gold cits total')\n",
+    "\n",
+    "\n",
+    "def encode_corpus_to_cpu(_model, batch_size=None, desc='corpus'):\n",
+    "    \"\"\"Encode laws corpus in chunks; return float32 numpy on CPU.\n",
+    "    Releases per-chunk GPU tensors aggressively to avoid 24GB OOM.\"\"\"\n",
+    "    bs = batch_size or CFG['ENCODE_BATCH']\n",
+    "    out = np.empty((len(laws_cits_arr), 1024), dtype=np.float32)\n",
+    "    chunk = 8192\n",
+    "    texts_all = laws['text'].astype(str).tolist()\n",
+    "    t0 = time.time()\n",
+    "    for s in range(0, len(texts_all), chunk):\n",
+    "        e = min(s + chunk, len(texts_all))\n",
+    "        emb = _model.encode(\n",
+    "            [P_PFX + t for t in texts_all[s:e]],\n",
+    "            batch_size=bs,\n",
+    "            convert_to_numpy=True, normalize_embeddings=True,\n",
+    "            show_progress_bar=False,\n",
+    "        )\n",
+    "        out[s:e] = emb\n",
+    "        if torch.cuda.is_available():\n",
+    "            torch.cuda.empty_cache()\n",
+    "        if s % (chunk * 4) == 0:\n",
+    "            print(f'  {desc} {e}/{len(texts_all)} [{time.time()-t0:.0f}s]', flush=True)\n",
+    "    return out\n",
+    "\n",
+    "\n",
+    "def eval_recall_and_rank(_model, k=40, label='', encode_bs=None):\n",
+    "    \"\"\"Manual val Recall@K + mean gold rank. Returns (recall_at_k, mean_rank, median_rank).\n",
+    "    Memory-conscious: encodes corpus in chunks, computes per-query similarity on CPU.\"\"\"\n",
+    "    print(f'\\n[eval] {label}: encoding corpus...', flush=True)\n",
+    "    cemb = encode_corpus_to_cpu(_model, batch_size=encode_bs, desc='laws')\n",
+    "\n",
+    "    print(f'[eval] {label}: encoding {len(val_qtexts)} val queries...', flush=True)\n",
+    "    qemb = _model.encode(\n",
+    "        val_qtexts, batch_size=CFG['ENCODE_BATCH'],\n",
+    "        convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False,\n",
+    "    ).astype(np.float32)\n",
+    "\n",
+    "    if torch.cuda.is_available():\n",
+    "        torch.cuda.empty_cache()\n",
+    "\n",
+    "    # Per-query similarity + rank\n",
+    "    recalls, ranks = [], []\n",
+    "    for i, gold_set in enumerate(val_gold_idx):\n",
+    "        if not gold_set:\n",
+    "            recalls.append(0.0); continue\n",
+    "        sims = qemb[i] @ cemb.T   # (N_corpus,)\n",
+    "        # Top-K ids\n",
+    "        topk = np.argpartition(-sims, k)[:k]\n",
+    "        topk_set = set(int(x) for x in topk)\n",
+    "        hit = len(gold_set & topk_set)\n",
+    "        recalls.append(hit / len(gold_set))\n",
+    "        # Ranks of gold\n",
+    "        order = np.argsort(-sims)\n",
+    "        rank_of = {int(x): r for r, x in enumerate(order)}\n",
+    "        for g in gold_set:\n",
+    "            ranks.append(rank_of[g])\n",
+    "\n",
+    "    mean_recall = float(np.mean(recalls)) if recalls else 0.0\n",
+    "    mean_rank   = float(np.mean(ranks))   if ranks   else float('inf')\n",
+    "    median_rank = float(np.median(ranks)) if ranks   else float('inf')\n",
+    "    print(f'[eval] {label}: Recall@{k}={mean_recall:.4f}  mean_rank={mean_rank:.0f}  median_rank={median_rank:.0f}')\n",
+    "    # Free corpus embedding immediately\n",
+    "    del cemb, qemb\n",
+    "    gc.collect()\n",
+    "    return mean_recall, mean_rank, median_rank"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 16 — SMOKE TEST 6: Baseline val Recall@40 (no fine-tune)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bf1a37ad",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print('Computing BASELINE val IR metrics (pre-fine-tune)...')\n",
+    "t0 = time.time()\n",
+    "baseline_recall40, mean_r_base, med_r_base = eval_recall_and_rank(\n",
+    "    model, k=CFG['VAL_RECALL_K'], label='baseline'\n",
+    ")\n",
+    "print(f'  total baseline eval: {time.time()-t0:.1f}s')\n",
+    "print(f'\\nBaseline (multilingual-e5-large, no fine-tune):')\n",
+    "print(f'  Recall@{CFG[\"VAL_RECALL_K\"]} = {baseline_recall40:.4f}')\n",
+    "print(f'  mean_rank   = {mean_r_base:.0f}')\n",
+    "print(f'  median_rank = {med_r_base:.0f}')\n",
+    "\n",
+    "# Free any GPU residue from corpus encoding before training starts\n",
+    "gc.collect()\n",
+    "if torch.cuda.is_available():\n",
+    "    torch.cuda.empty_cache()\n",
+    "    free_gb, total_gb = (x/1e9 for x in torch.cuda.mem_get_info())\n",
+    "    print(f'  GPU after baseline: free={free_gb:.2f} GB / {total_gb:.2f} GB')\n",
+    "print('\\nSMOKE 6 PASS — record this baseline. Fine-tune must improve it.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "af0558ec",
+   "metadata": {},
+   "source": [
+    "## Cell 17 — SMOKE TEST 7: Mini-train (50 steps), confirm loss decreases"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9c2f432c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Aggressive cleanup BEFORE fit. Even small leftovers (filter model fragments,\n",
+    "# baseline corpus tensors) push 24GB cards over the edge during MNR training,\n",
+    "# which materializes a (B, S, H) activation graph for 3 streams (q, p, n).\n",
+    "import gc as _gc\n",
+    "_gc.collect()\n",
+    "if torch.cuda.is_available():\n",
+    "    torch.cuda.empty_cache()\n",
+    "    torch.cuda.synchronize()\n",
+    "    torch.cuda.empty_cache()\n",
+    "    try: torch.cuda.reset_peak_memory_stats()\n",
+    "    except Exception: pass\n",
+    "    free_gb, total_gb = (x/1e9 for x in torch.cuda.mem_get_info())\n",
+    "    print(f'GPU before mini-fit: free={free_gb:.2f} GB / {total_gb:.2f} GB')\n",
+    "\n",
+    "_random.seed(11)\n",
+    "mini_examples = _random.sample(examples, min(2000, len(examples)))\n",
+    "mini_loader = DataLoader(mini_examples, batch_size=CFG['BATCH_SIZE'], shuffle=True, drop_last=True)\n",
+    "\n",
+    "loss_fn = losses.MultipleNegativesRankingLoss(model)\n",
+    "\n",
+    "import logging; logging.getLogger('sentence_transformers').setLevel(logging.WARNING)\n",
+    "\n",
+    "print(f'Mini-train (50 steps, batch={CFG[\"BATCH_SIZE\"]})...')\n",
+    "model.fit(\n",
+    "    train_objectives=[(mini_loader, loss_fn)],\n",
+    "    epochs=1,\n",
+    "    steps_per_epoch=50,\n",
+    "    warmup_steps=5,\n",
+    "    optimizer_params={'lr': CFG['LEARNING_RATE']},\n",
+    "    use_amp=True,           # critical: cuts activation memory ~2x\n",
+    "    show_progress_bar=True,\n",
+    ")\n",
+    "# Re-check the same triple — margin should INCREASE\n",
+    "_emb = model.encode([_q, _p, _n], convert_to_tensor=True, normalize_embeddings=True)\n",
+    "cos_qp_after = (_emb[0] @ _emb[1].T).item()\n",
+    "cos_qn_after = (_emb[0] @ _emb[2].T).item()\n",
+    "print(f'\\n  After 50 steps: cos(Q,P)={cos_qp_after:.4f}  cos(Q,N)={cos_qn_after:.4f}  margin={cos_qp_after - cos_qn_after:+.4f}')\n",
+    "print(f'  Δmargin = {(cos_qp_after - cos_qn_after) - (cos_qp - cos_qn):+.4f}  (expect positive)')\n",
+    "\n",
+    "# Cleanup post-mini-fit before full training\n",
+    "del mini_loader, mini_examples, _emb\n",
+    "_gc.collect()\n",
+    "if torch.cuda.is_available():\n",
+    "    torch.cuda.empty_cache()\n",
+    "print('SMOKE 7 PASS — model is learning the contrastive objective')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "af3e4ec7",
+   "metadata": {},
+   "source": [
+    "## Cell 17.5 — Custom batch sampler: unique qid per batch (MNR hygiene)\n",
+    "\n",
+    "MNR loss treats every other in-batch positive as a negative for the current anchor. If two examples in the same batch share a `qid` (different gold cits for the same query), one example's positive becomes the other's negative — punishing the model for correct retrieval.\n",
+    "\n",
+    "Fix: greedy batch sampler that drops the second occurrence of any qid within a batch. Cleaner loss signal at the cost of slightly fewer training steps."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cf930ce7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class NoDupQidBatchSampler(torch.utils.data.Sampler):\n",
+    "    \"\"\"Yield batches of indices such that no two examples share the same qid.\n",
+    "    Uses example_qids tracked alongside `examples` in Cell 12.\"\"\"\n",
+    "    def __init__(self, qids, batch_size, drop_last=True, seed=0):\n",
+    "        self.qids = list(qids)\n",
+    "        self.batch_size = batch_size\n",
+    "        self.drop_last = drop_last\n",
+    "        self.seed = seed\n",
+    "        self.epoch = 0\n",
+    "    def __iter__(self):\n",
+    "        rng = np.random.RandomState(self.seed + self.epoch)\n",
+    "        order = rng.permutation(len(self.qids))\n",
+    "        used = set()\n",
+    "        batch = []\n",
+    "        for i in order:\n",
+    "            qid = self.qids[i]\n",
+    "            if qid in used:\n",
+    "                continue\n",
+    "            batch.append(int(i)); used.add(qid)\n",
+    "            if len(batch) == self.batch_size:\n",
+    "                yield batch\n",
+    "                batch = []; used = set()\n",
+    "        if batch and not self.drop_last:\n",
+    "            yield batch\n",
+    "        self.epoch += 1\n",
+    "    def __len__(self):\n",
+    "        return len(self.qids) // self.batch_size\n",
+    "\n",
+    "# example_qids was already built in parallel with `examples` in Cell 12.\n",
+    "# No rebuild needed — that wastes memory and creates a stale list.\n",
+    "assert len(examples) == len(example_qids), (\n",
+    "    f'examples ({len(examples)}) and example_qids ({len(example_qids)}) drifted — rerun Cell 12.'\n",
+    ")\n",
+    "print(f'Sampler ready. {len(examples):,} examples, {len(set(example_qids)):,} unique qids.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9cabf1a8",
+   "metadata": {},
+   "source": [
+    "## Cell 18 — Full training run"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ddb6a4b6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Aggressive cleanup before full training\n",
+    "gc.collect()\n",
+    "if torch.cuda.is_available():\n",
+    "    torch.cuda.empty_cache()\n",
+    "    torch.cuda.synchronize()\n",
+    "    torch.cuda.empty_cache()\n",
+    "    try: torch.cuda.reset_peak_memory_stats()\n",
+    "    except Exception: pass\n",
+    "    free_gb, total_gb = (x/1e9 for x in torch.cuda.mem_get_info())\n",
+    "    print(f'GPU before full fit: free={free_gb:.2f} GB / {total_gb:.2f} GB')\n",
+    "\n",
+    "if CFG['BATCH_NO_DUP_QID']:\n",
+    "    sampler = NoDupQidBatchSampler(example_qids, batch_size=CFG['BATCH_SIZE'], drop_last=True, seed=CFG['SEED'])\n",
+    "    train_loader = DataLoader(examples, batch_sampler=sampler, collate_fn=model.smart_batching_collate)\n",
+    "    print(f'Using NoDupQidBatchSampler — {len(sampler)} batches/epoch')\n",
+    "else:\n",
+    "    _random.shuffle(examples)\n",
+    "    train_loader = DataLoader(examples, batch_size=CFG['BATCH_SIZE'], shuffle=True, drop_last=True)\n",
+    "\n",
+    "steps_per_epoch = len(train_loader)\n",
+    "warmup = int(CFG['WARMUP_RATIO'] * steps_per_epoch * CFG['EPOCHS'])\n",
+    "print(f'Full training: {CFG[\"EPOCHS\"]} epochs × {steps_per_epoch} steps  (warmup {warmup})')\n",
+    "\n",
+    "ckpt_dir = OUT / 'checkpoint_e5legal_v1'\n",
+    "# NOTE: we DO NOT pass evaluator=... to fit() because:\n",
+    "#   1. ST 2.7's IR evaluator returns 0.0000 (silently broken)\n",
+    "#   2. it re-encodes 175k corpus on GPU each eval step → OOM\n",
+    "# We do a single manual eval after training in Cell 19.\n",
+    "model.fit(\n",
+    "    train_objectives=[(train_loader, loss_fn)],\n",
+    "    epochs=CFG['EPOCHS'],\n",
+    "    output_path=str(ckpt_dir),\n",
+    "    save_best_model=False,   # no evaluator, so save final state at end\n",
+    "    warmup_steps=warmup,\n",
+    "    optimizer_params={'lr': CFG['LEARNING_RATE']},\n",
+    "    use_amp=True,\n",
+    "    show_progress_bar=True,\n",
+    ")\n",
+    "# Save explicitly — fit() with save_best_model=False doesn't always save a final checkpoint\n",
+    "model.save(str(ckpt_dir))\n",
+    "print(f'\\nFine-tuned model saved to: {ckpt_dir}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "aa377513",
+   "metadata": {},
+   "source": [
+    "## Cell 19 — SMOKE TEST 8: Final val Recall@40, compare to baseline"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "90530a21",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Move trained model off GPU before loading the saved checkpoint to free VRAM\n",
+    "print('Releasing training model from GPU...')\n",
+    "try:\n",
+    "    model.to('cpu')\n",
+    "except Exception: pass\n",
+    "gc.collect()\n",
+    "if torch.cuda.is_available():\n",
+    "    torch.cuda.empty_cache()\n",
+    "    torch.cuda.synchronize()\n",
+    "    torch.cuda.empty_cache()\n",
+    "    free_gb, total_gb = (x/1e9 for x in torch.cuda.mem_get_info())\n",
+    "    print(f'GPU after release: free={free_gb:.2f} GB / {total_gb:.2f} GB')\n",
+    "\n",
+    "# Reload from disk for a clean state\n",
+    "print(f'Loading fine-tuned checkpoint from {ckpt_dir}...')\n",
+    "best_model = SentenceTransformer(str(ckpt_dir), device='cuda')\n",
+    "\n",
+    "# Manual eval\n",
+    "finetuned_recall40, mean_r_ft, med_r_ft = eval_recall_and_rank(\n",
+    "    best_model, k=CFG['VAL_RECALL_K'], label='fine-tuned'\n",
+    ")\n",
+    "\n",
+    "print(f'\\nVal Recall@{CFG[\"VAL_RECALL_K\"]}:')\n",
+    "print(f'  baseline    = {baseline_recall40:.4f}')\n",
+    "print(f'  fine-tuned  = {finetuned_recall40:.4f}')\n",
+    "delta = finetuned_recall40 - baseline_recall40\n",
+    "print(f'  Δ           = {delta:+.4f}')\n",
+    "if delta < 0.01:\n",
+    "    print('  WARNING: improvement under 1pp on val Recall@40. Sanity-check hard negatives + epochs.')\n",
+    "else:\n",
+    "    print('  PASS — meaningful upstream improvement.')\n",
+    "print('\\nSMOKE 8 PASS')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "86bc4061",
+   "metadata": {},
+   "source": [
+    "## Cell 19.5 — GATE B: Avg gold rank improvement\n",
+    "\n",
+    "Second eval gate alongside ΔRecall@40. We require the **mean rank of gold cits** in the corpus retrieval to drop after fine-tuning. If Recall@40 improves but mean rank doesn't, the model just shuffled within the top-40 without lifting harder gold up the list.\n",
+    "\n",
+    "**Both gates must pass before re-encoding.**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c994bddd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Gate B uses mean gold rank — both baseline and fine-tuned values were ALREADY\n",
+    "# computed by eval_recall_and_rank() in Cells 16 and 19. No need to reload the\n",
+    "# baseline encoder or re-encode the corpus.\n",
+    "\n",
+    "print(f'\\n  Baseline:    mean_rank={mean_r_base:>8.0f}  median_rank={med_r_base:>6.0f}')\n",
+    "print(f'  Fine-tuned:  mean_rank={mean_r_ft:>8.0f}  median_rank={med_r_ft:>6.0f}')\n",
+    "print(f'  Δmean_rank: {mean_r_base - mean_r_ft:+.0f}  (positive = improvement)')\n",
+    "\n",
+    "gate_a_pass = (delta is not None and delta >= 0.01)\n",
+    "gate_b_pass = (mean_r_ft <= mean_r_base)\n",
+    "\n",
+    "print(f'\\n  Gate A (ΔRecall@{CFG[\"VAL_RECALL_K\"]} ≥ +0.01): {\"PASS\" if gate_a_pass else \"FAIL\"} ({delta:+.4f})')\n",
+    "print(f'  Gate B (mean gold rank improves): {\"PASS\" if gate_b_pass else \"FAIL\"}')\n",
+    "\n",
+    "if gate_a_pass and gate_b_pass:\n",
+    "    print('\\n>>> BOTH GATES PASS — proceed to corpus re-encoding cells.')\n",
+    "    PROCEED_TO_REENCODE = True\n",
+    "else:\n",
+    "    print('\\n>>> STOP — do NOT re-encode laws_de or queries yet.')\n",
+    "    print('    Iterate hyperparameters before continuing:')\n",
+    "    print('     - try LEARNING_RATE 5e-5 or 1e-5 instead of 1.5e-5')\n",
+    "    print('     - increase EPOCHS to 3')\n",
+    "    print('     - lower FALSE_NEG_SIM_RATIO to 0.92 (more aggressive filter)')\n",
+    "    print('     - increase N_HARD_NEGS_PER_POS to 6')\n",
+    "    PROCEED_TO_REENCODE = False"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 20 — Re-encode laws_de.csv with fine-tuned model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "passages = [P_PFX + str(t) for t in laws['text'].astype(str).tolist()]\n",
+    "print(f'Re-encoding {len(passages):,} law passages...')\n",
+    "t0 = time.time()\n",
+    "laws_emb = best_model.encode(\n",
+    "    passages,\n",
+    "    batch_size=CFG['ENCODE_BATCH'],\n",
+    "    show_progress_bar=True,\n",
+    "    convert_to_numpy=True,\n",
+    "    normalize_embeddings=True,\n",
+    "    precision='float32',\n",
+    ")\n",
+    "print(f'  done in {time.time()-t0:.1f}s  shape={laws_emb.shape}')\n",
+    "np.save(OUT / 'laws_emb_v1.npy', laws_emb.astype(np.float16))\n",
+    "print(f'  saved {OUT}/laws_emb_v1.npy ({laws_emb.nbytes/2/1e6:.1f} MB as fp16)')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 21 — Re-encode train/val/test queries"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def enc(texts, prefix=Q_PFX):\n",
+    "    arr = best_model.encode(\n",
+    "        [prefix + str(t) for t in texts],\n",
+    "        batch_size=CFG['ENCODE_BATCH'], show_progress_bar=True,\n",
+    "        convert_to_numpy=True, normalize_embeddings=True,\n",
+    "    ).astype(np.float32)\n",
+    "    return arr\n",
+    "\n",
+    "q_train = enc(train['query'].tolist())\n",
+    "q_val   = enc(val['query'].tolist())\n",
+    "q_test  = enc(test['query'].tolist())\n",
+    "print(f'shapes: train={q_train.shape} val={q_val.shape} test={q_test.shape}')\n",
+    "\n",
+    "np.savez_compressed(OUT / 'queries_v1.npz',\n",
+    "                    train=q_train.astype(np.float16),\n",
+    "                    val=q_val.astype(np.float16),\n",
+    "                    test=q_test.astype(np.float16))\n",
+    "print(f'saved {OUT}/queries_v1.npz')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 22 — Quick val F1 sanity (top-K cosine retrieval, K=18 to mimic v51 K)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "K_PRED = 18\n",
+    "rows = []\n",
+    "for i, r in val.iterrows():\n",
+    "    qe = q_val[i:i+1]\n",
+    "    sims = (qe @ laws_emb.T)[0]\n",
+    "    top = np.argsort(-sims)[:K_PRED]\n",
+    "    pred = set(laws_cits_arr[j] for j in top)\n",
+    "    gold = set(c.strip() for c in str(r['gold_citations']).split(';') if c.strip())\n",
+    "    gold_in_laws = gold & laws_cit_set\n",
+    "    if not gold_in_laws:\n",
+    "        rows.append({'qid': r['query_id'], 'tp': 0, 'precision': 0, 'recall': 0, 'f1': 0, 'gold_in_laws': 0})\n",
+    "        continue\n",
+    "    tp = len(pred & gold_in_laws)\n",
+    "    prec = tp / max(1, len(pred))\n",
+    "    rec  = tp / len(gold_in_laws)\n",
+    "    f1   = 2*prec*rec / max(1e-9, prec+rec)\n",
+    "    rows.append({'qid': r['query_id'], 'tp': tp, 'precision': prec, 'recall': rec, 'f1': f1, 'gold_in_laws': len(gold_in_laws)})\n",
+    "\n",
+    "ev = pd.DataFrame(rows)\n",
+    "print(ev)\n",
+    "print(f'\\nVal mean macro-F1 (laws-only retrieval @ K={K_PRED}): {ev[\"f1\"].mean():.4f}')\n",
+    "ev.to_csv(OUT / 'val_baseline_vs_finetuned.csv', index=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 23 — Push checkpoint + embeddings to Hugging Face (optional)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Uncomment + set token to push:\n",
+    "# from huggingface_hub import HfApi, login\n",
+    "# login(token='hf_xxx')   # or env var HF_TOKEN\n",
+    "# api = HfApi()\n",
+    "# repo = 'Dharun72/swiss-legal-checkpoints'\n",
+    "# api.create_repo(repo, repo_type='model', exist_ok=True)\n",
+    "# api.upload_folder(folder_path=str(ckpt_dir), repo_id=repo, repo_type='model', path_in_repo='e5legal_v1')\n",
+    "# api.upload_file(path_or_fileobj=str(OUT / 'laws_emb_v1.npy'), path_in_repo='laws_emb_v1.npy', repo_id=repo, repo_type='model')\n",
+    "# api.upload_file(path_or_fileobj=str(OUT / 'queries_v1.npz'), path_in_repo='queries_v1.npz', repo_id=repo, repo_type='model')\n",
+    "# print('uploaded')\n",
+    "print('skipped — uncomment when ready to push')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## What to look at while it runs\n",
+    "\n",
+    "1. **SMOKE 1** — GPU mem >= 16GB and loads tiny e5. If fails, you're CPU-only — abort.\n",
+    "2. **SMOKE 2** — train gold coverage in laws_de should be ≥80%. If lower, your CSV is corrupted or laws_de is stale.\n",
+    "3. **SMOKE 3** — pair count >= 1000. If much lower, citation strings don't match laws_de format — likely a normalization bug.\n",
+    "4. **SMOKE 4** — printed hard negs look topically similar to query but distinct from gold. If they look random, BM25 isn't being scored on the right tokens.\n",
+    "5. **SMOKE 5** — baseline cos(Q,Pos) should be > cos(Q,Neg). If reversed, prefix or pooling is broken.\n",
+    "6. **SMOKE 6** — record baseline Recall@40. This is what fine-tune must beat.\n",
+    "7. **SMOKE 7** — after 50 steps, margin (Q,Pos)-(Q,Neg) MUST increase. If not, lr too low or loss broken.\n",
+    "8. **SMOKE 8** — Δ Recall@40 ≥ +0.01 is a green light. < +0.005 means the contrastive signal is weak; iterate hyperparameters before re-encoding the corpus.\n",
+    "9. **GATE B** — mean gold rank must improve. If Recall@40 lifts but rank doesn't, the model is shuffling within top-40 instead of pulling harder gold up. Iterate before re-encoding.\n",
+    "\n",
+    "**Gate B (avg gold rank)** must ALSO pass alongside SMOKE 8.\n",
+    "\n",
+    "**If both gates pass**, the new `laws_emb_v1.npy` and `queries_v1.npz` go back into the Kaggle notebook to replace the existing precomputed embeddings — that's where the LB lift will show up."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}