hack/tools/scripts/llm-key-hunter/results/_cache/content/bhattadics_research-thesis-overthinking-boundary.json

{"repo": "bhattadics/research-thesis-overthinking-boundary", "files": {"research/real_trace_experiments.py": {"7400d3a6a855dcc850af07d70db520fd74adbd4a": {"t": 1785846947.0883954, "n": 124294, "h": "7c856b5c4c433075", "c": "from __future__ import annotations\n\nimport argparse\nimport gc\nimport hashlib\nimport importlib.util\nimport inspect\nimport json\nimport logging\nimport math\nimport random\nimport re\nimport shutil\nimport sys\nimport time\nfrom concurrent.futures import ThreadPoolExecutor\nfrom dataclasses import asdict, dataclass\nfrom fractions import Fraction\nfrom pathlib import Path\nfrom typing import Any, Iterable\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom tqdm.auto import tqdm\n\nTORCHVISION_STUB_ROOT = Path(__file__).resolve().parent / \"_vendor\" / \"torchvision_stub\"\nif TORCHVISION_STUB_ROOT.exists() and importlib.util.find_spec(\"torchvision\") is None:\n    sys.path.insert(0, str(TORCHVISION_STUB_ROOT))\n\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, set_seed\n\ntry:\n    from datasets import load_dataset\nexcept ImportError:\n    load_dataset = None\n\ntry:\n    from transformers import BitsAndBytesConfig\nexcept ImportError:\n    BitsAndBytesConfig = None\n\n\ndef ensure_bitsandbytes_params4bit_compatibility() -> None:\n    try:\n        from bitsandbytes.nn.modules import Params4bit\n    except Exception:\n        return\n\n    current_new = Params4bit.__new__\n    if getattr(current_new, \"_hf_compat_patched\", False):\n        return\n\n    try:\n        signature = inspect.signature(current_new)\n    except (TypeError, ValueError):\n        return\n\n    if \"_is_hf_initialized\" in signature.parameters:\n        return\n\n    def compat_new(cls, data=None, *args, **kwargs):\n        kwargs.pop(\"_is_hf_initialized\", None)\n        return current_new(cls, data, *args, **kwargs)\n\n    compat_new._hf_compat_patched = True\n    Params4bit.__new__ = compat_new\n    logging.info(\n        \"Applied Params4bit compatibility shim for Accelerate/Transformers quantized dispatch.\"\n    )\n\n\ndef ensure_accelerate_meta_state_dict_compatibility() -> None:\n    try:\n        import accelerate.hooks as accelerate_hooks\n        import accelerate.big_modeling as accelerate_big_modeling\n    except Exception:\n        return\n\n    current_hook = accelerate_hooks.attach_execution_device_hook\n    if getattr(current_hook, \"_hf_compat_patched\", False):\n        return\n\n    add_hook_to_module = accelerate_hooks.add_hook_to_module\n    align_devices_hook = accelerate_hooks.AlignDevicesHook\n\n    def compat_attach_execution_device_hook(\n        module: torch.nn.Module,\n        execution_device: int | str | torch.device,\n        skip_keys: str | list[str] | None = None,\n        preload_module_classes: list[str] | None = None,\n        tied_params_map: dict[int, dict[torch.device, torch.Tensor]] | None = None,\n    ) -> None:\n        if not hasattr(module, \"_hf_hook\"):\n            try:\n                has_state = len(module.state_dict()) > 0\n            except RuntimeError as exc:\n                if \"Tensor.item() cannot be called on meta tensors\" not in str(exc):\n                    raise\n                has_state = any(value is not None for value in module._parameters.values()) or any(\n                    value is not None for value in module._buffers.values()\n                )\n            if has_state:\n                add_hook_to_module(\n                    module,\n                    align_devices_hook(\n                        execution_device,\n                        skip_keys=skip_keys,\n                        tied_params_map=tied_params_map,\n                    ),\n                )\n\n        if preload_module_classes is not None and module.__class__.__name__ in preload_module_classes:\n            return\n\n        for child in module.children():\n            compat_attach_execution_device_hook(\n                child,\n                execution_device,\n                skip_keys=skip_keys,\n                preload_module_classes=preload_module_classes,\n                tied_params_map=tied_params_map,\n            )\n\n    compat_attach_execution_device_hook._hf_compat_patched = True\n    accelerate_hooks.attach_execution_device_hook = compat_attach_execution_device_hook\n    if hasattr(accelerate_big_modeling, \"attach_execution_device_hook\"):\n        accelerate_big_modeling.attach_execution_device_hook = compat_attach_execution_device_hook\n    logging.info(\n        \"Applied Accelerate meta state_dict compatibility shim for quantized dispatch.\"\n    )\n\n\ndef resolve_low_vram_device_map(\n    model_spec: ModelSpec,\n    actual_device: str,\n    target_precision: str,\n    device_map: str | None,\n) -> str | dict[str, int | str] | None:\n    if device_map is not None:\n        return device_map\n    if actual_device != \"cuda\" or target_precision != \"4bit\":\n        return None\n    # For small models (≤ 8B params), force single-GPU loading.\n    # device_map=\"auto\" estimates memory using the unquantized model size,\n    # which can cause accelerate to offload 4-bit layers to CPU as meta\n    # tensors -- bitsandbytes then crashes with \"Cannot copy out of meta\n    # tensor; no data!\" when generate() tries to move them back to CUDA.\n    try:\n        num_params = float(re.findall(r\"[\\d.]+\", model_spec.parameter_count)[0])\n        if num_params <= 8:\n            return {\"\": 0}\n    except (IndexError, ValueError):\n        pass\n    return None\n\n\nDEFAULT_OUTPUT_DIR = Path(__file__).resolve().parent / \"outputs\" / \"real_traces\"\nSTEP_COST = 0.05\n\nclass IOManager:\n    def __init__(self, max_threads: int = 4):\n        self.executor = ThreadPoolExecutor(max_workers=max_threads)\n        self.pending_futures = []\n\n    def save_npz_async(self, path: Path, data_dict: dict[str, np.ndarray]) -> None:\n        future = self.executor.submit(np.savez_compressed, path, **data_dict)\n        self.pending_futures.append(future)\n        self._cleanup()\n\n    def _cleanup(self) -> None:\n        self.pending_futures = [f for f in self.pending_futures if not f.done()]\n\n    def wait_all(self) -> None:\n        for f in self.pending_futures:\n            f.result()\n        self.pending_futures = []\n\n    def shutdown(self) -> None:\n        self.wait_all()\n        self.executor.shutdown()\n\nSYSTEM_PROMPT = (\n    \"You are participating in a research protocol about incremental reasoning. \"\n    \"Each step should be brief and update the current best answer rather than re-deriving the full solution. \"\n    \"If you already know the answer, do one short verification step and still return the requested format.\"\n)\nDAY_NAMES = [\n    \"monday\",\n    \"tuesday\",\n    \"wednesday\",\n    \"thursday\",\n    \"friday\",\n    \"saturday\",\n    \"sunday\",\n]\nNUMBER_WORDS = {\n    \"a\": 1,\n    \"an\": 1,\n    \"one\": 1,\n    \"two\": 2,\n    \"three\": 3,\n    \"four\": 4,\n    \"five\": 5,\n    \"six\": 6,\n    \"seven\": 7,\n    \"eight\": 8,\n    \"nine\": 9,\n    \"ten\": 10,\n}\nFRACTION_DENOMINATORS = {\n    \"half\": 2,\n    \"halves\": 2,\n    \"third\": 3,\n    \"thirds\": 3,\n    \"quarter\": 4,\n    \"quarters\": 4,\n    \"fourth\": 4,\n    \"fourths\": 4,\n    \"fifth\": 5,\n    \"fifths\": 5,\n    \"sixth\": 6,\n    \"sixths\": 6,\n    \"seventh\": 7,\n    \"sevenths\": 7,\n    \"eighth\": 8,\n    \"eighths\": 8,\n    \"ninth\": 9,\n    \"ninths\": 9,\n    \"tenth\": 10,\n    \"tenths\": 10,\n}\nPLURAL_FRACTION_WORDS = {\n    \"halves\",\n    \"thirds\",\n    \"quarters\",\n    \"fourths\",\n    \"fifths\",\n    \"sixths\",\n    \"sevenths\",\n    \"eighths\",\n    \"ninths\",\n    \"tenths\",\n}\nANSWER_CUE_PHRASES = (\n    \"answer\",\n    \"final\",\n    \"therefore\",\n    \"thus\",\n    \"hence\",\n    \"so\",\n    \"equals\",\n    \"equal to\",\n    \"gives\",\n    \"will be\",\n    \"would be\",\n    \"should be\",\n    \"left with\",\n    \"remaining\",\n)\n\n\n@dataclass(frozen=True)\nclass TaskSpec:\n    task_id: str\n    domain: str\n    difficulty: str\n    prompt: str\n    answer_type: str\n    expected_answer: str\n    notes: str\n    source: str = \"builtin\"\n    source_index: int = -1\n\n\n@dataclass(frozen=True)\nclass ModelSpec:\n    alias: str\n    hf_name: str\n    family: str\n    parameter_count: str\n\n\nBUILTIN_TASKS = [\n    TaskSpec(\n        task_id=\"calendar_offset_easy\",\n        domain=\"calendar\",\n        difficulty=\"easy\",\n        prompt=(\n            \"If today is Tuesday, what day of the week will it be 14 days later? \"\n            \"Answer with only the weekday.\"\n        ),\n        answer_type=\"day\",\n        expected_answer=\"tuesday\",\n        notes=\"Very easy task where extra verification can become redundant.\",\n    ),\n    TaskSpec(\n        task_id=\"fraction_remaining\",\n        domain=\"arithmetic\",\n        difficulty=\"medium\",\n        prompt=(\n            \"A tank is half full. Then one quarter of the whole tank is drained. \"\n            \"What fraction of the tank is full now?\"\n        ),\n        answer_type=\"fraction\",\n        expected_answer=\"1/4\",\n        notes=\"Easy fraction task with exact verification.\",\n    ),\n    TaskSpec(\n        task_id=\"linear_equation\",\n        domain=\"algebra\",\n        difficulty=\"medium\",\n        prompt=\"Solve 2x + 5 = 19. Answer with only the value of x.\",\n        answer_type=\"int\",\n        expected_answer=\"7\",\n        notes=\"Simple algebra task that should be solvable by small instruct models.\",\n    ),\n    TaskSpec(\n        task_id=\"python_trace\",\n        domain=\"code\",\n        difficulty=\"medium\",\n        prompt=(\n            \"What does this Python code print?\\n\"\n            \"total = 0\\n\"\n            \"for i in range(1, 5):\\n\"\n            \"    total += i\\n\"\n            \"print(total)\\n\"\n            \"Answer with only the printed value.\"\n        ),\n        answer_type=\"int\",\n        expected_answer=\"10\",\n        notes=\"Straightforward code tracing benchmark.\",\n    ),\n]\n\n\nMODEL_CATALOG = {\n    \"deepseek_r1_distill_1p5b\": ModelSpec(\n        alias=\"deepseek_r1_distill_1p5b\",\n        hf_name=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B\",\n        family=\"DeepSeek-R1 distill\",\n        parameter_count=\"1.5B\",\n    ),\n    \"deepseek_r1_distill_7b\": ModelSpec(\n        alias=\"deepseek_r1_distill_7b\",\n        hf_name=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\",\n        family=\"DeepSeek-R1 distill\",\n        parameter_count=\"7B\",\n    ),\n    \"qwen2p5_0p5b\": ModelSpec(\n        alias=\"qwen2p5_0p5b\",\n        hf_name=\"Qwen/Qwen2.5-0.5B-Instruct\",\n        family=\"Qwen2.5 instruct\",\n        parameter_count=\"0.5B\",\n    ),\n    \"qwen2p5_3b\": ModelSpec(\n        alias=\"qwen2p5_3b\",\n        hf_name=\"Qwen/Qwen2.5-3B-Instruct\",\n        family=\"Qwen2.5 instruct\",\n        parameter_count=\"3B\",\n    ),\n    \"qwen2p5_7b\": ModelSpec(\n        alias=\"qwen2p5_7b\",\n        hf_name=\"Qwen/Qwen2.5-7B-Instruct\",\n        family=\"Qwen2.5 instruct\",\n        parameter_count=\"7B\",\n    ),\n    \"qwen2p5_14b\": ModelSpec(\n        alias=\"qwen2p5_14b\",\n        hf_name=\"Qwen/Qwen2.5-14B-Instruct\",\n        family=\"Qwen2.5 instruct\",\n        parameter_count=\"14B\",\n    ),\n    \"qwen2p5_32b\": ModelSpec(\n        alias=\"qwen2p5_32b\",\n        hf_name=\"Qwen/Qwen2.5-32B-Instruct\",\n        family=\"Qwen2.5 instruct\",\n        parameter_count=\"32B\",\n    ),\n    \"mistral_7b_instruct_v0p3\": ModelSpec(\n        alias=\"mistral_7b_instruct_v0p3\",\n        hf_name=\"mistralai/Mistral-7B-Instruct-v0.3\",\n        family=\"Mistral instruct\",\n        parameter_count=\"7B\",\n    ),\n    \"gemma_4_31b_it\": ModelSpec(\n        alias=\"gemma_4_31b_it\",\n        hf_name=\"google/gemma-4-31B-it\",\n        family=\"Gemma 4\",\n        parameter_count=\"31B\",\n    ),\n    \"gemma_4_26b_moe_it\": ModelSpec(\n        alias=\"gemma_4_26b_moe_it\",\n        hf_name=\"google/gemma-4-26B-A4B-it\",\n        family=\"Gemma 4 MoE\",\n        parameter_count=\"26B\",\n    ),\n    \"gemma_4_e4b_it\": ModelSpec(\n        alias=\"gemma_4_e4b_it\",\n        hf_name=\"principled-intelligence/gemma-4-E4B-it-text-only\",\n        family=\"Gemma 4 Edge (text-only)\",\n        parameter_count=\"4B\",\n    ),\n    \"qwen_3p5_9b\": ModelSpec(\n        alias=\"qwen_3p5_9b\",\n        hf_name=\"Qwen/Qwen3.5-9B\",\n        family=\"Qwen3.5 instruct\",\n        parameter_count=\"9B\",\n    ),\n    \"qwen_3p5_35b_moe_it\": ModelSpec(\n        alias=\"qwen_3p5_35b_moe_it\",\n        hf_name=\"Qwen/Qwen3.5-35B-A3B\",\n        family=\"Qwen3.5 MoE\",\n        parameter_count=\"35B\",\n    ),\n    \"qwen_3p5_35b_gptq_int4\": ModelSpec(\n        alias=\"qwen_3p5_35b_gptq_int4\",\n        hf_name=\"Qwen/Qwen3.5-35B-A3B-GPTQ-Int4\",\n        family=\"Qwen3.5 MoE GPTQ\",\n        parameter_count=\"35B\",\n    ),\n    \"qwen_3p5_35b_awq\": ModelSpec(\n        alias=\"qwen_3p5_35b_awq\",\n        hf_name=\"QuantTrio/Qwen3.5-35B-A3B-AWQ\",\n        family=\"Qwen3.5 MoE AWQ\",\n        parameter_count=\"35B\",\n    ),\n    \"llama_3p1_8b_instruct\": ModelSpec(\n        alias=\"llama_3p1_8b_instruct\",\n        hf_name=\"meta-llama/Llama-3.1-8B-Instruct\",\n        family=\"Llama 3.1\",\n        parameter_count=\"8B\",\n    ),\n    \"phi_4_mini_instruct\": ModelSpec(\n        alias=\"phi_4_mini_instruct\",\n        hf_name=\"microsoft/Phi-4-mini-instruct\",\n        family=\"Phi 4\",\n        parameter_count=\"4B\",\n    ),\n    \"internlm3_8b_instruct\": ModelSpec(\n        alias=\"internlm3_8b_instruct\",\n        hf_name=\"internlm/InternLM3-8B-Instruct\",\n        family=\"InternLM3\",\n        parameter_count=\"8B\",\n    ),\n    \"yi_1p5_9b_chat\": ModelSpec(\n        alias=\"yi_1p5_9b_chat\",\n        hf_name=\"01-ai/Yi-1.5-9B-Chat\",\n        family=\"Yi 1.5\",\n        parameter_count=\"9B\",\n    ),\n    \"mistral_small_24b_2409\": ModelSpec(\n        # Text-only Mistral-Small (22B, MistralForCausalLM). Replaces the\n        # Mistral-Small-3.1-24B-2503 entry, which is a vision-language model\n        # (Mistral3ForConditionalGeneration) that this text-only harness cannot\n        # load. Gated -- needs HF_TOKEN + accepting the license on the HF page.\n        alias=\"mistral_small_24b_2409\",\n        hf_name=\"mistralai/Mistral-Small-Instruct-2409\",\n        family=\"Mistral Small\",\n        parameter_count=\"22B\",\n    ),\n}\n\n\ndef sanitize_split_name(dataset_split: str) -> str:\n    return re.sub(r\"[^a-z0-9]+\", \"_\", dataset_split.lower()).strip(\"_\") or \"split\"\n\n\ndef extract_word_fraction_values(text: str) -> list[str]:\n    values: list[str] = []\n    pattern = re.compile(\n        r\"\\b(?:(?P<numerator>a|an|one|two|three|four|five|six|seven|eight|nine|ten)\\s+)?\"\n        r\"(?P<denominator>half|halves|third|thirds|quarter|quarters|fourth|fourths|fifth|fifths|sixth|sixths|seventh|sevenths|eighth|eighths|ninth|ninths|tenth|tenths)\\b\",\n        flags=re.IGNORECASE,\n    )\n    for match in pattern.finditer(text):\n        # AUDIT FIX: require an EXPLICIT numerator. Previously a bare singular\n        # ordinal (\"third\", \"quarter\") was read as 1/3, 1/4 via an implicit\n        # \"one\", turning ordinary words into fraction answers and creating\n        # false negatives on GSM8K.\n        if match.group(\"numerator\") is None:\n            continue\n        numerator_token = match.group(\"numerator\").lower()\n        denominator_token = match.group(\"denominator\").lower()\n        numerator = NUMBER_WORDS.get(numerator_token)\n        denominator = FRACTION_DENOMINATORS.get(denominator_token)\n        if numerator is None or denominator is None:\n            continue\n        values.append(str(Fraction(numerator, denominator)))\n    return values\n\n\ndef _canonical_fraction(value: str) -> str:\n    try:\n        normalized = Fraction(value)\n    except (ValueError, ZeroDivisionError):\n        return value\n    return str(normalized.numerator) if normalized.denominator == 1 else str(normalized)\n\n\ndef extract_numeric_candidate(text: str) -> str:\n    # AUDIT FIX: the final answer is the numeric token that appears LAST by\n    # position, regardless of whether it is a fraction or an integer/decimal.\n    # The previous implementation scanned fractions first and returned a leading\n    # fraction even when the real answer was the trailing integer, deflating\n    # GSM8K correctness on ratio-flavoured problems.\n    stripped = text.strip().lower().replace(\",\", \"\")\n    if not stripped:\n        return \"\"\n\n    boxed_matches = re.findall(r\"\\\\boxed\\{([^{}]+)\\}\", stripped)\n    if boxed_matches:\n        stripped = boxed_matches[-1].strip().lower().replace(\",\", \"\")\n\n    candidates: list[tuple[int, str]] = []\n    fraction_spans: list[tuple[int, int]] = []\n    for match in re.finditer(r\"-?\\d+\\s*/\\s*-?\\d+\", stripped):\n        fraction_spans.append((match.start(), match.end()))\n        candidates.append((match.start(), _canonical_fraction(match.group().replace(\" \", \"\"))))\n    for match in re.finditer(r\"-?\\d+(?:\\.\\d+)?\", stripped):\n        if any(start <= match.start() < end for start, end in fraction_spans):\n            continue  # already covered by a fraction match\n        candidates.append((match.start(), _canonical_fraction(match.group())))\n\n    if candidates:\n        candidates.sort(key=lambda item: item[0])\n        return candidates[-1][1]\n\n    word_fraction_matches = extract_word_fraction_values(stripped)\n    if word_fraction_matches:\n        return word_fraction_matches[-1]\n\n    return \"\"\n\n\ndef _last_boxed_content(s: str) -> str | None:\n    \"\"\"Return the content of the last \\\\boxed{...} with balanced braces (handles\n    nesting like \\\\boxed{\\\\frac{1}{2}}, which a non-greedy regex truncates).\"\"\"\n    key = \"\\\\boxed{\"\n    idx = s.rfind(key)\n    if idx == -1:\n        return None\n    i = idx + len(key)\n    depth = 1\n    out: list[str] = []\n    while i < len(s) and depth > 0:\n        c = s[i]\n        if c == \"{\":\n            depth += 1\n        elif c == \"}\":\n            depth -= 1\n            if depth == 0:\n                break\n        out.append(c)\n        i += 1\n    return \"\".join(out)\n\n\ndef _strip_latex_math(text: str) -> str:\n    \"\"\"Reduce a LaTeX/competition-math answer to a compact comparable form.\"\"\"\n    s = str(text).strip()\n    # AUDIT FIX: only peel a leading single-variable assignment (\"x = ...\").\n    # The previous \"split on last =\" collapsed equations/systems/lines to a\n    # single token (e.g. \"5x-7y+11z+4=0\" -> \"0\"), accepting \"0\" as correct for\n    # a plane. Keep the full expression for everything else.\n    if (s.count(\"=\") == 1 and re.match(r\"^\\s*[A-Za-z]\\w*\\s*=\", s)\n            and not any(op in s for op in (\"<=\", \">=\", \"==\", \"!=\"))):\n        s = s.split(\"=\", 1)[1].strip()\n    s = re.sub(r\"^\\$+|\\$+$\", \"\", s).strip()\n    boxed_content = _last_boxed_content(s)\n    if boxed_content is not None:\n        s = boxed_content\n    s = re.sub(r\"\\\\(?:text|mathrm|mbox|mathbf|mathit|operatorname)\\s*\\{([^{}]*)\\}\", r\"\\1\", s)\n    s = s.replace(\"\\\\dfrac\", \"\\\\frac\").replace(\"\\\\tfrac\", \"\\\\frac\")\n    s = re.sub(r\"\\\\frac\\{([^{}]+)\\}\\{([^{}]+)\\}\", r\"((\\1)/(\\2))\", s)\n    s = re.sub(r\"\\\\frac\\s*(\\d)\\s*(\\d)\", r\"((\\1)/(\\2))\", s)\n    s = re.sub(r\"\\\\sqrt\\{([^{}]+)\\}\", r\"sqrt(\\1)\", s)\n    s = re.sub(r\"\\\\sqrt\\s*(\\w)\", r\"sqrt(\\1)\", s)\n    s = s.replace(\"\\\\cdot\", \"*\").replace(\"\\\\times\", \"*\").replace(\"\\\\div\", \"/\").replace(\"\\\\pi\", \"pi\")\n    for token in (\"\\\\left\", \"\\\\right\", \"\\\\!\", \"\\\\,\", \"\\\\;\", \"\\\\:\", \"\\\\ \", \"\\\\$\",\n                  \"\\\\(\", \"\\\\)\", \"\\\\[\", \"\\\\]\", \"^{\\\\circ}\", \"^\\\\circ\", \"\\\\circ\", \"\\\\%\", \"%\", \"$\"):\n        s = s.replace(token, \"\")\n    s = s.replace(\"{\", \"\").replace(\"}\", \"\").replace(\"\\\\\", \"\")\n    s = re.sub(r\"(?<=\\d),(?=\\d{3}\\b)\", \"\", s)  # 1,000 -> 1000\n    return s.replace(\" \", \"\").rstrip(\".\").strip()\n\n\ndef normalize_math_answer(raw_answer: str) -> str:\n    if raw_answer is None:\n        return \"\"\n    s = _strip_latex_math(raw_answer)\n    if s == \"\":\n        return \"\"\n    bare = s.replace(\"(\", \"\").replace(\")\", \"\")\n    if re.fullmatch(r\"-?\\d+(\\.\\d+)?\", bare) or re.fullmatch(r\"-?\\d+/-?\\d+\", bare):\n        try:\n            frac = Fraction(bare).limit_denominator(10**6)\n            return str(frac.numerator) if frac.denominator == 1 else f\"{frac.numerator}/{frac.denominator}\"\n        except (ValueError, ZeroDivisionError):\n            pass\n    return s.lower()\n\n\ndef _as_number(text: str | None) -> float | None:\n    if not text:\n        return None\n    try:\n        return float(Fraction(text))\n    except (ValueError, ZeroDivisionError):\n        try:\n            return float(text)\n        except (ValueError, TypeError):\n            return None\n\n\ndef math_answers_equivalent(candidate: str, expected: str) -> bool:\n    norm_candidate = normalize_math_answer(candidate)\n    norm_expected = normalize_math_answer(expected)\n    if not norm_expected:\n        return False\n    if norm_candidate == norm_expected:\n        return True\n    try:\n        if abs(float(norm_candidate) - float(norm_expected)) < 1e-6:\n            return True\n    except (ValueError, TypeError):\n        pass\n    # VALIDATION FIX: when the gold answer is purely numeric, the model often\n    # gives the right number wrapped in words/units (\"Savings: 550 gallons\",\n    # \"1.25 miles\"). Compare the candidate's trailing numeric token by value.\n    # Gated on a numeric gold answer so a verbose wrong answer that merely\n    # mentions a different number cannot become a false positive.\n    exp_num = _as_number(norm_expected)\n    candidate_str = str(candidate)\n    # Only number-grab from PROSE-wrapped numbers, never from symbolic answers\n    # (equations/expressions: '=', '^', '\\', or digit-letter adjacency like \"5r\"),\n    # so \"5x-7y+11z+4=0\" is never reduced to its trailing \"0\".\n    looks_symbolic = bool(re.search(r\"[=^\\\\]|\\d[a-zA-Z]|[a-zA-Z]\\d\", candidate_str))\n    if exp_num is not None and not looks_symbolic:\n        cand_num = _as_number(extract_numeric_candidate(candidate_str))\n        if cand_num is not None and abs(cand_num - exp_num) < 1e-6:\n            return True\n    try:  # symbolic fallback; sympy ships as an indirect torch dependency\n        import sympy\n        from sympy.parsing.sympy_parser import (\n            parse_expr, standard_transformations, implicit_multiplication_application,\n        )\n        import signal\n        transforms = standard_transformations + (implicit_multiplication_application,)\n        def _sympy_ready(expr: str) -> str:\n            # Digit-letter adjacency like '0x^2' tokenizes as a broken hex\n            # literal BEFORE sympy's implicit-multiplication transform runs\n            # (TokenError; rigor_audit/01 §4c) -- make the product explicit,\n            # leaving scientific notation (2e3) alone.\n            return re.sub(r\"(\\d)(?![eE][0-9])([a-zA-Z(])\", r\"\\1*\\2\", expr.replace(\"^\", \"**\"))\n\n        # Use signal-based alarm to prevent infinite loops in sympy's parse_expr or simplify\n        # which can get stuck on complex/recursive expressions from LLM outputs.\n        has_sigalrm = hasattr(signal, \"SIGALRM\")\n        if has_sigalrm:\n            def handler(signum, frame):\n                raise TimeoutError(\"SymPy evaluation timed out\")\n            old_handler = signal.signal(signal.SIGALRM, handler)\n            signal.alarm(5)  # 5-second timeout\n\n        try:\n            expr_a = parse_expr(_sympy_ready(norm_candidate), transformations=transforms)\n            expr_b = parse_expr(_sympy_ready(norm_expected), transformations=transforms)\n            is_equiv = sympy.simplify(expr_a - expr_b) == 0\n        finally:\n            if has_sigalrm:\n                signal.alarm(0)\n                signal.signal(signal.SIGALRM, old_handler)\n\n        if is_equiv:\n            return True\n    except Exception:  # noqa: BLE001 -- unparseable expressions or timeouts just fall through\n        pass\n    return False\n\n\n_MCQ_LETTERS = \"ABCDEFGH\"\n\n\ndef normalize_mcq_answer(raw_answer: str) -> str:\n    \"\"\"Extract the single multiple-choice letter (A-H) from a model answer.\n    Handles 'B', '(C)', 'Answer: D', 'The answer is A', boxed letters, etc.\"\"\"\n    if raw_answer is None:\n        return \"\"\n    s = str(raw_answer).strip()\n    boxed = _last_boxed_content(s)\n    if boxed is not None:\n        s = boxed\n    s = s.upper()\n    for pattern in (r\"(?:ANSWER|OPTION|CHOICE)\\s*(?:IS|:|=)?\\s*\\(?([A-H])\\)?\",\n                    r\"\\(([A-H])\\)\",\n                    # '/' guards: a bare letter adjacent to a slash is part of a\n                    # token like 'N/A', not a choice (rigor_audit/01 §4a).\n                    r\"(?<![A-Z/])([A-H])(?![A-Z/])\"):\n        matches = re.findall(pattern, s)\n        if matches:\n            return matches[-1]\n    return s[:1] if s[:1] in _MCQ_LETTERS else \"\"\n\n\ndef _format_mcq_prompt(question: str, labeled_options: list[tuple[str, str]]) -> str:\n    lines = [question.strip(), \"\", \"Options:\"]\n    lines += [f\"{label}. {text}\" for label, text in labeled_options]\n    lines += [\"\", \"Respond with the letter of the single correct option.\"]\n    return \"\\n\".join(lines)\n\n\ndef normalize_answer(raw_answer: str, answer_type: str) -> str:\n    if answer_type == \"math\":\n        return normalize_math_answer(raw_answer)\n    if answer_type == \"mcq\":\n        return normalize_mcq_answer(raw_answer)\n    text = raw_answer.strip().lower()\n    text = re.sub(r\"\\s+\", \" \", text)\n    boxed_matches = re.findall(r\"\\\\boxed\\{([^{}]+)\\}\", text)\n    if boxed_matches:\n        text = boxed_matches[-1].strip().lower()\n    if not text:\n        return \"\"\n\n    if answer_type == \"day\":\n        for day_name in DAY_NAMES:\n            if day_name in text:\n                return day_name\n        return text.split()[0]\n\n    if answer_type == \"int\":\n        match = re.findall(r\"-?\\d+\", text)\n        return match[-1] if match else text\n\n    if answer_type in {\"fraction\", \"number\"}:\n        numeric_candidate = extract_numeric_candidate(text)\n        return numeric_candidate or text\n\n    return text\n\n\ndef verify_answer(task: TaskSpec, candidate_answer: str) -> bool:\n    if task.answer_type == \"math\":\n        return math_answers_equivalent(candidate_answer, task.expected_answer)\n    if task.answer_type == \"mcq\":\n        gold = normalize_mcq_answer(task.expected_answer)\n        return gold != \"\" and normalize_mcq_answer(candidate_answer) == gold\n    return normalize_answer(candidate_answer, task.answer_type) == normalize_answer(task.expected_answer, task.answer_type)\n\n\ndef utility(correct: bool, step: int) -> float:\n    return float(correct) - STEP_COST * (step - 1)\n\n\ndef lexical_overlap(current_text: str, prior_text: str) -> float:\n    current_tokens = set(re.findall(r\"[a-z0-9]+\", current_text.lower()))\n    prior_tokens = set(re.findall(r\"[a-z0-9]+\", prior_text.lower()))\n    if not current_tokens or not prior_tokens:\n        return 0.0\n    return len(current_tokens & prior_tokens) / len(current_tokens | prior_tokens)\n\n\ndef conversation_prompt(task: TaskSpec, history: list[dict[str, Any]], step: int, max_steps: int, prompt_mode: str) -> str:\n    if history:\n        rendered_history = \"\\n\".join(\n            (\n                f\"Step {item['step']}: thought={item['thought']} | answer={item['answer']} | \"\n                f\"confidence={item['confidence']} | stop={item['model_stop_flag']}\"\n            )\n            for item in history\n        )\n    else:\n        rendered_history = \"No previous steps.\"\n\n    protocol_text = (\n        f\"Research protocol: you are at incremental reasoning step {step} of {max_steps}. \"\n        \"Even if you already know the answer, continue with one short reasoning or verification step. \"\n        \"Do not repeat the full solution. Update the answer field every time.\\n\\n\"\n        f\"Previous steps:\\n{rendered_history}\\n\\n\"\n    )\n\n    if prompt_mode == \"structured_four_line\":\n        format_instr = (\n            \"Return exactly four lines:\\n\"\n            \"THOUGHT: <one short sentence>\\n\"\n            \"ANSWER: <current best final answer only>\\n\"\n            \"CONFIDENCE: <integer 0-100>\\n\"\n            \"STOP: <yes or no>\"\n        )\n    elif prompt_mode == \"minimal_json\":\n        format_instr = (\n            \"Return a single JSON object with the following keys: 'thought' (string), 'answer' (string), \"\n            \"'confidence' (integer 0-100), and 'stop' (boolean).\"\n        )\n    else:\n        if task.answer_type in {\"int\", \"fraction\", \"number\"}:\n            format_instr = \"Provide only the final numeric answer.\"\n        else:\n            format_instr = \"Provide only the final answer.\"\n\n    return (\n        f\"Task id: {task.task_id}\\n\"\n        f\"Domain: {task.domain}\\n\"\n        f\"Difficulty band: {task.difficulty}\\n\"\n        f\"Task: {task.prompt}\\n\\n\"\n        f\"{protocol_text if prompt_mode != 'answer_only' else ''}\"\n        f\"{format_instr}\"\n    )\n\n\ndef render_prompt(tokenizer: Any, user_prompt: str, system_prompt_mode: str) -> str:\n    if system_prompt_mode == \"default\":\n        sys_msg = SYSTEM_PROMPT\n    elif system_prompt_mode == \"short\":\n        sys_msg = \"You are a logical reasoning assistant.\"\n    else:\n        sys_msg = \"\"\n\n    messages = []\n    if sys_msg:\n        messages.append({\"role\": \"system\", \"content\": sys_msg})\n    messages.append({\"role\": \"user\", \"content\": user_prompt})\n\n    if hasattr(tokenizer, \"apply_chat_template\"):\n        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n    if sys_msg:\n        return f\"System: {sys_msg}\\nUser: {user_prompt}\\nAssistant:\"\n    return f\"User: {user_prompt}\\nAssistant:\"\n\n\ndef extract_typed_answer(text: str, answer_type: str) -> str:\n    stripped = text.strip()\n    if not stripped:\n        return \"\"\n\n    boxed_matches = re.findall(r\"\\\\boxed\\{([^{}]+)\\}\", stripped)\n    if boxed_matches:\n        stripped = boxed_matches[-1].strip()\n\n    if answer_type == \"math\":\n        return stripped.strip()\n\n    if answer_type == \"mcq\":\n        return normalize_mcq_answer(stripped)\n\n    lowered = stripped.lower()\n    if answer_type == \"day\":\n        day_matches = re.findall(r\"\\b(monday|tuesday|wednesday|thursday|friday|saturday|sunday)\\b\", lowered)\n        return day_matches[-1] if day_matches else \"\"\n\n    if answer_type == \"int\":\n        int_matches = re.findall(r\"-?\\d+\", stripped)\n        return int_matches[-1] if int_matches else \"\"\n\n    if answer_type in {\"fraction\", \"number\"}:\n        return extract_numeric_candidate(stripped)\n\n    return stripped\n\n\ndef split_answer_segments(text: str) -> list[str]:\n    return [\n        segment.strip(\" \\t\\r\\n-:;,\")\n        for segment in re.split(r\"(?:\\n+|(?<=[.!?])\\s+)\", text)\n        if segment.strip()\n    ]\n\n\ndef segment_has_answer_cue(text: str) -> bool:\n    lowered = text.lower()\n    return any(phrase in lowered for phrase in ANSWER_CUE_PHRASES)\n\n\ndef extract_answer(raw_text: str, answer_type: str) -> tuple[str, str]:\n    cleaned = raw_text.strip()\n    if not cleaned:\n        return \"\", \"fallback\"\n\n    boxed_matches = re.findall(r\"\\\\boxed\\{([^{}]+)\\}\", cleaned)\n    if boxed_matches:\n        candidate = extract_typed_answer(boxed_matches[-1], answer_type)\n        if candidate:\n            return candidate, \"boxed\"\n\n    regions: list[tuple[str, str]] = []\n    think_parts = re.split(r\"</think>\", cleaned, flags=re.IGNORECASE)\n    if len(think_parts) > 1 and think_parts[-1].strip():\n        regions.append((\"after_think\", think_parts[-1].strip()))\n    regions.append((\"full_text\", cleaned))\n\n    for region_name, region_text in regions:\n        for segment in reversed(split_answer_segments(region_text)):\n            candidate = extract_typed_answer(segment, answer_type)\n            if candidate and segment_has_answer_cue(segment):\n                return candidate, f\"{region_name}_cue_segment\"\n\n    for region_name, region_text in regions:\n        candidate = extract_typed_answer(region_text, answer_type)\n        if candidate:\n            return candidate, region_name\n\n    return \"\", \"fallback\"\n\n\ndef parse_generation(raw_text: str, answer_type: str, prompt_mode: str) -> dict[str, Any]:\n    cleaned = raw_text.strip()\n    cleaned = re.sub(r\"^```\\w*\\n\", \"\", cleaned, flags=re.MULTILINE)\n    cleaned = re.sub(r\"\\n```$\", \"\", cleaned, flags=re.MULTILINE)\n    cleaned = cleaned.strip()\n\n    result = {\n        \"thought\": \"\",\n        \"answer\": \"\",\n        \"confidence\": 50,\n        \"stop\": False,\n        \"parse_success\": 0,\n        \"output_format_type\": \"fallback_freeform\",\n        \"answer_extraction_source\": \"fallback\",\n        \"stop_extraction_source\": \"default\",\n        \"confidence_extraction_source\": \"default\",\n    }\n\n    if prompt_mode == \"structured_four_line\":\n        line_thought = re.search(r\"^THOUGHT\\s*:\\s*(.*)$\", cleaned, flags=re.IGNORECASE | re.MULTILINE)\n        line_answer = re.search(r\"^ANSWER\\s*:\\s*(.*)$\", cleaned, flags=re.IGNORECASE | re.MULTILINE)\n        line_confidence = re.search(r\"^CONFIDENCE\\s*:\\s*(-?\\d+(?:\\.\\d+)?)\", cleaned, flags=re.IGNORECASE | re.MULTILINE)\n        line_stop = re.search(r\"^STOP\\s*:\\s*(yes|no|true|false)\", cleaned, flags=re.IGNORECASE | re.MULTILINE)\n\n        if line_thought and line_answer and line_confidence and line_stop:\n            result[\"parse_success\"] = 1\n            result[\"output_format_type\"] = \"structured_exact\"\n        elif line_thought or line_answer or line_confidence or line_stop:\n            result[\"output_format_type\"] = \"structured_partial\"\n\n        if line_thought:\n            result[\"thought\"] = line_thought.group(1).strip()\n        if line_answer:\n            result[\"answer\"] = line_answer.group(1).strip()\n            result[\"answer_extraction_source\"] = \"ANSWER_line\"\n        if line_confidence:\n            result[\"confidence\"] = int(float(line_confidence.group(1)))\n            result[\"confidence_extraction_source\"] = \"CONFIDENCE_line\"\n        if line_stop:\n            result[\"stop\"] = line_stop.group(1).lower() in {\"yes\", \"true\"}\n            result[\"stop_extraction_source\"] = \"STOP_line\"\n\n    elif prompt_mode == \"minimal_json\":\n        start = cleaned.find(\"{\")\n        end = cleaned.rfind(\"}\")\n        if start != -1 and end != -1 and end > start:\n            candidate = cleaned[start : end + 1]\n            try:\n                parsed = json.loads(candidate)\n                result[\"thought\"] = str(parsed.get(\"thought\", \"\")).strip()\n                result[\"answer\"] = str(parsed.get(\"answer\", \"\")).strip()\n                result[\"confidence\"] = int(float(parsed.get(\"confidence\", 50)))\n                result[\"stop\"] = bool(parsed.get(\"stop\", False))\n                result[\"parse_success\"] = 1\n                result[\"output_format_type\"] = \"json_exact\"\n                result[\"answer_extraction_source\"] = \"json_field\"\n                result[\"confidence_extraction_source\"] = \"json_field\"\n                result[\"stop_extraction_source\"] = \"json_field\"\n            except (json.JSONDecodeError, TypeError, ValueError):\n                result[\"output_format_type\"] = \"json_partial\"\n\n    if not result[\"answer\"]:\n        extracted_answer, extraction_source = extract_answer(cleaned, answer_type)\n        if extracted_answer:\n            result[\"answer\"] = extracted_answer\n            result[\"answer_extraction_source\"] = extraction_source\n\n    if not result[\"thought\"]:\n        think_match = re.search(r\"<think>(.*?)</think>\", cleaned, flags=re.IGNORECASE | re.DOTALL)\n        if think_match:\n            result[\"thought\"] = think_match.group(1).strip()\n        else:\n            lines = [line.strip() for line in cleaned.splitlines() if line.strip()]\n            if lines:\n                result[\"thought\"] = lines[0]\n\n    return result\n\n\ndef extract_gsm8k_reference_answer(answer_text: str) -> str:\n    match = re.search(r\"####\\s*([^\\n]+)\", answer_text)\n    if match:\n        return normalize_answer(match.group(1), \"number\")\n    return normalize_answer(answer_text, \"number\")\n\n\ndef load_gsm8k_tasks(\n    max_tasks: int,\n    dataset_split: str,\n    shuffle_seed: int | None,\n    dataset_revision: str | None = None,\n) -> list[TaskSpec]:\n    if load_dataset is None:\n        raise ImportError(\"datasets is required for --task-source gsm8k. Install it with pip install datasets evaluate bitsandbytes tqdm.\")\n\n    try:\n        load_kwargs: dict[str, Any] = {\"split\": dataset_split}\n        if dataset_revision:\n            load_kwargs[\"revision\"] = dataset_revision\n        dataset = load_dataset(\"openai/gsm8k\", \"main\", **load_kwargs)\n    except Exception as exc:\n        raise RuntimeError(\n            f\"Failed to download or load the GSM8K dataset from HuggingFace. \"\n            f\"Please check your internet connection on the VM and verify that you can access HuggingFace. \"\n            f\"Error details: {exc}\"\n        ) from exc\n    if shuffle_seed is not None:\n        dataset = dataset.shuffle(seed=shuffle_seed)\n    if max_tasks > 0 and len(dataset) > max_tasks:\n        dataset = dataset.select(range(max_tasks))\n\n    split_name = sanitize_split_name(dataset_split)\n    tasks: list[TaskSpec] = []\n    for index, example in enumerate(dataset):\n        question = str(example[\"question\"]).strip()\n        expected_answer = extract_gsm8k_reference_answer(str(example[\"answer\"]))\n        short_hash = hashlib.md5(question.encode(\"utf-8\")).hexdigest()[:8]\n        tasks.append(\n            TaskSpec(\n                task_id=f\"gsm8k_{split_name}_{index:05d}_{short_hash}\",\n                domain=\"gsm8k\",\n                difficulty=\"grade_school_math\",\n                prompt=question,\n                answer_type=\"number\",\n                expected_answer=expected_answer,\n                notes=f\"gsm8k main split={dataset_split}\",\n                source=\"gsm8k\",\n                source_index=index,\n            )\n        )\n    return tasks\n\n\ndef load_math_tasks(\n    max_tasks: int,\n    dataset_split: str,\n    shuffle_seed: int | None,\n    dataset_revision: str | None = None,\n) -> list[TaskSpec]:\n    \"\"\"Competition MATH. Uses HuggingFaceH4/MATH-500 (500 problems, openly\n    available); falls back to lighteval/MATH. Answers are graded with\n    math_answers_equivalent (LaTeX-normalized + numeric + sympy fallback).\"\"\"\n    if load_dataset is None:\n        raise ImportError(\"datasets is required for --task-source math. Install it with pip install datasets.\")\n    split = dataset_split if dataset_split and dataset_split not in {\"train\", \"main\"} else \"test\"\n    # AUDIT FIX: load ONLY the pinned corpus and FAIL HARD if unavailable. The\n    # previous silent fallback to lighteval/MATH (~5000 problems) could swap the\n    # whole corpus on a transient error, giving different models different\n    # problems and silently breaking cross-model comparability.\n    repo = \"HuggingFaceH4/MATH-500\"\n    try:\n        load_kwargs: dict[str, Any] = {\"split\": split}\n        if dataset_revision:\n            load_kwargs[\"revision\"] = dataset_revision\n        dataset = load_dataset(repo, **load_kwargs)\n    except Exception as exc:\n        raise RuntimeError(\n            f\"Failed to load the pinned MATH dataset {repo} (split={split}). Failing hard rather than \"\n            f\"silently substituting a different corpus. Error details: {exc}\"\n        ) from exc\n    if shuffle_seed is not None:\n        dataset = dataset.shuffle(seed=shuffle_seed)\n    if max_tasks > 0 and len(dataset) > max_tasks:\n        dataset = dataset.select(range(max_tasks))\n\n    split_name = sanitize_split_name(split)\n    tasks: list[TaskSpec] = []\n    for index, example in enumerate(dataset):\n        problem = str(example.get(\"problem\", \"\")).strip()\n        answer = example.get(\"answer\")\n        if answer is None or str(answer).strip() == \"\":\n            boxed = re.findall(r\"\\\\boxed\\{(.+?)\\}\", str(example.get(\"solution\", \"\")))\n            answer = boxed[-1] if boxed else str(example.get(\"solution\", \"\"))\n        level = str(example.get(\"level\", \"\")).strip()\n        subject = str(example.get(\"subject\", example.get(\"type\", \"\"))).strip()\n        short_hash = hashlib.md5(problem.encode(\"utf-8\")).hexdigest()[:8]\n        tasks.append(\n            TaskSpec(\n                task_id=f\"math_{split_name}_{index:05d}_{short_hash}\",\n                domain=\"math\",\n                difficulty=(f\"MATH {level}\".strip() or \"competition_math\"),\n                prompt=problem,\n                answer_type=\"math\",\n                expected_answer=str(answer).strip(),\n                notes=f\"MATH repo={repo} split={split} subject={subject} level={level}\",\n                source=\"math\",\n                source_index=index,\n            )\n        )\n    return tasks\n\n\ndef load_arc_tasks(\n    max_tasks: int,\n    dataset_split: str,\n    shuffle_seed: int | None,\n    dataset_revision: str | None = None,\n) -> list[TaskSpec]:\n    \"\"\"ARC-Challenge (allenai/ai2_arc, public). Multiple choice; all options are\n    relabeled A,B,C(,D,E) and the gold answerKey is mapped to the new letter.\"\"\"\n    if load_dataset is None:\n        raise ImportError(\"datasets is required for --task-source arc.\")\n    split = dataset_split if dataset_split and dataset_split not in {\"train\", \"main\"} else \"test\"\n    try:\n        load_kwargs: dict[str, Any] = {\"split\": split}\n        if dataset_revision:\n            load_kwargs[\"revision\"] = dataset_revision\n        dataset = load_dataset(\"allenai/ai2_arc\", \"ARC-Challenge\", **load_kwargs)\n    except Exception as exc:\n        raise RuntimeError(f\"Failed to load allenai/ai2_arc ARC-Challenge (split={split}). Error: {exc}\") from exc\n    if shuffle_seed is not None:\n        dataset = dataset.shuffle(seed=shuffle_seed)\n    if max_tasks > 0 and len(dataset) > max_tasks:\n        dataset = dataset.select(range(max_tasks))\n    split_name = sanitize_split_name(split)\n    tasks: list[TaskSpec] = []\n    for index, example in enumerate(dataset):\n        question = str(example[\"question\"]).strip()\n        texts = list(example[\"choices\"][\"text\"])\n        labels = [str(label) for label in example[\"choices\"][\"label\"]]\n        new_labels = [chr(ord(\"A\") + j) for j in range(len(texts))]\n        gold = dict(zip(labels, new_labels)).get(str(example[\"answerKey\"]).strip(), \"\")\n        if not gold or not texts:\n            continue\n        short_hash = hashlib.md5(question.encode(\"utf-8\")).hexdigest()[:8]\n        tasks.append(TaskSpec(\n            task_id=f\"arc_{split_name}_{index:05d}_{short_hash}\", domain=\"arc\",\n            difficulty=\"arc_challenge\", prompt=_format_mcq_prompt(question, list(zip(new_labels, texts))),\n            answer_type=\"mcq\", expected_answer=gold, notes=f\"ARC-Challenge split={split}\",\n            source=\"arc\", source_index=index))\n    return tasks\n\n\ndef load_gpqa_tasks(\n    max_tasks: int,\n    dataset_split: str,\n    shuffle_seed: int | None,\n    dataset_revision: str | None = None,\n) -> list[TaskSpec]:\n    \"\"\"GPQA main (Idavidrein/gpqa, GATED -- needs HF_TOKEN + accepted terms).\n    The correct + 3 incorrect answers are deterministically shuffled per item.\"\"\"\n    if load_dataset is None:\n        raise ImportError(\"datasets is required for --task-source gpqa.\")\n    try:\n        load_kwargs: dict[str, Any] = {\"split\": \"train\"}\n        if dataset_revision:\n            load_kwargs[\"revision\"] = dataset_revision\n        dataset = load_dataset(\"Idavidrein/gpqa\", \"gpqa_main\", **load_kwargs)\n    except Exception as exc:\n        raise RuntimeError(\n            \"Failed to load Idavidrein/gpqa (gpqa_main). It is GATED: accept the terms on the dataset \"\n            f\"page and set HF_TOKEN. Error: {exc}\") from exc\n    if shuffle_seed is not None:\n        dataset = dataset.shuffle(seed=shuffle_seed)\n    if max_tasks > 0 and len(dataset) > max_tasks:\n        dataset = dataset.select(range(max_tasks))\n    new_labels = [\"A\", \"B\", \"C\", \"D\"]\n    tasks: list[TaskSpec] = []\n    for index, example in enumerate(dataset):\n        question = str(example[\"Question\"]).strip()\n        options = [str(example[\"Correct Answer\"]).strip(),\n                   str(example[\"Incorrect Answer 1\"]).strip(),\n                   str(example[\"Incorrect Answer 2\"]).strip(),\n                   str(example[\"Incorrect Answer 3\"]).strip()]\n        rng = random.Random((shuffle_seed if shuffle_seed is not None else 0) * 100003 + index)\n        order = [0, 1, 2, 3]\n        rng.shuffle(order)\n        gold = new_labels[order.index(0)]  # where the correct option (index 0) landed\n        labeled = list(zip(new_labels, [options[k] for k in order]))\n        short_hash = hashlib.md5(question.encode(\"utf-8\")).hexdigest()[:8]\n        tasks.append(TaskSpec(\n            task_id=f\"gpqa_main_{index:05d}_{short_hash}\", domain=\"gpqa\", difficulty=\"gpqa_main\",\n            prompt=_format_mcq_prompt(question, labeled), answer_type=\"mcq\", expected_answer=gold,\n            notes=\"GPQA main\", source=\"gpqa\", source_index=index))\n    return tasks\n\n\ndef load_tasks(\n    task_source: str,\n    max_tasks: int,\n    dataset_split: str,\n    shuffle_seed: int | None,\n    dataset_revision: str | None = None,\n) -> list[TaskSpec]:\n    if task_source == \"builtin\":\n        return BUILTIN_TASKS[:max_tasks]\n    if task_source == \"gsm8k\":\n        return load_gsm8k_tasks(\n            max_tasks=max_tasks,\n            dataset_split=dataset_split,\n            shuffle_seed=shuffle_seed,\n            dataset_revision=dataset_revision,\n        )\n    if task_source == \"math\":\n        return load_math_tasks(\n            max_tasks=max_tasks,\n            dataset_split=dataset_split,\n            shuffle_seed=shuffle_seed,\n            dataset_revision=dataset_revision,\n        )\n    if task_source == \"arc\":\n        return load_arc_tasks(\n            max_tasks=max_tasks,\n            dataset_split=dataset_split,\n            shuffle_seed=shuffle_seed,\n            dataset_revision=dataset_revision,\n        )\n    if task_source == \"gpqa\":\n        return load_gpqa_tasks(\n            max_tasks=max_tasks,\n            dataset_split=dataset_split,\n            shuffle_seed=shuffle_seed,\n            dataset_revision=dataset_revision,\n        )\n    raise ValueError(f\"Unsupported task source: {task_source}\")\n\n\ndef run_id_for(model_alias: str, task: TaskSpec, temperature: float, seed: int) -> str:\n    return f\"{model_alias}__{task.task_id}__temp{temperature:.2f}__seed{seed}\"\n\n\ndef chunked(items: list[TaskSpec], size: int) -> Iterable[list[TaskSpec]]:\n    for start in range(0, len(items), size):\n        yield items[start : start + size]\n\n\ndef output_paths(output_dir: Path, is_baseline: bool) -> dict[str, Path]:\n    steps_name = \"baseline_steps.csv\" if is_baseline else \"trace_steps.csv\"\n    runs_name = \"baseline_runs.csv\" if is_baseline else \"trace_runs.csv\"\n    batch_metrics_name = \"baseline_batch_metrics.csv\" if is_baseline else \"batch_metrics.csv\"\n    return {\n        \"steps\": output_dir / steps_name,\n        \"runs\": output_dir / runs_name,\n        \"batch_metrics\": output_dir / batch_metrics_name,\n        \"hazard\": output_dir / \"hazard_by_step.csv\",\n        \"pilot\": output_dir / \"pilot_summary.csv\",\n        \"metadata\": output_dir / \"metadata.json\",\n    }\n\n\ndef expected_steps_per_run(is_baseline: bool, max_steps: int) -> int:\n    return 1 if is_baseline else max_steps\n\n\ndef _escape_control_chars(frame: pd.DataFrame) -> pd.DataFrame:\n    # Control characters (\\r, \\f from decoded LaTeX like \\rho, \\frac) inside\n    # text fields break CSV row framing when writes are interrupted\n    # (rigor_audit/01 Axis 5a root cause). Tabs and newlines survive pandas\n    # quoting; everything else in C0 gets replaced.\n    pattern = re.compile(r\"[\\x00-\\x08\\x0b-\\x1f]\")\n    for column in frame.columns:\n        if frame[column].dtype == object or pd.api.types.is_string_dtype(frame[column]):\n            frame[column] = frame[column].map(\n                lambda value: pattern.sub(\" \", value) if isinstance(value, str) else value\n            )\n    return frame\n\n\ndef append_records(path: Path, rows: list[dict[str, Any]]) -> None:\n    if not rows:\n        return\n    frame = _escape_control_chars(pd.DataFrame(rows))\n    if path.exists():\n        try:\n            existing_columns = pd.read_csv(path, nrows=0).columns.tolist()\n            for column in existing_columns:\n                if column not in frame.columns:\n                    frame[column] = np.nan\n            extra_columns = [column for column in frame.columns if column not in existing_columns]\n            frame = frame[existing_columns + extra_columns]\n            frame.to_csv(path, mode=\"a\", header=False, index=False)\n            return\n        except pd.errors.EmptyDataError:\n            pass\n    frame.to_csv(path, index=False)\n\n\ndef max_nan(values: Iterable[float]) -> float:\n    finite_values = [float(value) for value in values if not math.isnan(float(value))]\n    if not finite_values:\n        return float(\"nan\")\n    return max(finite_values)\n\n\ndef backup_for_reconciliation(path: Path) -> Path:\n    backup_path = path.with_name(f\"{path.stem}.pre_reconcile{path.suffix}\")\n    if path.exists() and not backup_path.exists():\n        shutil.copy2(path, backup_path)\n    return backup_path\n\n\ndef summarize_run_rows(run_rows: pd.DataFrame) -> dict[str, Any]:\n    ordered = run_rows.sort_values(\"step\").drop_duplicates(subset=[\"step\"], keep=\"last\").reset_index(drop=True)\n    first_row = ordered.iloc[0]\n    final_row = ordered.iloc[-1]\n    first_correct = ordered.loc[ordered[\"correct\"] == 1, \"step\"]\n    first_model_stop = ordered.loc[ordered[\"model_stop_flag\"] == 1, \"step\"]\n    oracle_row = ordered.loc[ordered[\"utility\"].idxmax()]\n\n    return {\n        \"run_id\": str(first_row[\"run_id\"]),\n        \"model_alias\": str(first_row[\"model_alias\"]),\n        \"model_name\": str(first_row[\"model_name\"]),\n        \"task_id\": str(first_row[\"task_id\"]),\n        \"domain\": str(first_row[\"domain\"]),\n        \"difficulty\": str(first_row[\"difficulty\"]),\n        \"task_source\": str(first_row[\"task_source\"]),\n        \"task_source_index\": int(first_row[\"task_source_index\"]),\n        \"temperature\": float(first_row[\"temperature\"]),\n        \"seed\": int(first_row[\"seed\"]),\n        \"prompt_mode\": str(first_row[\"prompt_mode\"]),\n        \"system_prompt_mode\": str(first_row[\"system_prompt_mode\"]),\n        \"is_baseline\": int(first_row[\"is_baseline\"]),\n        \"ever_correct\": int(ordered[\"correct\"].max()),\n        \"correct_at_step_1\": int(first_row[\"correct\"]),\n        \"oracle_stop\": int(oracle_row[\"step\"]),\n        \"first_correct_step\": int(first_correct.iloc[0]) if not first_correct.empty else -1,\n        \"first_model_stop_step\": int(first_model_stop.iloc[0]) if not first_model_stop.empty else -1,\n        \"revision_count\": int(ordered[\"answer_changed\"].sum()),\n        \"best_utility\": float(ordered[\"utility\"].max()),\n        \"final_correct\": int(final_row[\"correct\"]),\n        \"device\": str(first_row[\"device\"]),\n    }\n\n\ndef _drop_truncated_rows(frame: pd.DataFrame, key_cols: list[str]) -> tuple[pd.DataFrame, int]:\n    \"\"\"Drop rows whose key identifier columns are missing/unparseable -- the\n    signature of an append-only CSV line truncated by a hard kill mid-write\n    (e.g. the autostart watchdog SIGKILLing a hung run). Such a half-row would\n    otherwise poison numeric columns (object dtype) or blow up int(step) on\n    resume, crashing the whole cell. The affected run simply lacks a complete\n    row set and gets re-collected, which is exactly what we want.\"\"\"\n    if frame.empty:\n        return frame, 0\n    before = len(frame)\n    out = frame\n    for col in key_cols:\n        if col not in out.columns:\n            continue\n        out = out[out[col].notna()]\n        if col == \"step\":  # must be a clean integer; garbled value -> drop the row\n            out = out[pd.to_numeric(out[col], errors=\"coerce\").notna()]\n    out = out.reset_index(drop=True)\n    return out, before - len(out)\n\n\ndef reconcile_existing_outputs(\n    output_dir: Path,\n    hidden_dir: Path,\n    is_baseline: bool,\n    max_steps: int,\n) -> tuple[list[dict[str, Any]], list[dict[str, Any]], set[str], dict[str, Any]]:\n    paths = output_paths(output_dir, is_baseline)\n    expected_steps = expected_steps_per_run(is_baseline, max_steps)\n    expected_step_sequence = list(range(1, expected_steps + 1))\n\n    try:\n        steps_frame = pd.read_csv(paths[\"steps\"]) if paths[\"steps\"].exists() else pd.DataFrame()\n    except pd.errors.EmptyDataError:\n        steps_frame = pd.DataFrame()\n    try:\n        runs_frame = pd.read_csv(paths[\"runs\"]) if paths[\"runs\"].exists() else pd.DataFrame()\n    except pd.errors.EmptyDataError:\n        runs_frame = pd.DataFrame()\n    steps_frame, truncated_step_rows = _drop_truncated_rows(steps_frame, [\"run_id\", \"step\"])\n    runs_frame, truncated_run_rows = _drop_truncated_rows(runs_frame, [\"run_id\"])\n    if truncated_step_rows or truncated_run_rows:\n        print(\n            f\"[reconcile] dropped {truncated_step_rows} truncated step row(s) and \"\n            f\"{truncated_run_rows} truncated run row(s) from a prior hard kill; \"\n            \"affected runs will be re-collected.\",\n            flush=True,\n        )\n    hidden_run_ids = {path.stem for path in hidden_dir.glob(\"*.npz\")} if hidden_dir.exists() else set()\n\n    duplicate_step_rows = 0\n    duplicate_run_rows = 0\n    if not steps_frame.empty:\n        steps_frame = steps_frame.reset_index(drop=True)\n        steps_frame[\"_row_order\"] = np.arange(len(steps_frame))\n        duplicate_step_rows = int(steps_frame.duplicated(subset=[\"run_id\", \"step\"], keep=False).sum())\n        steps_frame = steps_frame.sort_values(\"_row_order\").drop_duplicates(subset=[\"run_id\", \"step\"], keep=\"last\")\n    if not runs_frame.empty:\n        runs_frame = runs_frame.reset_index(drop=True)\n        runs_frame[\"_row_order\"] = np.arange(len(runs_frame))\n        duplicate_run_rows = int(runs_frame.duplicated(subset=[\"run_id\"], keep=False).sum())\n        runs_frame = runs_frame.sort_values(\"_row_order\").drop_duplicates(subset=[\"run_id\"], keep=\"last\")\n\n    completed_from_steps: set[str] = set()\n    incomplete_step_run_ids: list[str] = []\n    reconstructed_runs: list[dict[str, Any]] = []\n    existing_run_ids = set(runs_frame[\"run_id\"].astype(str)) if not runs_frame.empty else set()\n\n    if not steps_frame.empty:\n        for run_id, group in steps_frame.groupby(\"run_id\", sort=False):\n            run_id_str = str(run_id)\n            ordered = group.sort_values(\"step\")\n            observed_steps = [int(value) for value in ordered[\"step\"].tolist()]\n            is_complete = observed_steps == expected_step_sequence\n            if is_complete:\n                completed_from_steps.add(run_id_str)\n                if run_id_str not in existing_run_ids:\n                    reconstructed_runs.append(summarize_run_rows(ordered))\n            else:\n                incomplete_step_run_ids.append(run_id_str)\n\n    runs_without_complete_steps = sorted(existing_run_ids - completed_from_steps)\n    hidden_without_complete_steps = sorted(hidden_run_ids - completed_from_steps)\n    completed_run_ids = set(completed_from_steps)\n\n    if not steps_frame.empty:\n        sanitized_steps = steps_frame[steps_frame[\"run_id\"].astype(str).isin(completed_run_ids)].copy()\n        sanitized_steps = sanitized_steps.sort_values([\"run_id\", \"step\", \"_row_order\"])\n        sanitized_steps = sanitized_steps.drop(columns=[\"_row_order\"], errors=\"ignore\")\n    else:\n        sanitized_steps = steps_frame\n\n    sanitized_runs_frames: list[pd.DataFrame] = []\n    if not runs_frame.empty:\n        sanitized_runs_frames.append(runs_frame[runs_frame[\"run_id\"].astype(str).isin(completed_run_ids)].copy())\n    if reconstructed_runs:\n        sanitized_runs_frames.append(pd.DataFrame(reconstructed_runs))\n    if sanitized_runs_frames:\n        sanitized_runs = pd.concat(sanitized_runs_frames, ignore_index=True, sort=False)\n        sanitized_runs[\"task_source_index\"] = pd.to_numeric(sanitized_runs[\"task_source_index\"], errors=\"coerce\").fillna(-1).astype(int)\n        sanitized_runs[\"temperature\"] = pd.to_numeric(sanitized_runs[\"temperature\"], errors=\"coerce\")\n        sanitized_runs[\"seed\"] = pd.to_numeric(sanitized_runs[\"seed\"], errors=\"coerce\").fillna(-1).astype(int)\n        sanitized_runs = sanitized_runs.sort_values([\"temperature\", \"seed\", \"task_source_index\", \"run_id\"])\n        sanitized_runs = sanitized_runs.drop(columns=[\"_row_order\"], errors=\"ignore\")\n    else:\n        sanitized_runs = pd.DataFrame(columns=runs_frame.columns.drop(\"_row_order\", errors=\"ignore\"))\n\n    anomalies_detected = any(\n        [\n            duplicate_step_rows,\n            duplicate_run_rows,\n            reconstructed_runs,\n            incomplete_step_run_ids,\n            runs_without_complete_steps,\n            hidden_without_complete_steps,\n        ]\n    )\n    reconciliation_report = {\n        \"expected_steps_per_run\": expected_steps,\n        \"completed_run_count\": len(completed_run_ids),\n        \"hidden_state_file_count\": len(hidden_run_ids),\n        \"duplicate_step_rows_removed\": duplicate_step_rows,\n        \"duplicate_run_rows_removed\": duplicate_run_rows,\n        \"reconstructed_run_summaries\": len(reconstructed_runs),\n        \"incomplete_step_run_ids\": incomplete_step_run_ids,\n        \"runs_without_complete_steps\": runs_without_complete_steps,\n        \"hidden_without_complete_steps\": hidden_without_complete_steps,\n        \"anomalies_detected\": bool(anomalies_detected),\n    }\n    if anomalies_detected:\n        if paths[\"steps\"].exists():\n            backup_for_reconciliation(paths[\"steps\"])\n        if paths[\"runs\"].exists():\n            backup_for_reconciliation(paths[\"runs\"])\n        sanitized_steps.to_csv(paths[\"steps\"], index=False)\n        sanitized_runs.to_csv(paths[\"runs\"], index=False)\n        with (output_dir / \"checkpoint_reconciliation.json\").open(\"w\", encoding=\"utf-8\") as handle:\n            json.dump(reconciliation_report, handle, indent=2)\n\n    return sanitized_steps.to_dict(\"records\"), sanitized_runs.to_dict(\"records\"), completed_run_ids, reconciliation_report\n\n\ndef load_existing_outputs(output_dir: Path, is_baseline: bool) -> tuple[list[dict[str, Any]], list[dict[str, Any]], set[str]]:\n    paths = output_paths(output_dir, is_baseline)\n    existing_steps = []\n    if paths[\"steps\"].exists():\n        try:\n            existing_steps = pd.read_csv(paths[\"steps\"]).to_dict(\"records\")\n        except pd.errors.EmptyDataError:\n            existing_steps = []\n    existing_runs = []\n    if paths[\"runs\"].exists():\n        try:\n            existing_runs = pd.read_csv(paths[\"runs\"]).to_dict(\"records\")\n        except pd.errors.EmptyDataError:\n            existing_runs = []\n    completed_run_ids = {str(row[\"run_id\"]) for row in existing_runs}\n    return existing_steps, existing_runs, completed_run_ids\n\n\ndef load_model(\n    model_spec: ModelSpec,\n    device: str,\n    quantization: str,\n    device_map: str | None,\n    attn_implementation: str,\n    model_path_override: str | None = None,\n    offload_folder: str | None = None,\n    model_revision: str | None = None,\n) -> tuple[Any, Any, str, str]:\n    actual_device = device if device != \"auto\" else (\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model_source = model_path_override or model_spec.hf_name\n    \n    # Precision Auto-Tuning for L4\n    target_precision = quantization\n    if quantization == \"auto\" and actual_device == \"cuda\":\n        # Extract numeric parameter count\n        try:\n            num_params = float(re.findall(r\"[\\d.]+\", model_spec.parameter_count)[0])\n            # If < 12B, use bfloat16 for speed on L4 24GB.\n            target_precision = \"bf16\" if num_params < 12 else \"4bit\"\n        except:\n            target_precision = \"4bit\"\n            \n    backend = \"transformers+torch(cpu)\"\n    load_kwargs: dict[str, Any] = {\"trust_remote_code\": True, \"low_cpu_mem_usage\": True}\n    if model_revision:\n        load_kwargs[\"revision\"] = model_revision\n    compute_dtype = torch.bfloat16 if actual_device == \"cuda\" and torch.cuda.is_bf16_supported() else torch.float16\n    dtype_label = \"bfloat16\" if compute_dtype == torch.bfloat16 else \"float16\"\n    logging.info(\"Precision auto-tuning: %s detected for %s model.\", dtype_label, model_spec.parameter_count)\n    \n    if actual_device == \"cuda\":\n        load_kwargs[\"torch_dtype\"] = compute_dtype\n        # Flash Attention 2 Enforcement\n        if attn_implementation == \"auto\" or attn_implementation == \"flash_attention_2\":\n            load_kwargs[\"attn_implementation\"] = \"flash_attention_2\"\n            logging.info(\"Flash Attention 2 explicitly enabled.\")\n        else:\n            load_kwargs[\"attn_implementation\"] = attn_implementation\n            logging.info(\"Attention implementation: %s (native FA2 kernels via SDPA).\", attn_implementation)\n\n    _MULTIMODAL_FAMILIES = [\"gemma 4\", \"llama 4\", \"qwen3.5\"]\n    _TEXT_ONLY_OVERRIDE = \"text-only\" in model_spec.family.lower()\n    is_multimodal = (\n        not _TEXT_ONLY_OVERRIDE\n        and any(f in model_spec.family.lower() for f in _MULTIMODAL_FAMILIES)\n    )\n    if is_multimodal:\n        from transformers import AutoModelForMultimodalLM\n        model_class = AutoModelForMultimodalLM\n    else:\n        model_class = AutoModelForCausalLM\n    \n    # Loading handle (Processor for multimodal, Tokenizer for vanilla)\n    if is_multimodal:\n        from transformers import AutoProcessor\n        processor_kwargs: dict[str, Any] = {\"trust_remote_code\": True}\n        if model_revision:\n            processor_kwargs[\"revision\"] = model_revision\n        processor = AutoProcessor.from_pretrained(model_source, **processor_kwargs)\n        # Use the inner tokenizer for text encoding; it also has apply_chat_template\n        tokenizer = getattr(processor, \"tokenizer\", processor)\n    else:\n        tokenizer_kwargs: dict[str, Any] = {\"trust_remote_code\": True}\n        if model_revision:\n            tokenizer_kwargs[\"revision\"] = model_revision\n        tokenizer = AutoTokenizer.from_pretrained(model_source, **tokenizer_kwargs)\n\n    if tokenizer.pad_token is None:\n        tokenizer.pad_token = tokenizer.eos_token\n    tokenizer.padding_side = \"left\"\n\n    resolved_device_map = resolve_low_vram_device_map(\n        model_spec=model_spec,\n        actual_device=actual_device,\n        target_precision=target_precision,\n        device_map=device_map,\n    )\n    if resolved_device_map is not None and resolved_device_map != device_map:\n        logging.info(\"Using explicit low-VRAM device map for %s: %s\", model_spec.alias, resolved_device_map)\n\n    try:\n        if actual_device == \"cuda\" and target_precision in {\"8bit\", \"4bit\"}:\n            if BitsAndBytesConfig is None:\n                raise ImportError(\"bitsandbytes support is unavailable in the installed transformers stack.\")\n            if target_precision == \"4bit\":\n                ensure_bitsandbytes_params4bit_compatibility()\n                ensure_accelerate_meta_state_dict_compatibility()\n            if target_precision == \"8bit\":\n                load_kwargs[\"quantization_config\"] = BitsAndBytesConfig(\n                    load_in_8bit=True,\n                    llm_int8_enable_fp32_cpu_offload=True,\n                )\n            else:\n                load_kwargs[\"quantization_config\"] = BitsAndBytesConfig(\n                    load_in_4bit=True,\n                    llm_int8_enable_fp32_cpu_offload=True,\n                    bnb_4bit_quant_type=\"nf4\",\n                    bnb_4bit_use_double_quant=True,\n                    bnb_4bit_compute_dtype=compute_dtype,\n                )\n            load_kwargs[\"device_map\"] = resolved_device_map or \"auto\"\n            if offload_folder:\n                load_kwargs[\"offload_folder\"] = offload_folder\n                load_kwargs[\"offload_state_dict\"] = True\n            try:\n                model = model_class.from_pretrained(model_source, **load_kwargs)\n            except ValueError as ve:\n                if \"flash_attention_2\" in str(ve):\n                    logging.warning(\"Flash Attention 2 not supported for this model. Falling back to SDPA.\")\n                    load_kwargs[\"attn_implementation\"] = \"sdpa\"\n                    model = model_class.from_pretrained(model_source, **load_kwargs)\n                else:\n                    raise\n            backend = f\"transformers+torch(cuda-{target_precision})\"\n        elif actual_device == \"cuda\" and resolved_device_map:\n            load_kwargs[\"device_map\"] = resolved_device_map\n            if offload_folder:\n                load_kwargs[\"offload_folder\"] = offload_folder\n                load_kwargs[\"offload_state_dict\"] = True\n            model = model_class.from_pretrained(model_source, **load_kwargs)\n            backend = f\"transformers+torch(cuda-device-map={resolved_device_map})\"\n        else:\n            model = model_class.from_pretrained(model_source, **load_kwargs)\n            model.to(actual_device)\n            backend = f\"transformers+torch({actual_device})\"\n    except Exception as exc:\n        logging.warning(\"Primary model load failed for %s: %s\", model_source, exc)\n        if (\n            actual_device == \"cuda\"\n            and target_precision == \"4bit\"\n            and \"Some modules are dispatched on the CPU or the disk\" in str(exc)\n        ):\n            logging.warning(\n                \"4-bit loading requested CPU offload but %s still does not fit cleanly on this GPU. \"\n                \"This usually means the card is too small for the model family even with quantization.\",\n                model_source,\n            )\n        if actual_device == \"cuda\":\n            raise\n        fallback_kwargs = {\"trust_remote_code\": True}\n        model = AutoModelForCausalLM.from_pretrained(model_source, **fallback_kwargs)\n        model.to(\"cpu\")\n        actual_device = \"cpu\"\n        backend = \"transformers+torch(cpu-fallback)\"\n\n    model.eval()\n    return model, tokenizer, actual_device, backend\n\n\ndef trim_completion_ids(completion_ids: torch.Tensor, pad_token_id: int | None, eos_token_id: int | None) -> torch.Tensor:\n    trimmed = completion_ids\n    trim_token_ids = {token_id for token_id in (pad_token_id, eos_token_id) if token_id is not None}\n    while len(trimmed) > 0 and int(trimmed[-1].item()) in trim_token_ids:\n        trimmed = trimmed[:-1]\n    return trimmed\n\n\ndef apply_output_head_with_model_postprocessing(model: Any, hidden_states: torch.Tensor) -> torch.Tensor:\n    output_embeddings = model.get_output_embeddings() if hasattr(model, \"get_output_embeddings\") else None\n    if output_embeddings is None:\n        raise AttributeError(\"Model does not expose output embeddings for low-memory scoring.\")\n\n    logits = output_embeddings(hidden_states)\n    config = getattr(model, \"config\", None)\n    if config is None:\n        return logits\n\n    text_config_getter = getattr(config, \"get_text_config\", None)\n    text_config = text_config_getter() if callable(text_config_getter) else config\n    final_logit_softcapping = getattr(text_config, \"final_logit_softcapping\", None)\n    if final_logit_softcapping is not None:\n        logits = logits / final_logit_softcapping\n        logits = torch.tanh(logits)\n        logits = logits * final_logit_softcapping\n    return logits\n\n\ndef low_memory_scoring_pass(\n    model: Any,\n    generated_ids: torch.Tensor,\n    generated_attention_mask: torch.Tensor,\n    prompt_width: int,\n    completion_width: int,\n    output_hidden_states: bool = False,\n) -> tuple[torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...] | None]:\n    base_model = getattr(model, \"model\", None)\n    if base_model is None:\n        base_model = getattr(model, \"base_model\", None)\n    if base_model is None:\n        raise AttributeError(\"Model does not expose a base model for low-memory scoring.\")\n\n    base_outputs = base_model(\n        input_ids=generated_ids,\n        attention_mask=generated_attention_mask,\n        use_cache=False,\n        return_dict=True,\n        output_hidden_states=output_hidden_states,\n    )\n    last_hidden_state = getattr(base_outputs, \"last_hidden_state\", None)\n    if last_hidden_state is None:\n        raise AttributeError(\"Base model did not return last_hidden_state for low-memory scoring.\")\n\n    scoring_start = prompt_width - 1\n    scoring_stop = scoring_start + completion_width\n    scoring_hidden_states = last_hidden_state[:, scoring_start:scoring_stop, :]\n    scoring_logits = apply_output_head_with_model_postprocessing(model, scoring_hidden_states)\n    \n    all_hidden_states = getattr(base_outputs, \"hidden_states\", None)\n    return scoring_logits, last_hidden_state, all_hidden_states\n\n\n_PROJECTION_MATRIX_CACHE = {}\n\ndef project_hidden_features(hidden_vector: Any, target_dim: int = 64, seed: int = 42) -> np.ndarray:\n    \"\"\"Projects a hidden state vector/tensor to a lower dimension using a cached CPU projection matrix.\"\"\"\n    if isinstance(hidden_vector, torch.Tensor):\n        hidden_vector = hidden_vector.float().cpu().numpy()\n        \n    D = hidden_vector.shape[0]\n    cache_key = (D, target_dim, seed)\n    \n    if cache_key not in _PROJECTION_MATRIX_CACHE:\n        rng = np.random.RandomState(seed)\n        _PROJECTION_MATRIX_CACHE[cache_key] = rng.normal(size=(D, target_dim)) / np.sqrt(target_dim)\n        \n    proj_matrix = _PROJECTION_MATRIX_CACHE[cache_key]\n    return np.dot(hidden_vector, proj_matrix)\n\n\ndef find_answer_token_indices(completion_ids: torch.Tensor, raw_text: str, answer_str: str, prompt_mode: str, tokenizer: Any) -> list[int]:\n    \"\"\"Finds the indices of tokens in completion_ids that correspond to the answer_str.\"\"\"\n    if not answer_str:\n        return []\n        \n    char_start = -1\n    if prompt_mode == \"structured_four_line\":\n        # Search for ANSWER: <answer_str>\n        match = re.search(r\"ANSWER\\s*:\\s*\" + re.escape(answer_str), raw_text, re.IGNORECASE)\n        if match:\n            char_start = match.end(0) - len(answer_str)\n    elif prompt_mode == \"minimal_json\":\n        # Search for \"answer\"\\s*:\\s*\"<answer_str>\"\n        match = re.search(r'\"answer\"\\s*:\\s*\"' + re.escape(answer_str) + r'\"', raw_text)\n        if match:\n            char_start = match.end(0) - len(answer_str) - 1 # exclude closing quote\n            \n    if char_start == -1:\n        char_start = raw_text.rfind(answer_str)\n        \n    if char_start == -1:\n        return []\n        \n    char_end = char_start + len(answer_str)\n    \n    token_indices = []\n    for idx in range(len(completion_ids)):\n        prefix_text_prev = tokenizer.decode(completion_ids[:idx], skip_special_tokens=True)\n        prefix_text_curr = tokenizer.decode(completion_ids[:idx+1], skip_special_tokens=True)\n        tok_start = len(prefix_text_prev)\n        tok_end = len(prefix_text_curr)\n        # Check overlap\n        if max(tok_start, char_start) < min(tok_end, char_end):\n            token_indices.append(idx)\n    return token_indices\n\n\ndef generate_batch_with_diagnostics(\n    model: Any,\n    tokenizer: Any,\n    prompt_texts: list[str],\n    actual_device: str,\n    temperature: float,\n    max_new_tokens: int,\n    answer_types: list[str] | None = None,\n    prompt_mode: str = \"structured_four_line\",\n    enable_extended_observables: bool = False,\n) -> tuple[list[dict[str, Any]], dict[str, Any]]:\n    if actual_device == \"cuda\" and torch.cuda.is_available():\n        torch.cuda.synchronize()\n        torch.cuda.reset_peak_memory_stats()\n\n    started_at = time.perf_counter()\n    encode_started_at = time.perf_counter()\n    encoded = tokenizer(\n        prompt_texts,\n        padding=True,\n        pad_to_multiple_of=8 if actual_device == \"cuda\" else None,\n        return_tensors=\"pt\",\n    )\n    if actual_device == \"cuda\" and torch.cuda.is_available():\n        torch.cuda.synchronize()\n    encode_seconds = time.perf_counter() - encode_started_at\n\n    prompt_tokens = int(encoded[\"attention_mask\"].sum().item())\n    input_ids = encoded[\"input_ids\"].to(actual_device)\n    attention_mask = encoded[\"attention_mask\"].to(actual_device)\n    prompt_width = input_ids.shape[1]\n\n    generation_kwargs = {\n        \"input_ids\": input_ids,\n        \"attention_mask\": attention_mask,\n        \"max_new_tokens\": max_new_tokens,\n        \"pad_token_id\": tokenizer.pad_token_id,\n        \"eos_token_id\": tokenizer.eos_token_id,\n        \"repetition_penalty\": 1.02,\n    }\n    if temperature > 0:\n        generation_kwargs.update({\"do_sample\": True, \"temperature\": temperature, \"top_p\": 0.95})\n    else:\n        generation_kwargs.update({\"do_sample\": False})\n\n    with torch.inference_mode():\n        generation_started_at = time.perf_counter()\n        generated_ids = model.generate(**generation_kwargs)\n        if actual_device == \"cuda\" and torch.cuda.is_available():\n            torch.cuda.synchronize()\n        generation_seconds = time.perf_counter() - generation_started_at\n\n        forward_started_at = time.perf_counter()\n        completion_width = generated_ids.shape[1] - prompt_width\n        generated_attention_mask = torch.cat(\n            [\n                attention_mask,\n                torch.ones((attention_mask.shape[0], completion_width), dtype=attention_mask.dtype, device=attention_mask.device),\n            ],\n            dim=1,\n        )\n        try:\n            scoring_logits, final_hidden_states, all_hidden_states = low_memory_scoring_pass(\n                model=model,\n                generated_ids=generated_ids,\n                generated_attention_mask=generated_attention_mask,\n                prompt_width=prompt_width,\n                completion_width=completion_width,\n                output_hidden_states=enable_extended_observables,\n            )\n        except Exception as exc:\n            logging.warning(\"Low-memory scoring path failed; falling back to wrapper forward: %s\", exc)\n            forward_outputs = model(\n                generated_ids,\n                attention_mask=generated_attention_mask,\n                output_hidden_states=True,\n                use_cache=False,\n            )\n            scoring_logits = forward_outputs.logits[:, prompt_width - 1 : prompt_width - 1 + completion_width, :]\n            final_hidden_states = forward_outputs.hidden_states[-1]\n            all_hidden_states = forward_outputs.hidden_states\n        if actual_device == \"cuda\" and torch.cuda.is_available():\n            torch.cuda.synchronize()\n        forward_seconds = time.perf_counter() - forward_started_at\n\n    # Postprocess on CPU to avoid expanding completion logits into softmax buffers on a constrained GPU.\n    generated_ids_cpu = generated_ids.detach().cpu()\n    scoring_logits = scoring_logits.float().cpu()\n    final_hidden_states = final_hidden_states.float().cpu()\n    \n    # Perform GPU-side mean pooling BEFORE clearing the cache.\n    # This reduces the CPU copy size from 1 GB (entire layer) to 1 MB (mean pooled vector).\n    L1_pooled_list = []\n    L2_pooled_list = []\n    if enable_extended_observables and all_hidden_states is not None:\n        num_layers = len(all_hidden_states) - 1\n        L1_idx = num_layers // 3\n        L2_idx = 2 * num_layers // 3\n        L1_layer = all_hidden_states[L1_idx + 1]\n        L2_layer = all_hidden_states[L2_idx + 1]\n        \n        batch_size = scoring_logits.shape[0]\n        for index in range(batch_size):\n            raw_completion_ids = generated_ids_cpu[index, prompt_width:]\n            completion_ids = trim_completion_ids(raw_completion_ids, tokenizer.pad_token_id, tokenizer.eos_token_id)\n            generated_length = int(completion_ids.shape[0])\n            \n            if generated_length > 0:\n                L1_pooled = L1_layer[index, prompt_width : prompt_width + generated_length, :].mean(dim=0).float().cpu().numpy()\n                L2_pooled = L2_layer[index, prompt_width : prompt_width + generated_length, :].mean(dim=0).float().cpu().numpy()\n            else:\n                L1_pooled = np.zeros((model.config.hidden_size,), dtype=np.float32)\n                L2_pooled = np.zeros((model.config.hidden_size,), dtype=np.float32)\n                \n            L1_pooled_list.append(L1_pooled)\n            L2_pooled_list.append(L2_pooled)\n            \n    if actual_device == \"cuda\":\n        del input_ids\n        del attention_mask\n        del generated_ids\n        del generated_attention_mask\n        if \"all_hidden_states\" in locals():\n            del all_hidden_states\n        release_cuda_memory()\n\n    results: list[dict[str, Any]] = []\n    postprocess_started_at = time.perf_counter()\n    batch_size = scoring_logits.shape[0]\n    for index in range(batch_size):\n        raw_completion_ids = generated_ids_cpu[index, prompt_width:]\n        completion_ids = trim_completion_ids(raw_completion_ids, tokenizer.pad_token_id, tokenizer.eos_token_id)\n        generated_length = int(completion_ids.shape[0])\n        raw_text = tokenizer.decode(completion_ids, skip_special_tokens=True).strip()\n\n        # Default extended observables\n        answer_span_mean_logprob = float(\"nan\")\n        answer_span_min_logprob = float(\"nan\")\n        answer_span_mean_entropy = float(\"nan\")\n        answer_span_std_entropy = 0.0\n        mid_hidden_1_proj_str = \"\"\n        mid_hidden_2_proj_str = \"\"\n\n        if generated_length > 0:\n            completion_scoring_logits = scoring_logits[index, :generated_length]\n            log_probs = torch.log_softmax(completion_scoring_logits, dim=-1)\n            probs = torch.softmax(completion_scoring_logits, dim=-1)\n            token_logprobs = log_probs.gather(1, completion_ids.unsqueeze(1)).squeeze(1)\n            token_entropies = -(probs * log_probs).sum(dim=-1)\n            pooled_hidden = final_hidden_states[index, prompt_width : prompt_width + generated_length, :].mean(dim=0).float().cpu().numpy()\n            \n            if enable_extended_observables and answer_types is not None:\n                # N8a: Answer-span token diagnostics\n                ans_type = answer_types[index]\n                parsed = parse_generation(raw_text, ans_type, prompt_mode)\n                ans_str = parsed[\"answer\"]\n                ans_tokens = find_answer_token_indices(completion_ids, raw_text, ans_str, prompt_mode, tokenizer)\n                \n                if ans_tokens:\n                    ans_logprobs = token_logprobs[ans_tokens]\n                    ans_entropies = token_entropies[ans_tokens]\n                    \n                    answer_span_mean_logprob = float(ans_logprobs.mean().item())\n                    answer_span_min_logprob = float(ans_logprobs.min().item())\n                    answer_span_mean_entropy = float(ans_entropies.mean().item())\n                    answer_span_std_entropy = float(ans_entropies.std().item()) if len(ans_entropies) > 1 else 0.0\n                \n                # N8b: Pooled mid-layer hidden states projection\n                if L1_pooled_list and L2_pooled_list:\n                    L1_tensor = L1_pooled_list[index]\n                    L2_tensor = L2_pooled_list[index]\n                    \n                    L1_proj = project_hidden_features(L1_tensor, target_dim=64, seed=42)\n                    L2_proj = project_hidden_features(L2_tensor, target_dim=64, seed=42)\n                    \n                    mid_hidden_1_proj_str = \",\".join(f\"{x:.6f}\" for x in L1_proj)\n                    mid_hidden_2_proj_str = \",\".join(f\"{x:.6f}\" for x in L2_proj)\n        else:\n            token_logprobs = torch.empty(0)\n            token_entropies = torch.empty(0)\n            pooled_hidden = np.zeros((model.config.hidden_size,), dtype=np.float32)\n\n        results.append(\n            {\n                \"raw_text\": raw_text,\n                \"generated_tokens\": generated_length,\n                \"mean_token_logprob\": float(token_logprobs.mean().item()) if len(token_logprobs) else float(\"nan\"),\n                \"mean_entropy\": float(token_entropies.mean().item()) if len(token_entropies) else float(\"nan\"),\n                \"entropy_std\": float(token_entropies.std().item()) if len(token_entropies) > 1 else 0.0,\n                \"pooled_hidden\": pooled_hidden,\n                \"answer_span_mean_logprob\": answer_span_mean_logprob,\n                \"answer_span_min_logprob\": answer_span_min_logprob,\n                \"answer_span_mean_entropy\": answer_span_mean_entropy,\n                \"answer_span_std_entropy\": answer_span_std_entropy,\n                \"mid_hidden_1_proj\": mid_hidden_1_proj_str,\n                \"mid_hidden_2_proj\": mid_hidden_2_proj_str,\n            }\n        )\n    if actual_device == \"cuda\" and torch.cuda.is_available():\n        torch.cuda.synchronize()\n    postprocess_seconds = time.perf_counter() - postprocess_started_at\n    elapsed_seconds = time.perf_counter() - started_at\n    generated_tokens = int(sum(item[\"generated_tokens\"] for item in results))\n    batch_metrics = {\n        \"batch_size\": len(prompt_texts),\n        \"prompt_tokens\": prompt_tokens,\n        \"generated_tokens\": generated_tokens,\n        \"total_seconds\": elapsed_seconds,\n        \"tokenize_seconds\": encode_seconds,\n        \"generation_seconds\": generation_seconds,\n        \"forward_seconds\": forward_seconds,\n        \"postprocess_seconds\": postprocess_seconds,\n        \"gpu_memory_allocated_gb\": gpu_memory_allocated_gb(actual_device),\n        \"gpu_memory_reserved_gb\": gpu_memory_reserved_gb(actual_device),\n        \"gpu_max_memory_allocated_gb\": gpu_max_memory_allocated_gb(actual_device),\n        \"gpu_max_memory_reserved_gb\": gpu_max_memory_reserved_gb(actual_device),\n        \"split_count\": 1,\n        \"oom_retry_count\": 0,\n    }\n    return results, batch_metrics\n\n\ndef release_cuda_memory() -> None:\n    gc.collect()\n    if not torch.cuda.is_available():\n        return\n    torch.cuda.empty_cache()\n    if hasattr(torch.cuda, \"ipc_collect\"):\n        torch.cuda.ipc_collect()\n\n\ndef safe_generate_batch_with_diagnostics(\n    model: Any,\n    tokenizer: Any,\n    prompt_texts: list[str],\n    actual_device: str,\n    temperature: float,\n    max_new_tokens: int,\n    allow_single_retry: bool = True,\n    answer_types: list[str] | None = None,\n    prompt_mode: str = \"structured_four_line\",\n    enable_extended_observables: bool = False,\n) -> tuple[list[dict[str, Any]], dict[str, Any]]:\n    oom_occurred = False\n    try:\n        return generate_batch_with_diagnostics(\n            model=model,\n            tokenizer=tokenizer,\n            prompt_texts=prompt_texts,\n            actual_device=actual_device,\n            temperature=temperature,\n            max_new_tokens=max_new_tokens,\n            answer_types=answer_types,\n            prompt_mode=prompt_mode,\n            enable_extended_observables=enable_extended_observables,\n        )\n    except RuntimeError as exc:\n        if actual_device != \"cuda\" or \"out of memory\" not in str(exc).lower():\n            raise\n        oom_occurred = True\n\n    if oom_occurred:\n        import gc\n        release_cuda_memory()\n        gc.collect()\n\n        if len(prompt_texts) == 1:\n            if not allow_single_retry:\n                raise RuntimeError(\"CUDA out of memory for single prompt; retry disabled or already attempted.\")\n            logging.warning(\"CUDA OOM for single prompt. Retrying once after cache release.\")\n            retry_results, retry_metrics = safe_generate_batch_with_diagnostics(\n                model=model,\n                tokenizer=tokenizer,\n                prompt_texts=prompt_texts,\n                actual_device=actual_device,\n                temperature=temperature,\n                max_new_tokens=max_new_tokens,\n                allow_single_retry=False,\n                answer_types=answer_types,\n                prompt_mode=prompt_mode,\n                enable_extended_observables=enable_extended_observables,\n            )\n            retry_metrics[\"oom_retry_count\"] = int(retry_metrics.get(\"oom_retry_count\", 0)) + 1\n            return retry_results, retry_metrics\n\n        logging.warning(\"CUDA OOM for microbatch size %d. Retrying with smaller splits.\", len(prompt_texts))\n        midpoint = max(1, len(prompt_texts) // 2)\n        first_results, first_metrics = safe_generate_batch_with_diagnostics(\n            model=model,\n            tokenizer=tokenizer,\n            prompt_texts=prompt_texts[:midpoint],\n            actual_device=actual_device,\n            temperature=temperature,\n            max_new_tokens=max_new_tokens,\n            allow_single_retry=allow_single_retry,\n            answer_types=answer_types[:midpoint] if answer_types else None,\n            prompt_mode=prompt_mode,\n            enable_extended_observables=enable_extended_observables,\n        )\n        release_cuda_memory()\n        gc.collect()\n\n        second_results, second_metrics = safe_generate_batch_with_diagnostics(\n            model=model,\n            tokenizer=tokenizer,\n            prompt_texts=prompt_texts[midpoint:],\n            actual_device=actual_device,\n            temperature=temperature,\n            max_new_tokens=max_new_tokens,\n            allow_single_retry=allow_single_retry,\n            answer_types=answer_types[midpoint:] if answer_types else None,\n            prompt_mode=prompt_mode,\n            enable_extended_observables=enable_extended_observables,\n        )\n        merged_metrics = {\n            \"batch_size\": int(first_metrics[\"batch_size\"]) + int(second_metrics[\"batch_size\"]),\n            \"prompt_tokens\": int(first_metrics[\"prompt_tokens\"]) + int(second_metrics[\"prompt_tokens\"]),\n            \"generated_tokens\": int(first_metrics[\"generated_tokens\"]) + int(second_metrics[\"generated_tokens\"]),\n            \"total_seconds\": float(first_metrics[\"total_seconds\"]) + float(second_metrics[\"total_seconds\"]),\n            \"tokenize_seconds\": float(first_metrics[\"tokenize_seconds\"]) + float(second_metrics[\"tokenize_seconds\"]),\n            \"generation_seconds\": float(first_metrics[\"generation_seconds\"]) + float(second_metrics[\"generation_seconds\"]),\n            \"forward_seconds\": float(first_metrics[\"forward_seconds\"]) + float(second_metrics[\"forward_seconds\"]),\n            \"postprocess_seconds\": float(first_metrics[\"postprocess_seconds\"]) + float(second_metrics[\"postprocess_seconds\"]),\n            \"gpu_memory_allocated_gb\": max_nan([\n                float(first_metrics[\"gpu_memory_allocated_gb\"]),\n                float(second_metrics[\"gpu_memory_allocated_gb\"]),\n            ]),\n            \"gpu_memory_reserved_gb\": max_nan([\n                float(first_metrics[\"gpu_memory_reserved_gb\"]),\n                float(second_metrics[\"gpu_memory_reserved_gb\"]),\n            ]),\n            \"gpu_max_memory_allocated_gb\": max_nan([\n                float(first_metrics[\"gpu_max_memory_allocated_gb\"]),\n                float(second_metrics[\"gpu_max_memory_allocated_gb\"]),\n            ]),\n            \"gpu_max_memory_reserved_gb\": max_nan([\n                float(first_metrics[\"gpu_max_memory_reserved_gb\"]),\n                float(second_metrics[\"gpu_max_memory_reserved_gb\"]),\n            ]),\n            \"split_count\": int(first_metrics.get(\"split_count\", 1)) + int(second_metrics.get(\"split_count\", 1)),\n            \"oom_retry_count\": int(first_metrics.get(\"oom_retry_count\", 0)) + int(second_metrics.get(\"oom_retry_count\", 0)) + 1,\n        }\n        return first_results + second_results, merged_metrics\n\n\n\ndef gpu_memory_allocated_gb(actual_device: str) -> float:\n    if actual_device != \"cuda\" or not torch.cuda.is_available():\n        return float(\"nan\")\n    return torch.cuda.memory_allocated() / (1024**3)\n\n\ndef gpu_memory_reserved_gb(actual_device: str) -> float:\n    if actual_device != \"cuda\" or not torch.cuda.is_available():\n        return float(\"nan\")\n    return torch.cuda.memory_reserved() / (1024**3)\n\n\ndef gpu_max_memory_allocated_gb(actual_device: str) -> float:\n    if actual_device != \"cuda\" or not torch.cuda.is_available():\n        return float(\"nan\")\n    return torch.cuda.max_memory_allocated() / (1024**3)\n\n\ndef gpu_max_memory_reserved_gb(actual_device: str) -> float:\n    if actual_device != \"cuda\" or not torch.cuda.is_available():\n        return float(\"nan\")\n    return torch.cuda.max_memory_reserved() / (1024**3)\n\n\ndef first_pending_run(\n    model_spec: ModelSpec,\n    tasks: list[TaskSpec],\n    temperatures: list[float],\n    seeds: list[int],\n    completed_run_ids: set[str],\n) -> dict[str, Any] | None:\n    for temperature in temperatures:\n        for seed in seeds:\n            for task in tasks:\n                run_id = run_id_for(model_spec.alias, task, temperature, seed)\n                if run_id in completed_run_ids:\n                    continue\n                return {\n                    \"run_id\": run_id,\n                    \"task_id\": task.task_id,\n                    \"task_source_index\": task.source_index,\n                    \"temperature\": temperature,\n                    \"seed\": seed,\n                }\n    return None\n\n\ndef write_runtime_metadata(\n    *,\n    metadata_path: Path,\n    model_spec: ModelSpec,\n    backend: str,\n    actual_device: str,\n    quantization: str,\n    device_map: str | None,\n    attn_implementation: str,\n    temperatures: list[float],\n    seeds: list[int],\n    max_steps: int,\n    max_new_tokens: int,\n    max_tasks: int,\n    task_source: str,\n    dataset_split: str,\n    dataset_shuffle_seed: int,\n    batch_size: int,\n    prompt_mode: str,\n    system_prompt_mode: str,\n    resume_enabled: bool,\n    completed_run_ids: set[str],\n    pending_run_count: int,\n    next_pending_run: dict[str, Any] | None,\n    reconciliation_report: dict[str, Any],\n    tasks: list[TaskSpec],\n    model_source: str | None,\n) -> None:\n    metadata = {\n        \"model\": asdict(model_spec),\n        \"backend\": backend,\n        \"model_source\": model_source or model_spec.hf_name,\n        \"device\": actual_device,\n        \"quantization\": quantization,\n        \"device_map\": device_map,\n        \"attn_implementation\": attn_implementation,\n        \"temperatures\": temperatures,\n        \"seeds\": seeds,\n        \"max_steps\": max_steps,\n        \"max_new_tokens\": max_new_tokens,\n        \"max_tasks\": max_tasks,\n        \"task_source\": task_source,\n        \"dataset_split\": dataset_split,\n        \"dataset_shuffle_seed\": dataset_shuffle_seed,\n        \"batch_size\": batch_size,\n        \"step_cost\": STEP_COST,\n        \"prompt_mode\": prompt_mode,\n        \"system_prompt_mode\": system_prompt_mode,\n        \"resume_enabled\": resume_enabled,\n        \"completed_run_count\": len(completed_run_ids),\n        \"pending_run_count\": pending_run_count,\n        \"next_pending_run\": next_pending_run,\n        \"checkpoint_reconciliation\": reconciliation_report,\n        \"tasks\": [asdict(task) for task in tasks],\n        \"hidden_state_accessible\": True,\n        \"token_logprobs_accessible\": True,\n        \"reasoning_traces_accessible\": True,\n        \"batch_metrics_accessible\": True,\n        \"runtime_limit_notes\": \"Batch-first GSM8K runner with checkpoint reconciliation, append-only CSV outputs, incremental metadata snapshots, and automatic microbatch splitting on CUDA OOM.\",\n    }\n    with metadata_path.open(\"w\", encoding=\"utf-8\") as handle:\n        json.dump(metadata, handle, indent=2)\n\n\ndef run_batch_traces(\n    model: Any,\n    tokenizer: Any,\n    model_spec: ModelSpec,\n    tasks: list[TaskSpec],\n    actual_device: str,\n    temperature: float,\n    seed: int,\n    max_steps: int,\n    max_new_tokens: int,\n    prompt_mode: str,\n    system_prompt_mode: str,\n    is_baseline: bool,\n    hidden_dir: Path,\n    io_manager: IOManager | None = None,\n    enable_k2_agreement: bool = False,\n    enable_extended_observables: bool = False,\n) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]]]:\n    if not tasks:\n        return [], [], []\n\n    actual_prompt_mode = \"answer_only\" if is_baseline else prompt_mode\n    steps_to_run = 1 if is_baseline else max_steps\n    batch_metric_rows: list[dict[str, Any]] = []\n    contexts = [\n        {\n            \"task\": task,\n            \"run_id\": run_id_for(model_spec.alias, task, temperature, seed),\n            \"history\": [],\n            \"hidden_vectors\": [],\n            \"rows\": [],\n        }\n        for task in tasks\n    ]\n\n    for step in range(1, steps_to_run + 1):\n        prompts = [\n            render_prompt(\n                tokenizer,\n                conversation_prompt(\n                    task=context[\"task\"],\n                    history=context[\"history\"],\n                    step=step,\n                    max_steps=max_steps,\n                    prompt_mode=actual_prompt_mode,\n                ),\n                system_prompt_mode,\n            )\n            for context in contexts\n        ]\n\n        # First path generation (main trajectory)\n        generated_batch, batch_metrics = safe_generate_batch_with_diagnostics(\n            model=model,\n            tokenizer=tokenizer,\n            prompt_texts=prompts,\n            actual_device=actual_device,\n            temperature=temperature,\n            max_new_tokens=max_new_tokens,\n            answer_types=[context[\"task\"].answer_type for context in contexts],\n            prompt_mode=actual_prompt_mode,\n            enable_extended_observables=enable_extended_observables,\n        )\n        \n        # Second path generation (N7: k=2 agreement)\n        k2_batch = None\n        if enable_k2_agreement and not is_baseline:\n            k2_batch, _ = safe_generate_batch_with_diagnostics(\n                model=model,\n                tokenizer=tokenizer,\n                prompt_texts=prompts,\n                actual_device=actual_device,\n                temperature=temperature,\n                max_new_tokens=max_new_tokens,\n                answer_types=None,\n                prompt_mode=actual_prompt_mode,\n                enable_extended_observables=False,\n            )\n\n        batch_elapsed_seconds = float(batch_metrics[\"total_seconds\"])\n        batch_tokens = sum(item[\"generated_tokens\"] for item in generated_batch)\n        if k2_batch is not None:\n            batch_tokens += sum(item[\"generated_tokens\"] for item in k2_batch)\n            \n        batch_examples_per_second = len(contexts) / max(batch_elapsed_seconds, 1e-6)\n        batch_tokens_per_second = batch_tokens / max(batch_elapsed_seconds, 1e-6)\n        batch_metric_rows.append(\n            {\n                \"phase\": \"generate\",\n                \"temperature\": temperature,\n                \"seed\": seed,\n                \"step\": step,\n                \"requested_batch_size\": len(contexts),\n                \"realized_batch_size\": int(batch_metrics[\"batch_size\"]),\n                \"prompt_tokens\": int(batch_metrics[\"prompt_tokens\"]),\n                \"generated_tokens\": batch_tokens,\n                \"wall_clock_seconds\": batch_elapsed_seconds,\n                \"examples_per_second\": batch_examples_per_second,\n                \"tokens_per_second\": batch_tokens_per_second,\n                \"tokenize_seconds\": float(batch_metrics[\"tokenize_seconds\"]),\n                \"generation_seconds\": float(batch_metrics[\"generation_seconds\"]),\n                \"forward_seconds\": float(batch_metrics[\"forward_seconds\"]),\n                \"postprocess_seconds\": float(batch_metrics[\"postprocess_seconds\"]),\n                \"hidden_state_write_seconds\": 0.0,\n                \"gpu_memory_allocated_gb\": float(batch_metrics[\"gpu_memory_allocated_gb\"]),\n                \"gpu_memory_reserved_gb\": float(batch_metrics[\"gpu_memory_reserved_gb\"]),\n                \"gpu_max_memory_allocated_gb\": float(batch_metrics[\"gpu_max_memory_allocated_gb\"]),\n                \"gpu_max_memory_reserved_gb\": float(batch_metrics[\"gpu_max_memory_reserved_gb\"]),\n                \"split_count\": int(batch_metrics.get(\"split_count\", 1)),\n                \"oom_retry_count\": int(batch_metrics.get(\"oom_retry_count\", 0)),\n            }\n        )\n        logging.info(\n            \"temp=%.2f seed=%d step=%d/%d | batch=%d | generated_tokens=%d | ex_s=%.2f | tok_s=%.2f | gpu_peak_alloc_gb=%.2f | gpu_peak_reserved_gb=%.2f | splits=%d | oom_retries=%d\",\n            temperature,\n            seed,\n            step,\n            steps_to_run,\n            len(contexts),\n            batch_tokens,\n            batch_examples_per_second,\n            batch_tokens_per_second,\n            float(batch_metrics[\"gpu_max_memory_allocated_gb\"]),\n            float(batch_metrics[\"gpu_max_memory_reserved_gb\"]),\n            int(batch_metrics.get(\"split_count\", 1)),\n            int(batch_metrics.get(\"oom_retry_count\", 0)),\n        )\n\n        for idx, (context, generated) in enumerate(zip(contexts, generated_batch, strict=True)):\n            history = context[\"history\"]\n            parsed = parse_generation(generated[\"raw_text\"], context[\"task\"].answer_type, actual_prompt_mode)\n            normalized_answer = normalize_answer(parsed[\"answer\"], context[\"task\"].answer_type)\n            is_correct = verify_answer(context[\"task\"], parsed[\"answer\"])\n            prior_answer = history[-1][\"answer_normalized\"] if history else \"\"\n            answer_changed = bool(history) and normalized_answer != prior_answer and normalized_answer != \"\"\n            prior_thought = \" \".join(item[\"thought\"] for item in history[-2:])\n\n            hidden_vector = generated[\"pooled_hidden\"]\n            if context[\"hidden_vectors\"]:\n                previous_hidden = context[\"hidden_vectors\"][-1]\n                denominator = float(np.linalg.norm(hidden_vector) * np.linalg.norm(previous_hidden))\n                hidden_cosine_shift = 1.0 - float(np.dot(hidden_vector, previous_hidden) / denominator) if denominator else 0.0\n                hidden_l2_shift = float(np.linalg.norm(hidden_vector - previous_hidden))\n            else:\n                hidden_cosine_shift = 0.0\n                hidden_l2_shift = 0.0\n            context[\"hidden_vectors\"].append(hidden_vector)\n\n            hit_max_new_tokens = int(generated[\"generated_tokens\"] == max_new_tokens)\n            truncated_output_suspected = int(hit_max_new_tokens and not parsed[\"parse_success\"])\n            tokens_per_second = generated[\"generated_tokens\"] / max(batch_elapsed_seconds, 1e-6)\n\n            # N7: k=2 agreement extraction\n            k2_agreement = -1\n            k2_tokens = 0\n            if k2_batch is not None:\n                k2_gen = k2_batch[idx]\n                parsed_k2 = parse_generation(k2_gen[\"raw_text\"], context[\"task\"].answer_type, actual_prompt_mode)\n                normalized_k2 = normalize_answer(parsed_k2[\"answer\"], context[\"task\"].answer_type)\n                if context[\"task\"].answer_type == \"math\":\n                    k2_agreement = int(math_answers_equivalent(normalized_answer, normalized_k2))\n                else:\n                    k2_agreement = int(normalized_answer == normalized_k2)\n                k2_tokens = k2_gen[\"generated_tokens\"]\n\n            row = {\n                \"run_id\": context[\"run_id\"],\n                \"model_alias\": model_spec.alias,\n                \"model_name\": model_spec.hf_name,\n                \"task_id\": context[\"task\"].task_id,\n                \"domain\": context[\"task\"].domain,\n                \"difficulty\": context[\"task\"].difficulty,\n                \"task_source\": context[\"task\"].source,\n                \"task_source_index\": context[\"task\"].source_index,\n                \"expected_answer\": context[\"task\"].expected_answer,\n                \"step\": step,\n                \"thought\": parsed[\"thought\"],\n                \"answer\": parsed[\"answer\"],\n                \"answer_normalized\": normalized_answer,\n                \"correct\": int(is_correct),\n                \"confidence\": int(max(0, min(100, parsed[\"confidence\"]))),\n                \"model_stop_flag\": int(bool(parsed[\"stop\"])),\n                \"answer_changed\": int(answer_changed),\n                \"thought_token_count\": int(len(re.findall(r\"\\w+\", parsed[\"thought\"]))),\n                \"raw_generation_tokens\": generated[\"generated_tokens\"],\n                \"mean_token_logprob\": generated[\"mean_token_logprob\"],\n                \"entropy_mean\": generated[\"mean_entropy\"],\n                \"entropy_std\": generated[\"entropy_std\"],\n                \"hidden_norm\": float(np.linalg.norm(hidden_vector)),\n                \"hidden_l2_shift\": hidden_l2_shift,\n                \"hidden_cosine_shift\": hidden_cosine_shift,\n                \"lexical_echo\": lexical_overlap(parsed[\"thought\"], prior_thought),\n                \"verbose_confidence_proxy\": float(parsed[\"confidence\"]) / 100.0 + 0.01 * generated[\"generated_tokens\"],\n                \"utility\": utility(correct=is_correct, step=step),\n                \"elapsed_seconds\": batch_elapsed_seconds,\n                \"tokens_per_second\": tokens_per_second,\n                \"gpu_memory_allocated_gb\": gpu_memory_allocated_gb(actual_device),\n                \"seed\": seed,\n                \"temperature\": temperature,\n                \"device\": actual_device,\n                \"prompt_mode\": actual_prompt_mode,\n                \"system_prompt_mode\": system_prompt_mode,\n                \"is_baseline\": int(is_baseline),\n                \"parse_success\": parsed[\"parse_success\"],\n                \"output_format_type\": parsed[\"output_format_type\"],\n                \"answer_extraction_source\": parsed[\"answer_extraction_source\"],\n                \"stop_extraction_source\": parsed[\"stop_extraction_source\"],\n                \"confidence_extraction_source\": parsed[\"confidence_extraction_source\"],\n                \"hit_max_new_tokens\": hit_max_new_tokens,\n                \"truncated_output_suspected\": truncated_output_suspected,\n                \"raw_text_length_chars\": len(generated[\"raw_text\"]),\n                \"raw_text_length_tokens\": generated[\"generated_tokens\"],\n                \"raw_text\": generated[\"raw_text\"],\n                # N7 Columns\n                \"k2_agreement\": k2_agreement,\n                \"k2_raw_generation_tokens\": k2_tokens,\n                # N8 Columns\n                \"answer_span_mean_logprob\": generated.get(\"answer_span_mean_logprob\", float(\"nan\")),\n                \"answer_span_min_logprob\": generated.get(\"answer_span_min_logprob\", float(\"nan\")),\n                \"answer_span_mean_entropy\": generated.get(\"answer_span_mean_entropy\", float(\"nan\")),\n                \"answer_span_std_entropy\": generated.get(\"answer_span_std_entropy\", 0.0),\n                \"mid_hidden_1_proj\": generated.get(\"mid_hidden_1_proj\", \"\"),\n                \"mid_hidden_2_proj\": generated.get(\"mid_hidden_2_proj\", \"\"),\n            }\n            history.append(row)\n            context[\"rows\"].append(row)\n\n    hidden_dir.mkdir(parents=True, exist_ok=True)\n    all_rows: list[dict[str, Any]] = []\n    all_runs: list[dict[str, Any]] = []\n    hidden_state_write_seconds = 0.0\n    for context in contexts:\n        run_rows = context[\"rows\"]\n        hidden_vectors = context[\"hidden_vectors\"]\n        if not run_rows:\n            continue\n\n        hidden_write_started_at = time.perf_counter()\n        npz_path = hidden_dir / f\"{context['run_id']}.npz\"\n        if io_manager:\n            io_manager.save_npz_async(npz_path, {\"hidden_states\": np.stack(hidden_vectors)})\n            logging.info(\"IOManager: Background saving started for %s\", npz_path.name)\n        else:\n            np.savez_compressed(npz_path, hidden_states=np.stack(hidden_vectors))\n        hidden_state_write_seconds += time.perf_counter() - hidden_write_started_at\n        logging.info(\"batch_metrics: hidden_state_write_seconds = %.4f\", hidden_state_write_seconds)\n        first_correct_step = next((row[\"step\"] for row in run_rows if row[\"correct\"] == 1), -1)\n        first_model_stop_step = next((row[\"step\"] for row in run_rows if row[\"model_stop_flag\"] == 1), -1)\n        oracle_stop = max(run_rows, key=lambda row: row[\"utility\"])[\"step\"]\n        revision_count = int(sum(row[\"answer_changed\"] for row in run_rows))\n\n        run_summary = {\n            \"run_id\": context[\"run_id\"],\n            \"model_alias\": model_spec.alias,\n            \"model_name\": model_spec.hf_name,\n            \"task_id\": context[\"task\"].task_id,\n            \"domain\": context[\"task\"].domain,\n            \"difficulty\": context[\"task\"].difficulty,\n            \"task_source\": context[\"task\"].source,\n            \"task_source_index\": context[\"task\"].source_index,\n            \"temperature\": temperature,\n            \"seed\": seed,\n            \"prompt_mode\": actual_prompt_mode,\n            \"system_prompt_mode\": system_prompt_mode,\n            \"is_baseline\": int(is_baseline),\n            \"ever_correct\": int(first_correct_step != -1),\n            \"correct_at_step_1\": int(run_rows[0][\"correct\"] == 1),\n            \"oracle_stop\": oracle_stop,\n            \"first_correct_step\": first_correct_step if first_correct_step != -1 else max_steps,\n            \"first_model_stop_step\": first_model_stop_step if first_model_stop_step != -1 else max_steps,\n            \"revision_count\": revision_count,\n            \"best_utility\": max(row[\"utility\"] for row in run_rows),\n            \"final_correct\": run_rows[-1][\"correct\"],\n            \"device\": actual_device,\n        }\n        all_rows.extend(run_rows)\n        all_runs.append(run_summary)\n        logging.info(\n            \"completed %s | ever_correct=%d | first_correct=%s | final_correct=%d | revisions=%d\",\n            context[\"run_id\"],\n            run_summary[\"ever_correct\"],\n            run_summary[\"first_correct_step\"],\n            run_summary[\"final_correct\"],\n            revision_count,\n        )\n\n    batch_metric_rows.append(\n        {\n            \"phase\": \"hidden_state_write\",\n            \"temperature\": temperature,\n            \"seed\": seed,\n            \"step\": 0,\n            \"requested_batch_size\": len(contexts),\n            \"realized_batch_size\": len(contexts),\n            \"prompt_tokens\": 0,\n            \"generated_tokens\": 0,\n            \"wall_clock_seconds\": hidden_state_write_seconds,\n            \"examples_per_second\": float(\"nan\"),\n            \"tokens_per_second\": float(\"nan\"),\n            \"tokenize_seconds\": 0.0,\n            \"generation_seconds\": 0.0,\n            \"forward_seconds\": 0.0,\n            \"postprocess_seconds\": 0.0,\n            \"hidden_state_write_seconds\": hidden_state_write_seconds,\n            \"gpu_memory_allocated_gb\": gpu_memory_allocated_gb(actual_device),\n            \"gpu_memory_reserved_gb\": gpu_memory_reserved_gb(actual_device),\n            \"gpu_max_memory_allocated_gb\": gpu_max_memory_allocated_gb(actual_device),\n            \"gpu_max_memory_reserved_gb\": gpu_max_memory_reserved_gb(actual_device),\n            \"split_count\": 0,\n            \"oom_retry_count\": 0,\n        }\n    )\n\n    return all_rows, all_runs, batch_metric_rows\n\n\ndef _safe_mean(series: pd.Series) -> float:\n    \"\"\"Mean that tolerates a string-contaminated numeric column: coerce\n    non-numeric entries to NaN and average the rest. AUDIT/STABILITY FIX -- a\n    single stray string in a summary column was doing str+float inside pandas\n    .mean() and crashing a cell's end-of-run summary, throwing away hours of\n    completed collection (e.g. yi x math failed after 5 h).\"\"\"\n    return float(pd.to_numeric(series, errors=\"coerce\").mean())\n\n\ndef summarize_transitions(step_frame: pd.DataFrame) -> pd.DataFrame:\n    ordered = step_frame.sort_values([\"run_id\", \"step\"]).copy()\n    ordered[\"next_correct\"] = ordered.groupby(\"run_id\")[\"correct\"].shift(-1)\n    ordered = ordered.dropna(subset=[\"next_correct\"])\n    ordered[\"next_correct\"] = ordered[\"next_correct\"].astype(int)\n    ordered[\"repair\"] = ((ordered[\"correct\"] == 0) & (ordered[\"next_correct\"] == 1)).astype(int)\n    ordered[\"corruption\"] = ((ordered[\"correct\"] == 1) & (ordered[\"next_correct\"] == 0)).astype(int)\n    rows: list[dict[str, Any]] = []\n    for step, group in ordered.groupby(\"step\"):\n        wrong_group = group[group[\"correct\"] == 0]\n        correct_group = group[group[\"correct\"] == 1]\n        n_repairs = wrong_group[\"repair\"].sum()\n        n_corruptions = correct_group[\"corruption\"].sum()\n        repair_rate = _safe_mean(wrong_group[\"repair\"]) if len(wrong_group) else float(\"nan\")\n        corruption_rate = _safe_mean(correct_group[\"corruption\"]) if len(correct_group) else float(\"nan\")\n        q_t = _safe_mean(group[\"correct\"])\n        if (n_repairs + n_corruptions) >= 3 and not math.isnan(repair_rate) and not math.isnan(corruption_rate):\n            hazard_mu = (1.0 - q_t) * repair_rate - q_t * corruption_rate - STEP_COST\n        else:\n            hazard_mu = float(\"nan\")\n        rows.append(\n            {\n                \"step\": int(step),\n                \"q_t\": q_t,\n                \"n_repairs\": int(n_repairs),\n                \"n_corruptions\": int(n_corruptions),\n                \"repair_rate\": repair_rate,\n                \"corruption_rate\": corruption_rate,\n                \"hazard_mu\": hazard_mu,\n                \"entropy_mean\": _safe_mean(group[\"entropy_mean\"]),\n                \"confidence_mean\": _safe_mean(group[\"confidence\"]),\n                \"answer_changed_rate\": _safe_mean(group[\"answer_changed\"]),\n                \"hidden_shift_mean\": _safe_mean(group[\"hidden_l2_shift\"]),\n                \"n_transitions\": int(len(group)),\n            }\n        )\n    return pd.DataFrame(rows)\n\n\ndef build_pilot_summary(\n    step_frame: pd.DataFrame,\n    run_frame: pd.DataFrame,\n    transition_frame: pd.DataFrame,\n    model_spec: ModelSpec,\n    backend: str,\n    actual_device: str,\n) -> pd.DataFrame:\n    return pd.DataFrame(\n        [\n            {\n                \"model_alias\": model_spec.alias,\n                \"model_name\": model_spec.hf_name,\n                \"family\": model_spec.family,\n                \"parameter_count\": model_spec.parameter_count,\n                \"n_runs\": int(len(run_frame)),\n                \"n_tasks\": int(step_frame[\"task_id\"].nunique()),\n                \"mean_oracle_stop\": _safe_mean(run_frame[\"oracle_stop\"]),\n                \"mean_first_correct_step\": _safe_mean(run_frame[\"first_correct_step\"]),\n                \"mean_model_stop_step\": _safe_mean(run_frame[\"first_model_stop_step\"]),\n                \"mean_revision_count\": _safe_mean(run_frame[\"revision_count\"]),\n                \"mean_entropy\": _safe_mean(step_frame[\"entropy_mean\"]),\n                \"mean_hidden_shift\": _safe_mean(step_frame[\"hidden_l2_shift\"]),\n                \"repair_rate_overall\": _safe_mean(transition_frame[\"repair_rate\"]) if not transition_frame.empty else float(\"nan\"),\n                \"corruption_rate_overall\": _safe_mean(transition_frame[\"corruption_rate\"]) if not transition_frame.empty else float(\"nan\"),\n                \"runs_ever_correct\": int(run_frame[\"ever_correct\"].sum()),\n                \"backend\": backend,\n                \"device\": actual_device,\n            }\n        ]\n    )\n\n\ndef infer_existing_runtime_context(paths: dict[str, Path], requested_device: str) -> tuple[str, str]:\n    backend = \"transformers+torch(uninitialized)\"\n    if requested_device == \"auto\":\n        actual_device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    else:\n        actual_device = requested_device\n\n    metadata_path = paths[\"metadata\"]\n    if metadata_path.exists():\n        try:\n            with metadata_path.open(\"r\", encoding=\"utf-8\") as handle:\n                metadata = json.load(handle)\n            backend = str(metadata.get(\"backend\", backend))\n            actual_device = str(metadata.get(\"device\", actual_device))\n            return backend, actual_device\n        except (json.JSONDecodeError, OSError, TypeError, ValueError):\n            logging.warning(\"Failed to read existing metadata from %s for runtime context recovery.\", metadata_path)\n\n    pilot_path = paths[\"pilot\"]\n    if pilot_path.exists():\n        try:\n            pilot_summary = pd.read_csv(pilot_path)\n            if not pilot_summary.empty:\n                backend = str(pilot_summary.iloc[0].get(\"backend\", backend))\n                actual_device = str(pilot_summary.iloc[0].get(\"device\", actual_device))\n        except (OSError, ValueError, pd.errors.EmptyDataError):\n            logging.warning(\"Failed to read existing pilot summary from %s for runtime context recovery.\", pilot_path)\n\n    return backend, actual_device\n\n\ndef main() -> None:\n    parser = argparse.ArgumentParser(description=\"Run real-trace reasoning experiments on builtin tasks or GSM8K.\")\n    parser.add_argument(\"--model\", default=\"deepseek_r1_distill_1p5b\", choices=sorted(MODEL_CATALOG.keys()))\n    parser.add_argument(\"--device\", default=\"auto\", choices=[\"auto\", \"cpu\", \"cuda\"])\n    parser.add_argument(\"--quantization\", default=\"none\", choices=[\"none\", \"8bit\", \"4bit\"])\n    parser.add_argument(\"--device-map\", default=None)\n    parser.add_argument(\"--model-path-override\", default=None)\n    parser.add_argument(\"--attn-implementation\", default=\"auto\", choices=[\"auto\", \"sdpa\", \"flash_attention_2\", \"eager\"])\n    parser.add_argument(\"--temperatures\", nargs=\"+\", type=float, default=[0.6])\n    parser.add_argument(\"--seeds\", nargs=\"+\", type=int, default=[7])\n    parser.add_argument(\"--max-steps\", type=int, default=5)\n    parser.add_argument(\"--max-new-tokens\", type=int, default=256)\n    parser.add_argument(\"--max-tasks\", type=int, default=len(BUILTIN_TASKS))\n    parser.add_argument(\"--task-source\", default=\"builtin\", choices=[\"builtin\", \"gsm8k\", \"math\", \"arc\", \"gpqa\"])\n    parser.add_argument(\"--dataset-split\", default=\"train\")\n    parser.add_argument(\"--dataset-shuffle-seed\", type=int, default=17)\n    parser.add_argument(\"--batch-size\", type=int, default=4)\n    parser.add_argument(\"--prompt-mode\", default=\"structured_four_line\", choices=[\"structured_four_line\", \"minimal_json\", \"answer_only\"])\n    parser.add_argument(\"--system-prompt-mode\", default=\"default\", choices=[\"default\", \"short\", \"none\"])\n    parser.add_argument(\"--run-baseline\", action=\"store_true\", help=\"Run competence baseline instead of iterative reasoning\")\n    parser.add_argument(\"--output-dir\", default=str(DEFAULT_OUTPUT_DIR))\n    parser.add_argument(\"--io-threads\", type=int, default=4, help=\"Number of background threads for IO operations\")\n    parser.add_argument(\"--no-resume\", dest=\"resume\", action=\"store_false\")\n    parser.add_argument(\"--enable-k2-agreement\", action=\"store_true\", help=\"Enable k=2 self-consistency agreement recording (N7)\")\n    parser.add_argument(\"--enable-extended-observables\", action=\"store_true\", help=\"Enable answer-span metrics and mid-layer hidden projections (N8)\")\n    parser.set_defaults(resume=True)\n    args = parser.parse_args()\n\n    logging.basicConfig(level=logging.INFO, format=\"[%(levelname)s] %(message)s\")\n    if torch.cuda.is_available():\n        torch.backends.cuda.matmul.allow_tf32 = True\n\n    output_dir = Path(args.output_dir)\n    output_dir.mkdir(parents=True, exist_ok=True)\n    hidden_dir = output_dir / (\"hidden_states_baseline\" if args.run_baseline else \"hidden_states\")\n    paths = output_paths(output_dir, args.run_baseline)\n    model_spec = MODEL_CATALOG[args.model]\n\n    tasks = load_tasks(\n        task_source=args.task_source,\n        max_tasks=args.max_tasks,\n        dataset_split=args.dataset_split,\n        shuffle_seed=args.dataset_shuffle_seed,\n    )\n    logging.info(\"Loaded %d tasks from %s.\", len(tasks), args.task_source)\n\n    existing_steps, existing_runs, completed_run_ids, reconciliation_report = reconcile_existing_outputs(\n        output_dir=output_dir,\n        hidden_dir=hidden_dir,\n        is_baseline=args.run_baseline,\n        max_steps=args.max_steps,\n    )\n    if reconciliation_report[\"anomalies_detected\"]:\n        logging.warning(\"Checkpoint reconciliation adjusted on-disk artifacts: %s\", reconciliation_report)\n    if args.resume and completed_run_ids:\n        logging.info(\"Resuming with %d completed runs already on disk.\", len(completed_run_ids))\n\n    total_requested_runs = len(tasks) * len(args.temperatures) * len(args.seeds)\n    pending_requested_runs = 0\n    for temperature in args.temperatures:\n        for seed in args.seeds:\n            for task in tasks:\n                run_id = run_id_for(model_spec.alias, task, temperature, seed)\n                if not args.resume or run_id not in completed_run_ids:\n                    pending_requested_runs += 1\n    logging.info(\"Requested runs=%d | pending runs=%d\", total_requested_runs, pending_requested_runs)\n    next_pending_run = first_pending_run(\n        model_spec=model_spec,\n        tasks=tasks,\n        temperatures=args.temperatures,\n        seeds=args.seeds,\n        completed_run_ids=completed_run_ids if args.resume else set(),\n    )\n    if next_pending_run is not None:\n        logging.info(\n            \"Next pending run: task_index=%s | temperature=%.2f | seed=%d | run_id=%s\",\n            next_pending_run[\"task_source_index\"],\n            next_pending_run[\"temperature\"],\n            next_pending_run[\"seed\"],\n            next_pending_run[\"run_id\"],\n        )\n\n    if pending_requested_runs == 0 and existing_runs:\n        logging.info(\"No pending runs detected. Rebuilding summaries from existing artifacts.\")\n    elif pending_requested_runs == 0:\n        logging.info(\"No runs to execute.\")\n\n    model = None\n    tokenizer = None\n    backend, actual_device = infer_existing_runtime_context(paths, args.device)\n    if pending_requested_runs > 0:\n        offload_folder = output_dir / \"model_offload\"\n        offload_folder.mkdir(parents=True, exist_ok=True)\n        model, tokenizer, actual_device, backend = load_model(\n            model_spec=model_spec,\n            device=args.device,\n            quantization=args.quantization,\n            device_map=args.device_map,\n            attn_implementation=args.attn_implementation,\n            model_path_override=args.model_path_override,\n            offload_folder=str(offload_folder),\n        )\n        logging.info(\"Model loaded: %s | backend=%s | device=%s\", model_spec.hf_name, backend, actual_device)\n\n    io_manager = IOManager(max_threads=args.io_threads)\n    try:\n        write_runtime_metadata(\n            metadata_path=paths[\"metadata\"],\n            model_spec=model_spec,\n            backend=backend,\n            actual_device=actual_device,\n            quantization=args.quantization,\n            device_map=args.device_map,\n            attn_implementation=args.attn_implementation,\n            temperatures=args.temperatures,\n            seeds=args.seeds,\n            max_steps=args.max_steps,\n            max_new_tokens=args.max_new_tokens,\n            max_tasks=args.max_tasks,\n            task_source=args.task_source,\n            dataset_split=args.dataset_split,\n            dataset_shuffle_seed=args.dataset_shuffle_seed,\n            batch_size=args.batch_size,\n            prompt_mode=args.prompt_mode,\n            system_prompt_mode=args.system_prompt_mode,\n            resume_enabled=args.resume,\n            completed_run_ids=completed_run_ids,\n            pending_run_count=pending_requested_runs,\n            next_pending_run=next_pending_run,\n            reconciliation_report=reconciliation_report,\n            tasks=tasks,\n            model_source=args.model_path_override,\n        )\n\n        all_rows: list[dict[str, Any]] = list(existing_steps)\n        all_runs: list[dict[str, Any]] = list(existing_runs)\n\n        if model is not None and tokenizer is not None:\n            for temperature in args.temperatures:\n                for seed in args.seeds:\n                    pending_tasks = [\n                        task\n                        for task in tasks\n                        if not args.resume or run_id_for(model_spec.alias, task, temperature, seed) not in completed_run_ids\n                    ]\n                    if not pending_tasks:\n                        logging.info(\"Skipping temp=%.2f seed=%d because all runs already exist.\", temperature, seed)\n                        continue\n\n                    set_seed(seed)\n                    batches = list(chunked(pending_tasks, max(1, args.batch_size)))\n                    progress = tqdm(batches, desc=f\"temp={temperature:.2f} seed={seed}\", unit=\"batch\")\n                    for task_batch in progress:\n                        batch_rows, batch_runs, batch_metrics = run_batch_traces(\n                            model=model,\n                            tokenizer=tokenizer,\n                            model_spec=model_spec,\n                            tasks=task_batch,\n                            actual_device=actual_device,\n                            temperature=temperature,\n                            seed=seed,\n                            max_steps=args.max_steps,\n                            max_new_tokens=args.max_new_tokens,\n                            prompt_mode=args.prompt_mode,\n                            system_prompt_mode=args.system_prompt_mode,\n                            is_baseline=args.run_baseline,\n                            hidden_dir=hidden_dir,\n                            io_manager=io_manager,\n                            enable_k2_agreement=args.enable_k2_agreement,\n                            enable_extended_observables=args.enable_extended_observables,\n                        )\n                        append_records(paths[\"steps\"], batch_rows)\n                        append_records(paths[\"runs\"], batch_runs)\n                        append_records(paths[\"batch_metrics\"], batch_metrics)\n                        all_rows.extend(batch_rows)\n                        all_runs.extend(batch_runs)\n                        completed_run_ids.update(run[\"run_id\"] for run in batch_runs)\n                        next_pending_run = first_pending_run(\n                            model_spec=model_spec,\n                            tasks=tasks,\n                            temperatures=args.temperatures,\n                            seeds=args.seeds,\n                            completed_run_ids=completed_run_ids if args.resume else set(),\n                        )\n                        write_runtime_metadata(\n                            metadata_path=paths[\"metadata\"],\n                            model_spec=model_spec,\n                            backend=backend,\n                            actual_device=actual_device,\n                            quantization=args.quantization,\n                            device_map=args.device_map,\n                            attn_implementation=args.attn_implementation,\n                            temperatures=args.temperatures,\n                            seeds=args.seeds,\n                            max_steps=args.max_steps,\n                            max_new_tokens=args.max_new_tokens,\n                            max_tasks=args.max_tasks,\n                            task_source=args.task_source,\n                            dataset_split=args.dataset_split,\n                            dataset_shuffle_seed=args.dataset_shuffle_seed,\n                            batch_size=args.batch_size,\n                            prompt_mode=args.prompt_mode,\n                            system_prompt_mode=args.system_prompt_mode,\n                            resume_enabled=args.resume,\n                            completed_run_ids=completed_run_ids,\n                            pending_run_count=max(total_requested_runs - len(completed_run_ids), 0),\n                            next_pending_run=next_pending_run,\n                            reconciliation_report=reconciliation_report,\n                            tasks=tasks,\n                            model_source=args.model_path_override,\n                        )\n                        progress.set_postfix(\n                            runs=len(batch_runs),\n                            correct=sum(int(run[\"ever_correct\"]) for run in batch_runs),\n                        )\n\n        if not all_rows or not all_runs:\n            raise RuntimeError(\"No trace data is available to summarize. The experiment produced no rows.\")\n\n        if args.run_baseline:\n            logging.info(\"Wrote baseline artifacts to: %s\", output_dir)\n            return\n\n        step_frame = pd.DataFrame(all_rows)\n        run_frame = pd.DataFrame(all_runs)\n        transition_frame = summarize_transitions(step_frame)\n        pilot_summary = build_pilot_summary(step_frame, run_frame, transition_frame, model_spec, backend, actual_device)\n    finally:\n        if io_manager:\n            io_manager.shutdown()\n\n    transition_frame.to_csv(paths[\"hazard\"], index=False)\n    pilot_summary.to_csv(paths[\"pilot\"], index=False)\n\n    next_pending_run = first_pending_run(\n        model_spec=model_spec,\n        tasks=tasks,\n        temperatures=args.temperatures,\n        seeds=args.seeds,\n        completed_run_ids=completed_run_ids if args.resume else set(),\n    )\n    write_runtime_metadata(\n        metadata_path=paths[\"metadata\"],\n        model_spec=model_spec,\n        backend=backend,\n        actual_device=actual_device,\n        quantization=args.quantization,\n        device_map=args.device_map,\n        attn_implementation=args.attn_implementation,\n        temperatures=args.temperatures,\n        seeds=args.seeds,\n        max_steps=args.max_steps,\n        max_new_tokens=args.max_new_tokens,\n        max_tasks=args.max_tasks,\n        task_source=args.task_source,\n        dataset_split=args.dataset_split,\n        dataset_shuffle_seed=args.dataset_shuffle_seed,\n        batch_size=args.batch_size,\n        prompt_mode=args.prompt_mode,\n        system_prompt_mode=args.system_prompt_mode,\n        resume_enabled=args.resume,\n        completed_run_ids=completed_run_ids,\n        pending_run_count=max(total_requested_runs - len(completed_run_ids), 0),\n        next_pending_run=next_pending_run,\n        reconciliation_report=reconciliation_report,\n        tasks=tasks,\n        model_source=args.model_path_override,\n    )\n\n    for _, row in pilot_summary.iterrows():\n        print(\n            f\"{row['model_alias']}: runs={int(row['n_runs'])}, \"\n            f\"ever_correct={int(row['runs_ever_correct'])}, \"\n            f\"oracle={row['mean_oracle_stop']:.2f}, first_correct={row['mean_first_correct_step']:.2f}, \"\n            f\"repair={row['repair_rate_overall']:.3f}, corruption={row['corruption_rate_overall']:.3f}, \"\n            f\"device={row['device']}\"\n        )\n    print(f\"Wrote real-trace artifacts to: {output_dir}\")\n\n\nif __name__ == \"__main__\":\n    main()\n"}}}}