Blame for TokenCal_Q1_Notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ac6d6be7",
   "metadata": {},
   "source": [
    "# TokenCal Q1 — Complete GPU Experiment Notebook\n",
    "\n",
    "**Three modes: Preflight → Benchmark → Full Training**  \n",
    "**Backbones:** Qwen2.5-1.5B, Gemma3-4B-IT, Llama3.2-3B-Instruct  \n",
    "**Seeds:** 42, 123, 456  \n",
    "**Comparators:** M0, M1, M2, M3, Focal loss, R-Drop, Temperature Scaling, Constant and Shuffled controls  \n",
    "**Datasets:** BEnQA (train/val/test), B-REASO (Bengali OOD), MILU (Bengali + GU/ML/TE)\n",
    "\n",
    "**Private HF repo created:** https://huggingface.co/Tanvir106/TokenCal-Q1-Research\n",
    "\n",
    "This is a **complete source notebook**, not completed model-training results. It must be executed with a CUDA GPU, compatible dependencies, a Hugging Face login, and the source datasets. The module's CPU tests passed in the authoring environment; GPU preflight, benchmark, full training, and remote uploads have **not been executed**. Run one mode at a time.\n",
    "\n",
    "**Before you start:** Choose `Runtime → Change runtime type → T4 GPU` in Colab or a GPU accelerator in Kaggle. Gemma3 4B might require more memory than one T4; the preflight should report fit/compatibility issues rather than beginning full training."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9df808b9",
   "metadata": {},
   "source": [
    "## 0 — Set up a persistent project directory\n",
    "\n",
    "All outputs save under `tokencal_workspace`. For a long run on Colab, mount Google Drive or persistent storage **if available**; session disks are temporary. The `--sync-hf checkpoints` option also copies completed checkpoints/logs to your private Hugging Face repo for recovery. **Make sure your account has write access to `Tanvir106/TokenCal-Q1-Research`.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70196340",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys\n",
    "from pathlib import Path\n",
    "PROJECT_ROOT = Path('/content/TokenCal-Q1-Research') if Path('/content').exists() else Path.cwd() / 'TokenCal-Q1-Research'\n",
    "PROJECT_ROOT.mkdir(parents=True, exist_ok=True)\n",
    "os.chdir(PROJECT_ROOT)\n",
    "os.environ['TOKENCAL_ROOT'] = str(PROJECT_ROOT / 'tokencal_workspace')\n",
    "print('Project:', PROJECT_ROOT)\n",
    "print('Workspace:', os.environ['TOKENCAL_ROOT'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "484f5495",
   "metadata": {},
   "source": [
    "## 1 — Write version-aware dependencies and install"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "519a8fea",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile requirements.txt\n",
    "# Match the training environment recorded in the original TokenCal ZIP.\n",
    "# CUDA-enabled torch should be provided by Colab/Kaggle (do not replace torch blindly).\n",
    "transformers==5.16.1\n",
    "peft==0.20.0\n",
    "bitsandbytes==0.50.2\n",
    "accelerate==1.14.0\n",
    "datasets==4.8.5\n",
    "huggingface_hub==1.29.0\n",
    "pandas>=2.2\n",
    "numpy>=2.1\n",
    "scipy>=1.16\n",
    "scikit-learn>=1.6\n",
    "matplotlib>=3.10\n",
    "pytest>=8.0\n",
    "safetensors>=0.5\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5da94207",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Internet required. Torch itself is provided by the active GPU runtime.\n",
    "%pip install -q -r requirements.txt\n",
    "import torch\n",
    "print('PyTorch:', torch.__version__, 'CUDA available:', torch.cuda.is_available())\n",
    "print('GPU(s):', [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])\n",
    "if not torch.cuda.is_available():\n",
    "    print('WARNING: switch to a GPU runtime before running preflight/benchmark/full training.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e457e93c",
   "metadata": {},
   "source": [
    "## 2 — Generate the complete research engine\n",
    "\n",
    "The next cell writes the **complete** Python implementation. It includes the original TokenCal M0–M3 losses (with normalized weighted JS), QLoRA adapters, deterministic seeds, original BEnQA split locking, three modes, full evaluation, confidence metrics, bootstrap, automatic checkpoints, save/resume, and Hugging Face uploads."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bfdc838",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile tokencal_q1.py\n",
    "\"\"\"TokenCal Q1: controlled, multi-backbone experiments.\n",
    "\n",
    "Runs on a Linux NVIDIA CUDA runtime after installing requirements.txt.\n",
    "Source fidelity: BEnQA original fixed 5,077 split; exact A/B/C/D next-token\n",
    "scoring, full-vocabulary task CE, Brier/2, JS/ln2 and normalized weighted JS.\n",
    "\n",
    "Modes:\n",
    "  preflight : forward/backward, optimizer, save/reload/resume on tiny data\n",
    "  benchmark : short measured training, never counted as experimental result\n",
    "  full      : complete fixed-dataset experiments, selection, external evaluation\n",
    "\n",
    "Strictly no CUDA/full-model execution occurs merely by importing this module.\n",
    "\"\"\"\n",
    "from __future__ import annotations\n",
    "\n",
    "import argparse\n",
    "import csv\n",
    "import dataclasses\n",
    "import gc\n",
    "import hashlib\n",
    "import json\n",
    "import math\n",
    "import os\n",
    "import random\n",
    "import re\n",
    "import shutil\n",
    "import subprocess\n",
    "import sys\n",
    "import time\n",
    "import traceback\n",
    "import uuid\n",
    "import zipfile\n",
    "from pathlib import Path\n",
    "from typing import Any, Optional\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "SEEDS = (42, 123, 456)\n",
    "BASE_REVISIONS = {\n",
    "    \"qwen\": (\"Qwen/Qwen2.5-1.5B-Instruct\", \"989aa7980e4cf806f80c7fef2b1adb7bc71aa306\"),\n",
    "    \"gemma\": (\"google/gemma-3-4b-it\", None),\n",
    "    \"llama\": (\"meta-llama/Llama-3.2-3B-Instruct\", None),\n",
    "}\n",
    "VARIANTS = (\"M0\", \"M1\", \"M2\", \"M3\", \"Focal\", \"RDrop\", \"Const\", \"Shuffle\", \"TS\")\n",
    "LANGUAGES = {\"bn\": \"Bengali\", \"gu\": \"Gujarati\", \"ml\": \"Malayalam\", \"te\": \"Telugu\"}\n",
    "BENQA_SHA = \"b8cd623abbe78b5e5d4dac422cfe4026163d0e68\"\n",
    "BENQA_ORIGIN = \"https://github.com/sheikhshafayat/BEnQA.git\"\n",
    "ORIGINAL_HF_RESULTS = \"hr1018/tokencal-results\"\n",
    "EXPECTED_SPLITS = {\"train\": 3553, \"validation\": 762, \"test\": 762}\n",
    "EXPECTED_MILU = {\"bn\": 6636, \"gu\": 4825, \"ml\": 4321, \"te\": 7304}\n",
    "HF_REPO = \"Tanvir106/TokenCal-Q1-Research\"\n",
    "PROJECT = Path(os.environ.get(\"TOKENCAL_ROOT\", \"./tokencal_workspace\")).expanduser().resolve()\n",
    "CHOICES = \"ABCD\"\n",
    "\n",
    "\n",
    "def canonical_text(x: Any) -> str:\n",
    "    if x is None or (isinstance(x, (float, np.floating)) and np.isnan(x)):\n",
    "        return \"\"\n",
    "    return re.sub(r\"\\s+\", \" \", str(x).replace(\"\\u00a0\", \" \").replace(\"\\u200b\", \"\").replace(\"\\ufeff\", \"\")).strip()\n",
    "\n",
    "\n",
    "def canonical_answer(x: Any) -> Optional[str]:\n",
    "    s = canonical_text(x).upper()\n",
    "    if s in CHOICES:\n",
    "        return s\n",
    "    # Exactly like original BEnQA answer normalization; not used for MILU numeric labels.\n",
    "    match = re.search(r\"\\b([ABCD])\\b\", s)\n",
    "    return match.group(1) if match else None\n",
    "\n",
    "\n",
    "def milu_answer(x: Any) -> Optional[str]:\n",
    "    s = canonical_text(x).upper()\n",
    "    if s in CHOICES:\n",
    "        return s\n",
    "    if s in {\"1\", \"2\", \"3\", \"4\"}:\n",
    "        return CHOICES[int(s) - 1]\n",
    "    if s.startswith(\"OPTION\"):\n",
    "        return canonical_answer(s)\n",
    "    return None\n",
    "\n",
    "\n",
    "def sha_file(path: Path) -> str:\n",
    "    h = hashlib.sha256()\n",
    "    with open(path, \"rb\") as f:\n",
    "        for chunk in iter(lambda: f.read(1024 * 1024), b\"\"):\n",
    "            h.update(chunk)\n",
    "    return h.hexdigest()\n",
    "\n",
    "\n",
    "def json_dump_atomic(obj: Any, path: Path) -> None:\n",
    "    path.parent.mkdir(parents=True, exist_ok=True)\n",
    "    temp = path.with_name(path.name + \".tmp.\" + uuid.uuid4().hex[:8])\n",
    "    with open(temp, \"w\", encoding=\"utf-8\") as f:\n",
    "        json.dump(obj, f, ensure_ascii=False, indent=2, default=str)\n",
    "        f.flush(); os.fsync(f.fileno())\n",
    "    os.replace(temp, path)\n",
    "\n",
    "\n",
    "def replace_dir_atomic(staged: Path, destination: Path) -> None:\n",
    "    \"\"\"Keep previous checkpoint until next checkpoint is completely written.\"\"\"\n",
    "    old = destination.with_name(destination.name + \".old\")\n",
    "    if old.exists(): shutil.rmtree(old)\n",
    "    if destination.exists(): os.replace(destination, old)\n",
    "    try:\n",
    "        os.replace(staged, destination)\n",
    "    except Exception:\n",
    "        if old.exists() and not destination.exists(): os.replace(old, destination)\n",
    "        raise\n",
    "    if old.exists(): shutil.rmtree(old)\n",
    "\n",
    "\n",
    "def seed_all(seed: int) -> None:\n",
    "    import torch\n",
    "    random.seed(seed); np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)\n",
    "\n",
    "\n",
    "def rng_snapshot() -> dict:\n",
    "    import torch\n",
    "    return {\"python\": random.getstate(), \"numpy\": np.random.get_state(),\n",
    "            \"torch\": torch.get_rng_state(),\n",
    "            \"cuda\": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None}\n",
    "\n",
    "\n",
    "def restore_rng(state: dict) -> None:\n",
    "    import torch\n",
    "    random.setstate(state[\"python\"]); np.random.set_state(state[\"numpy\"])\n",
    "    torch.set_rng_state(state[\"torch\"])\n",
    "    if torch.cuda.is_available() and state.get(\"cuda\"):\n",
    "        torch.cuda.set_rng_state_all(state[\"cuda\"])\n",
    "\n",
    "\n",
    "def hf_token() -> Optional[str]:\n",
    "    from huggingface_hub import get_token\n",
    "    return os.environ.get(\"HF_TOKEN\") or get_token()\n",
    "\n",
    "\n",
    "def ensure_model_revision(model_key: str, work_root: Path) -> tuple[str, str]:\n",
    "    from huggingface_hub import HfApi\n",
    "    model_id, pinned = BASE_REVISIONS[model_key]\n",
    "    lock = work_root / \"locks\" / f\"model_{model_key}.json\"\n",
    "    if lock.exists():\n",
    "        old = json.loads(lock.read_text());\n",
    "        if old[\"model_id\"] != model_id: raise RuntimeError(\"Model lock mismatch\")\n",
    "        return model_id, old[\"revision\"]\n",
    "    info = HfApi(token=hf_token()).model_info(model_id, revision=pinned or \"main\")\n",
    "    revision = info.sha\n",
    "    json_dump_atomic({\"model_id\": model_id, \"revision\": revision}, lock)\n",
    "    return model_id, revision\n",
    "\n",
    "\n",
    "def source_reference(path_name: str, root: Path, source_zip: Optional[Path]) -> Path:\n",
    "    \"\"\"Never silently create a new train/test split.\"\"\"\n",
    "    target = root / \"reference\" / path_name\n",
    "    if target.exists(): return target\n",
    "    target.parent.mkdir(parents=True, exist_ok=True)\n",
    "    if source_zip and source_zip.is_file():\n",
    "        with zipfile.ZipFile(source_zip) as z:\n",
    "            key = next((k for k in z.namelist() if k.endswith(\"/datasets/\" + path_name)), None)\n",
    "            if key:\n",
    "                with open(target, \"wb\") as f: f.write(z.read(key))\n",
    "                return target\n",
    "    # Original results repo is referenced READ-ONLY; raw source data are not re-uploaded.\n",
    "    try:\n",
    "        from huggingface_hub import hf_hub_download\n",
    "        f = hf_hub_download(repo_id=ORIGINAL_HF_RESULTS, repo_type=\"dataset\",\n",
    "                            filename=\"datasets/\" + path_name, token=hf_token())\n",
    "        shutil.copy2(f, target)\n",
    "        return target\n",
    "    except Exception as e:\n",
    "        raise FileNotFoundError(f\"Missing {path_name}. Provide --source-zip path/to/tokencal-results.zip \"\n",
    "                                f\"(uploaded original archive). Original Hub fallback failed: {e}\") from e\n",
    "\n",
    "\n",
    "def load_benqa(root: Path, source_zip: Optional[Path]) -> pd.DataFrame:\n",
    "    \"\"\"Rebuild 5,077 source examples, then reuse ORIGINAL split assignments exactly.\"\"\"\n",
    "    cache = root / \"private_data\" / \"benqa_canonical.csv\"\n",
    "    if cache.exists():\n",
    "        df = pd.read_csv(cache, dtype=str, keep_default_na=False)\n",
    "        if df[\"split\"].value_counts().to_dict() == EXPECTED_SPLITS: return df\n",
    "        raise RuntimeError(\"BEnQA cache split counts mismatch; refusing to reuse\")\n",
    "    repo = root / \"private_data\" / \"BEnQA_source\"\n",
    "    if not (repo / \".git\").exists():\n",
    "        repo.parent.mkdir(parents=True, exist_ok=True)\n",
    "        subprocess.run([\"git\", \"clone\", BENQA_ORIGIN, str(repo)], check=True)\n",
    "    subprocess.run([\"git\", \"checkout\", \"--detach\", BENQA_SHA], cwd=repo, check=True)\n",
    "    sha = subprocess.check_output([\"git\", \"rev-parse\", \"HEAD\"], cwd=repo, text=True).strip()\n",
    "    assert sha == BENQA_SHA, (sha, BENQA_SHA)\n",
    "    folder = repo / \"data\" / \"BEnQA\"\n",
    "    csvs = sorted(folder.glob(\"*.csv\"))\n",
    "    if len(csvs) != 15: raise RuntimeError(f\"BEnQA source file count {len(csvs)} != 15\")\n",
    "    records = []\n",
    "    keys = [\"English Question\", \"Bengali Question\", \"A\", \"B\", \"C\", \"D\",\n",
    "            \"A Bn\", \"B Bn\", \"C Bn\", \"D Bn\", \"Correct Answer\"]\n",
    "    for path in csvs:\n",
    "        raw = pd.read_csv(path, low_memory=False)\n",
    "        if not set(keys).issubset(raw.columns): raise ValueError(f\"BEnQA schema changed: {path.name}\")\n",
    "        for idx, row in raw.iterrows():\n",
    "            pair_id = \"benqa_\" + hashlib.sha256(\n",
    "                f\"{BENQA_SHA}|{path.name}|{idx + 1}\".encode()).hexdigest()[:16]\n",
    "            rec = {\"pair_id\": pair_id, \"source_file\": path.name, \"source_row\": idx + 1,\n",
    "                   \"english_question\": canonical_text(row[\"English Question\"]),\n",
    "                   \"bengali_question\": canonical_text(row[\"Bengali Question\"]),\n",
    "                   \"correct_answer\": canonical_answer(row[\"Correct Answer\"])}\n",
    "            for k, suffix in [(\"en\", \"\"), (\"bn\", \" Bn\")]:\n",
    "                for c in CHOICES: rec[f\"option_{c.lower()}_{k}\"] = canonical_text(row[c + suffix])\n",
    "            records.append(rec)\n",
    "    if len(records) != 5087: raise RuntimeError(f\"Unexpected raw BEnQA rows: {len(records)}\")\n",
    "    df = pd.DataFrame(records)\n",
    "    splits = pd.read_csv(source_reference(\"benqa_split_assignments.csv\", root, source_zip), dtype=str)\n",
    "    if len(splits) != 5077 or not splits[\"pair_id\"].is_unique:\n",
    "        raise RuntimeError(\"Original split assignment source invalid\")\n",
    "    df = df.merge(splits[[\"pair_id\", \"question_group_id\", \"record_fingerprint\", \"split\"]],\n",
    "                  how=\"inner\", on=\"pair_id\", validate=\"one_to_one\")\n",
    "    if len(df) != 5077: raise RuntimeError(\"BEnQA joins differ from original; check source revision\")\n",
    "    if not df[\"correct_answer\"].isin(list(CHOICES)).all(): raise RuntimeError(\"Invalid BEnQA answer\")\n",
    "    if df[\"split\"].value_counts().to_dict() != EXPECTED_SPLITS: raise RuntimeError(\"Split size drift\")\n",
    "    if df.groupby(\"question_group_id\")[\"split\"].nunique().max() != 1: raise RuntimeError(\"Split leakage\")\n",
    "    needed = [\"english_question\", \"bengali_question\"] + [f\"option_{c.lower()}_{lang}\" for lang in (\"en\", \"bn\") for c in CHOICES]\n",
    "    if df[needed].eq(\"\").any().any(): raise RuntimeError(\"Unexpected empty BEnQA content\")\n",
    "    cache.parent.mkdir(parents=True, exist_ok=True)\n",
    "    df.to_csv(cache, index=False)\n",
    "    json_dump_atomic({\"sha\": BENQA_SHA, \"counts\": df.split.value_counts().to_dict(),\n",
    "                      \"split_assignments_sha256\": sha_file(root / \"reference\" / \"benqa_split_assignments.csv\")},\n",
    "                     root / \"locks\" / \"benqa.json\")\n",
    "    return df\n",
    "\n",
    "\n",
    "def load_breaso(root: Path, source_zip: Optional[Path]) -> pd.DataFrame:\n",
    "    f = source_reference(\"b_reaso_val_ood_canonical.csv\", root, source_zip)\n",
    "    d = pd.read_csv(f, dtype=str, keep_default_na=False)\n",
    "    if len(d) != 1215: raise RuntimeError(f\"B-REASO count changed: {len(d)} != 1215\")\n",
    "    if not d[\"answer_label\"].isin(list(CHOICES)).all(): raise RuntimeError(\"B-REASO invalid labels\")\n",
    "    records = []\n",
    "    for _, row in d.iterrows():\n",
    "        records.append({\"item_id\": row[\"b_reaso_ood_id\"], \"language\": \"bn\",\n",
    "                        \"question\": row[\"question\"], \"answer\": row[\"answer_label\"],\n",
    "                        **{f\"option_{c}\": row[f\"option_{c}\"] for c in CHOICES}})\n",
    "    return pd.DataFrame(records)\n",
    "\n",
    "\n",
    "def load_milu(root: Path) -> pd.DataFrame:\n",
    "    \"\"\"The official ai4bharat/MILU dataset is gated; never substitute a sample.\"\"\"\n",
    "    from datasets import load_dataset\n",
    "    from huggingface_hub import HfApi\n",
    "    lock_path = root / \"locks\" / \"milu.json\"\n",
    "    sha = json.loads(lock_path.read_text())[\"revision\"] if lock_path.exists() else HfApi(token=hf_token()).dataset_info(\"ai4bharat/MILU\").sha\n",
    "    records = []\n",
    "    for code, language in LANGUAGES.items():\n",
    "        try:\n",
    "            raw = load_dataset(\"ai4bharat/MILU\", data_dir=language, split=\"test\", revision=sha, token=hf_token())\n",
    "        except Exception:\n",
    "            try:\n",
    "                raw = load_dataset(\"ai4bharat/MILU\", name=language, split=\"test\", revision=sha, token=hf_token())\n",
    "            except Exception as e:\n",
    "                raise RuntimeError(f\"MILU {language} unavailable. Accept gated access in the official \"\n",
    "                                   f\"dataset repo and rerun. {e}\") from e\n",
    "        for i, r in enumerate(raw):\n",
    "            options = [canonical_text(r[f\"option{j}\"]) for j in range(1,5)]\n",
    "            ans = milu_answer(r.get(\"target\"))\n",
    "            if ans is None or not canonical_text(r.get(\"question\")) or not all(options):\n",
    "                raise RuntimeError(f\"MILU unexpected schema/invalid row {code}/{i}; no silent skipping\")\n",
    "            rec = {\"item_id\": f\"MILU_{code}_{i}\", \"language\": code, \"question\": canonical_text(r[\"question\"]),\n",
    "                   \"answer\": ans, \"is_translated\": r.get(\"is_translated\"), \"subject\": r.get(\"subject\"),\n",
    "                   **{f\"option_{c}\": options[j] for j,c in enumerate(CHOICES)}}\n",
    "            rec[\"_content_hash\"] = hashlib.sha256(json.dumps([code,rec[\"question\"],*options,ans],ensure_ascii=False).encode()).hexdigest()\n",
    "            records.append(rec)\n",
    "    df = pd.DataFrame(records)\n",
    "    df = df.drop_duplicates(subset=[\"language\", \"_content_hash\"], keep=\"first\").drop(columns=\"_content_hash\")\n",
    "    observed = df.language.value_counts().to_dict()\n",
    "    if observed != EXPECTED_MILU:\n",
    "        raise RuntimeError(f\"MILU counts differ from original locked evaluation {observed} != {EXPECTED_MILU}. \"\n",
    "                           \"Check source revision and dedup policy; do not silently change benchmark.\")\n",
    "    json_dump_atomic({\"repo\": \"ai4bharat/MILU\", \"revision\": sha, \"counts\": observed},lock_path)\n",
    "    return df\n",
    "\n",
    "\n",
    "def render_messages(row: Any, lang: str, external: bool = False) -> list:\n",
    "    if external:\n",
    "        question = canonical_text(row[\"question\"])\n",
    "        options = {c: canonical_text(row[f\"option_{c}\"]) for c in CHOICES}\n",
    "    else:\n",
    "        question = canonical_text(row[\"english_question\" if lang == \"en\" else \"bengali_question\"])\n",
    "        options = {c: canonical_text(row[f\"option_{c.lower()}_{lang}\"]) for c in CHOICES}\n",
    "    if lang == \"en\":\n",
    "        system = \"You are a multiple-choice question answering system. Return only one label: A, B, C, or D.\"\n",
    "        user = (f\"Question:\\n{question}\\n\\nOptions:\\n\"\n",
    "                + \"\\n\".join(f\"{c}. {options[c]}\" for c in CHOICES)\n",
    "                + \"\\n\\nAnswer with only A, B, C, or D.\")\n",
    "    elif lang == \"bn\":\n",
    "        system = \"তুমি একটি বহুনির্বাচনী প্রশ্নের উত্তরদাতা সিস্টেম। শুধু একটি লেবেল দাও: A, B, C, অথবা D।\"\n",
    "        user = (f\"প্রশ্ন:\\n{question}\\n\\nবিকল্প:\\n\"\n",
    "                + \"\\n\".join(f\"{c}. {options[c]}\" for c in CHOICES)\n",
    "                + \"\\n\\nশুধু A, B, C, অথবা D-এর মধ্যে একটি দিয়ে উত্তর দাও।\")\n",
    "    else:\n",
    "        # Applies the same MCQ task to previously unseen languages; preserve question/options verbatim.\n",
    "        system = \"Answer the multiple-choice question. Return only one label: A, B, C, or D.\"\n",
    "        user = (f\"Question:\\n{question}\\n\\nOptions:\\n\"\n",
    "                + \"\\n\".join(f\"{c}. {options[c]}\" for c in CHOICES)\n",
    "                + \"\\n\\nAnswer with only A, B, C, or D.\")\n",
    "    return [{\"role\":\"system\", \"content\":system},{\"role\":\"user\", \"content\":user}]\n",
    "\n",
    "\n",
    "def tokenize_mcq(tok: Any, row: Any, lang: str, context: int, external: bool = False) -> list[int]:\n",
    "    rendered = tok.apply_chat_template(render_messages(row, lang, external), tokenize=False, add_generation_prompt=True)\n",
    "    ids = tok.encode(rendered, add_special_tokens=False)\n",
    "    if not ids: raise RuntimeError(\"Empty rendered MCQ prompt\")\n",
    "    if len(ids) + 1 > context:\n",
    "        raise RuntimeError(f\"Input too long ({len(ids)}+target>{context}), no truncation allowed. \"\n",
    "                           \"Increase context per model or document an explicit exclusion policy.\")\n",
    "    return ids\n",
    "\n",
    "\n",
    "def choice_token_ids(tok: Any) -> list[int]:\n",
    "    # Identical to the original Qwen next-token choice scoring; never hardcode 32–35 for other backbones.\n",
    "    ids = [tok.encode(c, add_special_tokens=False) for c in CHOICES]\n",
    "    if not all(len(x) == 1 for x in ids) or len({x[0] for x in ids}) != 4:\n",
    "        raise RuntimeError(f\"This backbone cannot support the original single-token A/B/C/D protocol: {ids}. \"\n",
    "                           \"Implement and validate a sequence-logprob protocol before comparing it.\")\n",
    "    return [x[0] for x in ids]\n",
    "\n",
    "\n",
    "def content_fertility(tok: Any, row: Any, lang: str) -> float:\n",
    "    # Identical MCQ content framing used by original CELL 10 (not the system/chat prompt).\n",
    "    question = canonical_text(row[\"english_question\" if lang == \"en\" else \"bengali_question\"])\n",
    "    header = \"Question:\" if lang == \"en\" else \"প্রশ্ন:\"\n",
    "    text = canonical_text(f\"{header} {question}\\n\" + \"\\n\".join(\n",
    "        f\"{c}. {canonical_text(row[f'option_{c.lower()}_{lang}'])}\" for c in CHOICES))\n",
    "    words = len(text.split())\n",
    "    if words < 1: raise ValueError(\"Empty content for fertility\")\n",
    "    return len(tok.encode(text, add_special_tokens=False)) / words\n",
    "\n",
    "\n",
    "def prepare_train_records(df: pd.DataFrame, tok: Any, root: Path, model_key: str,\n",
    "                          model_rev: str, limit: Optional[int] = None) -> tuple[list[dict],list[dict],list[dict]]:\n",
    "    # The q05/q95 bounds MUST be learned from full original TRAIN split only.\n",
    "    cache = root / \"private_data\" / f\"tokenized_{model_key}_{model_rev[:12]}.json\"\n",
    "    if limit is None and cache.exists():\n",
    "        saved=json.loads(cache.read_text())\n",
    "        if [len(saved[k]) for k in (\"train\",\"validation\",\"test\")] == [3553,762,762]:\n",
    "            return saved[\"train\"],saved[\"validation\"],saved[\"test\"]\n",
    "        raise RuntimeError(\"Corrupt token cache; remove it manually after auditing\")\n",
    "    tr = df[df.split == \"train\"].copy().reset_index(drop=True)\n",
    "    va = df[df.split == \"validation\"].copy().reset_index(drop=True)\n",
    "    te = df[df.split == \"test\"].copy().reset_index(drop=True)\n",
    "    if len(tr) != 3553 or len(va) != 762 or len(te) != 762: raise RuntimeError(\"Split drift\")\n",
    "    deltas = np.array([content_fertility(tok,r,\"bn\") - content_fertility(tok,r,\"en\") for _,r in tr.iterrows()])\n",
    "    q05,q95 = np.quantile(deltas,[.05,.95])\n",
    "    if not q95 > q05: raise RuntimeError(\"Bad fragmentation normalization range\")\n",
    "    json_dump_atomic({\"model\":model_key,\"revision\":model_rev,\"train_q05\":float(q05),\"train_q95\":float(q95),\n",
    "                      \"train_count\":len(tr)}, root / \"locks\" / f\"fragmentation_{model_key}.json\")\n",
    "    if limit: tr=tr.iloc[:limit];va=va.iloc[:max(8,min(24,limit))]\n",
    "    def encode(d: pd.DataFrame) -> list[dict]:\n",
    "        out=[]\n",
    "        for _, r in d.iterrows():\n",
    "            delta=content_fertility(tok,r,\"bn\")-content_fertility(tok,r,\"en\")\n",
    "            out.append({\"id\":r[\"pair_id\"],\"en\":tokenize_mcq(tok,r,\"en\",1024),\n",
    "                        \"bn\":tokenize_mcq(tok,r,\"bn\",1024),\"target\":CHOICES.index(r[\"correct_answer\"]),\n",
    "                        \"weight\":float(np.clip((delta-q05)/(q95-q05),0,1)),\"delta_tf\":float(delta)})\n",
    "        return out\n",
    "    trains,valids,tests=encode(tr),encode(va),encode(te) if limit is None else []\n",
    "    if limit is None:\n",
    "        json_dump_atomic({\"train\":trains,\"validation\":valids,\"test\":tests},cache)\n",
    "    return trains,valids,tests\n",
    "\n",
    "\n",
    "def format_external_records(df: pd.DataFrame, tok: Any, max_context: int) -> list[dict]:\n",
    "    return [{\"id\":r[\"item_id\"], \"lang\":r[\"language\"],\"tokens\":tokenize_mcq(tok,r,r[\"language\"],max_context,external=True),\n",
    "             \"target\":CHOICES.index(r[\"answer\"]),\"is_translated\":r.get(\"is_translated\",None)} for _,r in df.iterrows()]\n",
    "\n",
    "# --------------------- Model, objectives, metrics ---------------------\n",
    "def load_backbone(model_key: str, root: Path, seed: int, trainable: bool,\n",
    "                  adapter_checkpoint: Optional[Path] = None):\n",
    "    import torch\n",
    "    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n",
    "    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel\n",
    "    if not torch.cuda.is_available():\n",
    "        raise RuntimeError(\"QLoRA needs CUDA; choose a GPU-enabled T4/H100 runtime\")\n",
    "    seed_all(seed)\n",
    "    model_id,rev = ensure_model_revision(model_key,root)\n",
    "    tok = AutoTokenizer.from_pretrained(model_id, revision=rev, token=hf_token(), trust_remote_code=False)\n",
    "    if tok.chat_template is None: raise RuntimeError(f\"No chat template: {model_id}\")\n",
    "    tok.padding_side=\"right\"\n",
    "    if tok.pad_token_id is None: tok.pad_token=tok.eos_token\n",
    "    cids=choice_token_ids(tok)\n",
    "    bnb=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type=\"nf4\",\n",
    "                           bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.float16)\n",
    "    num=torch.cuda.device_count()\n",
    "    device_map=\"auto\" if num>1 else {\"\":0}\n",
    "    opts=dict(revision=rev, token=hf_token(), quantization_config=bnb,\n",
    "              device_map=device_map, dtype=torch.float16, trust_remote_code=False,\n",
    "              low_cpu_mem_usage=True)\n",
    "    try:\n",
    "        if model_key==\"gemma\":\n",
    "            from transformers import Gemma3ForConditionalGeneration\n",
    "            base=Gemma3ForConditionalGeneration.from_pretrained(model_id,**opts)\n",
    "        else:\n",
    "            base=AutoModelForCausalLM.from_pretrained(model_id,**opts)\n",
    "    except Exception as e:\n",
    "        raise RuntimeError(f\"Model loading failed for {model_id}. Check gated access, Gemma3 support, \"\n",
    "                           f\"quantization and 2xT4 sharding. Cause: {e}\") from e\n",
    "    if hasattr(base,\"hf_device_map\"):\n",
    "        bad={str(v) for v in base.hf_device_map.values() if str(v).startswith(\"cpu\") or str(v)==\"disk\"}\n",
    "        if bad: raise RuntimeError(f\"Model offloaded to CPU/disk: {bad}; refusing slow/unsafe training\")\n",
    "    base.config.use_cache=False\n",
    "    if trainable:\n",
    "        base=prepare_model_for_kbit_training(base, use_gradient_checkpointing=True,\n",
    "                                              gradient_checkpointing_kwargs={\"use_reentrant\":False})\n",
    "        if adapter_checkpoint:\n",
    "            model=PeftModel.from_pretrained(base,adapter_checkpoint,is_trainable=True)\n",
    "        else:\n",
    "            # Re-seed just before LoRA initialization for paired, matched seed runs.\n",
    "            seed_all(seed)\n",
    "            target=[\"q_proj\",\"k_proj\",\"v_proj\",\"o_proj\",\"gate_proj\",\"up_proj\",\"down_proj\"]\n",
    "            model=get_peft_model(base,LoraConfig(r=16,lora_alpha=32,lora_dropout=.05,bias=\"none\",\n",
    "                                    task_type=\"CAUSAL_LM\",target_modules=target))\n",
    "        trainable_params=[p for p in model.parameters() if p.requires_grad]\n",
    "        if not trainable_params:raise RuntimeError(\"No trainable adapter parameters\")\n",
    "    else:\n",
    "        model=PeftModel.from_pretrained(base, adapter_checkpoint, is_trainable=False) if adapter_checkpoint else base\n",
    "    model.train(trainable)\n",
    "    return model,tok,cids,model_id,rev\n",
    "\n",
    "\n",
    "def input_device(model) -> Any:\n",
    "    # Model may be sharded across two GPUs; send IDs to embedding device.\n",
    "    return model.get_input_embeddings().weight.device\n",
    "\n",
    "\n",
    "def batch_view(rows: list[dict], key: str, tok: Any, device: Any) -> dict:\n",
    "    import torch\n",
    "    samples=[r[key] for r in rows]\n",
    "    maxlen=max(map(len,samples))\n",
    "    pad=tok.pad_token_id\n",
    "    ids=torch.full((len(rows),maxlen),pad,dtype=torch.long,device=device)\n",
    "    mask=torch.zeros((len(rows),maxlen),dtype=torch.long,device=device)\n",
    "    for i,seq in enumerate(samples):\n",
    "        ids[i,:len(seq)]=torch.tensor(seq,dtype=torch.long,device=device)\n",
    "        mask[i,:len(seq)]=1\n",
    "    return {\"input_ids\":ids,\"attention_mask\":mask}\n",
    "\n",
    "\n",
    "def model_choice_logits(model, inputs: dict, choice_ids: list[int]):\n",
    "    import torch\n",
    "    out=model(**inputs, use_cache=False)\n",
    "    positions=inputs[\"attention_mask\"].sum(-1)-1\n",
    "    batch_idx=torch.arange(inputs[\"input_ids\"].size(0),device=positions.device)\n",
    "    # Full-vocabulary logits must remain available for exact task CE.\n",
    "    logits=out.logits[batch_idx,positions,:].float()\n",
    "    return logits,logits[:,choice_ids].softmax(dim=-1)\n",
    "\n",
    "\n",
    "def multiclass_brier(p,targets):\n",
    "    import torch.nn.functional as F\n",
    "    return ((p-F.one_hot(targets,num_classes=4).float())**2).sum(-1)\n",
    "\n",
    "\n",
    "def js_per_pair(p,q,eps=1e-8):\n",
    "    import torch\n",
    "    p=p.clamp_min(eps);q=q.clamp_min(eps);m=(p+q)*.5\n",
    "    return .5*(p*(p.log()-m.log())).sum(-1)+.5*(q*(q.log()-m.log())).sum(-1)\n",
    "\n",
    "\n",
    "def objectives(en_full,bn_full,en_p,bn_p,targets,weights,variant:str,\n",
    "               focal_gamma:float=2.0,rdrop_second=None):\n",
    "    \"\"\"Exactly reproduce original M0--M3 objective including weighted mean JS.\n",
    "\n",
    "    Ltask: mean full-vocab cross entropy of EN/BN next-token answer IDs.\n",
    "    Lcal: mean multiclass Brier divided by 2.\n",
    "    Lpair: mean JS/ln(2).\n",
    "    Lfrag: sum_i(w_i JS_i) / sum_i(w_i) / ln(2), with zero-safe fallback.\n",
    "    \"\"\"\n",
    "    import torch\n",
    "    import torch.nn.functional as F\n",
    "    # `targets` is integer choice indices; target token ids are explicit caller input.\n",
    "    if isinstance(targets,tuple):choice_idx,target_token_ids=targets\n",
    "    else:raise ValueError(\"Pass (choice_index_tensor, answer_token_id_tensor)\")\n",
    "    loss_ce_en=F.cross_entropy(en_full,target_token_ids,reduction=\"none\")\n",
    "    loss_ce_bn=F.cross_entropy(bn_full,target_token_ids,reduction=\"none\")\n",
    "    if variant==\"Focal\":\n",
    "        a=((1-(-loss_ce_en).exp())**focal_gamma)*loss_ce_en\n",
    "        b=((1-(-loss_ce_bn).exp())**focal_gamma)*loss_ce_bn\n",
    "        ltask=(a.mean()+b.mean())*.5\n",
    "    else:\n",
    "        ltask=(loss_ce_en.mean()+loss_ce_bn.mean())*.5\n",
    "    lcal=(multiclass_brier(en_p,choice_idx).mean()+multiclass_brier(bn_p,choice_idx).mean())*.25\n",
    "    js=js_per_pair(en_p,bn_p)\n",
    "    lpair=js.mean()/math.log(2.)\n",
    "    w=weights.float().clamp(0,1)\n",
    "    lfrag=torch.where(w.sum()>1e-8,(w*js).sum()/w.sum().clamp_min(1e-8),js.sum()*0.)/math.log(2.)\n",
    "    if variant==\"M0\" or variant==\"Focal\": total=ltask\n",
    "    elif variant==\"M1\": total=ltask+lcal\n",
    "    elif variant==\"M2\": total=ltask+lcal+lpair\n",
    "    elif variant in {\"M3\",\"Shuffle\"}:total=ltask+lcal+lpair+lfrag\n",
    "    elif variant==\"Const\":total=ltask+lcal+2*lpair\n",
    "    elif variant==\"RDrop\":\n",
    "        if rdrop_second is None: raise ValueError(\"RDrop requires a second stochastic forward pass\")\n",
    "        en2,bn2,pen2,pbn2=rdrop_second\n",
    "        ce2=(F.cross_entropy(en2,target_token_ids)+F.cross_entropy(bn2,target_token_ids))*.5\n",
    "        def bidirectional_kl(a,b):\n",
    "            a=a.clamp_min(1e-8); b=b.clamp_min(1e-8)\n",
    "            return ((a*(a.log()-b.log())).sum(-1)+(b*(b.log()-a.log())).sum(-1)).mean()*.5\n",
    "        total=(ltask+ce2)*.5 + (bidirectional_kl(en_p,pen2)+bidirectional_kl(bn_p,pbn2))*.5\n",
    "    else:raise ValueError(variant)\n",
    "    return total,{\"task\":ltask,\"cal\":lcal,\"pair\":lpair,\"frag\":lfrag}\n",
    "\n",
    "\n",
    "def forward_pair(model,tok,cids,rows:list[dict],variant:str,device):\n",
    "    import torch\n",
    "    en_full,en_p=model_choice_logits(model,batch_view(rows,\"en\",tok,device),cids)\n",
    "    bn_full,bn_p=model_choice_logits(model,batch_view(rows,\"bn\",tok,device),cids)\n",
    "    ci=torch.tensor([r[\"target\"] for r in rows],device=device,dtype=torch.long)\n",
    "    ti=torch.tensor([cids[r[\"target\"]] for r in rows],device=device,dtype=torch.long)\n",
    "    w=torch.tensor([r[\"weight\"] for r in rows],device=device,dtype=torch.float32)\n",
    "    if variant==\"Shuffle\":w=torch.tensor([r[\"shuffled_weight\"] for r in rows],device=device,dtype=torch.float32)\n",
    "    second=None\n",
    "    if variant==\"RDrop\":\n",
    "        # Two *independent* dropout draws on the same rows (EN and BN separately).\n",
    "        en2,pen2=model_choice_logits(model,batch_view(rows,\"en\",tok,device),cids)\n",
    "        bn2,pbn2=model_choice_logits(model,batch_view(rows,\"bn\",tok,device),cids)\n",
    "        second=(en2,bn2,pen2,pbn2)\n",
    "    return objectives(en_full,bn_full,en_p,bn_p,(ci,ti),w,variant,rdrop_second=second)\n",
    "\n",
    "\n",
    "def metrics_array(probs:np.ndarray, y:np.ndarray, ece_bins:int=15) -> dict:\n",
    "    from sklearn.metrics import roc_auc_score\n",
    "    if len(y)!=len(probs) or probs.ndim!=2 or probs.shape[1]!=4:raise ValueError(\"Invalid prediction shapes\")\n",
    "    pred=probs.argmax(axis=1);correct=(pred==y).astype(float)\n",
    "    conf=probs.max(axis=1);onehot=np.eye(4)[y]\n",
    "    brier=((probs-onehot)**2).sum(axis=1)\n",
    "    nll=-np.log(np.clip(probs[np.arange(len(y)),y],1e-12,1))\n",
    "    ece=0.\n",
    "    for j in range(ece_bins):\n",
    "        lo,hi=j/ece_bins,(j+1)/ece_bins\n",
    "        mask=(conf>=lo)&((conf<=hi) if j==ece_bins-1 else (conf<hi))\n",
    "        if mask.any():ece+=mask.mean()*abs(correct[mask].mean()-conf[mask].mean())\n",
    "    order=np.argsort(-conf,kind=\"stable\");errors=1-correct[order]\n",
    "    risk=np.cumsum(errors)/np.arange(1,len(errors)+1)\n",
    "    aurc=np.trapz(risk,x=np.arange(1,len(errors)+1)/len(errors)) if len(y)>1 else float(risk.mean())\n",
    "    try: err_auroc=float(roc_auc_score(1-correct,1-conf))\n",
    "    except ValueError:err_auroc=float(\"nan\")\n",
    "    return {\"n\":len(y),\"accuracy\":float(correct.mean()),\"mean_confidence\":float(conf.mean()),\n",
    "            \"brier\":float(brier.mean()),\"nll\":float(nll.mean()),\"ece\":float(ece),\n",
    "            \"aurc\":float(aurc),\"error_auroc\":err_auroc}\n",
    "\n",
    "\n",
    "def pair_metrics(pe:np.ndarray,pb:np.ndarray)->dict:\n",
    "    import scipy.special as sp\n",
    "    m=(pe+pb)*.5\n",
    "    js=.5*(sp.rel_entr(pe,m).sum(axis=1)+sp.rel_entr(pb,m).sum(axis=1))\n",
    "    return {\"mean_js\":float(js.mean()),\"disagreement\":float((pe.argmax(1)!=pb.argmax(1)).mean())}\n",
    "\n",
    "\n",
    "def predict_pair_dataset(model,tok,cids,rows:list[dict],batch_size:int=4)->pd.DataFrame:\n",
    "    import torch\n",
    "    model.eval();device=input_device(model)\n",
    "    out=[]\n",
    "    with torch.inference_mode():\n",
    "        for start in range(0,len(rows),batch_size):\n",
    "            batch=rows[start:start+batch_size]\n",
    "            en_logits,en_p=model_choice_logits(model,batch_view(batch,\"en\",tok,device),cids)\n",
    "            bn_logits,bn_p=model_choice_logits(model,batch_view(batch,\"bn\",tok,device),cids)\n",
    "            for i,r in enumerate(batch):\n",
    "                rec={\"item_id\":r[\"id\"],\"target\":r[\"target\"]}\n",
    "                for lang,z,p in ((\"en\",en_logits,en_p),(\"bn\",bn_logits,bn_p)):\n",
    "                    for j,c in enumerate(CHOICES):\n",
    "                        rec[f\"{lang}_logit_{c}\"]=float(z[i,cids[j]].item())\n",
    "                        rec[f\"{lang}_p_{c}\"]=float(p[i,j].item())\n",
    "                out.append(rec)\n",
    "    return pd.DataFrame(out)\n",
    "\n",
    "\n",
    "def predict_external(model,tok,cids,rows:list[dict],batch_size:int=4)->pd.DataFrame:\n",
    "    import torch\n",
    "    model.eval();device=input_device(model);out=[]\n",
    "    with torch.inference_mode():\n",
    "        for start in range(0,len(rows),batch_size):\n",
    "            chunk=rows[start:start+batch_size]\n",
    "            inputs=batch_view([{\"tokens\":r[\"tokens\"]} for r in chunk],\"tokens\",tok,device)\n",
    "            logits,p=model_choice_logits(model,inputs,cids)\n",
    "            for i,r in enumerate(chunk):\n",
    "                rec={\"item_id\":r[\"id\"],\"language\":r[\"lang\"],\"target\":r[\"target\"],\n",
    "                     \"is_translated\":r.get(\"is_translated\")}\n",
    "                for j,c in enumerate(CHOICES):\n",
    "                    rec[f\"logit_{c}\"]=float(logits[i,cids[j]].item())\n",
    "                    rec[f\"p_{c}\"]=float(p[i,j].item())\n",
    "                out.append(rec)\n",
    "    return pd.DataFrame(out)\n",
    "\n",
    "\n",
    "def score_predictions(df:pd.DataFrame,prefix:Optional[str]=None)->dict:\n",
    "    if prefix:\n",
    "        p=df[[f\"{prefix}_p_{c}\" for c in CHOICES]].to_numpy(float)\n",
    "    else:p=df[[f\"p_{c}\" for c in CHOICES]].to_numpy(float)\n",
    "    return metrics_array(p,df[\"target\"].to_numpy(int))\n",
    "\n",
    "\n",
    "def pair_scores(df:pd.DataFrame)->dict:\n",
    "    en=df[[f\"en_p_{c}\" for c in CHOICES]].to_numpy(float)\n",
    "    bn=df[[f\"bn_p_{c}\" for c in CHOICES]].to_numpy(float)\n",
    "    return {\"en\":metrics_array(en,df.target.to_numpy(int)),\n",
    "            \"bn\":metrics_array(bn,df.target.to_numpy(int)),\"pair\":pair_metrics(en,bn)}\n",
    "\n",
    "\n",
    "def fit_temperature(val_predictions:pd.DataFrame)->float:\n",
    "    from scipy.optimize import minimize_scalar\n",
    "    from scipy.special import logsumexp\n",
    "    logits=np.concatenate([val_predictions[[f\"{lang}_logit_{c}\" for c in CHOICES]].to_numpy(float)\n",
    "                           for lang in (\"en\",\"bn\")],axis=0)\n",
    "    y=np.tile(val_predictions.target.to_numpy(int),2)\n",
    "    def nll(logt):\n",
    "        z=logits/np.exp(logt)\n",
    "        return float(np.mean(logsumexp(z,axis=1)-z[np.arange(len(z)),y]))\n",
    "    res=minimize_scalar(nll,bounds=(np.log(.05),np.log(20.)),method=\"bounded\")\n",
    "    if not res.success:raise RuntimeError(\"Temperature fit failed\")\n",
    "    return float(np.exp(res.x))\n",
    "\n",
    "\n",
    "def add_temperature_predictions(df:pd.DataFrame,T:float,paired:bool)->pd.DataFrame:\n",
    "    from scipy.special import softmax\n",
    "    d=df.copy()\n",
    "    for lang in ([\"en\",\"bn\"] if paired else [None]):\n",
    "        pfx=f\"{lang}_\" if lang else \"\"\n",
    "        z=d[[f\"{pfx}logit_{c}\" for c in CHOICES]].to_numpy(float)\n",
    "        probs=softmax(z/T,axis=1)\n",
    "        for i,c in enumerate(CHOICES):d[f\"{pfx}p_{c}\"]=probs[:,i]\n",
    "    return d\n",
    "\n",
    "\n",
    "def paired_bootstrap_difference(a:pd.DataFrame,b:pd.DataFrame, metric:str=\"accuracy\",\n",
    "                                n_boot:int=10000,seed:int=2026)->dict:\n",
    "    \"\"\"Paired item bootstrap for accuracy or Brier; separate from training-seed CI.\"\"\"\n",
    "    keys=[\"item_id\",\"target\"]\n",
    "    if a[keys].duplicated().any() or b[keys].duplicated().any():raise ValueError(\"Duplicate IDs\")\n",
    "    joined=a.merge(b,on=keys,validate=\"one_to_one\",suffixes=(\"_a\",\"_b\"))\n",
    "    if len(joined)!=len(a) or len(joined)!=len(b):raise ValueError(\"Paired item mismatch\")\n",
    "    y=joined.target.to_numpy(int);arr=[]\n",
    "    for suffix in (\"a\",\"b\"):\n",
    "        p=joined[[f\"p_{c}_{suffix}\" for c in CHOICES]].to_numpy(float)\n",
    "        if metric==\"accuracy\":arr.append((p.argmax(1)==y).astype(float))\n",
    "        elif metric==\"brier\":arr.append(((p-np.eye(4)[y])**2).sum(1))\n",
    "        else:raise ValueError(metric)\n",
    "    diff=arr[0]-arr[1];rng=np.random.default_rng(seed)\n",
    "    means=np.array([diff[rng.integers(0,len(diff),len(diff))].mean() for _ in range(n_boot)])\n",
    "    return {\"n\":len(diff),\"metric\":metric,\"delta\":float(diff.mean()),\n",
    "            \"ci_low\":float(np.quantile(means,.025)),\"ci_high\":float(np.quantile(means,.975)),\"bootstrap_n\":n_boot}\n",
    "\n",
    "# --------------------- Training, crash recovery and Hub sync ---------------------\n",
    "@dataclasses.dataclass\n",
    "class Settings:\n",
    "    model: str\n",
    "    variant: str\n",
    "    seed: int\n",
    "    mode: str\n",
    "    epochs: int = 3\n",
    "    lr: float = 2e-4\n",
    "    weight_decay: float = 0.01\n",
    "    pair_batch: int = 2\n",
    "    accumulation: int = 8\n",
    "    val_batch: int = 4\n",
    "    max_steps: Optional[int] = None\n",
    "    ckpt_steps: int = 50\n",
    "    ckpt_minutes: int = 30\n",
    "    sync_hf: str = \"none\"\n",
    "    hf_repo: str = HF_REPO\n",
    "\n",
    "    @property\n",
    "    def run_id(self):\n",
    "        return f\"{self.model}/{self.variant}/seed_{self.seed}/{self.mode}\"\n",
    "\n",
    "\n",
    "def stable_epoch_batches(records:list[dict],epoch:int,seed:int,batch_size:int)->list[list[int]]:\n",
    "    \"\"\"Seed-specific reproducible length-aware buckets, independent of global RNG state.\"\"\"\n",
    "    rng=np.random.default_rng(seed + 1000003*epoch)\n",
    "    shuffled=rng.permutation(len(records)).tolist();batches=[]\n",
    "    for i in range(0,len(shuffled),128):\n",
    "        chunk=shuffled[i:i+128]\n",
    "        chunk.sort(key=lambda j:max(len(records[j][\"en\"]),len(records[j][\"bn\"])))\n",
    "        for j in range(0,len(chunk),batch_size):batches.append(chunk[j:j+batch_size])\n",
    "    # Shuffle batches while preserving each length-sorted microbatch.\n",
    "    return [batches[i] for i in rng.permutation(len(batches))]\n",
    "\n",
    "\n",
    "def adamw_scheduler(model, settings:Settings, steps_per_epoch:int):\n",
    "    import torch\n",
    "    params=[p for p in model.parameters() if p.requires_grad]\n",
    "    optimizer=torch.optim.AdamW(params,lr=settings.lr,weight_decay=settings.weight_decay,\n",
    "                                betas=(.9,.999),eps=1e-8)\n",
    "    total=steps_per_epoch*settings.epochs\n",
    "    warmup=max(1,round(total*.03))\n",
    "    def schedule(step):\n",
    "        if step < warmup:return max(step/warmup,1e-8)\n",
    "        return max(0.,(total-step)/max(1,total-warmup))\n",
    "    scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer,schedule)\n",
    "    scaler=torch.amp.GradScaler(\"cuda\",enabled=True)\n",
    "    return optimizer,scheduler,scaler,total\n",
    "\n",
    "\n",
    "def adapter_sha256(model)->str:\n",
    "    import torch\n",
    "    sha=hashlib.sha256()\n",
    "    for n,p in sorted(model.named_parameters()):\n",
    "        if p.requires_grad:\n",
    "            sha.update(n.encode())\n",
    "            sha.update(p.detach().float().cpu().contiguous().numpy().tobytes())\n",
    "    return sha.hexdigest()\n",
    "\n",
    "\n",
    "def run_fingerprint(settings:Settings,revision:str,train_count:int)->str:\n",
    "    identity=dataclasses.asdict(settings)\n",
    "    for key in (\"sync_hf\",\"hf_repo\",\"ckpt_steps\",\"ckpt_minutes\"):\n",
    "        identity.pop(key,None)  # transport / checkpoint cadence do not change the objective\n",
    "    text=json.dumps({\"settings\":identity,\"revision\":revision,\n",
    "                     \"train_count\":train_count,\"benqa_source\":BENQA_SHA},sort_keys=True)\n",
    "    return hashlib.sha256(text.encode()).hexdigest()\n",
    "\n",
    "\n",
    "def write_checkpoint(model,optimizer,scheduler,scaler,run_dir:Path,location:str,\n",
    "                     state:dict,settings:Settings)->Path:\n",
    "    import torch\n",
    "    target=run_dir/location;target.parent.mkdir(parents=True,exist_ok=True)\n",
    "    stage=run_dir/(\".staging_\"+location+\"_\"+uuid.uuid4().hex[:8]);stage.mkdir(parents=True)\n",
    "    try:\n",
    "        adapter=stage/\"adapter\";model.save_pretrained(adapter,safe_serialization=True)\n",
    "        torch.save({\"optimizer\":optimizer.state_dict(),\"scheduler\":scheduler.state_dict(),\n",
    "                    \"scaler\":scaler.state_dict(),\"rng\":rng_snapshot()},stage/\"trainer.pt\")\n",
    "        json_dump_atomic(state,stage/\"state.json\")\n",
    "        replace_dir_atomic(stage,target)\n",
    "    finally:\n",
    "        if stage.exists():shutil.rmtree(stage)\n",
    "    return target\n",
    "\n",
    "\n",
    "def save_best_adapter(model,run_dir:Path,epoch:int,metric:float)->Path:\n",
    "    stage=run_dir/(\".best_staging_\"+uuid.uuid4().hex[:8]);stage.mkdir(parents=True)\n",
    "    try:\n",
    "        model.save_pretrained(stage/\"adapter\",safe_serialization=True)\n",
    "        json_dump_atomic({\"epoch\":epoch,\"validation_mean_brier\":metric},stage/\"selection.json\")\n",
    "        replace_dir_atomic(stage,run_dir/\"best\")\n",
    "    finally:\n",
    "        if stage.exists():shutil.rmtree(stage)\n",
    "    return run_dir/\"best\"\n",
    "\n",
    "\n",
    "def restore_checkpoint_state(path:Path,optimizer,scheduler,scaler,expected_fingerprint:str)->dict:\n",
    "    import torch\n",
    "    state=json.loads((path/\"state.json\").read_text())\n",
    "    if state[\"fingerprint\"]!=expected_fingerprint:\n",
    "        raise RuntimeError(\"Checkpoint config/model/data fingerprint mismatch. Never mix different runs.\")\n",
    "    # Only load your OWN checkpoint: trainer.pt contains serialized Python state.\n",
    "    blob=torch.load(path/\"trainer.pt\",map_location=\"cpu\",weights_only=False)\n",
    "    optimizer.load_state_dict(blob[\"optimizer\"])\n",
    "    scheduler.load_state_dict(blob[\"scheduler\"])\n",
    "    scaler.load_state_dict(blob[\"scaler\"])\n",
    "    restore_rng(blob[\"rng\"])\n",
    "    return state\n",
    "\n",
    "\n",
    "def sync_run_to_hub(run_dir:Path,settings:Settings,phase:str=\"update\"):\n",
    "    if settings.sync_hf==\"none\":return\n",
    "    from huggingface_hub import HfApi\n",
    "    api=HfApi(token=hf_token())\n",
    "    if not hf_token(): raise RuntimeError(\"HF_TOKEN missing; login before Hub sync\")\n",
    "    api.upload_folder(folder_path=str(run_dir),repo_id=settings.hf_repo,repo_type=\"model\",\n",
    "                      path_in_repo=\"runs/\"+settings.run_id,\n",
    "                      commit_message=f\"TokenCal {settings.run_id}: {phase}\")\n",
    "\n",
    "\n",
    "def restore_run_from_hub(run_dir:Path,settings:Settings)->None:\n",
    "    if run_dir.exists() and (run_dir/\"latest\"/\"state.json\").exists():return\n",
    "    from huggingface_hub import snapshot_download\n",
    "    snapshot=snapshot_download(repo_id=settings.hf_repo,repo_type=\"model\",token=hf_token(),\n",
    "                               allow_patterns=[\"runs/\"+settings.run_id+\"/**\"])\n",
    "    origin=Path(snapshot)/\"runs\"/settings.run_id\n",
    "    if not (origin/\"latest\"/\"state.json\").exists():\n",
    "        raise RuntimeError(f\"No resumable remote checkpoint for {settings.run_id}\")\n",
    "    run_dir.parent.mkdir(parents=True,exist_ok=True)\n",
    "    shutil.copytree(origin,run_dir,dirs_exist_ok=True)\n",
    "\n",
    "\n",
    "def serialize_epoch_metrics(stats:dict,run_dir:Path):\n",
    "    p=run_dir/\"validation_history.json\"\n",
    "    old=json.loads(p.read_text()) if p.exists() else []\n",
    "    old=[o for o in old if o.get(\"epoch\")!=stats[\"epoch\"]]\n",
    "    old.append(stats);json_dump_atomic(old,p)\n",
    "\n",
    "\n",
    "def train_one(settings:Settings,root:Path,benqa:pd.DataFrame,source_zip:Optional[Path],\n",
    "              include_external:bool=True,resume_hub:bool=False)->Path:\n",
    "    import torch\n",
    "    if settings.variant==\"TS\":raise ValueError(\"TS is post-hoc; use temperature_eval() on trained M0\")\n",
    "    if settings.seed not in SEEDS:raise ValueError(\"Use only seeds 42,123,456\")\n",
    "    run_dir=root/\"outputs\"/\"runs\"/settings.run_id\n",
    "    if resume_hub:restore_run_from_hub(run_dir,settings)\n",
    "    run_dir.mkdir(parents=True,exist_ok=True)\n",
    "    latest=run_dir/\"latest\"\n",
    "    adapter_ck=latest/\"adapter\" if (latest/\"state.json\").exists() else None\n",
    "    model,tok,cids,mid,rev=load_backbone(settings.model,root,settings.seed,trainable=True,\n",
    "                                         adapter_checkpoint=adapter_ck)\n",
    "    device=input_device(model)\n",
    "    subset=16 if settings.mode==\"preflight\" else None\n",
    "    train_rows,val_rows,test_rows=prepare_train_records(benqa,tok,root,settings.model,rev,limit=subset)\n",
    "    # Training-only control: one fixed permutation per seed, not per minibatch.\n",
    "    if settings.mode==\"preflight\" and include_external:\n",
    "        # Full token-length audit once per model revision, before costly runs.\n",
    "        audit_path=root/\"locks\"/f\"external_token_audit_{settings.model}_{rev[:12]}.json\"\n",
    "        if not audit_path.exists():\n",
    "            external_data=[(\"breaso\",load_breaso(root,source_zip),2048),\n",
    "                           (\"milu\",load_milu(root),4096)]\n",
    "            audits={}\n",
    "            for label,external_df,ctx in external_data:\n",
    "                lengths=[]\n",
    "                for _,entry in external_df.iterrows():\n",
    "                    ids=tokenize_mcq(tok,entry,entry[\"language\"],ctx,external=True)\n",
    "                    lengths.append(len(ids))\n",
    "                audits[label]={\"items\":len(lengths),\"min_tokens\":min(lengths),\n",
    "                               \"max_tokens\":max(lengths),\"context\":ctx,\"no_truncation\":True}\n",
    "            json_dump_atomic(audits,audit_path)\n",
    "    if settings.variant==\"Shuffle\":\n",
    "        g=np.random.default_rng(settings.seed+811)\n",
    "        perm=g.permutation([r[\"weight\"] for r in train_rows])\n",
    "        for r,w in zip(train_rows,perm):r[\"shuffled_weight\"]=float(w)\n",
    "    fprint=run_fingerprint(settings,rev,len(train_rows))\n",
    "    init_sha=adapter_sha256(model) if adapter_ck is None else None\n",
    "    batches_per_epoch=len(stable_epoch_batches(train_rows,0,settings.seed,settings.pair_batch))\n",
    "    steps_per_epoch=math.ceil(batches_per_epoch/settings.accumulation)\n",
    "    optimizer,scheduler,scaler,total_steps=adamw_scheduler(model,settings,steps_per_epoch)\n",
    "    first_epoch=0;first_batch=0;global_step=0;best=float(\"inf\")\n",
    "    if adapter_ck is not None:\n",
    "        state=restore_checkpoint_state(latest,optimizer,scheduler,scaler,fprint)\n",
    "        first_epoch=int(state[\"next_epoch\"]);first_batch=int(state[\"next_batch\"])\n",
    "        global_step=int(state[\"global_step\"]);best=float(state[\"best_validation_brier\"])\n",
    "    else:\n",
    "        json_dump_atomic({\"model\":mid,\"revision\":rev,\"choice_token_ids\":cids,\n",
    "                          \"initial_adapter_sha256\":init_sha,\"settings\":dataclasses.asdict(settings),\n",
    "                          \"fingerprint\":fprint,\"status\":\"started\"},run_dir/\"manifest.json\")\n",
    "    torch.cuda.reset_peak_memory_stats()\n",
    "    log_path=run_dir/\"step_log.csv\";t0=time.monotonic();last_save=t0\n",
    "    interrupted=False\n",
    "    for epoch in range(first_epoch,settings.epochs):\n",
    "        model.train();optimizer.zero_grad(set_to_none=True)\n",
    "        all_batches=stable_epoch_batches(train_rows,epoch,settings.seed,settings.pair_batch)\n",
    "        begin=first_batch if epoch==first_epoch else 0\n",
    "        accumulated=0\n",
    "        for bi in range(begin,len(all_batches)):\n",
    "            batch=[train_rows[i] for i in all_batches[bi]]\n",
    "            with torch.autocast(device_type=\"cuda\",dtype=torch.float16):\n",
    "                loss,parts=forward_pair(model,tok,cids,batch,settings.variant,device)\n",
    "            if not torch.isfinite(loss):raise FloatingPointError(f\"Nonfinite loss: {settings.run_id} e{epoch} b{bi}\")\n",
    "            # Scale by actual number of accumulation microbatches, including partial final group.\n",
    "            left=(bi//settings.accumulation)*settings.accumulation\n",
    "            group_len=min(settings.accumulation,len(all_batches)-left)\n",
    "            scaler.scale(loss/group_len).backward();accumulated+=1\n",
    "            if (bi+1)%settings.accumulation==0 or bi+1==len(all_batches):\n",
    "                scaler.unscale_(optimizer)\n",
    "                norm=torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad],1.)\n",
    "                if not torch.isfinite(norm):raise FloatingPointError(\"NaN/Inf in gradients\")\n",
    "                before=scaler.get_scale();scaler.step(optimizer);scaler.update()\n",
    "                if scaler.get_scale()>=before: scheduler.step()\n",
    "                optimizer.zero_grad(set_to_none=True)\n",
    "                global_step+=1;accumulated=0\n",
    "                row={\"epoch\":epoch+1,\"step\":global_step,\"loss\":float(loss.item()),\n",
    "                     \"grad_norm\":float(norm.item()),\"lr\":optimizer.param_groups[0][\"lr\"],\n",
    "                     \"time_sec\":round(time.monotonic()-t0,1)}\n",
    "                with open(log_path,\"a\",newline=\"\") as f:\n",
    "                    w=csv.DictWriter(f,fieldnames=list(row))\n",
    "                    if f.tell()==0:w.writeheader()\n",
    "                    w.writerow(row)\n",
    "                if global_step%10==0:print(settings.run_id,\"step\",global_step,\"loss\",round(row[\"loss\"],4),flush=True)\n",
    "                if settings.mode==\"full\" and bi+1<len(all_batches) and (\n",
    "                      global_step%settings.ckpt_steps==0 or time.monotonic()-last_save>settings.ckpt_minutes*60):\n",
    "                    state={\"fingerprint\":fprint,\"next_epoch\":epoch,\"next_batch\":bi+1,\n",
    "                           \"global_step\":global_step,\"best_validation_brier\":best}\n",
    "                    write_checkpoint(model,optimizer,scheduler,scaler,run_dir,\"latest\",state,settings)\n",
    "                    if settings.sync_hf==\"checkpoints\":sync_run_to_hub(run_dir,settings,\"checkpoint\")\n",
    "                    last_save=time.monotonic()\n",
    "                if settings.max_steps is not None and global_step>=settings.max_steps:\n",
    "                    interrupted=True;break\n",
    "        if interrupted and settings.mode!=\"full\":\n",
    "            # On preflight / benchmark, small step count is intentional and NOT a final experiment.\n",
    "            state={\"fingerprint\":fprint,\"next_epoch\":epoch,\"next_batch\":min(bi+1,len(all_batches)),\n",
    "                   \"global_step\":global_step,\"best_validation_brier\":best}\n",
    "            write_checkpoint(model,optimizer,scheduler,scaler,run_dir,\"latest\",state,settings)\n",
    "            break\n",
    "        if interrupted:\n",
    "            state={\"fingerprint\":fprint,\"next_epoch\":epoch,\"next_batch\":min(bi+1,len(all_batches)),\n",
    "                   \"global_step\":global_step,\"best_validation_brier\":best}\n",
    "            write_checkpoint(model,optimizer,scheduler,scaler,run_dir,\"latest\",state,settings)\n",
    "            break\n",
    "        vpred=predict_pair_dataset(model,tok,cids,val_rows,settings.val_batch)\n",
    "        vm=pair_scores(vpred)\n",
    "        v_brier=(vm[\"en\"][\"brier\"]+vm[\"bn\"][\"brier\"])*.5\n",
    "        info={\"epoch\":epoch+1,\"global_step\":global_step,\"validation_mean_brier\":v_brier,\n",
    "              \"english_accuracy\":vm[\"en\"][\"accuracy\"],\"bengali_accuracy\":vm[\"bn\"][\"accuracy\"]}\n",
    "        serialize_epoch_metrics(info,run_dir)\n",
    "        if v_brier < best:\n",
    "            best=v_brier;save_best_adapter(model,run_dir,epoch+1,best)\n",
    "        state={\"fingerprint\":fprint,\"next_epoch\":epoch+1,\"next_batch\":0,\n",
    "               \"global_step\":global_step,\"best_validation_brier\":best}\n",
    "        write_checkpoint(model,optimizer,scheduler,scaler,run_dir,\"latest\",state,settings)\n",
    "        if settings.sync_hf==\"checkpoints\":sync_run_to_hub(run_dir,settings,f\"epoch {epoch+1}\")\n",
    "    peak=max(torch.cuda.max_memory_allocated(d) for d in range(torch.cuda.device_count()))/2**30\n",
    "    status={\"status\":\"short_run\" if settings.mode!=\"full\" else \"trained\",\n",
    "            \"train_steps\":global_step,\"elapsed_seconds\":round(time.monotonic()-t0,2),\n",
    "            \"peak_gpu_gib\":round(peak,3),\"best_validation_brier\":best}\n",
    "    json_dump_atomic(status,run_dir/\"training_status.json\")\n",
    "    del model,optimizer,scheduler,scaler;gc.collect();torch.cuda.empty_cache()\n",
    "    if settings.mode==\"preflight\":\n",
    "        # Genuine GPU reload, optimizer/scheduler state restore and one forward pass.\n",
    "        ck=run_dir/\"latest\"\n",
    "        probe,tok2,cids2,_,_=load_backbone(settings.model,root,settings.seed,True,ck/\"adapter\")\n",
    "        opt,sch,scale,_=adamw_scheduler(probe,settings,steps_per_epoch)\n",
    "        restored=restore_checkpoint_state(ck,opt,sch,scale,fprint)\n",
    "        # Actual resumed optimizer UPDATE, not merely loading and running inference.\n",
    "        probe.train();opt.zero_grad(set_to_none=True)\n",
    "        with torch.autocast(device_type=\"cuda\",dtype=torch.float16):\n",
    "            recovered_loss,_=forward_pair(probe,tok2,cids2,train_rows[:2],settings.variant,input_device(probe))\n",
    "        if not torch.isfinite(recovered_loss):raise RuntimeError(\"Nonfinite resumed loss\")\n",
    "        scale.scale(recovered_loss).backward()\n",
    "        scale.unscale_(opt)\n",
    "        if not torch.isfinite(torch.nn.utils.clip_grad_norm_([p for p in probe.parameters() if p.requires_grad],1.)):\n",
    "            raise RuntimeError(\"Nonfinite resumed gradients\")\n",
    "        scale.step(opt);scale.update();sch.step()\n",
    "        opt.zero_grad(set_to_none=True)\n",
    "        with torch.no_grad():\n",
    "            probe.eval()\n",
    "            pp=predict_pair_dataset(probe,tok2,cids2,val_rows[:2],1)\n",
    "        if pp.isnull().any().any() or len(pp)!=2:raise RuntimeError(\"Resume preflight failed\")\n",
    "        if include_external:\n",
    "            ood_sample=load_breaso(root,source_zip).head(2)\n",
    "            milu_data=load_milu(root)\n",
    "            milu_sample=pd.concat([x.head(1) for _,x in milu_data.groupby(\"language\")],ignore_index=True)\n",
    "            for ext,ctx in ((ood_sample,2048),(milu_sample,4096)):\n",
    "                rows=format_external_records(ext,tok2,ctx)\n",
    "                pred=predict_external(probe,tok2,cids2,rows,1)\n",
    "                if len(pred)!=len(rows) or pred.isnull().loc[:,[f\"p_{c}\" for c in CHOICES]].any().any():\n",
    "                    raise RuntimeError(\"External evaluation preflight failed\")\n",
    "        json_dump_atomic({\"status\":\"passed\",\"resumed_step\":restored[\"global_step\"],\n",
    "                          \"two_pair_prediction\":\"passed\",\"external_evaluation\":bool(include_external)},\n",
    "                         run_dir/\"preflight.json\")\n",
    "        del probe,opt,sch,scale;gc.collect();torch.cuda.empty_cache()\n",
    "    if settings.mode==\"full\" and not interrupted:\n",
    "        best_path=run_dir/\"best\"/\"adapter\"\n",
    "        if not best_path.exists():raise RuntimeError(\"Full training ended without selected checkpoint\")\n",
    "        frozen,tok,cids,_,_=load_backbone(settings.model,root,settings.seed,False,best_path)\n",
    "        val_pred=predict_pair_dataset(frozen,tok,cids,val_rows,settings.val_batch)\n",
    "        val_pred.to_csv(run_dir/\"validation_predictions.csv\",index=False)\n",
    "        test_pred=predict_pair_dataset(frozen,tok,cids,test_rows,settings.val_batch)\n",
    "        test_pred.to_csv(run_dir/\"benqa_predictions.csv\",index=False)\n",
    "        summary={\"benqa\":pair_scores(test_pred)}\n",
    "        if include_external:\n",
    "            breaso=load_breaso(root,source_zip)\n",
    "            breaso_rows=format_external_records(breaso,tok,2048)\n",
    "            breaso_pred=predict_external(frozen,tok,cids,breaso_rows,settings.val_batch)\n",
    "            breaso_pred.to_csv(run_dir/\"breaso_predictions.csv\",index=False)\n",
    "            summary[\"breaso\"]=score_predictions(breaso_pred)\n",
    "            milu=load_milu(root)\n",
    "            milu_rows=format_external_records(milu,tok,4096)\n",
    "            milu_pred=predict_external(frozen,tok,cids,milu_rows,settings.val_batch)\n",
    "            milu_pred.to_csv(run_dir/\"milu_predictions.csv\",index=False)\n",
    "            summary[\"milu\"]={lang:score_predictions(d) for lang,d in milu_pred.groupby(\"language\")}\n",
    "            summary[\"milu\"][\"unseen_pooled\"]=score_predictions(milu_pred[milu_pred.language.isin([\"gu\",\"ml\",\"te\"])])\n",
    "            if \"is_translated\" in milu_pred:\n",
    "                summary[\"milu\"][\"native_translated\"]={str(v):score_predictions(d) for v,d in milu_pred.groupby(\"is_translated\")}\n",
    "        json_dump_atomic(summary,run_dir/\"results_summary.json\")\n",
    "        status[\"status\"]=\"completed\";json_dump_atomic(status,run_dir/\"training_status.json\")\n",
    "        del frozen;gc.collect();torch.cuda.empty_cache()\n",
    "    if settings.sync_hf!=\"none\":sync_run_to_hub(run_dir,settings,\"completed or short run\")\n",
    "    return run_dir\n",
    "\n",
    "# --------------------- Temperature scaling, summaries, orchestration ---------------------\n",
    "def temperature_eval(model_key:str,seed:int,root:Path,repo:str=HF_REPO,sync_hf:str=\"none\") ->Path:\n",
    "    src=root/\"outputs\"/\"runs\"/model_key/\"M0\"/f\"seed_{seed}\"/\"full\"\n",
    "    if not (src/\"validation_predictions.csv\").exists():\n",
    "        raise RuntimeError(f\"Temperature scaling requires completed M0 validation predictions: {src}\")\n",
    "    val=pd.read_csv(src/\"validation_predictions.csv\")\n",
    "    T=fit_temperature(val)\n",
    "    dest=root/\"outputs\"/\"runs\"/model_key/\"TS\"/f\"seed_{seed}\"/\"full\"\n",
    "    dest.mkdir(parents=True,exist_ok=True)\n",
    "    summary={\"temperature\":T,\"selection\":\"validation NLL only\"}\n",
    "    pairs={\"benqa_predictions.csv\":True,\"breaso_predictions.csv\":False,\"milu_predictions.csv\":False}\n",
    "    for name,paired in pairs.items():\n",
    "        if not (src/name).exists():continue\n",
    "        frame=add_temperature_predictions(pd.read_csv(src/name),T,paired)\n",
    "        frame.to_csv(dest/name,index=False)\n",
    "        if paired:summary[\"benqa\"]=pair_scores(frame)\n",
    "        elif name.startswith(\"breaso\"):summary[\"breaso\"]=score_predictions(frame)\n",
    "        else:\n",
    "            summary[\"milu\"]={k:score_predictions(g) for k,g in frame.groupby(\"language\")}\n",
    "            summary[\"milu\"][\"unseen_pooled\"]=score_predictions(frame[frame.language.isin([\"gu\",\"ml\",\"te\"])])\n",
    "    json_dump_atomic(summary,dest/\"results_summary.json\")\n",
    "    json_dump_atomic({\"status\":\"completed\",\"model\":model_key,\"seed\":seed,\"method\":\"TS\",\"source\":\"M0\"},\n",
    "                     dest/\"training_status.json\")\n",
    "    if sync_hf!=\"none\":\n",
    "        settings=Settings(model_key,\"TS\",seed,\"full\",sync_hf=sync_hf,hf_repo=repo)\n",
    "        sync_run_to_hub(dest,settings,\"posthoc temperature scaling\")\n",
    "    return dest\n",
    "\n",
    "\n",
    "def aggregate_results(root:Path,bootstrap_n:int=10000)->Path:\n",
    "    main=root/\"outputs\"/\"runs\";out=root/\"outputs\"/\"summary\"\n",
    "    out.mkdir(parents=True,exist_ok=True);rows=[]\n",
    "    for file in sorted(main.glob(\"*/*/seed_*/full/results_summary.json\")):\n",
    "        d=json.loads(file.read_text());bits=file.relative_to(main).parts\n",
    "        model,variant,seed=bits[:3]\n",
    "        def add(dataset,language,m):\n",
    "            rows.append({\"model\":model,\"method\":variant,\"seed\":int(seed.replace(\"seed_\",\"\")),\n",
    "                         \"dataset\":dataset,\"language\":language,**m})\n",
    "        if \"benqa\" in d:\n",
    "            for lang in (\"en\",\"bn\"):add(\"BEnQA\",lang,d[\"benqa\"][lang])\n",
    "        if \"breaso\" in d:add(\"B-REASO\",\"bn\",d[\"breaso\"])\n",
    "        if \"milu\" in d:\n",
    "            for lang,m in d[\"milu\"].items():\n",
    "                if lang!=\"native_translated\":add(\"MILU\",lang,m)\n",
    "    if not rows:return out\n",
    "    frame=pd.DataFrame(rows).sort_values([\"model\",\"method\",\"seed\",\"dataset\",\"language\"])\n",
    "    frame.to_csv(out/\"all_results.csv\",index=False)\n",
    "    metrics=[\"accuracy\",\"brier\",\"nll\",\"ece\",\"aurc\"]\n",
    "    agg=frame.groupby([\"model\",\"method\",\"dataset\",\"language\"])[metrics].agg([\"mean\",\"std\",\"count\"])\n",
    "    agg.columns=[f\"{m}_{stat}\" for m,stat in agg.columns]\n",
    "    agg.reset_index().to_csv(out/\"multi_seed_summary.csv\",index=False)\n",
    "    comparisons=[]\n",
    "    # Paired BEnQA EN/BN and OOD comparisons per seed: M3 vs M2.\n",
    "    for model in BASE_REVISIONS:\n",
    "        for seed in SEEDS:\n",
    "            a=main/model/\"M3\"/f\"seed_{seed}\"/\"full\"\n",
    "            b=main/model/\"M2\"/f\"seed_{seed}\"/\"full\"\n",
    "            for csv_name in (\"benqa_predictions.csv\",\"breaso_predictions.csv\",\"milu_predictions.csv\"):\n",
    "                if not (a/csv_name).exists() or not (b/csv_name).exists():continue\n",
    "                aa=pd.read_csv(a/csv_name);bb=pd.read_csv(b/csv_name)\n",
    "                if csv_name.startswith(\"benqa\"):\n",
    "                    for lang in (\"en\",\"bn\"):\n",
    "                        def normalize(d):\n",
    "                            return d.rename(columns={f\"{lang}_p_{c}\":f\"p_{c}\" for c in CHOICES})\n",
    "                        ci=paired_bootstrap_difference(normalize(aa),normalize(bb),n_boot=bootstrap_n)\n",
    "                        comparisons.append({\"backbone\":model,\"seed\":seed,\"dataset\":\"BEnQA\",\"language\":lang,**ci})\n",
    "                else:\n",
    "                    if csv_name.startswith(\"milu\"):\n",
    "                        aa=aa[aa.language.isin([\"gu\",\"ml\",\"te\"])];bb=bb[bb.language.isin([\"gu\",\"ml\",\"te\"])]\n",
    "                    ci=paired_bootstrap_difference(aa,bb,n_boot=bootstrap_n)\n",
    "                    comparisons.append({\"backbone\":model,\"seed\":seed,\n",
    "                                        \"dataset\":\"MILU\" if csv_name.startswith(\"milu\") else \"B-REASO\",\n",
    "                                        \"language\":\"unseen_pooled\" if csv_name.startswith(\"milu\") else \"bn\",**ci})\n",
    "    if comparisons:pd.DataFrame(comparisons).to_csv(out/\"m3_vs_m2_bootstrap_by_seed.csv\",index=False)\n",
    "    # Publication-grade figure uses available final data only (never simulated).\n",
    "    try:\n",
    "        import matplotlib.pyplot as plt\n",
    "        subset=frame[(frame.dataset==\"MILU\")&(frame.language==\"unseen_pooled\")]\n",
    "        if not subset.empty:\n",
    "            fig,ax=plt.subplots(figsize=(8,5))\n",
    "            for (model,method),g in subset.groupby([\"model\",\"method\"]):\n",
    "                ax.scatter(g.accuracy.mean(),g.brier.mean(),label=f\"{model} / {method}\")\n",
    "            ax.set(xlabel=\"Accuracy (higher is better)\",ylabel=\"Brier (lower is better)\",\n",
    "                   title=\"TokenCal: unseen-language accuracy and probability quality\")\n",
    "            ax.legend(fontsize=7,bbox_to_anchor=(1.01,1),loc=\"upper left\")\n",
    "            fig.tight_layout();fig.savefig(out/\"accuracy_brier.png\",dpi=180,bbox_inches=\"tight\");plt.close(fig)\n",
    "    except ImportError:pass\n",
    "    return out\n",
    "\n",
    "\n",
    "def source_sync(root:Path,repo:str):\n",
    "    \"\"\"Upload ONLY code/documentation; never third-party source datasets or tokens.\"\"\"\n",
    "    from huggingface_hub import HfApi\n",
    "    api=HfApi(token=hf_token());\n",
    "    if not hf_token():raise RuntimeError(\"HF_TOKEN missing, cannot sync source\")\n",
    "    for filename in (\"tokencal_q1.py\",\"TokenCal_Q1_Notebook.ipynb\",\"README.md\",\"requirements.txt\",\"test_tokencal_q1.py\"):\n",
    "        source=root/filename\n",
    "        if source.exists():\n",
    "            api.upload_file(path_or_fileobj=str(source),path_in_repo=filename,\n",
    "                            repo_id=repo,repo_type=\"model\",commit_message=f\"Upload TokenCal source: {filename}\")\n",
    "\n",
    "\n",
    "def valid_combinations(models:list[str],variants:list[str],mode:str):\n",
    "    for model in models:\n",
    "        for method in variants:\n",
    "            if method in {\"Focal\",\"RDrop\"} and model not in (\"qwen\",\"gemma\"):continue\n",
    "            if method in {\"Const\",\"Shuffle\"} and model!=\"qwen\":continue\n",
    "            if method==\"TS\" and mode!=\"full\":continue\n",
    "            yield model,method\n",
    "\n",
    "\n",
    "def main(argv: Optional[list[str]]=None)->None:\n",
    "    parser=argparse.ArgumentParser(description=\"TokenCal Q1 reproducible research runner\")\n",
    "    parser.add_argument(\"--mode\",choices=[\"preflight\",\"benchmark\",\"full\"],default=\"preflight\")\n",
    "    parser.add_argument(\"--models\",default=\"qwen,gemma,llama\")\n",
    "    parser.add_argument(\"--variants\",default=\"M0,M1,M2,M3,Focal,RDrop,Const,Shuffle,TS\")\n",
    "    parser.add_argument(\"--seeds\",default=\"42,123,456\")\n",
    "    parser.add_argument(\"--root\",type=Path,default=PROJECT)\n",
    "    parser.add_argument(\"--source-zip\",type=Path,default=None)\n",
    "    parser.add_argument(\"--sync-hf\",choices=[\"none\",\"final\",\"checkpoints\"],default=\"none\")\n",
    "    parser.add_argument(\"--hf-repo\",default=HF_REPO)\n",
    "    parser.add_argument(\"--resume-hub\",action=\"store_true\")\n",
    "    parser.add_argument(\"--no-external-eval\",action=\"store_true\")\n",
    "    parser.add_argument(\"--benchmark-steps\",type=int,default=50)\n",
    "    parser.add_argument(\"--preflight-steps\",type=int,default=2)\n",
    "    parser.add_argument(\"--bootstrap-resamples\",type=int,default=10000)\n",
    "    parser.add_argument(\"--source-upload\",action=\"store_true\")\n",
    "    parser.add_argument(\"--force\",action=\"store_true\",help=\"Rerun even if status completed (will resume where possible)\")\n",
    "    args=parser.parse_args(argv)\n",
    "    root=args.root.expanduser().resolve();root.mkdir(parents=True,exist_ok=True)\n",
    "    model_list=[x.strip().lower() for x in args.models.split(\",\") if x.strip()]\n",
    "    variant_list=[x.strip() for x in args.variants.split(\",\") if x.strip()]\n",
    "    seeds=[int(x) for x in args.seeds.split(\",\")]\n",
    "    if any(x not in BASE_REVISIONS for x in model_list):raise ValueError(\"Unknown backbone\")\n",
    "    if any(x not in VARIANTS for x in variant_list):raise ValueError(\"Unknown method\")\n",
    "    if any(x not in SEEDS for x in seeds):raise ValueError(\"Allowed seeds are 42,123,456 only\")\n",
    "    if args.mode!=\"full\":seeds=[seeds[0]]\n",
    "    # Source upload (after notebook writes source into root) is explicit.\n",
    "    if args.source_upload:source_sync(Path(__file__).resolve().parent,args.hf_repo)\n",
    "    print(f\"MODE={args.mode} MODELS={model_list} VARIANTS={variant_list} SEEDS={seeds}\",flush=True)\n",
    "    print(f\"OUTPUT={root} HF_REPO={args.hf_repo} SYNC={args.sync_hf}\",flush=True)\n",
    "    source_zip=args.source_zip.expanduser().resolve() if args.source_zip else None\n",
    "    benqa=load_benqa(root,source_zip)\n",
    "    # Dataset access must be checked before lengthy training; external datasets are not used to tune hyperparameters.\n",
    "    if args.mode in (\"full\",\"preflight\") and not args.no_external_eval:\n",
    "        print(\"Checking external dataset access before running...\",flush=True)\n",
    "        _=load_breaso(root,source_zip)\n",
    "        _=load_milu(root)  # Gated access and count check. Does not train on MILU.\n",
    "    combinations=list(valid_combinations(model_list,variant_list,args.mode))\n",
    "    if args.mode==\"full\":\n",
    "        # Ensure M0 runs before post-hoc TS by stable ordering.\n",
    "        order={k:i for i,k in enumerate(VARIANTS)}\n",
    "        combinations.sort(key=lambda pair:(model_list.index(pair[0]),order[pair[1]]))\n",
    "    print(f\"Planned {'training' if args.mode=='full' else args.mode} combinations: {len(combinations)}\",flush=True)\n",
    "    for model,variant in combinations:\n",
    "        for seed in seeds:\n",
    "            if variant==\"TS\":\n",
    "                temperature_eval(model,seed,root,args.hf_repo,args.sync_hf);continue\n",
    "            settings=Settings(model,variant,seed,args.mode,sync_hf=args.sync_hf,hf_repo=args.hf_repo,\n",
    "                              max_steps=(args.preflight_steps if args.mode==\"preflight\" else\n",
    "                                         args.benchmark_steps if args.mode==\"benchmark\" else None))\n",
    "            location=root/\"outputs\"/\"runs\"/settings.run_id\n",
    "            status_path=location/\"training_status.json\"\n",
    "            if status_path.exists() and not args.force:\n",
    "                previous=json.loads(status_path.read_text())\n",
    "                if previous.get(\"status\")==\"completed\" and args.mode==\"full\":\n",
    "                    print(\"SKIP completed\",settings.run_id,flush=True);continue\n",
    "                if args.mode==\"preflight\" and (location/\"preflight.json\").exists():\n",
    "                    print(\"SKIP preflight passed\",settings.run_id,flush=True);continue\n",
    "            print(\"START\",settings.run_id,flush=True)\n",
    "            try:\n",
    "                dest=train_one(settings,root,benqa,source_zip,\n",
    "                               include_external=(args.mode in (\"full\",\"preflight\") and not args.no_external_eval),\n",
    "                               resume_hub=args.resume_hub)\n",
    "                print(\"DONE\",dest,flush=True)\n",
    "            except Exception as e:\n",
    "                fail={\"run\":settings.run_id,\"failure\":repr(e),\"traceback\":traceback.format_exc()}\n",
    "                json_dump_atomic(fail,root/\"outputs\"/\"failures\"/(model+\"_\"+variant+\"_\"+str(seed)+\".json\"))\n",
    "                print(\"FAILED\",settings.run_id,repr(e),flush=True)\n",
    "                raise  # fail fast; do not silently continue after a broken scientific run\n",
    "    if args.mode==\"full\":\n",
    "        out=aggregate_results(root,bootstrap_n=args.bootstrap_resamples)\n",
    "        print(\"Final comparison tables:\",out,flush=True)\n",
    "        if args.sync_hf!=\"none\":\n",
    "            from huggingface_hub import HfApi\n",
    "            HfApi(token=hf_token()).upload_folder(folder_path=str(out),repo_id=args.hf_repo,\n",
    "                                                    repo_type=\"model\",path_in_repo=\"summary\",\n",
    "                                                    commit_message=\"TokenCal multi-seed tables\")\n",
    "    print(\"All selected runs finished for mode\",args.mode,flush=True)\n",
    "\n",
    "\n",
    "if __name__==\"__main__\":main()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22bca5bf",
   "metadata": {},
   "source": [
    "## 3 — Generate automated tests and run them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beba286b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile test_tokencal_q1.py\n",
    "\"\"\"CPU-only mathematical and integrity tests; real CUDA tests run in preflight mode.\"\"\"\n",
    "import math\n",
    "import zipfile\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pytest\n",
    "import torch\n",
    "\n",
    "import tokencal_q1 as t\n",
    "\n",
    "\n",
    "def toy_logits():\n",
    "    torch.manual_seed(1)\n",
    "    en=torch.randn(5,8,requires_grad=True)\n",
    "    bn=torch.randn(5,8,requires_grad=True)\n",
    "    cids=[0,2,4,6]\n",
    "    return en,bn,en[:,cids].softmax(-1),bn[:,cids].softmax(-1),cids\n",
    "\n",
    "\n",
    "def test_loss_equations_and_gradients():\n",
    "    en,bn,pe,pb,cids=toy_logits()\n",
    "    idx=torch.tensor([0,1,2,3,0]);tok=torch.tensor([cids[i] for i in idx])\n",
    "    w=torch.tensor([.1,.25,.4,.8,.9])\n",
    "    import torch.nn.functional as F\n",
    "    task=(F.cross_entropy(en,tok)+F.cross_entropy(bn,tok))/2\n",
    "    cal=(t.multiclass_brier(pe,idx).mean()+t.multiclass_brier(pb,idx).mean())/4\n",
    "    js=t.js_per_pair(pe,pb)\n",
    "    pair=js.mean()/math.log(2)\n",
    "    frag=(js*w).sum()/w.sum()/math.log(2)\n",
    "    reference={\"M0\":task,\"M1\":task+cal,\"M2\":task+cal+pair,\n",
    "               \"M3\":task+cal+pair+frag,\"Const\":task+cal+2*pair}\n",
    "    for variant,expected in reference.items():\n",
    "        actual,parts=t.objectives(en,bn,pe,pb,(idx,tok),w,variant)\n",
    "        assert torch.allclose(actual,expected,atol=2e-6),(variant,actual.item(),expected.item())\n",
    "        assert torch.isfinite(actual)\n",
    "    total,_=t.objectives(en,bn,pe,pb,(idx,tok),w,\"M3\")\n",
    "    total.backward()\n",
    "    assert torch.isfinite(en.grad).all() and torch.isfinite(bn.grad).all()\n",
    "    assert en.grad.abs().sum()>0\n",
    "\n",
    "\n",
    "def test_zero_fragmentation_weights():\n",
    "    en,bn,pe,pb,cids=toy_logits();idx=torch.tensor([0,1,2,3,0]);tok=torch.tensor([cids[i] for i in idx])\n",
    "    total,parts=t.objectives(en,bn,pe,pb,(idx,tok),torch.zeros(5),\"M3\")\n",
    "    assert parts[\"frag\"].item()==0\n",
    "    assert math.isfinite(total.item())\n",
    "\n",
    "\n",
    "def test_external_loss_variants():\n",
    "    en,bn,pe,pb,cids=toy_logits();idx=torch.tensor([0,1,2,3,0]);tok=torch.tensor([cids[i] for i in idx]);w=torch.rand(5)\n",
    "    focal,_=t.objectives(en,bn,pe,pb,(idx,tok),w,\"Focal\")\n",
    "    assert math.isfinite(focal.item())\n",
    "    rdrop,_=t.objectives(en,bn,pe,pb,(idx,tok),w,\"RDrop\",rdrop_second=(en,bn,pe,pb))\n",
    "    assert math.isfinite(rdrop.item())\n",
    "    assert torch.allclose(rdrop,(torch.nn.functional.cross_entropy(en,tok)+torch.nn.functional.cross_entropy(bn,tok))/2,atol=1e-5)\n",
    "\n",
    "\n",
    "def test_choice_token_ids_and_no_truncation():\n",
    "    class Tok:\n",
    "        chat_template=\"test\"\n",
    "        def encode(self,txt,add_special_tokens=False):return {\"A\":[5],\"B\":[7],\"C\":[9],\"D\":[11]}.get(txt,[1,2,3,4])\n",
    "        def apply_chat_template(self,*args,**kwargs):return \"prompt\"\n",
    "    tok=Tok()\n",
    "    assert t.choice_token_ids(tok)==[5,7,9,11]\n",
    "    row={\"question\":\"Where?\",\"option_A\":\"a\",\"option_B\":\"b\",\"option_C\":\"c\",\"option_D\":\"d\"}\n",
    "    assert t.tokenize_mcq(tok,row,\"en\",5,external=True)==[1,2,3,4]\n",
    "    with pytest.raises(RuntimeError,match=\"no truncation\"):\n",
    "        t.tokenize_mcq(tok,row,\"en\",4,external=True)\n",
    "\n",
    "\n",
    "def test_metrics_and_calibration():\n",
    "    p=np.array([[.8,.1,.05,.05],[.1,.7,.1,.1],[.1,.1,.7,.1],[.1,.1,.1,.7]])\n",
    "    y=np.array([0,1,2,3]);m=t.metrics_array(p,y)\n",
    "    assert m[\"accuracy\"]==1.\n",
    "    assert np.isclose(m[\"mean_confidence\"],.725)\n",
    "    assert m[\"brier\"]>0\n",
    "    assert t.pair_metrics(p,p)[\"mean_js\"]==0\n",
    "    frame=pd.DataFrame({\"target\":y,**{f\"en_logit_{c}\":np.log(p[:,i]) for i,c in enumerate(t.CHOICES)},\n",
    "                        **{f\"bn_logit_{c}\":np.log(p[:,i]) for i,c in enumerate(t.CHOICES)}})\n",
    "    temp=t.fit_temperature(frame)\n",
    "    assert temp>0\n",
    "    out=t.add_temperature_predictions(frame,temp,paired=True)\n",
    "    assert np.allclose(out[[f\"en_p_{c}\" for c in t.CHOICES]].sum(axis=1),1.)\n",
    "\n",
    "\n",
    "def test_paired_bootstrap():\n",
    "    p=np.array([[.7,.1,.1,.1],[.1,.7,.1,.1],[.1,.1,.7,.1],[.1,.1,.1,.7]])\n",
    "    df=pd.DataFrame({\"item_id\":[\"a\",\"b\",\"c\",\"d\"],\"target\":[0,1,2,3],\n",
    "                     **{f\"p_{c}\":p[:,i] for i,c in enumerate(t.CHOICES)}})\n",
    "    ci=t.paired_bootstrap_difference(df,df,n_boot=200)\n",
    "    assert ci[\"delta\"]==ci[\"ci_low\"]==ci[\"ci_high\"]==0\n",
    "    with pytest.raises(ValueError,match=\"Paired item mismatch\"):\n",
    "        t.paired_bootstrap_difference(df.iloc[:3],df,n_boot=20)\n",
    "\n",
    "\n",
    "def test_epoch_batches_reproducible_no_losses():\n",
    "    records=[{\"en\":[1]*i,\"bn\":[2]*(i+1)} for i in range(1,45)]\n",
    "    a=t.stable_epoch_batches(records,0,42,2)\n",
    "    b=t.stable_epoch_batches(records,0,42,2)\n",
    "    assert a==b\n",
    "    assert sorted(sum(a,[]))==list(range(44))\n",
    "    assert a!=t.stable_epoch_batches(records,1,42,2)\n",
    "\n",
    "\n",
    "def test_source_zip_retrieval(tmp_path):\n",
    "    archive=tmp_path/\"results.zip\"\n",
    "    with zipfile.ZipFile(archive,\"w\") as z:z.writestr(\"tokencal-results/datasets/example.csv\",\"a,b\\n1,2\\n\")\n",
    "    got=t.source_reference(\"example.csv\",tmp_path,archive)\n",
    "    assert got.read_text()==\"a,b\\n1,2\\n\"\n",
    "\n",
    "\n",
    "def test_checkpoint_fingerprint():\n",
    "    s=t.Settings(\"qwen\",\"M3\",42,\"full\")\n",
    "    assert s.run_id==\"qwen/M3/seed_42/full\"\n",
    "    assert t.run_fingerprint(s,\"revision\",3553)==t.run_fingerprint(s,\"revision\",3553)\n",
    "    assert t.run_fingerprint(s,\"revision\",3553)!=t.run_fingerprint(s,\"revision\",16)\n",
    "\n",
    "\n",
    "def test_planned_runs_exactly_54():\n",
    "    combos=list(t.valid_combinations([\"qwen\",\"gemma\",\"llama\"],list(t.VARIANTS),\"full\"))\n",
    "    trains=[(m,v) for m,v in combos if v!=\"TS\"]\n",
    "    assert len(trains)==18  # qwen:8, gemma:6, llama:4\n",
    "    assert len(trains)*len(t.SEEDS)==54\n",
    "    assert len([x for x in combos if x[1]==\"TS\"])==3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a97a1c4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile README.md\n",
    "---\n",
    "library_name: peft\n",
    "tags:\n",
    "- tokenization\n",
    "- bengali\n",
    "- calibration\n",
    "- qlora\n",
    "- cross-lingual\n",
    "- research\n",
    "license: other\n",
    "---\n",
    "\n",
    "# TokenCal Q1 — Three modes, three model families, three training seeds\n",
    "\n",
    "**Research repository:** https://huggingface.co/Tanvir106/TokenCal-Q1-Research (private)\n",
    "\n",
    "Research reproduction/extension code for *Tokenization Fragmentation and Predictive Reliability in Low-Resource LLMs: A Bengali-Centric Cross-Lingual Study*.\n",
    "\n",
    "> **Status:** The code was syntax-checked and its CPU mathematical/unit tests passed. Full CUDA training, external dataset downloads, and Hugging Face push functionality have **not** been executed in the authoring environment. Do not report planned results as completed results. Preflight must run successfully on your GPU before full training.\n",
    "\n",
    "## Quick start in Google Colab or Kaggle\n",
    "\n",
    "1. Open `TokenCal_Q1_Notebook.ipynb` using a GPU runtime. The notebook is standalone: its cells create the Python runner, requirements and tests.\n",
    "2. Log in to the **Hugging Face account that owns this private repository** (`Tanvir106`). Your access to gated Gemma/Llama checkpoints and `ai4bharat/MILU` must be enabled separately on Hugging Face.\n",
    "3. The notebook will rebuild BEnQA from the **official** GitHub commit. The original exact split assignments and cleaned B-REASO validation set are available from your uploaded `tokencal-results.zip`; upload the ZIP in the notebook when asked, or use the included `reference/` files when running from this full project ZIP. Original data are **not** re-uploaded to Hugging Face.\n",
    "4. Choose `preflight`, then `benchmark`, then `full`. Select models and variants in the editable options cell.\n",
    "5. Enable `SYNC_HF=checkpoints` for remote crash-recovery checkpoints. The code will upload the **source code, manifests, checkpoint artifacts, metrics, predictions, and summary outputs** to this private model repository. This can consume significant storage, and uploads may add runtime. Ensure access rights and storage quota are sufficient.\n",
    "6. Read `outputs/summary/` after full training. If interrupted, rerun the same configuration and optionally set `--resume-hub`.\n",
    "\n",
    "## Experiment matrix (54 *full training* runs)\n",
    "\n",
    "| Backbone | Variants trained | Seeds | Runs |\n",
    "|---|---|---|---:|\n",
    "| Qwen2.5-1.5B-Instruct | M0 M1 M2 M3 Focal RDrop Const Shuffle | 42,123,456 | 24 |\n",
    "| Gemma 3 4B IT | M0 M1 M2 M3 Focal RDrop | 42,123,456 | 18 |\n",
    "| Llama 3.2 3B Instruct | M0 M1 M2 M3 | 42,123,456 | 12 |\n",
    "| **Total** | | | **54** |\n",
    "\n",
    "Temperature Scaling (`TS`) is fitted using the **M0 validation predictions** for each backbone and seed; it requires no additional LLM training. Sarvam-1 is not included.\n",
    "\n",
    "### Dataset protocol\n",
    "\n",
    "| Dataset | Role | Population | Access |\n",
    "|---|---|---|---|\n",
    "| BEnQA | train/validation/held-out test | 3,553 / 762 / 762 **pairs** | original official GitHub repo; original split metadata |\n",
    "| B-REASO | Bengali OOD evaluation | 1,215 labeled validation examples | canonical cleaned reference from original ZIP |\n",
    "| MILU | multilingual **test only** | BN 6,636; GU 4,825; ML 4,321; TE 7,304 | gated official `ai4bharat/MILU` |\n",
    "\n",
    "- **Never tune on BEnQA test, B-REASO, or MILU.**\n",
    "- All three backbones recompute their own tokenizer-specific fragmentation differences using training data.\n",
    "- Locked BEnQA split assignment metadata prevents a silently new split.\n",
    "- A/B/C/D next-token scoring is validated separately per tokenizer. If any label takes multiple tokens, the runner fails instead of silently using an incorrect scoring rule.\n",
    "- No overlength prompt is truncated, even on new backbones. You must revise the context policy **before** training if an entire dataset does not fit.\n",
    "\n",
    "### Core M0–M3 losses (faithful to original notebook)\n",
    "\n",
    "- `M0`: mean EN/BN full-vocabulary answer-token cross-entropy.\n",
    "- `M1`: M0 + normalized four-class Brier (divide raw score by 2).\n",
    "- `M2`: M1 + normalized EN–BN Jensen–Shannon divergence (divide by ln 2).\n",
    "- `M3`: M2 + `sum_i(weight_i * JS_i) / sum_i(weight_i) / ln(2)`.\n",
    "- Fragmentation weight is the 0–1 clipped percentile normalization of EN–BN content fertility difference, using 5th/95th **training-only** percentiles.\n",
    "- `Const`: M2 plus one extra unweighted pair-JS loss. `Shuffle`: M3 with train-only weights shuffled within the training population.\n",
    "- `Focal`: full-vocabulary answer-token focal loss, gamma=2. `RDrop`: stochastic dual-pass bidirectional KL on four-choice distributions, plus ordinary task loss. These are adaptations for the fixed four-choice experimental protocol, and are **not exact replicas** of every published implementation.\n",
    "- `TS`: validation-only temperature adjustment of M0's four-choice logits (NLL objective).\n",
    "\n",
    "### Crash-resume and validation\n",
    "\n",
    "- A run ID uniquely identifies backbone/variant/seed/mode.\n",
    "- Latest checkpoint includes LoRA adapter, optimizer, scheduler, gradient scaler, Python/NumPy/PyTorch/CUDA RNG, epoch, batch cursor, selected validation-Brier best score, and configuration fingerprint.\n",
    "- Checkpoints are saved at optimizer-step boundaries; staged writes are atomically swapped locally.\n",
    "- **Checkpoint integrity is checked on resume.** Do not load `trainer.pt` from untrusted sources (`torch.load(weights_only=False)` contains Python pickle objects).\n",
    "- `preflight` loads each selected GPU model, executes forward/backward, saves a checkpoint, reloads optimizer/scheduler/RNG, continues an optimizer update, and checks inference.\n",
    "- `benchmark` measures short training performance and peak GPU VRAM; it never contributes accuracy to research results.\n",
    "- `full` trains for three epochs per selected combination, chooses checkpoint by the average validation Brier, and evaluates the frozen selected checkpoint on held-out BEnQA, B-REASO, and MILU.\n",
    "- CPU mathematical tests: `python -m pytest -q test_tokencal_q1.py`.\n",
    "\n",
    "### Limitations and deployment warnings\n",
    "\n",
    "1. **The three non-identical backbones have different parameter sizes and tokenizer vocabularies.** Compare *within-backbone relative gains*, not unadjusted raw model accuracy as evidence of a superior training method.\n",
    "2. **Gemma3 4B IT is multimodal and Llama 3.2 / MILU may require separate gated access.** Gemma3 quantized LoRA can be memory-intensive on 2×T4. Two separate T4 cards do not provide a single contiguous 32GB device.\n",
    "3. Despite successful mathematical/unit tests, no full GPU pipeline was executed here. If any real GPU preflight fails, correct the implementation first and rerun preflight; do not launch all 54 experiments.\n",
    "4. The current implementation **extends** the original study rather than claiming to have exactly reproduced the original training bit-for-bit. Different tokenizers, optimizer policy (strict Torch AdamW), environment, seed and data processing decisions can change numbers.\n",
    "5. The notebook's first upload sends source files, not copyrighted third-party dataset content. The private run repo will store model adapters and evaluation artifacts. Inspect dataset licensing before sharing any derived raw text publicly.\n",
    "6. Temperature scaling is a post-hoc baseline that generally does not change argmax accuracy. Focal/R-Drop can affect both accuracy and calibration; evaluate all metrics.\n",
    "7. Original observed strict configuration comparisons were not independent multi-seed evidence. Future results must be statistically audited before Q1 submission.\n",
    "\n",
    "## CLI commands (optional)\n",
    "\n",
    "```bash\n",
    "python tokencal_q1.py --mode preflight --models qwen --variants M0,M3 --seeds 42 --source-zip /content/tokencal-results.zip\n",
    "python tokencal_q1.py --mode benchmark --models qwen,gemma,llama --variants M0,M3,RDrop --seeds 42 --source-zip /content/tokencal-results.zip --benchmark-steps 50\n",
    "python tokencal_q1.py --mode full --models qwen,gemma,llama --variants M0,M1,M2,M3,Focal,RDrop,Const,Shuffle,TS --seeds 42,123,456 --source-zip /content/tokencal-results.zip --sync-hf checkpoints --source-upload\n",
    "```\n",
    "\n",
    "When you restart a Colab runtime, execute the notebook's source-writing cells again, select the same model/seed/method, and set `--resume-hub` if the latest checkpoint is stored in the HF repo.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0460dff4",
   "metadata": {},
   "source": [
    "### Restore a local copy of this notebook for Hugging Face source upload"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f076b81",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile build_notebook.py\n",
    "\"\"\"Create the self-contained notebook from source; not a runtime dependency.\"\"\"\n",
    "import nbformat as nbf\n",
    "from pathlib import Path\n",
    "root=Path(__file__).parent\n",
    "source=(root/'tokencal_q1.py').read_text()\n",
    "tests=(root/'test_tokencal_q1.py').read_text()\n",
    "requirements=(root/'requirements.txt').read_text()\n",
    "readme=(root/'README.md').read_text()\n",
    "builder=(root/'build_notebook.py').read_text()\n",
    "N=nbf.v4.new_notebook()\n",
    "N.metadata={\n",
    " 'kernelspec':{'display_name':'Python 3','language':'python','name':'python3'},\n",
    " 'language_info':{'name':'python'},\n",
    " 'colab':{'provenance':[],'gpuType':'T4'},\n",
    "}\n",
    "c=[]\n",
    "def md(s):c.append(nbf.v4.new_markdown_cell(s))\n",
    "def code(s):c.append(nbf.v4.new_code_cell(s))\n",
    "md('''# TokenCal Q1 — Complete GPU Experiment Notebook\n",
    "\n",
    "**Three modes: Preflight → Benchmark → Full Training**  \n",
    "**Backbones:** Qwen2.5-1.5B, Gemma3-4B-IT, Llama3.2-3B-Instruct  \n",
    "**Seeds:** 42, 123, 456  \n",
    "**Comparators:** M0, M1, M2, M3, Focal loss, R-Drop, Temperature Scaling, Constant and Shuffled controls  \n",
    "**Datasets:** BEnQA (train/val/test), B-REASO (Bengali OOD), MILU (Bengali + GU/ML/TE)\n",
    "\n",
    "**Private HF repo created:** https://huggingface.co/Tanvir106/TokenCal-Q1-Research\n",
    "\n",
    "This is a **complete source notebook**, not completed model-training results. It must be executed with a CUDA GPU, compatible dependencies, a Hugging Face login, and the source datasets. The module's CPU tests passed in the authoring environment; GPU preflight, benchmark, full training, and remote uploads have **not been executed**. Run one mode at a time.\n",
    "\n",
    "**Before you start:** Choose `Runtime → Change runtime type → T4 GPU` in Colab or a GPU accelerator in Kaggle. Gemma3 4B might require more memory than one T4; the preflight should report fit/compatibility issues rather than beginning full training.''')\n",
    "md('''## 0 — Set up a persistent project directory\n",
    "\n",
    "All outputs save under `tokencal_workspace`. For a long run on Colab, mount Google Drive or persistent storage **if available**; session disks are temporary. The `--sync-hf checkpoints` option also copies completed checkpoints/logs to your private Hugging Face repo for recovery. **Make sure your account has write access to `Tanvir106/TokenCal-Q1-Research`.**''')\n",
    "code('''import os, sys\n",
    "from pathlib import Path\n",
    "PROJECT_ROOT = Path('/content/TokenCal-Q1-Research') if Path('/content').exists() else Path.cwd() / 'TokenCal-Q1-Research'\n",
    "PROJECT_ROOT.mkdir(parents=True, exist_ok=True)\n",
    "os.chdir(PROJECT_ROOT)\n",
    "os.environ['TOKENCAL_ROOT'] = str(PROJECT_ROOT / 'tokencal_workspace')\n",
    "print('Project:', PROJECT_ROOT)\n",
    "print('Workspace:', os.environ['TOKENCAL_ROOT'])\n",
    "''')\n",
    "md('## 1 — Write version-aware dependencies and install')\n",
    "code('%%writefile requirements.txt\\n'+requirements)\n",
    "code('''# Internet required. Torch itself is provided by the active GPU runtime.\n",
    "%pip install -q -r requirements.txt\n",
    "import torch\n",
    "print('PyTorch:', torch.__version__, 'CUDA available:', torch.cuda.is_available())\n",
    "print('GPU(s):', [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])\n",
    "if not torch.cuda.is_available():\n",
    "    print('WARNING: switch to a GPU runtime before running preflight/benchmark/full training.')''')\n",
    "md('''## 2 — Generate the complete research engine\n",
    "\n",
    "The next cell writes the **complete** Python implementation. It includes the original TokenCal M0–M3 losses (with normalized weighted JS), QLoRA adapters, deterministic seeds, original BEnQA split locking, three modes, full evaluation, confidence metrics, bootstrap, automatic checkpoints, save/resume, and Hugging Face uploads.''')\n",
    "code('%%writefile tokencal_q1.py\\n'+source)\n",
    "md('## 3 — Generate automated tests and run them')\n",
    "code('%%writefile test_tokencal_q1.py\\n'+tests)\n",
    "code('%%writefile README.md\\n'+readme)\n",
    "md('### Restore a local copy of this notebook for Hugging Face source upload')\n",
    "code('%%writefile build_notebook.py\\n'+builder)\n",
    "code('''import subprocess, sys\n",
    "subprocess.run([sys.executable, 'build_notebook.py'], check=True)\n",
    "print('Local notebook is ready for Hugging Face source upload.')''')\n",
    "code('''import subprocess, sys\n",
    "subprocess.run([sys.executable, '-m', 'pytest', '-q', 'test_tokencal_q1.py'], check=True)\n",
    "print('CPU logic passed. Real GPU loading/training remains to be checked by PRE-FLIGHT.')''')\n",
    "md('''## 4 — Authenticate Hugging Face (securely)\n",
    "\n",
    "This is a **private** model repository. Do **not** paste your access token into a notebook cell or upload `.env` files. This cell asks for secure interactive login if necessary. You must also accept access conditions separately for Llama 3.2 and the gated MILU dataset (and any gated checkpoints).''')\n",
    "code('''from huggingface_hub import get_token, notebook_login, HfApi\n",
    "if not get_token():\n",
    "    print('Log in with the Hugging Face account Tanvir106, which owns the new private repo:')\n",
    "    notebook_login()\n",
    "account=HfApi(token=get_token()).whoami()['name']\n",
    "print('Authenticated HF account:', account)\n",
    "if account != 'Tanvir106':\n",
    "    print('WARNING: The private repository was created under Tanvir106. Uploads will fail without permission.')''')\n",
    "md('''## 5 — Source dataset access\n",
    "\n",
    "- **BEnQA:** cloned from the original GitHub commit and matched to the original **5,077** split assignment CSV. No new data split is generated.\n",
    "- **B-REASO:** uses the cleaned 1,215-item official labeled validation subset from the uploaded original `tokencal-results.zip`.\n",
    "- **MILU:** downloads the official test split, four languages; requires the Hugging Face user to have accepted gated access.\n",
    "\n",
    "The runner first checks for `reference/` metadata locally; if absent, it tries the original public results repository read-only. If that fails, upload your original `tokencal-results.zip` below. Do not upload third-party raw data to your new research repo without permission.''')\n",
    "code('''#@title Optional — upload the original results ZIP to this session\n",
    "UPLOAD_ORIGINAL_ZIP = False #@param {type:\"boolean\"}\n",
    "SOURCE_ZIP_PATH = '' #@param {type:\"string\"}\n",
    "if UPLOAD_ORIGINAL_ZIP:\n",
    "    from google.colab import files\n",
    "    uploaded=files.upload()\n",
    "    names=[n for n in uploaded if n.endswith('.zip')]\n",
    "    if not names: raise ValueError('Upload tokencal-results.zip')\n",
    "    SOURCE_ZIP_PATH=str(PROJECT_ROOT / names[0])\n",
    "if SOURCE_ZIP_PATH and not Path(SOURCE_ZIP_PATH).is_file():\n",
    "    raise FileNotFoundError(SOURCE_ZIP_PATH)\n",
    "print('Original results ZIP:', SOURCE_ZIP_PATH or '(try original Hub repository reference files)')''')\n",
    "md('''## 6 — Choose one run mode and experiment options\n",
    "\n",
    "**PRE-FLIGHT:** 2 optimizer steps on a tiny training subset, then save/reload optimizer, scheduler, RNG and adapter, run another optimizer step, and score predictions. Checks external dataset access before GPU tests.  \n",
    "**BENCHMARK:** 50 short optimizer steps (adjustable) to measure time/step and peak VRAM; outputs are not scientific results.  \n",
    "**FULL:** fixed 3-epoch training and validation-Brier checkpoint selection across all selected model×method×seed combinations, then full BEnQA/B-REASO/MILU evaluations. `FULL` with the default selection entails **54 complete training runs**; run sequentially, consider starting with Qwen only, and keep checkpoints remotely synced.\n",
    "\n",
    "Preflight/benchmark always use **the first selected seed**. Full uses all selected seeds. `TS` is post-hoc on M0 validation logits (no extra model training).''')\n",
    "code('''#@title Experiment controls\n",
    "RUN_MODE = 'preflight' #@param ['preflight', 'benchmark', 'full']\n",
    "MODELS = 'qwen,gemma,llama' #@param {type:\"string\"}\n",
    "VARIANTS = 'M0,M1,M2,M3,Focal,RDrop,Const,Shuffle,TS' #@param {type:\"string\"}\n",
    "SEEDS = '42,123,456' #@param {type:\"string\"}\n",
    "BENCHMARK_STEPS = 50 #@param {type:\"integer\"}\n",
    "PREFLIGHT_STEPS = 2 #@param {type:\"integer\"}\n",
    "SYNC_HF = 'none' #@param ['none','final','checkpoints']\n",
    "EXTERNAL_EVAL = True #@param {type:\"boolean\"}\n",
    "RESTORE_FROM_HF = False #@param {type:\"boolean\"}\n",
    "UPLOAD_CODE_TO_HF = True #@param {type:\"boolean\"}\n",
    "HF_REPOSITORY = 'Tanvir106/TokenCal-Q1-Research' #@param {type:\"string\"}\n",
    "\n",
    "from tokencal_q1 import valid_combinations\n",
    "models=[m.strip() for m in MODELS.split(',') if m.strip()]\n",
    "methods=[m.strip() for m in VARIANTS.split(',') if m.strip()]\n",
    "selected=list(valid_combinations(models,methods,RUN_MODE))\n",
    "train_methods=[(m,v) for m,v in selected if v!='TS']\n",
    "training_seeds=len(SEEDS.split(',')) if RUN_MODE=='full' else 1\n",
    "print('Mode:', RUN_MODE, 'backbones:', models)\n",
    "print('Selected model × method pairs:', len(selected))\n",
    "print('Planned training runs:', len(train_methods)*training_seeds)\n",
    "if RUN_MODE=='full' and SYNC_HF=='none':\n",
    "    print('CAUTION: SYNC_HF=none; checkpoints exist only on the current runtime disk.')''')\n",
    "md('''## 7 — Run selected phase (this is the launch cell)\n",
    "\n",
    "The command **fails fast** and saves diagnostic files in `outputs/failures/` on exceptions. Resume the same configuration after correcting the issue; do not use experimental test/OOD data to tune hyperparameters. Hugging Face checkpoint sync requires a valid token and the private repo write permission.\n",
    "\n",
    "**Storage note:** a 54-run study with frequent model/optimizer checkpoints can consume many gigabytes of local and remote storage. Monitor quota.''')\n",
    "code('''import importlib\n",
    "import tokencal_q1\n",
    "importlib.reload(tokencal_q1)\n",
    "args=['--mode',RUN_MODE,'--models',MODELS,'--variants',VARIANTS,'--seeds',SEEDS,\n",
    "      '--root',str(PROJECT_ROOT/'tokencal_workspace'), '--hf-repo',HF_REPOSITORY,\n",
    "      '--sync-hf',SYNC_HF,'--benchmark-steps',str(BENCHMARK_STEPS),\n",
    "      '--preflight-steps',str(PREFLIGHT_STEPS)]\n",
    "if SOURCE_ZIP_PATH: args+=['--source-zip',SOURCE_ZIP_PATH]\n",
    "if not EXTERNAL_EVAL: args+=['--no-external-eval']\n",
    "if RESTORE_FROM_HF: args+=['--resume-hub']\n",
    "if UPLOAD_CODE_TO_HF: args+=['--source-upload']\n",
    "print('Launching mode:',RUN_MODE)\n",
    "tokencal_q1.main(args)''')\n",
    "md('''## 8 — Read progress and published results\n",
    "\n",
    "Use this after preflight/benchmark/full. Notebook results are read from files, not invented. Preflight status is saved under `outputs/runs/.../preflight/preflight.json`, benchmark speed under `.../benchmark/training_status.json`, and full final tables under `outputs/summary`.''')\n",
    "code('''from pathlib import Path\n",
    "import pandas as pd\n",
    "from IPython.display import display\n",
    "out=PROJECT_ROOT/'tokencal_workspace'/'outputs'\n",
    "summary=out/'summary'/'multi_seed_summary.csv'\n",
    "if summary.exists():\n",
    "    display(pd.read_csv(summary).head(40))\n",
    "else:\n",
    "    print('No full-study summary yet. Existing output directories:')\n",
    "    for d in sorted((out/'runs').glob('*/*/seed_*/*/training_status.json'))[:30]:\n",
    "        print(' -',d.relative_to(out))\n",
    "print('Failures:',len(list((out/'failures').glob('*.json'))) if (out/'failures').exists() else 0)''')\n",
    "md('''## 9 — Optional: view repo and source files\n",
    "\n",
    "The private repository was created ahead of the notebook. The runner uploads source files using `--source-upload`, and uploads experiment files only when `SYNC_HF` is `final` or `checkpoints`. The original raw BEnQA/B-REASO/MILU examples stay local. **Notebook source file is uploaded only if it exists as a local file under the project directory;** in some Colab launch modes the currently open notebook is not available to the Python filesystem. In that case, upload the downloaded `.ipynb` through the Hugging Face website to store the notebook itself.''')\n",
    "code('''print('Hugging Face repo: https://huggingface.co/'+HF_REPOSITORY)\n",
    "print('Local source:',PROJECT_ROOT/'tokencal_q1.py')\n",
    "print('Checkpoints:',PROJECT_ROOT/'tokencal_workspace'/'outputs'/'runs')\n",
    "print('Summary:',PROJECT_ROOT/'tokencal_workspace'/'outputs'/'summary')''')\n",
    "N.cells=c\n",
    "p=root/'TokenCal_Q1_Notebook.ipynb'\n",
    "nbf.validate(N)\n",
    "nbf.write(N,p)\n",
    "print('Wrote notebook',p,p.stat().st_size,'bytes','cells',len(c))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84cf9534",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess, sys\n",
    "subprocess.run([sys.executable, 'build_notebook.py'], check=True)\n",
    "print('Local notebook is ready for Hugging Face source upload.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d46fc5d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess, sys\n",
    "subprocess.run([sys.executable, '-m', 'pytest', '-q', 'test_tokencal_q1.py'], check=True)\n",
    "print('CPU logic passed. Real GPU loading/training remains to be checked by PRE-FLIGHT.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45516898",
   "metadata": {},
   "source": [
    "## 4 — Authenticate Hugging Face (securely)\n",
    "\n",
    "This is a **private** model repository. Do **not** paste your access token into a notebook cell or upload `.env` files. This cell asks for secure interactive login if necessary. You must also accept access conditions separately for Llama 3.2 and the gated MILU dataset (and any gated checkpoints)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9b0b5f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import get_token, notebook_login, HfApi\n",
    "if not get_token():\n",
    "    print('Log in with the Hugging Face account Tanvir106, which owns the new private repo:')\n",
    "    notebook_login()\n",
    "account=HfApi(token=get_token()).whoami()['name']\n",
    "print('Authenticated HF account:', account)\n",
    "if account != 'Tanvir106':\n",
    "    print('WARNING: The private repository was created under Tanvir106. Uploads will fail without permission.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25b1c8c0",
   "metadata": {},
   "source": [
    "## 5 — Source dataset access\n",
    "\n",
    "- **BEnQA:** cloned from the original GitHub commit and matched to the original **5,077** split assignment CSV. No new data split is generated.\n",
    "- **B-REASO:** uses the cleaned 1,215-item official labeled validation subset from the uploaded original `tokencal-results.zip`.\n",
    "- **MILU:** downloads the official test split, four languages; requires the Hugging Face user to have accepted gated access.\n",
    "\n",
    "The runner first checks for `reference/` metadata locally; if absent, it tries the original public results repository read-only. If that fails, upload your original `tokencal-results.zip` below. Do not upload third-party raw data to your new research repo without permission."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "666a1603",
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Optional — upload the original results ZIP to this session\n",
    "UPLOAD_ORIGINAL_ZIP = False #@param {type:\"boolean\"}\n",
    "SOURCE_ZIP_PATH = '' #@param {type:\"string\"}\n",
    "if UPLOAD_ORIGINAL_ZIP:\n",
    "    from google.colab import files\n",
    "    uploaded=files.upload()\n",
    "    names=[n for n in uploaded if n.endswith('.zip')]\n",
    "    if not names: raise ValueError('Upload tokencal-results.zip')\n",
    "    SOURCE_ZIP_PATH=str(PROJECT_ROOT / names[0])\n",
    "if SOURCE_ZIP_PATH and not Path(SOURCE_ZIP_PATH).is_file():\n",
    "    raise FileNotFoundError(SOURCE_ZIP_PATH)\n",
    "print('Original results ZIP:', SOURCE_ZIP_PATH or '(try original Hub repository reference files)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a29837c6",
   "metadata": {},
   "source": [
    "## 6 — Choose one run mode and experiment options\n",
    "\n",
    "**PRE-FLIGHT:** 2 optimizer steps on a tiny training subset, then save/reload optimizer, scheduler, RNG and adapter, run another optimizer step, and score predictions. Checks external dataset access before GPU tests.  \n",
    "**BENCHMARK:** 50 short optimizer steps (adjustable) to measure time/step and peak VRAM; outputs are not scientific results.  \n",
    "**FULL:** fixed 3-epoch training and validation-Brier checkpoint selection across all selected model×method×seed combinations, then full BEnQA/B-REASO/MILU evaluations. `FULL` with the default selection entails **54 complete training runs**; run sequentially, consider starting with Qwen only, and keep checkpoints remotely synced.\n",
    "\n",
    "Preflight/benchmark always use **the first selected seed**. Full uses all selected seeds. `TS` is post-hoc on M0 validation logits (no extra model training)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "397a0360",
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Experiment controls\n",
    "RUN_MODE = 'preflight' #@param ['preflight', 'benchmark', 'full']\n",
    "MODELS = 'qwen,gemma,llama' #@param {type:\"string\"}\n",
    "VARIANTS = 'M0,M1,M2,M3,Focal,RDrop,Const,Shuffle,TS' #@param {type:\"string\"}\n",
    "SEEDS = '42,123,456' #@param {type:\"string\"}\n",
    "BENCHMARK_STEPS = 50 #@param {type:\"integer\"}\n",
    "PREFLIGHT_STEPS = 2 #@param {type:\"integer\"}\n",
    "SYNC_HF = 'none' #@param ['none','final','checkpoints']\n",
    "EXTERNAL_EVAL = True #@param {type:\"boolean\"}\n",
    "RESTORE_FROM_HF = False #@param {type:\"boolean\"}\n",
    "UPLOAD_CODE_TO_HF = True #@param {type:\"boolean\"}\n",
    "HF_REPOSITORY = 'Tanvir106/TokenCal-Q1-Research' #@param {type:\"string\"}\n",
    "\n",
    "from tokencal_q1 import valid_combinations\n",
    "models=[m.strip() for m in MODELS.split(',') if m.strip()]\n",
    "methods=[m.strip() for m in VARIANTS.split(',') if m.strip()]\n",
    "selected=list(valid_combinations(models,methods,RUN_MODE))\n",
    "train_methods=[(m,v) for m,v in selected if v!='TS']\n",
    "training_seeds=len(SEEDS.split(',')) if RUN_MODE=='full' else 1\n",
    "print('Mode:', RUN_MODE, 'backbones:', models)\n",
    "print('Selected model × method pairs:', len(selected))\n",
    "print('Planned training runs:', len(train_methods)*training_seeds)\n",
    "if RUN_MODE=='full' and SYNC_HF=='none':\n",
    "    print('CAUTION: SYNC_HF=none; checkpoints exist only on the current runtime disk.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed2354a4",
   "metadata": {},
   "source": [
    "## 7 — Run selected phase (this is the launch cell)\n",
    "\n",
    "The command **fails fast** and saves diagnostic files in `outputs/failures/` on exceptions. Resume the same configuration after correcting the issue; do not use experimental test/OOD data to tune hyperparameters. Hugging Face checkpoint sync requires a valid token and the private repo write permission.\n",
    "\n",
    "**Storage note:** a 54-run study with frequent model/optimizer checkpoints can consume many gigabytes of local and remote storage. Monitor quota."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2fec794",
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "import tokencal_q1\n",
    "importlib.reload(tokencal_q1)\n",
    "args=['--mode',RUN_MODE,'--models',MODELS,'--variants',VARIANTS,'--seeds',SEEDS,\n",
    "      '--root',str(PROJECT_ROOT/'tokencal_workspace'), '--hf-repo',HF_REPOSITORY,\n",
    "      '--sync-hf',SYNC_HF,'--benchmark-steps',str(BENCHMARK_STEPS),\n",
    "      '--preflight-steps',str(PREFLIGHT_STEPS)]\n",
    "if SOURCE_ZIP_PATH: args+=['--source-zip',SOURCE_ZIP_PATH]\n",
    "if not EXTERNAL_EVAL: args+=['--no-external-eval']\n",
    "if RESTORE_FROM_HF: args+=['--resume-hub']\n",
    "if UPLOAD_CODE_TO_HF: args+=['--source-upload']\n",
    "print('Launching mode:',RUN_MODE)\n",
    "tokencal_q1.main(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66b18360",
   "metadata": {},
   "source": [
    "## 8 — Read progress and published results\n",
    "\n",
    "Use this after preflight/benchmark/full. Notebook results are read from files, not invented. Preflight status is saved under `outputs/runs/.../preflight/preflight.json`, benchmark speed under `.../benchmark/training_status.json`, and full final tables under `outputs/summary`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61f85bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "from IPython.display import display\n",
    "out=PROJECT_ROOT/'tokencal_workspace'/'outputs'\n",
    "summary=out/'summary'/'multi_seed_summary.csv'\n",
    "if summary.exists():\n",
    "    display(pd.read_csv(summary).head(40))\n",
    "else:\n",
    "    print('No full-study summary yet. Existing output directories:')\n",
    "    for d in sorted((out/'runs').glob('*/*/seed_*/*/training_status.json'))[:30]:\n",
    "        print(' -',d.relative_to(out))\n",
    "print('Failures:',len(list((out/'failures').glob('*.json'))) if (out/'failures').exists() else 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3417b43d",
   "metadata": {},
   "source": [
    "## 9 — Optional: view repo and source files\n",
    "\n",
    "The private repository was created ahead of the notebook. The runner uploads source files using `--source-upload`, and uploads experiment files only when `SYNC_HF` is `final` or `checkpoints`. The original raw BEnQA/B-REASO/MILU examples stay local. **Notebook source file is uploaded only if it exists as a local file under the project directory;** in some Colab launch modes the currently open notebook is not available to the Python filesystem. In that case, upload the downloaded `.ipynb` through the Hugging Face website to store the notebook itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b185ac17",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Hugging Face repo: https://huggingface.co/'+HF_REPOSITORY)\n",
    "print('Local source:',PROJECT_ROOT/'tokencal_q1.py')\n",
    "print('Checkpoints:',PROJECT_ROOT/'tokencal_workspace'/'outputs'/'runs')\n",
    "print('Summary:',PROJECT_ROOT/'tokencal_workspace'/'outputs'/'summary')"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}