from __future__ import annotations

import os
import json
import time
import hashlib
import subprocess
import sys
from typing import Dict, List, Tuple, Optional, Set

from ..code_graph import CodeGraph
from .prompts import build_prompts_for_module
from .verify import verify_with_tests, extract_citations
from modules.embedding import build_subgraph_embedding_from_graph
from modules.adapter import (
    save_npz,
    generate_lora_from_embedding,
)
from model.inspect import detect_target_shapes_from_model_full, detect_target_shapes_from_model

def _run_repo_adapter_cli(
    model: str,
    adapters_dir: str,
    repo: str,
    prompt: str,
    *,
    cache_dir: Optional[str] = None,
    device: str = "cpu",
    gpu_ids: Optional[str] = None,
    context_tokens: int = 5000,
    timeout_sec: Optional[int] = None,
) -> Tuple[int, str, str]:
    """Minimal wrapper around the generic run CLI to obtain an answer text."""
    cmd = [
        sys.executable,
        "-m",
        "run",
        "--model", model,
        "--sources", repo,
        "--adapters-dir", adapters_dir,
        "--prompt", prompt,
        "--context-tokens", str(int(context_tokens)),
    ]
    if cache_dir:
        cmd += ["--cache-dir", str(cache_dir)]
    if device and device != "cpu":
        cmd += ["--device-map", "auto"]
    if gpu_ids:
        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_ids)
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout_sec)
        return int(proc.returncode), proc.stdout or "", proc.stderr or ""
    except subprocess.TimeoutExpired as te:
        return 124, "", str(te)
    except Exception as e:
        return 1, "", str(e)

def distill_repo(
    repo: str,
    model: str,
    adapters: str,
    out_dir: str,
    *,
    ignore: Optional[List[str]] = None,
    max_prompts: int = 3,
    cache_dir: Optional[str] = None,
    device: str = "cpu",
    gpu_ids: Optional[str] = None,
    context_tokens: int = 5000,
    timeout_sec: Optional[int] = None,
    resume: bool = False,
    log_every: int = 25,
    citations_per_paragraph: bool = False,
) -> Tuple[Set[str], str]:
    out_dir = os.path.abspath(os.path.expanduser(os.path.expandvars(out_dir)))
    os.makedirs(out_dir, exist_ok=True)
    g = CodeGraph.load_or_build(repo, ignore=[s for s in (ignore or []) if s])

    buffer_path = os.path.join(out_dir, "distill.jsonl")
    verified_modules: Dict[str, bool] = {}
    processed: Dict[Tuple[str, str], bool] = {}
    if resume and os.path.isfile(buffer_path):
        try:
            with open(buffer_path, "r", encoding="utf-8") as fh:
                for line in fh:
                    try:
                        row = json.loads(line)
                        k = (str(row.get("module", "")), str(row.get("prompt", "")))
                        processed[k] = True
                    except Exception:
                        continue
        except Exception:
            processed = {}

    mode = ("a" if resume and os.path.isfile(buffer_path) else "w")
    t0 = time.time()
    logs_dir = os.path.join(out_dir, "logs")
    try:
        os.makedirs(logs_dir, exist_ok=True)
    except Exception:
        pass

    env = os.environ.copy()
    with open(buffer_path, mode, encoding="utf-8") as fh:
        for idx, module in enumerate(sorted(g.modules.keys()), start=1):
            mi = g.modules[module]
            if mi.is_test:
                continue
            prompts = build_prompts_for_module(g, module, max_q=int(max_prompts))
            module_verified = False
            for pidx, q in enumerate(prompts, start=1):
                if resume and ((module, q) in processed):
                    continue
                try:
                    print(json.dumps({
                        "event": "start_prompt",
                        "module": module,
                        "prompt_index": pidx,
                        "context_tokens": int(context_tokens),
                    }), flush=True)
                except Exception:
                    pass
                rc, ans_text, err_text = _run_repo_adapter_cli(
                    model,
                    adapters,
                    repo,
                    q,
                    cache_dir=cache_dir,
                    device=device,
                    gpu_ids=gpu_ids,
                    context_tokens=int(context_tokens),
                    timeout_sec=(int(timeout_sec) if (timeout_sec and int(timeout_sec) > 0) else None),
                )
                ok = (rc == 0)
                if ok:
                    cites = extract_citations(ans_text)
                    ok = bool(cites)
                    if ok:
                        ok = verify_with_tests(g, module, repo_root=repo, env=env)
                try:
                    h = hashlib.sha1((module + "\n" + q).encode("utf-8", errors="ignore")).hexdigest()[:12]
                    log_fp = os.path.join(logs_dir, f"{module.replace('/', '_')}.{pidx}.{h}.log")
                    with open(log_fp, "w", encoding="utf-8") as lf:
                        lf.write(ans_text)
                        lf.write("\n\n[stderr]\n")
                        lf.write(err_text or "")
                except Exception:
                    pass
                fh.write(json.dumps({
                    "module": module,
                    "prompt": q,
                    "answer": ans_text,
                    "verified": bool(ok),
                    "citations": extract_citations(ans_text),
                    "ignore": [s for s in (ignore or []) if s],
                    "device": device,
                    "context_tokens": int(context_tokens),
                    "rc": int(rc),
                }) + "\n")
                module_verified = module_verified or bool(ok)
            if module_verified:
                verified_modules[module] = True
            if (idx % max(1, int(log_every))) == 0:
                try:
                    elapsed = time.time() - t0
                except Exception:
                    elapsed = -1.0
                print(json.dumps({"progress": {"done": idx, "total": len(g.modules), "elapsed_sec": elapsed}}))

    return set(verified_modules.keys()), buffer_path

def export_tuned_adapters(
    repo: str,
    model: str,
    verified_modules: List[str],
    out_dir: str,
    *,
    per_module_adapters: bool = False,
    include_deps: bool = False,
    max_deps: int = 4,
    rank: int = 8,
    cache_dir: Optional[str] = None,
) -> None:
    os.makedirs(out_dir, exist_ok=True)
    g = CodeGraph.load_or_build(repo)
    emb = build_subgraph_embedding_from_graph(
        g,
        dim=1536,
        seed=0,
        include_modules=sorted(list(verified_modules)),
        include_files=None,
        include_text=True,
        text_max_bytes=250000,
        max_text_tokens=0,
        text_weight=0.25,
        graph_prop_hops=0,
        graph_prop_damp=0.85,
    )
    targets = ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"]
    adapters = {"layers": [], "rank": 8, "d_model": 0, "targets": targets, "gates": []}
    manifest = {
        "repo": os.path.abspath(repo),
        "verified_modules": sorted(list(verified_modules)),
        "buffer": os.path.join(out_dir, "distill.jsonl"),
        "schema_version": 1,
        "note": "Embedding from verified modules; adapter layers empty for downstream mixing.",
    }
    save_npz(out_dir, embedding=emb, adapters=adapters, manifest=manifest)

    if not per_module_adapters:
        return

    # Try to infer model dims and target shapes for per-module export
    num_layers: int = 0
    d_model: int = 0
    target_shapes: Optional[Dict[str, Tuple[int, int]]] = None
    try:
        target_shapes = detect_target_shapes_from_model_full(model, target_regex=None)
    except Exception:
        target_shapes = None
    if not target_shapes:
        try:
            target_shapes = detect_target_shapes_from_model(model)
        except Exception:
            target_shapes = None
    try:
        from transformers import AutoConfig  # type: ignore
        cfg = AutoConfig.from_pretrained(model, cache_dir=cache_dir)
        num_layers = int(getattr(cfg, "num_hidden_layers", 0) or 0)
        d_model = int(getattr(cfg, "hidden_size", 0) or 0)
    except Exception:
        num_layers, d_model = 0, 0

    sub_root = os.path.join(out_dir, "sub_adapters")
    os.makedirs(sub_root, exist_ok=True)
    for module in sorted(list(verified_modules)):
        mods = [module]
        if include_deps:
            deps = list(g.module_imports.get(module, []) or [])
            if int(max_deps) > 0:
                deps = deps[: int(max_deps)]
            mods.extend([m for m in deps if m])
        mods_unique = sorted({m for m in mods})
        emb_m = build_subgraph_embedding_from_graph(
            g,
            dim=1536,
            seed=0,
            include_modules=mods_unique,
            include_files=None,
            include_text=True,
            text_max_bytes=250000,
            max_text_tokens=0,
            text_weight=0.25,
            graph_prop_hops=0,
            graph_prop_damp=0.85,
        )
        if (num_layers <= 0) or (d_model <= 0):
            sub_dir = os.path.join(sub_root, module.replace("/", "_"))
            os.makedirs(sub_dir, exist_ok=True)
            save_npz(
                sub_dir,
                embedding=emb_m,
                adapters={"layers": [], "rank": rank, "d_model": 0, "targets": [], "gates": []},
                manifest={
                    "module": module,
                    "include_modules": mods_unique,
                    "note": "Embedding-only; shapes not inferred",
                },
            )
            continue
        try:
            adapters_m = generate_lora_from_embedding(
                emb_m["z"],
                d_model=int(d_model),
                num_layers=int(num_layers),
                rank=int(rank),
                seed=0,
                targets=list((target_shapes or {}).keys()) if target_shapes else None,
                target_shapes=target_shapes,
                layer_gate="zmean",
                target_weights=None,
            )
            sub_dir = os.path.join(sub_root, module.replace("/", "_"))
            os.makedirs(sub_dir, exist_ok=True)
            save_npz(
                sub_dir,
                embedding=emb_m,
                adapters=adapters_m,
                manifest={
                    "module": module,
                    "include_modules": mods_unique,
                    "rank": int(rank),
                    "d_model": int(d_model),
                    "layers": int(num_layers),
                },
            )
        except Exception:
            continue


import os
import json
import time
import hashlib
import subprocess
import sys
from typing import Dict, List, Tuple, Optional, Set

from ..code_graph import CodeGraph
from .prompts import build_prompts_for_module
from .verify import verify_with_tests, extract_citations
from modules.embedding import build_subgraph_embedding_from_graph
from modules.adapter import (
    save_npz,
    generate_lora_from_embedding,
)
from model.inspect import detect_target_shapes_from_model_full, detect_target_shapes_from_model

def _run_repo_adapter_cli(
    model: str,
    adapters_dir: str,
    repo: str,
    prompt: str,
    *,
    cache_dir: Optional[str] = None,
    device: str = "cpu",
    gpu_ids: Optional[str] = None,
    context_tokens: int = 5000,
    timeout_sec: Optional[int] = None,
) -> Tuple[int, str, str]:
    """Minimal wrapper around the generic run CLI to obtain an answer text."""
    cmd = [
        sys.executable,
        "-m",
        "run",
        "--model", model,
        "--sources", repo,
        "--adapters-dir", adapters_dir,
        "--prompt", prompt,
        "--context-tokens", str(int(context_tokens)),
    ]
    if cache_dir:
        cmd += ["--cache-dir", str(cache_dir)]
    if device and device != "cpu":
        cmd += ["--device-map", "auto"]
    if gpu_ids:
        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_ids)
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout_sec)
        return int(proc.returncode), proc.stdout or "", proc.stderr or ""
    except subprocess.TimeoutExpired as te:
        return 124, "", str(te)
    except Exception as e:
        return 1, "", str(e)

def distill_repo(
    repo: str,
    model: str,
    adapters: str,
    out_dir: str,
    *,
    ignore: Optional[List[str]] = None,
    max_prompts: int = 3,
    cache_dir: Optional[str] = None,
    device: str = "cpu",
    gpu_ids: Optional[str] = None,
    context_tokens: int = 5000,
    timeout_sec: Optional[int] = None,
    resume: bool = False,
    log_every: int = 25,
    citations_per_paragraph: bool = False,
) -> Tuple[Set[str], str]:
    out_dir = os.path.abspath(os.path.expanduser(os.path.expandvars(out_dir)))


        "--sources", repo,
        "--adapters-dir", adapters_dir,
        "--prompt", prompt,
        "--context-tokens", str(int(context_tokens)),
    ]
    if cache_dir:
        cmd += ["--cache-dir", str(cache_dir)]
    if device and device != "cpu":
        cmd += ["--device-map", "auto"]
    if gpu_ids:
        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_ids)
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout_sec)
        return int(proc.returncode), proc.stdout or "", proc.stderr or ""
    except subprocess.TimeoutExpired as te:
        return 124, "", str(te)
    except Exception as e:
        return 1, "", str(e)

def distill_repo(
    repo: str,
    model: str,
    adapters: str,
    out_dir: str,
    *,
    ignore: Optional[List[str]] = None,
    max_prompts: int = 3,
    cache_dir: Optional[str] = None,
    device: str = "cpu",
    gpu_ids: Optional[str] = None,
    context_tokens: int = 5000,
    timeout_sec: Optional[int] = None,
    resume: bool = False,
    log_every: int = 25,
    citations_per_paragraph: bool = False,
) -> Tuple[Set[str], str]:
    out_dir = os.path.abspath(os.path.expanduser(os.path.expandvars(out_dir)))
    os.makedirs(out_dir, exist_ok=True)
    g = CodeGraph.load_or_build(repo, ignore=[s for s in (ignore or []) if s])

    buffer_path = os.path.join(out_dir, "distill.jsonl")
    verified_modules: Dict[str, bool] = {}
    processed: Dict[Tuple[str, str], bool] = {}
    if resume and os.path.isfile(buffer_path):
        try:
            with open(buffer_path, "r", encoding="utf-8") as fh:
                for line in fh:
                    try:
                        row = json.loads(line)
                        k = (str(row.get("module", "")), str(row.get("prompt", "")))
                        processed[k] = True
                    except Exception:
                        continue
        except Exception:
            processed = {}

    mode = ("a" if resume and os.path.isfile(buffer_path) else "w")
    t0 = time.time()
    logs_dir = os.path.join(out_dir, "logs")
    try:
        os.makedirs(logs_dir, exist_ok=True)
    except Exception:
        pass

    env = os.environ.copy()
    with open(buffer_path, mode, encoding="utf-8") as fh:
        for idx, module in enumerate(sorted(g.modules.keys()), start=1):
            mi = g.modules[module]
            if mi.is_test:
                continue
            prompts = build_prompts_for_module(g, module, max_q=int(max_prompts))
            module_verified = False
            for pidx, q in enumerate(prompts, start=1):
                if resume and ((module, q) in processed):
                    continue
                try:
                    print(json.dumps({
                        "event": "start_prompt",
                        "module": module,
                        "prompt_index": pidx,
                        "context_tokens": int(context_tokens),
                    }), flush=True)
                except Exception:
                    pass
                rc, ans_text, err_text = _run_repo_adapter_cli(
                    model,
                    adapters,
                    repo,
                    q,
                    cache_dir=cache_dir,
                    device=device,
                    gpu_ids=gpu_ids,
                    context_tokens=int(context_tokens),
                    timeout_sec=(int(timeout_sec) if (timeout_sec and int(timeout_sec) > 0) else None),
                )
                ok = (rc == 0)
                if ok:
                    cites = extract_citations(ans_text)
                    ok = bool(cites)
                    if ok:
                        ok = verify_with_tests(g, module, repo_root=repo, env=env)
                try:
                    h = hashlib.sha1((module + "\n" + q).encode("utf-8", errors="ignore")).hexdigest()[:12]
                    log_fp = os.path.join(logs_dir, f"{module.replace('/', '_')}.{pidx}.{h}.log")
                    with open(log_fp, "w", encoding="utf-8") as lf:
                        lf.write(ans_text)
                        lf.write("\n\n[stderr]\n")
                        lf.write(err_text or "")
                except Exception:
                    pass
                fh.write(json.dumps({
                    "module": module,
                    "prompt": q,
                    "answer": ans_text,
                    "verified": bool(ok),
                    "citations": extract_citations(ans_text),
                    "ignore": [s for s in (ignore or []) if s],
                    "device": device,
                    "context_tokens": int(context_tokens),
                    "rc": int(rc),
                }) + "\n")
                module_verified = module_verified or bool(ok)
            if module_verified:
                verified_modules[module] = True
            if (idx % max(1, int(log_every))) == 0:
                try:
                    elapsed = time.time() - t0
                except Exception:
                    elapsed = -1.0
                print(json.dumps({"progress": {"done": idx, "total": len(g.modules), "elapsed_sec": elapsed}}))

    return set(verified_modules.keys()), buffer_path

def export_tuned_adapters(
    repo: str,
    model: str,
    verified_modules: List[str],
    out_dir: str,
    *,
    per_module_adapters: bool = False,
    include_deps: bool = False,
    max_deps: int = 4,
    rank: int = 8,
    cache_dir: Optional[str] = None,
) -> None:
    os.makedirs(out_dir, exist_ok=True)
    g = CodeGraph.load_or_build(repo)
    emb = build_subgraph_embedding_from_graph(
        g,
        dim=1536,
        seed=0,


                    "verified": bool(ok),
                    "citations": extract_citations(ans_text),
                    "ignore": [s for s in (ignore or []) if s],
                    "device": device,
                    "context_tokens": int(context_tokens),
                    "rc": int(rc),
                }) + "\n")
                module_verified = module_verified or bool(ok)
            if module_verified:
                verified_modules[module] = True
            if (idx % max(1, int(log_every))) == 0:
                try:
                    elapsed = time.time() - t0
                except Exception:
                    elapsed = -1.0
                print(json.dumps({"progress": {"done": idx, "total": len(g.modules), "elapsed_sec": elapsed}}))

    return set(verified_modules.keys()), buffer_path

def export_tuned_adapters(
    repo: str,
    model: str,
    verified_modules: List[str],
    out_dir: str,
    *,
    per_module_adapters: bool = False,
    include_deps: bool = False,
    max_deps: int = 4,
    rank: int = 8,
    cache_dir: Optional[str] = None,
) -> None:
    os.makedirs(out_dir, exist_ok=True)
    g = CodeGraph.load_or_build(repo)
    emb = build_subgraph_embedding_from_graph(
        g,
        dim=1536,
        seed=0,
        include_modules=sorted(list(verified_modules)),
        include_files=None,
        include_text=True,
        text_max_bytes=250000,
        max_text_tokens=0,
        text_weight=0.25,
        graph_prop_hops=0,
        graph_prop_damp=0.85,
    )
    targets = ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"]
    adapters = {"layers": [], "rank": 8, "d_model": 0, "targets": targets, "gates": []}
    manifest = {
        "repo": os.path.abspath(repo),
        "verified_modules": sorted(list(verified_modules)),
        "buffer": os.path.join(out_dir, "distill.jsonl"),
        "schema_version": 1,
        "note": "Embedding from verified modules; adapter layers empty for downstream mixing.",
    }
    save_npz(out_dir, embedding=emb, adapters=adapters, manifest=manifest)

    if not per_module_adapters:
        return

    # Try to infer model dims and target shapes for per-module export
    num_layers: int = 0
    d_model: int = 0
    target_shapes: Optional[Dict[str, Tuple[int, int]]] = None
    try:
        target_shapes = detect_target_shapes_from_model_full(model, target_regex=None)
    except Exception:
        target_shapes = None
    if not target_shapes:
        try:
            target_shapes = detect_target_shapes_from_model(model)
        except Exception:
            target_shapes = None
    try:
        from transformers import AutoConfig  # type: ignore
        cfg = AutoConfig.from_pretrained(model, cache_dir=cache_dir)
        num_layers = int(getattr(cfg, "num_hidden_layers", 0) or 0)
        d_model = int(getattr(cfg, "hidden_size", 0) or 0)
    except Exception:
        num_layers, d_model = 0, 0

    sub_root = os.path.join(out_dir, "sub_adapters")
    os.makedirs(sub_root, exist_ok=True)
    for module in sorted(list(verified_modules)):
        mods = [module]
        if include_deps:
            deps = list(g.module_imports.get(module, []) or [])
            if int(max_deps) > 0:
                deps = deps[: int(max_deps)]
            mods.extend([m for m in deps if m])
        mods_unique = sorted({m for m in mods})
        emb_m = build_subgraph_embedding_from_graph(
            g,
            dim=1536,
            seed=0,
            include_modules=mods_unique,
            include_files=None,
            include_text=True,
            text_max_bytes=250000,
            max_text_tokens=0,
            text_weight=0.25,
            graph_prop_hops=0,
            graph_prop_damp=0.85,
        )
        if (num_layers <= 0) or (d_model <= 0):
            sub_dir = os.path.join(sub_root, module.replace("/", "_"))
            os.makedirs(sub_dir, exist_ok=True)
            save_npz(
                sub_dir,
                embedding=emb_m,
                adapters={"layers": [], "rank": rank, "d_model": 0, "targets": [], "gates": []},
                manifest={
                    "module": module,
                    "include_modules": mods_unique,
                    "note": "Embedding-only; shapes not inferred",
                },
            )
            continue
        try:
            adapters_m = generate_lora_from_embedding(
                emb_m["z"],
                d_model=int(d_model),
                num_layers=int(num_layers),
                rank=int(rank),
                seed=0,
                targets=list((target_shapes or {}).keys()) if target_shapes else None,
                target_shapes=target_shapes,
                layer_gate="zmean",
                target_weights=None,
            )
            sub_dir = os.path.join(sub_root, module.replace("/", "_"))
            os.makedirs(sub_dir, exist_ok=True)
            save_npz(
                sub_dir,
                embedding=emb_m,
                adapters=adapters_m,
                manifest={
                    "module": module,
                    "include_modules": mods_unique,
                    "rank": int(rank),
                    "d_model": int(d_model),
                    "layers": int(num_layers),
                },
            )
        except Exception:
            continue


import os
import json
import argparse
from typing import Dict, Any, List

import torch

@torch.no_grad()
def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="meta-llama/Llama-3.1-8B-Instruct")
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints")
    p.add_argument("--prompt", default="Hello")
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--dtype", default="bfloat16")
    p.add_argument("--max-layers", type=int, default=8, help="Limit layers to trace for brevity")
    p.add_argument("--save-logits", action="store_true", help="Save HF/local logits to disk to minimize RAM/VRAM use")
    p.add_argument("--out-dir", default="", help="Directory to write saved tensors (defaults to cache-dir/trace)")
    args = p.parse_args()

    os.makedirs(args.cache_dir, exist_ok=True)
    os.environ.setdefault("ATTN_BACKEND", "torch")

    from transformers import AutoTokenizer, AutoModelForCausalLM  # type: ignore
    tok = AutoTokenizer.from_pretrained(args.model, use_fast=True, cache_dir=args.cache_dir)
    x = tok(args.prompt, return_tensors="pt")
    x = {k: v.to(args.device) for k, v in x.items()}

    # HF model (match dtype)
    hf_dtype = getattr(torch, args.dtype) if hasattr(torch, args.dtype) else torch.float32
    hf = AutoModelForCausalLM.from_pretrained(
        args.model,
        torch_dtype=hf_dtype,
        device_map={"": args.device},
        low_cpu_mem_usage=True,
        cache_dir=args.cache_dir,
    ).eval()

    # Build masks and positions exactly once (no cache), compute HF RoPE and per-layer traces first
    input_ids = x["input_ids"]
    attention_mask = x.get("attention_mask", None)
    B, T = int(input_ids.shape[0]), int(input_ids.shape[1])

    # HF position_ids and mask
    position_ids_hf = torch.arange(T, device=args.device).unsqueeze(0)
    from transformers.models.llama.modeling_llama import LlamaModel  # type: ignore
    # Prepare rotary cos/sin via HF module
    llama_model: LlamaModel = hf.model  # type: ignore[attr-defined]
    cos_hf, sin_hf = llama_model.rotary_emb(
        llama_model.embed_tokens(input_ids), position_ids=position_ids_hf
    )

    # Run HF forward with hooks to collect per-layer hidden (store on CPU to save VRAM)
    hf_layers_out: List[torch.Tensor] = []

    def _hf_layer_hook(module, inputs, output):
        hs = output[0] if isinstance(output, tuple) else output
        hf_layers_out.append(hs.detach().to(device="cpu", dtype=torch.float32))

    hooks = []
    try:
        for i, layer in enumerate(llama_model.layers[: int(args.max_layers) ]):
            hooks.append(layer.register_forward_hook(_hf_layer_hook))
        _ = hf(input_ids=input_ids, attention_mask=attention_mask)
    finally:
        for h in hooks:
            try:
                h.remove()
            except Exception:
                pass

    # Final HF logits (CPU copy)
    out_hf = hf(input_ids=input_ids, attention_mask=attention_mask)
    logits_hf = out_hf.logits[:, -1, :].to("cpu", torch.float32)
    hf_logits_path = None
    if args.save_logits:
        out_dir = args.out_dir or os.path.join(args.cache_dir, "trace")
        os.makedirs(out_dir, exist_ok=True)
        hf_logits_path = os.path.join(out_dir, "hf_logits.pt")
        torch.save(logits_hf, hf_logits_path)
        del logits_hf
        logits_hf = None

    # Free HF before building local to avoid OOM
    del out_hf
    del llama_model
    del hf
    try:
        torch.cuda.empty_cache()
    except Exception:
        pass

    # Local model (now that HF is freed)
    from model.hf_snapshot import ensure_snapshot
    ckpt_dir = ensure_snapshot(args.model, args.cache_dir)
    from specs.config import ModelConfig
    from model.factory import build_causal_lm
    from model.hf_llama_loader import load_hf_llama_weights_into_local
    cfg = json.load(open(os.path.join(ckpt_dir, "config.json"), "r", encoding="utf-8"))
    mc = ModelConfig(
        d_model=int(cfg.get("hidden_size")),
        n_heads=int(cfg.get("num_attention_heads")),
        n_layers=int(cfg.get("num_hidden_layers")),
        d_ff=int(cfg.get("intermediate_size")),
        vocab_size=int(cfg.get("vocab_size")),
        head_dim=int(cfg.get("head_dim", int(cfg.get("hidden_size")) // int(cfg.get("num_attention_heads")))),
        rope_theta=float(cfg.get("rope_parameters", {}).get("rope_theta", cfg.get("rope_theta", 1e6))),
        dtype=(args.dtype if args.device.startswith("cuda") else "float32"),
        attn_impl="sdpa",
        rms_norm_eps=float(cfg.get("rms_norm_eps", 1e-6)),
        rope_scaling_type=(cfg.get("rope_scaling", {}) or {}).get("type"),
        rope_scaling_factor=(cfg.get("rope_scaling", {}) or {}).get("factor"),
        rope_scaling_original_max_position_embeddings=(cfg.get("rope_scaling", {}) or {}).get("original_max_position_embeddings"),
        rope_scaling_low_freq_factor=(cfg.get("rope_scaling", {}) or {}).get("low_freq_factor"),
        rope_scaling_high_freq_factor=(cfg.get("rope_scaling", {}) or {}).get("high_freq_factor"),
    )
    n_kv_heads = int(cfg.get("num_key_value_heads", mc.n_heads))
    tie_we = bool(cfg.get("tie_word_embeddings", True))
    local = build_causal_lm(mc, block="llama", n_kv_heads=n_kv_heads, tie_weights=tie_we)
    load_hf_llama_weights_into_local(local, ckpt_dir)
    local = local.to(device=args.device, dtype=hf_dtype).eval()

    # Local position_ids and mask/cos/sin
    from tensor.masking import create_causal_mask
    from tensor.positional import RotaryEmbeddingHF as LocalRotary
    position_ids_local = torch.arange(T, device=args.device).unsqueeze(0)
    add_mask_local = create_causal_mask(
        input_embeds=local.embed(input_ids),
        attention_mask=attention_mask,
        cache_position=None,
        position_ids=position_ids_local,
        past_key_values=None,
    )  # (B,1,T,S)
    rope = LocalRotary(
        head_dim=int(mc.head_dim or mc.d_model // mc.n_heads),
        base_theta=float(mc.rope_theta),
        attention_scaling=float(getattr(mc, "rope_attention_scaling", 1.0) or 1.0),
        device=input_ids.device,
        scaling_type=getattr(mc, "rope_scaling_type", None),
        scaling_factor=getattr(mc, "rope_scaling_factor", None),
        original_max_position_embeddings=getattr(mc, "rope_scaling_original_max_position_embeddings", None),
        low_freq_factor=getattr(mc, "rope_scaling_low_freq_factor", None),
        high_freq_factor=getattr(mc, "rope_scaling_high_freq_factor", None),
    )
    cos_loc, sin_loc = rope.forward(local.embed(input_ids), position_ids=position_ids_local)

    def _stats(a: torch.Tensor, b: torch.Tensor) -> Dict[str, float]:
        a32 = a.detach().to(torch.float32, copy=False).view(-1)
        b32 = b.detach().to(torch.float32, copy=False).view(-1)
        diff = (a32 - b32)
        l2 = float(diff.norm().item())
        max_abs = float(diff.abs().max().item())
        mean_abs = float(diff.abs().mean().item())
        denom = float(max(1e-12, a32.norm().item()))
        rel = float(l2 / denom)
        cos = float(torch.nn.functional.cosine_similarity(a32.view(1, -1), b32.view(1, -1)).item())
        return {"l2": l2, "max_abs": max_abs, "mean_abs": mean_abs, "rel_l2": rel, "cos": cos}

    report: Dict[str, Any] = {"layers": []}

    # Compare masks and RoPE
    report["mask"] = {
        "local_shape": tuple(int(s) for s in add_mask_local.shape),
        "local_inf_counts": int(torch.isinf(add_mask_local).sum().item()),
    }
    report["rope"] = {
        "cos": _stats(cos_loc, cos_hf),
        "sin": _stats(sin_loc, sin_hf),
    }

    # Run local forward step-by-step to collect per-layer hidden
    local_layers_out: List[torch.Tensor] = []
    x_loc = local.embed(input_ids)
    mask_loc = add_mask_local
    pos_ids = position_ids_local
    for i, blk in enumerate(local.blocks[: int(args.max_layers) ]):
        x_loc = blk(x_loc, mask_loc, None, (cos_loc, sin_loc), pos_ids)
        local_layers_out.append(x_loc.detach().to(device="cpu", dtype=torch.float32))

    # Compare per-layer outputs
    for i in range(min(len(hf_layers_out), len(local_layers_out))):
        stats = _stats(local_layers_out[i], hf_layers_out[i])
        report["layers"].append({"index": i, **stats})

    # Final logits comparison (one step)
    out_loc = local(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
    logits_loc = out_loc["logits"][:, -1, :].to("cpu", torch.float32)
    loc_logits_path = None
    if args.save_logits:
        out_dir = args.out_dir or os.path.join(args.cache_dir, "trace")
        os.makedirs(out_dir, exist_ok=True)
        loc_logits_path = os.path.join(out_dir, "local_logits.pt")
        torch.save(logits_loc, loc_logits_path)
    # Load from disk for stats if we freed memory
    if logits_hf is None and hf_logits_path:
        logits_hf = torch.load(hf_logits_path, map_location="cpu").to(torch.float32)
    if args.save_logits and loc_logits_path is not None:
        # ensure we use the same on-disk tensor for consistency
        logits_loc = torch.load(loc_logits_path, map_location="cpu").to(torch.float32)
    report["logits"] = _stats(logits_loc, logits_hf)
    if args.save_logits:
        report["logits_paths"] = {"hf": hf_logits_path, "local": loc_logits_path}

    print(json.dumps(report, indent=2))

if __name__ == "__main__":
    main()


        cache_dir=args.cache_dir,
    ).eval()

    # Build masks and positions exactly once (no cache), compute HF RoPE and per-layer traces first
    input_ids = x["input_ids"]
    attention_mask = x.get("attention_mask", None)
    B, T = int(input_ids.shape[0]), int(input_ids.shape[1])

    # HF position_ids and mask
    position_ids_hf = torch.arange(T, device=args.device).unsqueeze(0)
    from transformers.models.llama.modeling_llama import LlamaModel  # type: ignore
    # Prepare rotary cos/sin via HF module
    llama_model: LlamaModel = hf.model  # type: ignore[attr-defined]
    cos_hf, sin_hf = llama_model.rotary_emb(
        llama_model.embed_tokens(input_ids), position_ids=position_ids_hf
    )

    # Run HF forward with hooks to collect per-layer hidden (store on CPU to save VRAM)
    hf_layers_out: List[torch.Tensor] = []

    def _hf_layer_hook(module, inputs, output):
        hs = output[0] if isinstance(output, tuple) else output
        hf_layers_out.append(hs.detach().to(device="cpu", dtype=torch.float32))

    hooks = []
    try:
        for i, layer in enumerate(llama_model.layers[: int(args.max_layers) ]):
            hooks.append(layer.register_forward_hook(_hf_layer_hook))
        _ = hf(input_ids=input_ids, attention_mask=attention_mask)
    finally:
        for h in hooks:
            try:
                h.remove()
            except Exception:
                pass

    # Final HF logits (CPU copy)
    out_hf = hf(input_ids=input_ids, attention_mask=attention_mask)
    logits_hf = out_hf.logits[:, -1, :].to("cpu", torch.float32)
    hf_logits_path = None
    if args.save_logits:
        out_dir = args.out_dir or os.path.join(args.cache_dir, "trace")
        os.makedirs(out_dir, exist_ok=True)


    add_mask_local = create_causal_mask(
        input_embeds=local.embed(input_ids),
        attention_mask=attention_mask,
        cache_position=None,
        position_ids=position_ids_local,
        past_key_values=None,
    )  # (B,1,T,S)
    rope = LocalRotary(
        head_dim=int(mc.head_dim or mc.d_model // mc.n_heads),
        base_theta=float(mc.rope_theta),
        attention_scaling=float(getattr(mc, "rope_attention_scaling", 1.0) or 1.0),
        device=input_ids.device,
        scaling_type=getattr(mc, "rope_scaling_type", None),
        scaling_factor=getattr(mc, "rope_scaling_factor", None),
        original_max_position_embeddings=getattr(mc, "rope_scaling_original_max_position_embeddings", None),
        low_freq_factor=getattr(mc, "rope_scaling_low_freq_factor", None),
        high_freq_factor=getattr(mc, "rope_scaling_high_freq_factor", None),
    )
    cos_loc, sin_loc = rope.forward(local.embed(input_ids), position_ids=position_ids_local)

    def _stats(a: torch.Tensor, b: torch.Tensor) -> Dict[str, float]:
        a32 = a.detach().to(torch.float32, copy=False).view(-1)
        b32 = b.detach().to(torch.float32, copy=False).view(-1)
        diff = (a32 - b32)
        l2 = float(diff.norm().item())
        max_abs = float(diff.abs().max().item())
        mean_abs = float(diff.abs().mean().item())
        denom = float(max(1e-12, a32.norm().item()))
        rel = float(l2 / denom)
        cos = float(torch.nn.functional.cosine_similarity(a32.view(1, -1), b32.view(1, -1)).item())
        return {"l2": l2, "max_abs": max_abs, "mean_abs": mean_abs, "rel_l2": rel, "cos": cos}

    report: Dict[str, Any] = {"layers": []}

    # Compare masks and RoPE
    report["mask"] = {
        "local_shape": tuple(int(s) for s in add_mask_local.shape),
        "local_inf_counts": int(torch.isinf(add_mask_local).sum().item()),
    }
    report["rope"] = {
        "cos": _stats(cos_loc, cos_hf),
        "sin": _stats(sin_loc, sin_hf),
    }

    # Run local forward step-by-step to collect per-layer hidden
    local_layers_out: List[torch.Tensor] = []
    x_loc = local.embed(input_ids)
    mask_loc = add_mask_local
    pos_ids = position_ids_local
    for i, blk in enumerate(local.blocks[: int(args.max_layers) ]):
        x_loc = blk(x_loc, mask_loc, None, (cos_loc, sin_loc), pos_ids)


import os
import json
from typing import List, Dict

def main() -> None:
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True, help="HF model id or local snapshot path")
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints")
    p.add_argument("--samples", type=int, default=0, help="Sample count from a default suite (0 uses all)")
    p.add_argument("--backend", default="local", choices=["hf","local","pure"], help="Which local tokenizer backend to compare against HF")
    args = p.parse_args()

    from model.hf_snapshot import ensure_snapshot
    snap = ensure_snapshot(args.model, args.cache_dir)

    # HF reference tokenizer
    try:
        from transformers import AutoTokenizer  # type: ignore
    except Exception as e:
        raise RuntimeError("Install transformers to run the comparison script") from e
    hf_tok = AutoTokenizer.from_pretrained(args.model, use_fast=True, cache_dir=args.cache_dir)
    if getattr(hf_tok, "pad_token", None) is None:
        try:
            hf_tok.pad_token = hf_tok.eos_token
        except Exception:
            pass

    # Local tokenizer
    if args.backend == "pure":
        from data.tokenizer import PureLlamaTokenizer
        local_tok = PureLlamaTokenizer(snap)
    else:
        from data.tokenizer import LocalLlamaTokenizer
        local_tok = LocalLlamaTokenizer(snap)

    tests: List[str] = [
        "Hello, world",
        "The quick brown fox jumps over 13 lazy dogs.",
        "  leading and  multiple   spaces  ",
        "tabs\tand\nnewlines\n",
        "emoji 😀😃😄 😁🖖🏽",
        "中文測試 — mixed-width punctuation、！",
        "العَرَبِيَّةُ لغةٌ جميلةٌ",
        "कृत्रिम बुद्धिमत्ता से नमस्ते",
        "geschichten999",
        "<|begin_of_text|> special token check <|end_of_text|>",
    ]
    if args.samples and args.samples > 0:
        tests = tests[: int(args.samples)]

    results: Dict[str, Dict] = {}
    ok = 0
    for t in tests:
        hf_ids = list(hf_tok.encode(t, add_special_tokens=False))
        loc_ids = local_tok.encode(t)
        enc_eq = (hf_ids == loc_ids)
        # Decode both from HF ids for stable baseline
        hf_dec = hf_tok.decode(hf_ids)
        loc_dec_from_hf = local_tok.decode(hf_ids)
        # Decode both from local ids
        loc_dec = local_tok.decode(loc_ids)
        hf_dec_from_loc = hf_tok.decode(loc_ids)
        results[t] = {
            "enc_equal": enc_eq,
            "hf_ids": hf_ids[:64],
            "loc_ids": loc_ids[:64],
            "hf_dec": hf_dec,
            "loc_dec_from_hf": loc_dec_from_hf,
            "loc_dec": loc_dec,
            "hf_dec_from_loc": hf_dec_from_loc,
        }
        ok += int(enc_eq)

    out = {
        "snapshot": snap,
        "hf_special": {
            "bos": getattr(hf_tok, "bos_token_id", None),
            "eos": getattr(hf_tok, "eos_token_id", None),
            "unk": getattr(hf_tok, "unk_token_id", None),
            "pad": getattr(hf_tok, "pad_token_id", None),
        },
        "local_special": {
            "bos": local_tok.bos_token_id,
            "eos": local_tok.eos_token_id,
            "unk": local_tok.unk_token_id,
            "pad": local_tok.pad_token_id,
        },
        "tests": results,
        "enc_equal_count": ok,
        "enc_total": len(tests),
    }
    print(json.dumps(out, ensure_ascii=False, indent=2))

if __name__ == "__main__":
    main()


import os
import json
import argparse
from typing import Dict, List, Tuple

import torch

def _ensure_snapshot(model_id: str, cache_dir: str) -> str:
    if os.path.isdir(model_id) and os.path.isfile(os.path.join(model_id, "config.json")):
        return model_id
    org_name = model_id.strip().split("/")[-2:]
    if len(org_name) == 2:
        org, name = org_name
        dir1 = os.path.join(cache_dir, f"models--{org}--{name}", "snapshots")
        cands = []
        if os.path.isdir(dir1):
            cands.extend([os.path.join(dir1, d) for d in os.listdir(dir1)])
        cands = [p for p in cands if os.path.isfile(os.path.join(p, "model.safetensors.index.json"))]
        if cands:
            cands.sort(key=lambda p: os.path.getmtime(p), reverse=True)
            return cands[0]
    from huggingface_hub import snapshot_download  # type: ignore
    return snapshot_download(repo_id=model_id, cache_dir=cache_dir)

def _index_shards(model_dir: str) -> Dict[str, List[str]]:
    index_fp = os.path.join(model_dir, "model.safetensors.index.json")
    if not os.path.isfile(index_fp):
        raise FileNotFoundError(f"Missing index file: {index_fp}")
    with open(index_fp, "r", encoding="utf-8") as fh:
        index_obj = json.load(fh)
    weight_map: Dict[str, str] = index_obj.get("weight_map", {})
    by_file: Dict[str, List[str]] = {}
    for k, fn in weight_map.items():
        by_file.setdefault(fn, []).append(k)
    return by_file

def _read_hf_tensor(model_dir: str, key: str) -> torch.Tensor:
    from safetensors import safe_open  # type: ignore
    by_file = _index_shards(model_dir)
    for shard_rel, keys in by_file.items():
        if key not in keys:
            continue
        fp = os.path.join(model_dir, shard_rel)
        with safe_open(fp, framework="pt", device="cpu") as f:  # type: ignore
            return f.get_tensor(key)  # type: ignore[attr-defined]
    raise KeyError(f"tensor {key} not found in shards")

def _metrics(a: torch.Tensor, b: torch.Tensor) -> Dict[str, float]:
    a32 = a.detach().to(torch.float32, copy=False)
    b32 = b.detach().to(torch.float32, copy=False)
    diff = (a32 - b32).view(-1)
    l2 = float(diff.norm().item())
    max_abs = float(diff.abs().max().item())
    denom = float(max(1e-12, a32.norm().item()))
    rel = float(l2 / denom)
    cos = float((torch.nn.functional.cosine_similarity(a32.view(1, -1), b32.view(1, -1))).item())
    return {"l2": l2, "max_abs": max_abs, "rel_l2": rel, "cos": cos}

def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="meta-llama/Llama-3.1-8B-Instruct")
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints")
    p.add_argument("--layers", default="0", help="Comma-separated layer indices to compare")
    args = p.parse_args()

    os.makedirs(args.cache_dir, exist_ok=True)
    ckpt = _ensure_snapshot(args.model, args.cache_dir)

    # Build local model and load HF weights via our loader
    from specs.config import ModelConfig
    from model.factory import build_causal_lm
    from model.hf_llama_loader import load_hf_llama_weights_into_local

    cfg = json.load(open(os.path.join(ckpt, "config.json"), "r", encoding="utf-8"))
    use_cuda = torch.cuda.is_available()
    mc = ModelConfig(
        d_model=int(cfg.get("hidden_size")),
        n_heads=int(cfg.get("num_attention_heads")),
        n_layers=int(cfg.get("num_hidden_layers")),
        d_ff=int(cfg.get("intermediate_size")),
        vocab_size=int(cfg.get("vocab_size")),
        head_dim=int(cfg.get("head_dim", int(cfg.get("hidden_size")) // int(cfg.get("num_attention_heads")))),
        rope_theta=float(cfg.get("rope_parameters", {}).get("rope_theta", cfg.get("rope_theta", 1e6))),
        dtype=("bfloat16" if use_cuda else "float32"),
        attn_impl="sdpa",
        rms_norm_eps=float(cfg.get("rms_norm_eps", 1e-6)),
        rope_scaling_type=(cfg.get("rope_scaling", {}) or {}).get("type"),
        rope_scaling_factor=(cfg.get("rope_scaling", {}) or {}).get("factor"),
        rope_scaling_original_max_position_embeddings=(cfg.get("rope_scaling", {}) or {}).get("original_max_position_embeddings"),
        rope_scaling_low_freq_factor=(cfg.get("rope_scaling", {}) or {}).get("low_freq_factor"),
        rope_scaling_high_freq_factor=(cfg.get("rope_scaling", {}) or {}).get("high_freq_factor"),
    )
    n_kv_heads = int(cfg.get("num_key_value_heads", mc.n_heads))
    tie_we = bool(cfg.get("tie_word_embeddings", True))
    local = build_causal_lm(mc, block="llama", n_kv_heads=n_kv_heads, tie_weights=tie_we)
    load_hf_llama_weights_into_local(local, ckpt)
    local = (local.to(device="cuda", dtype=torch.bfloat16) if use_cuda else local.to(device="cpu", dtype=torch.float32)).eval()

    report: Dict[str, Dict[str, float]] = {}

    # Embedding and lm_head
    try:
        hf_embed = _read_hf_tensor(ckpt, "model.embed_tokens.weight").to("cpu")
        report["embed.weight"] = _metrics(local.embed.weight.data.detach().to("cpu"), hf_embed)
    except Exception:
        report["embed.weight"] = {"error": 1.0}  # type: ignore
    try:
        hf_lm = _read_hf_tensor(ckpt, "lm_head.weight").to("cpu")
        report["lm_head.weight"] = _metrics(local.lm_head.weight.data.detach().to("cpu"), hf_lm)
    except Exception:
        report["lm_head.weight"] = {"error": 1.0}  # type: ignore

    layers = [int(x.strip()) for x in str(args.layers).split(",") if x.strip()]
    for li in layers:
        key = f"layers.{li}"
        try:
            blk = local.blocks[li]
        except Exception:
            report[key] = {"error": 1.0}  # type: ignore
            continue

        # Attention projections
        for short, hf_name, local_w in (
            ("q_proj", f"model.layers.{li}.self_attn.q_proj.weight", blk.attn.w_q.weight),
            ("k_proj", f"model.layers.{li}.self_attn.k_proj.weight", blk.attn.w_k.weight),
            ("v_proj", f"model.layers.{li}.self_attn.v_proj.weight", blk.attn.w_v.weight),
            ("o_proj", f"model.layers.{li}.self_attn.o_proj.weight", blk.attn.w_o.weight),
        ):
            try:
                hf_t = _read_hf_tensor(ckpt, hf_name).to("cpu")
                report[f"{key}.{short}"] = _metrics(local_w.data.detach().to("cpu"), hf_t)
            except Exception:
                report[f"{key}.{short}"] = {"error": 1.0}  # type: ignore

        # Layer norms
        for short, hf_name, local_w in (
            ("n1", f"model.layers.{li}.input_layernorm.weight", blk.n1.weight),
            ("n2", f"model.layers.{li}.post_attention_layernorm.weight", blk.n2.weight),
        ):
            try:
                hf_t = _read_hf_tensor(ckpt, hf_name).to("cpu")
                report[f"{key}.{short}"] = _metrics(local_w.data.detach().to("cpu"), hf_t)
            except Exception:
                report[f"{key}.{short}"] = {"error": 1.0}  # type: ignore

        # MLP down (w_out)
        try:
            hf_down = _read_hf_tensor(ckpt, f"model.layers.{li}.mlp.down_proj.weight").to("cpu")
            report[f"{key}.mlp.down"] = _metrics(blk.mlp.w_out.weight.data.detach().to("cpu"), hf_down)
        except Exception:
            report[f"{key}.mlp.down"] = {"error": 1.0}  # type: ignore

        # MLP fused w_in vs concat(gate, up)
        try:
            hf_gate = _read_hf_tensor(ckpt, f"model.layers.{li}.mlp.gate_proj.weight").to("cpu")
            hf_up = _read_hf_tensor(ckpt, f"model.layers.{li}.mlp.up_proj.weight").to("cpu")
            fused = torch.cat([hf_gate, hf_up], dim=0).contiguous()
            report[f"{key}.mlp.w_in"] = _metrics(blk.mlp.w_in.weight.data.detach().to("cpu"), fused)
        except Exception:
            report[f"{key}.mlp.w_in"] = {"error": 1.0}  # type: ignore

    print(json.dumps(report, indent=2))

if __name__ == "__main__":
    main()


import os
import json
import argparse
from typing import Dict, List, Tuple

import torch

def _ensure_snapshot(model_id: str, cache_dir: str) -> str:
    if os.path.isdir(model_id) and os.path.isfile(os.path.join(model_id, "config.json")):
        return model_id
    org_name = model_id.strip().split("/")[-2:]
    if len(org_name) == 2:
        org, name = org_name
        dir1 = os.path.join(cache_dir, f"models--{org}--{name}", "snapshots")
        cands = []
        if os.path.isdir(dir1):
            cands.extend([os.path.join(dir1, d) for d in os.listdir(dir1)])
        cands = [p for p in cands if os.path.isfile(os.path.join(p, "model.safetensors.index.json"))]
        if cands:
            cands.sort(key=lambda p: os.path.getmtime(p), reverse=True)
            return cands[0]
    from huggingface_hub import snapshot_download  # type: ignore
    return snapshot_download(repo_id=model_id, cache_dir=cache_dir)

def _index_shards(model_dir: str) -> Dict[str, List[str]]:
    index_fp = os.path.join(model_dir, "model.safetensors.index.json")
    if not os.path.isfile(index_fp):
        raise FileNotFoundError(f"Missing index file: {index_fp}")
    with open(index_fp, "r", encoding="utf-8") as fh:
        index_obj = json.load(fh)
    weight_map: Dict[str, str] = index_obj.get("weight_map", {})
    by_file: Dict[str, List[str]] = {}
    for k, fn in weight_map.items():
        by_file.setdefault(fn, []).append(k)
    return by_file

def _read_hf_tensor(model_dir: str, key: str) -> torch.Tensor:
    from safetensors import safe_open  # type: ignore
    by_file = _index_shards(model_dir)
    for shard_rel, keys in by_file.items():
        if key not in keys:


def _ensure_snapshot(model_id: str, cache_dir: str) -> str:
    if os.path.isdir(model_id) and os.path.isfile(os.path.join(model_id, "config.json")):
        return model_id
    org_name = model_id.strip().split("/")[-2:]
    if len(org_name) == 2:
        org, name = org_name
        dir1 = os.path.join(cache_dir, f"models--{org}--{name}", "snapshots")
        cands = []
        if os.path.isdir(dir1):
            cands.extend([os.path.join(dir1, d) for d in os.listdir(dir1)])
        cands = [p for p in cands if os.path.isfile(os.path.join(p, "model.safetensors.index.json"))]
        if cands:
            cands.sort(key=lambda p: os.path.getmtime(p), reverse=True)
            return cands[0]
    from huggingface_hub import snapshot_download  # type: ignore
    return snapshot_download(repo_id=model_id, cache_dir=cache_dir)

def _index_shards(model_dir: str) -> Dict[str, List[str]]:
    index_fp = os.path.join(model_dir, "model.safetensors.index.json")
    if not os.path.isfile(index_fp):
        raise FileNotFoundError(f"Missing index file: {index_fp}")
    with open(index_fp, "r", encoding="utf-8") as fh:
        index_obj = json.load(fh)
    weight_map: Dict[str, str] = index_obj.get("weight_map", {})
    by_file: Dict[str, List[str]] = {}
    for k, fn in weight_map.items():
        by_file.setdefault(fn, []).append(k)
    return by_file

def _read_hf_tensor(model_dir: str, key: str) -> torch.Tensor:
    from safetensors import safe_open  # type: ignore
    by_file = _index_shards(model_dir)
    for shard_rel, keys in by_file.items():
        if key not in keys:
            continue
        fp = os.path.join(model_dir, shard_rel)
        with safe_open(fp, framework="pt", device="cpu") as f:  # type: ignore
            return f.get_tensor(key)  # type: ignore[attr-defined]
    raise KeyError(f"tensor {key} not found in shards")

def _metrics(a: torch.Tensor, b: torch.Tensor) -> Dict[str, float]:
    a32 = a.detach().to(torch.float32, copy=False)
    b32 = b.detach().to(torch.float32, copy=False)
    diff = (a32 - b32).view(-1)
    l2 = float(diff.norm().item())
    max_abs = float(diff.abs().max().item())


        if cands:
            cands.sort(key=lambda p: os.path.getmtime(p), reverse=True)
            return cands[0]
    from huggingface_hub import snapshot_download  # type: ignore
    return snapshot_download(repo_id=model_id, cache_dir=cache_dir)

def _index_shards(model_dir: str) -> Dict[str, List[str]]:
    index_fp = os.path.join(model_dir, "model.safetensors.index.json")
    if not os.path.isfile(index_fp):
        raise FileNotFoundError(f"Missing index file: {index_fp}")
    with open(index_fp, "r", encoding="utf-8") as fh:
        index_obj = json.load(fh)
    weight_map: Dict[str, str] = index_obj.get("weight_map", {})
    by_file: Dict[str, List[str]] = {}
    for k, fn in weight_map.items():
        by_file.setdefault(fn, []).append(k)
    return by_file

def _read_hf_tensor(model_dir: str, key: str) -> torch.Tensor:
    from safetensors import safe_open  # type: ignore
    by_file = _index_shards(model_dir)
    for shard_rel, keys in by_file.items():
        if key not in keys:
            continue
        fp = os.path.join(model_dir, shard_rel)
        with safe_open(fp, framework="pt", device="cpu") as f:  # type: ignore
            return f.get_tensor(key)  # type: ignore[attr-defined]
    raise KeyError(f"tensor {key} not found in shards")

def _metrics(a: torch.Tensor, b: torch.Tensor) -> Dict[str, float]:
    a32 = a.detach().to(torch.float32, copy=False)
    b32 = b.detach().to(torch.float32, copy=False)
    diff = (a32 - b32).view(-1)
    l2 = float(diff.norm().item())
    max_abs = float(diff.abs().max().item())
    denom = float(max(1e-12, a32.norm().item()))
    rel = float(l2 / denom)
    cos = float((torch.nn.functional.cosine_similarity(a32.view(1, -1), b32.view(1, -1))).item())
    return {"l2": l2, "max_abs": max_abs, "rel_l2": rel, "cos": cos}

def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="meta-llama/Llama-3.1-8B-Instruct")
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints")
    p.add_argument("--layers", default="0", help="Comma-separated layer indices to compare")
    args = p.parse_args()


        index_obj = json.load(fh)
    weight_map: Dict[str, str] = index_obj.get("weight_map", {})
    by_file: Dict[str, List[str]] = {}
    for k, fn in weight_map.items():
        by_file.setdefault(fn, []).append(k)
    return by_file

def _read_hf_tensor(model_dir: str, key: str) -> torch.Tensor:
    from safetensors import safe_open  # type: ignore
    by_file = _index_shards(model_dir)
    for shard_rel, keys in by_file.items():
        if key not in keys:
            continue
        fp = os.path.join(model_dir, shard_rel)
        with safe_open(fp, framework="pt", device="cpu") as f:  # type: ignore
            return f.get_tensor(key)  # type: ignore[attr-defined]
    raise KeyError(f"tensor {key} not found in shards")

def _metrics(a: torch.Tensor, b: torch.Tensor) -> Dict[str, float]:
    a32 = a.detach().to(torch.float32, copy=False)
    b32 = b.detach().to(torch.float32, copy=False)
    diff = (a32 - b32).view(-1)
    l2 = float(diff.norm().item())
    max_abs = float(diff.abs().max().item())
    denom = float(max(1e-12, a32.norm().item()))
    rel = float(l2 / denom)
    cos = float((torch.nn.functional.cosine_similarity(a32.view(1, -1), b32.view(1, -1))).item())
    return {"l2": l2, "max_abs": max_abs, "rel_l2": rel, "cos": cos}

def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="meta-llama/Llama-3.1-8B-Instruct")
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints")
    p.add_argument("--layers", default="0", help="Comma-separated layer indices to compare")
    args = p.parse_args()

    os.makedirs(args.cache_dir, exist_ok=True)
    ckpt = _ensure_snapshot(args.model, args.cache_dir)

    # Build local model and load HF weights via our loader
    from specs.config import ModelConfig
    from model.factory import build_causal_lm
    from model.hf_llama_loader import load_hf_llama_weights_into_local

    cfg = json.load(open(os.path.join(ckpt, "config.json"), "r", encoding="utf-8"))
    use_cuda = torch.cuda.is_available()
    mc = ModelConfig(


main

        if key not in keys:
            continue
        fp = os.path.join(model_dir, shard_rel)
        with safe_open(fp, framework="pt", device="cpu") as f:  # type: ignore
            return f.get_tensor(key)  # type: ignore[attr-defined]
    raise KeyError(f"tensor {key} not found in shards")

def _metrics(a: torch.Tensor, b: torch.Tensor) -> Dict[str, float]:
    a32 = a.detach().to(torch.float32, copy=False)
    b32 = b.detach().to(torch.float32, copy=False)
    diff = (a32 - b32).view(-1)
    l2 = float(diff.norm().item())
    max_abs = float(diff.abs().max().item())
    denom = float(max(1e-12, a32.norm().item()))
    rel = float(l2 / denom)
    cos = float((torch.nn.functional.cosine_similarity(a32.view(1, -1), b32.view(1, -1))).item())
    return {"l2": l2, "max_abs": max_abs, "rel_l2": rel, "cos": cos}

def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="meta-llama/Llama-3.1-8B-Instruct")
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints")
    p.add_argument("--layers", default="0", help="Comma-separated layer indices to compare")
    args = p.parse_args()

    os.makedirs(args.cache_dir, exist_ok=True)
    ckpt = _ensure_snapshot(args.model, args.cache_dir)

    # Build local model and load HF weights via our loader
    from specs.config import ModelConfig
    from model.factory import build_causal_lm
    from model.hf_llama_loader import load_hf_llama_weights_into_local

    cfg = json.load(open(os.path.join(ckpt, "config.json"), "r", encoding="utf-8"))
    use_cuda = torch.cuda.is_available()
    mc = ModelConfig(
        d_model=int(cfg.get("hidden_size")),
        n_heads=int(cfg.get("num_attention_heads")),
        n_layers=int(cfg.get("num_hidden_layers")),
        d_ff=int(cfg.get("intermediate_size")),
        vocab_size=int(cfg.get("vocab_size")),
        head_dim=int(cfg.get("head_dim", int(cfg.get("hidden_size")) // int(cfg.get("num_attention_heads")))),
        rope_theta=float(cfg.get("rope_parameters", {}).get("rope_theta", cfg.get("rope_theta", 1e6))),
        dtype=("bfloat16" if use_cuda else "float32"),
        attn_impl="sdpa",
        rms_norm_eps=float(cfg.get("rms_norm_eps", 1e-6)),
        rope_scaling_type=(cfg.get("rope_scaling", {}) or {}).get("type"),
        rope_scaling_factor=(cfg.get("rope_scaling", {}) or {}).get("factor"),
        rope_scaling_original_max_position_embeddings=(cfg.get("rope_scaling", {}) or {}).get("original_max_position_embeddings"),
        rope_scaling_low_freq_factor=(cfg.get("rope_scaling", {}) or {}).get("low_freq_factor"),
        rope_scaling_high_freq_factor=(cfg.get("rope_scaling", {}) or {}).get("high_freq_factor"),
    )
    n_kv_heads = int(cfg.get("num_key_value_heads", mc.n_heads))
    tie_we = bool(cfg.get("tie_word_embeddings", True))
    local = build_causal_lm(mc, block="llama", n_kv_heads=n_kv_heads, tie_weights=tie_we)
    load_hf_llama_weights_into_local(local, ckpt)
    local = (local.to(device="cuda", dtype=torch.bfloat16) if use_cuda else local.to(device="cpu", dtype=torch.float32)).eval()

    report: Dict[str, Dict[str, float]] = {}

    # Embedding and lm_head
    try:
        hf_embed = _read_hf_tensor(ckpt, "model.embed_tokens.weight").to("cpu")
        report["embed.weight"] = _metrics(local.embed.weight.data.detach().to("cpu"), hf_embed)
    except Exception:
        report["embed.weight"] = {"error": 1.0}  # type: ignore
    try:
        hf_lm = _read_hf_tensor(ckpt, "lm_head.weight").to("cpu")
        report["lm_head.weight"] = _metrics(local.lm_head.weight.data.detach().to("cpu"), hf_lm)
    except Exception:
        report["lm_head.weight"] = {"error": 1.0}  # type: ignore

    layers = [int(x.strip()) for x in str(args.layers).split(",") if x.strip()]
    for li in layers:
        key = f"layers.{li}"
        try:
            blk = local.blocks[li]
        except Exception:
            report[key] = {"error": 1.0}  # type: ignore
            continue

        # Attention projections
        for short, hf_name, local_w in (
            ("q_proj", f"model.layers.{li}.self_attn.q_proj.weight", blk.attn.w_q.weight),
            ("k_proj", f"model.layers.{li}.self_attn.k_proj.weight", blk.attn.w_k.weight),
            ("v_proj", f"model.layers.{li}.self_attn.v_proj.weight", blk.attn.w_v.weight),
            ("o_proj", f"model.layers.{li}.self_attn.o_proj.weight", blk.attn.w_o.weight),
        ):
            try:
                hf_t = _read_hf_tensor(ckpt, hf_name).to("cpu")
                report[f"{key}.{short}"] = _metrics(local_w.data.detach().to("cpu"), hf_t)
            except Exception:
                report[f"{key}.{short}"] = {"error": 1.0}  # type: ignore

        # Layer norms
        for short, hf_name, local_w in (
            ("n1", f"model.layers.{li}.input_layernorm.weight", blk.n1.weight),
            ("n2", f"model.layers.{li}.post_attention_layernorm.weight", blk.n2.weight),
        ):
            try:
                hf_t = _read_hf_tensor(ckpt, hf_name).to("cpu")
                report[f"{key}.{short}"] = _metrics(local_w.data.detach().to("cpu"), hf_t)
            except Exception:
                report[f"{key}.{short}"] = {"error": 1.0}  # type: ignore

        # MLP down (w_out)
        try:
            hf_down = _read_hf_tensor(ckpt, f"model.layers.{li}.mlp.down_proj.weight").to("cpu")
            report[f"{key}.mlp.down"] = _metrics(blk.mlp.w_out.weight.data.detach().to("cpu"), hf_down)
        except Exception:
            report[f"{key}.mlp.down"] = {"error": 1.0}  # type: ignore

        # MLP fused w_in vs concat(gate, up)
        try:
            hf_gate = _read_hf_tensor(ckpt, f"model.layers.{li}.mlp.gate_proj.weight").to("cpu")
            hf_up = _read_hf_tensor(ckpt, f"model.layers.{li}.mlp.up_proj.weight").to("cpu")
            fused = torch.cat([hf_gate, hf_up], dim=0).contiguous()
            report[f"{key}.mlp.w_in"] = _metrics(blk.mlp.w_in.weight.data.detach().to("cpu"), fused)
        except Exception:
            report[f"{key}.mlp.w_in"] = {"error": 1.0}  # type: ignore

    print(json.dumps(report, indent=2))

if __name__ == "__main__":
    main()


import os
import sys
import json
import gc
import torch
from typing import Tuple
from model.hf_snapshot import ensure_snapshot
from model.runtime_utils import local_logits_last

def main() -> None:
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="meta-llama/Llama-3.1-8B-Instruct")
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints")
    p.add_argument("--prompt", default="Hello")
    p.add_argument("--max-tokens", type=int, default=64)
    p.add_argument("--gen-new", type=int, default=32, help="Generate this many new tokens for HF/local side-by-side")
    p.add_argument("--mode", default="hf_guided", choices=["independent", "hf_guided", "single_step", "independent_exact", "two_pass_exact"], help="Verification mode")
    p.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"], help="Computation dtype for both models during verify")
    p.add_argument("--save-logits", action="store_true", help="Save HF/local last-step logits to disk to minimize RAM/VRAM use")
    p.add_argument("--parity-exact", action="store_true", help="Use explicit matmul+float32 softmax attention for exact HF parity")
    p.add_argument("--out-dir", default="", help="Directory to write saved tensors (defaults to cache-dir/trace)")
    args = p.parse_args()

    os.makedirs(args.cache_dir, exist_ok=True)
    # Force stable attention backend for parity
    os.environ.setdefault("ATTN_BACKEND", "torch")
    if args.parity_exact:
        os.environ["ATTN_PARITY_EXACT"] = "1"

    from transformers import AutoTokenizer, AutoModelForCausalLM  # type: ignore
    tok = AutoTokenizer.from_pretrained(args.model, use_fast=True, cache_dir=args.cache_dir)
    x = tok(args.prompt, return_tensors="pt")

    use_cuda = torch.cuda.is_available()
    req_dtype = torch.bfloat16 if str(args.dtype) == "bfloat16" and use_cuda else torch.float32
    hf_device_map = "auto" if use_cuda else None

    # HF reference (load in lower precision on GPU to reduce RAM)
    hf = AutoModelForCausalLM.from_pretrained(
        args.model,
        torch_dtype=req_dtype,
        device_map=hf_device_map,
        low_cpu_mem_usage=True,
        cache_dir=args.cache_dir,
    )
    hf.eval()
    with torch.no_grad():
        # Ensure inputs are on the same device as the HF model
        try:
            hf_dev = next(hf.parameters()).device
        except Exception:
            hf_dev = torch.device("cpu")
        x_hf = {k: v.to(hf_dev) for k, v in x.items()}
        out_hf = hf(**x_hf)
        logits_hf = out_hf.logits[:, -1, :].to(device="cpu", dtype=torch.float32)
        hf_logits_path = None
        if args.save_logits:
            out_dir = args.out_dir or os.path.join(args.cache_dir, "trace")
            try:
                os.makedirs(out_dir, exist_ok=True)
            except Exception:
                pass
            hf_logits_path = os.path.join(out_dir, "hf_logits.pt")
            try:
                torch.save(logits_hf, hf_logits_path)
            except Exception:
                hf_logits_path = None

    # Also capture a short HF generation for comparison (force pure greedy, no suppression)
    try:
        try:
            gen_cfg = hf.generation_config  # type: ignore[attr-defined]
            gen_cfg = gen_cfg.clone()
            # Force deterministic greedy without suppression/penalties
            gen_cfg.do_sample = False
            gen_cfg.temperature = 1.0
            gen_cfg.top_p = 1.0
            gen_cfg.top_k = None
            gen_cfg.repetition_penalty = 1.0
            gen_cfg.no_repeat_ngram_size = 0
            # Neutralize suppression lists if present
            if hasattr(gen_cfg, "suppress_tokens"):
                try:
                    gen_cfg.suppress_tokens = None  # type: ignore[attr-defined]
                except Exception:
                    pass
            if hasattr(gen_cfg, "begin_suppress_tokens"):
                try:
                    gen_cfg.begin_suppress_tokens = None  # type: ignore[attr-defined]
                except Exception:
                    pass
        except Exception:
            gen_cfg = None  # best-effort
        gen_hf = hf.generate(
            **x_hf,
            generation_config=gen_cfg,
            max_new_tokens=int(args.gen_new),
            do_sample=False,
            eos_token_id=tok.eos_token_id,
        )
        try:
            in_len = int(x_hf["input_ids"].shape[1])
            seq = gen_hf[0]
            new = seq[in_len:] if (hasattr(seq, "shape") and seq.shape[0] > in_len) else seq
            text_hf = tok.decode(new, skip_special_tokens=True)
        except Exception:
            text_hf = tok.decode(gen_hf[0], skip_special_tokens=True)
    except Exception:
        text_hf = "[error] hf generate failed"

    # Keep HF loaded when it's needed later for stepwise checks
    keep_hf_loaded = args.mode in ("hf_guided", "single_step", "independent_exact", "two_pass_exact")

    # If not needed later, free HF before building local to avoid peak GPU memory
    if not keep_hf_loaded:
        try:
            del out_hf
        except Exception:
            pass
        try:
            del hf
        except Exception:
            pass
        gc.collect()
        try:
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        except Exception:
            pass

    # Local
    ckpt_dir = ensure_snapshot(args.model, args.cache_dir)
    from specs.config import ModelConfig
    from model.factory import build_causal_lm
    from model.hf_llama_loader import load_hf_llama_weights_into_local

    cfg = json.load(open(os.path.join(ckpt_dir, "config.json"), "r", encoding="utf-8"))
    # Derive pad token id from tokenizer or config
    pad_id = tok.pad_token_id
    if pad_id is None:
        try:
            pad_id = int(cfg.get("pad_token_id")) if (cfg.get("pad_token_id") is not None) else None
        except Exception:
            pad_id = None
    mc = ModelConfig(
        d_model=int(cfg.get("hidden_size")),
        n_heads=int(cfg.get("num_attention_heads")),
        n_layers=int(cfg.get("num_hidden_layers")),
        d_ff=int(cfg.get("intermediate_size")),
        vocab_size=int(cfg.get("vocab_size")),
        head_dim=int(cfg.get("head_dim", int(cfg.get("hidden_size")) // int(cfg.get("num_attention_heads")))),
        rope_theta=float(cfg.get("rope_parameters", {}).get("rope_theta", cfg.get("rope_theta", 1e6))),
        dtype=("bfloat16" if (req_dtype == torch.bfloat16) else "float32"),
        attn_impl="sdpa",
        rms_norm_eps=float(cfg.get("rms_norm_eps", 1e-6)),
        rope_scaling_type=(cfg.get("rope_scaling", {}) or {}).get("type"),
        rope_scaling_factor=(cfg.get("rope_scaling", {}) or {}).get("factor"),
        rope_scaling_original_max_position_embeddings=(cfg.get("rope_scaling", {}) or {}).get("original_max_position_embeddings"),
        rope_scaling_low_freq_factor=(cfg.get("rope_scaling", {}) or {}).get("low_freq_factor"),
        rope_scaling_high_freq_factor=(cfg.get("rope_scaling", {}) or {}).get("high_freq_factor"),
        pad_token_id=(int(pad_id) if pad_id is not None else None),
    )
    n_kv_heads = int(cfg.get("num_key_value_heads", mc.n_heads))
    tie_we = bool(cfg.get("tie_word_embeddings", True))
    local = build_causal_lm(mc, block="llama", n_kv_heads=n_kv_heads, tie_weights=tie_we)
    load_hf_llama_weights_into_local(local, ckpt_dir)
    if use_cuda:
        local = local.to(device="cuda", dtype=req_dtype).eval()
        x_cuda = {k: v.to("cuda") for k, v in x.items()}
        with torch.no_grad():
            out_loc_fwd = local(input_ids=x_cuda["input_ids"], attention_mask=x_cuda.get("attention_mask", None), return_dict=True)
            logits_local = out_loc_fwd["logits"][:, -1, :].to(device="cpu", dtype=torch.float32)
            loc_logits_path = None
            if args.save_logits:
                out_dir = args.out_dir or os.path.join(args.cache_dir, "trace")
                try:
                    os.makedirs(out_dir, exist_ok=True)
                except Exception:
                    pass
                loc_logits_path = os.path.join(out_dir, "local_logits.pt")
                try:
                    torch.save(logits_local, loc_logits_path)
                except Exception:
                    loc_logits_path = None
        # Local generation (greedy); exact parity can use stepwise no-cache recompute
        try:
            if args.parity_exact:
                seq = x_cuda["input_ids"].clone()
                attn = x_cuda.get("attention_mask", None)
                for _ in range(int(args.gen_new)):
                    out_step = local(input_ids=seq[:, -1:], attention_mask=attn, return_dict=True)
                    logits = out_step["logits"][:, -1, :]
                    nxt = torch.argmax(logits, dim=-1).view(1, 1)
                    seq = torch.cat([seq, nxt], dim=1)
                    if attn is not None:
                        attn = torch.cat([attn, torch.ones_like(nxt, device=attn.device, dtype=attn.dtype)], dim=1)
                    if tok.eos_token_id is not None and int(nxt[0,0].item()) == int(tok.eos_token_id):
                        break
                s0 = seq[0].to("cpu")
                in_len = int(x_cuda["input_ids"].shape[1])
                new = s0[in_len:] if (hasattr(s0, "shape") and s0.shape[0] > in_len) else s0
                text_local = tok.decode(new, skip_special_tokens=True)
            else:
                out_loc = local.generate(
                    input_ids=x_cuda["input_ids"],
                    attention_mask=x_cuda.get("attention_mask", None),
                    max_new_tokens=int(args.gen_new),
                    do_sample=False,
                    eos_token_id=tok.eos_token_id,
                    return_dict=True,
                )
                seq = out_loc.get("sequences", None)
                if seq is None:
                    seq = out_loc  # best-effort
                in_len = int(x_cuda["input_ids"].shape[1])
                s0 = seq[0].to("cpu")
                new = s0[in_len:] if (hasattr(s0, "shape") and s0.shape[0] > in_len) else s0
                text_local = tok.decode(new, skip_special_tokens=True)
        except Exception:
            text_local = "[error] local generate failed"
    else:
        local = local.to(device="cpu", dtype=torch.float32).eval()
        with torch.no_grad():
            out_loc_fwd = local(input_ids=x["input_ids"], attention_mask=x.get("attention_mask", None), return_dict=True)
            logits_local = out_loc_fwd["logits"][:, -1, :]
            loc_logits_path = None
            if args.save_logits:
                out_dir = args.out_dir or os.path.join(args.cache_dir, "trace")
                try:
                    os.makedirs(out_dir, exist_ok=True)
                except Exception:
                    pass
                loc_logits_path = os.path.join(out_dir, "local_logits.pt")
                try:
                    torch.save(logits_local.to(device="cpu", dtype=torch.float32), loc_logits_path)
                except Exception:
                    loc_logits_path = None
        try:
            if args.parity_exact:
                seq = x["input_ids"].clone()
                attn = x.get("attention_mask", None)
                for _ in range(int(args.gen_new)):
                    out_step = local(input_ids=seq[:, -1:], attention_mask=attn, return_dict=True)
                    logits = out_step["logits"][:, -1, :]
                    nxt = torch.argmax(logits, dim=-1).view(1, 1)
                    seq = torch.cat([seq, nxt], dim=1)
                    if attn is not None:
                        attn = torch.cat([attn, torch.ones_like(nxt, dtype=attn.dtype)], dim=1)
                    if tok.eos_token_id is not None and int(nxt[0,0].item()) == int(tok.eos_token_id):
                        break
                s0 = seq[0]
                in_len = int(x["input_ids"].shape[1])
                new = s0[in_len:] if (hasattr(s0, "shape") and s0.shape[0] > in_len) else s0
                text_local = tok.decode(new, skip_special_tokens=True)
            else:
                out_loc = local.generate(
                    input_ids=x["input_ids"],
                    attention_mask=x.get("attention_mask", None),
                    max_new_tokens=int(args.gen_new),
                    do_sample=False,
                    eos_token_id=tok.eos_token_id,
                    return_dict=True,
                )
                seq = out_loc.get("sequences", None)
                if seq is None:
                    seq = out_loc
                in_len = int(x["input_ids"].shape[1])
                s0 = seq[0]
                new = s0[in_len:] if (hasattr(s0, "shape") and s0.shape[0] > in_len) else s0
                text_local = tok.decode(new, skip_special_tokens=True)
        except Exception:
            text_local = "[error] local generate failed"

    # Guided/single-step parity and two-pass exact (single model in memory)
    guided_report = {}
    if args.mode in ("hf_guided", "single_step", "independent_exact"):
        import os as _os
        _trace = (_os.getenv("VERIFY_TRACE", "0") == "1")
        with torch.no_grad():
            # Devices for HF/local
            try:
                hf_dev = next(hf.parameters()).device
            except Exception:
                hf_dev = torch.device("cpu")
            loc_dev = (torch.device("cuda") if use_cuda else torch.device("cpu"))
            # Initialize sequences on devices
            seq_cpu = x["input_ids"].to(torch.long)
            seq_hf = seq_cpu.to(hf_dev)
            seq_loc = seq_cpu.to(loc_dev)
            attn_hf = x.get("attention_mask", None)
            attn_hf = attn_hf.to(hf_dev) if attn_hf is not None else None
            attn_loc = x.get("attention_mask", None)
            attn_loc = attn_loc.to(loc_dev) if attn_loc is not None else None
            steps = int(args.gen_new if args.mode in ("hf_guided", "independent_exact") else 1)
            mismatches = []
            for t in range(steps):
                # HF one-step (no cache)
                out_h = hf(input_ids=seq_hf, attention_mask=attn_hf)
                logits_h = out_h.logits[:, -1, :]
                next_h = torch.argmax(logits_h, dim=-1)
                # Local one-step (no cache, parity_exact attention if requested)
                out_l = local(input_ids=seq_loc, attention_mask=attn_loc, return_dict=True)
                logits_l = out_l["logits"][:, -1, :]
                next_l = torch.argmax(logits_l, dim=-1)
                if _trace:
                    try:
                        print(f"[verify] step={t} indep_exact next_h={int(next_h[0].item())} next_l={int(next_l[0].item())}")
                    except Exception:
                        pass
                if int(next_h[0].item()) != int(next_l[0].item()):
                    mismatches.append({
                        "step": t,
                        "hf": int(next_h[0].item()),
                        "local": int(next_l[0].item()),
                    })
                    if args.mode == "independent_exact":
                        break
                # Advance sequences
                if args.mode == "hf_guided":
                    # append HF token to both
                    seq_hf = torch.cat([seq_hf, next_h.view(1, 1)], dim=1)
                    seq_loc = torch.cat([seq_loc, next_h.to(device=seq_loc.device).view(1, 1)], dim=1)
                    if attn_hf is not None:
                        attn_hf = torch.cat([attn_hf, torch.ones_like(next_h).view(1, 1)], dim=1)
                    if attn_loc is not None:
                        attn_loc = torch.cat([attn_loc, torch.ones_like(next_h).view(1, 1).to(attn_loc.device)], dim=1)
                else:
                    # independent exact: each model appends its own token
                    seq_hf = torch.cat([seq_hf, next_h.view(1, 1)], dim=1)
                    seq_loc = torch.cat([seq_loc, next_l.view(1, 1)], dim=1)
                    if attn_hf is not None:
                        attn_hf = torch.cat([attn_hf, torch.ones_like(next_h).view(1, 1)], dim=1)
                    if attn_loc is not None:
                        attn_loc = torch.cat([attn_loc, torch.ones_like(next_l).view(1, 1)], dim=1)
                # Early stop on eos
                eos_id = tok.eos_token_id
                if eos_id is not None and int(next_h[0].item()) == int(eos_id):
                    break
            # Decode final guided sequence
            new_tokens = seq_cpu.shape[1]
            out_seq = seq_hf[0].to("cpu")
            text_guided = tok.decode(out_seq[new_tokens:], skip_special_tokens=True)
            guided_report = {"guided_text": text_guided, "mismatches": mismatches}

    if args.mode == "two_pass_exact":
        # Pass 1: HF stepwise on current device; record tokens
        import os as _os
        _trace = (_os.getenv("VERIFY_TRACE", "0") == "1")
        with torch.no_grad():
            try:
                hf_dev = next(hf.parameters()).device
            except Exception:
                hf_dev = torch.device("cpu")
            seq_cpu = x["input_ids"].to(torch.long)
            seq_hf = seq_cpu.to(hf_dev)
            attn_hf = x.get("attention_mask", None)
            attn_hf = attn_hf.to(hf_dev) if attn_hf is not None else None
            hf_tokens: list[int] = []
            for t in range(int(args.gen_new)):
                out_h = hf(input_ids=seq_hf, attention_mask=attn_hf)
                logits_h = out_h.logits[:, -1, :]
                next_h = int(torch.argmax(logits_h, dim=-1)[0].item())
                hf_tokens.append(next_h)
                if _trace:
                    print(f"[verify] pass1 step={t} hf_token={next_h}")
                seq_hf = torch.cat([seq_hf, torch.tensor([[next_h]], device=seq_hf.device)], dim=1)
                if attn_hf is not None:
                    attn_hf = torch.cat([attn_hf, torch.ones((1,1), device=attn_hf.device, dtype=attn_hf.dtype)], dim=1)
                eos_id = tok.eos_token_id
                if eos_id is not None and next_h == int(eos_id):
                    break
            # Decode HF text
            out_seq = seq_hf[0].to("cpu")
            text_guided = tok.decode(out_seq[seq_cpu.shape[1]:], skip_special_tokens=True)
            guided_report = {"guided_text": text_guided, "mismatches": []}
        # Free HF completely before local
        try:
            del hf
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        except Exception:
            pass
        # Pass 2: Local stepwise on GPU/CPU; compare tokens at each step
        # Rebuild local (already built above as `local`); if freed earlier, ensure available
        with torch.no_grad():
            # Initialize from original prompt
            if use_cuda:
                seq_loc = x["input_ids"].to("cuda")
                attn_loc = x.get("attention_mask", None)
                attn_loc = attn_loc.to("cuda") if attn_loc is not None else None
            else:
                seq_loc = x["input_ids"].clone()
                attn_loc = x.get("attention_mask", None)
            mismatches: list[dict] = []
            for step, tok_id in enumerate(hf_tokens):
                out_l = local(input_ids=seq_loc, attention_mask=attn_loc, return_dict=True)
                logits_l = out_l["logits"][:, -1, :]
                next_l = int(torch.argmax(logits_l, dim=-1)[0].item())
                if _trace:
                    print(f"[verify] pass2 step={step} hf_token={int(tok_id)} local_next={next_l}")
                if next_l != int(tok_id):
                    mismatches.append({"step": step, "hf": int(tok_id), "local": next_l})
                    break
                # Append HF token (teacher-forced) to keep prefixes identical
                t =
# ... truncated ...


run.py

from __future__ import annotations

import os
import sys
import argparse
from pathlib import Path
import subprocess
from typing import Dict, Any
import json
import os as _os
import importlib

import numpy as np
from model.hf_snapshot import ensure_snapshot
from modules.retrieval_policy import RetrievalPolicy  # type: ignore
from modules.runner_core import select_region, prepare_citations  # type: ignore

def main() -> None:


build.py

import os
import argparse
import json
import hashlib
import subprocess
from datetime import datetime, timezone
from typing import Optional, Dict, Tuple, List
import sys
import random
import platform
import importlib

import numpy as np

from modules.embedding import (  # type: ignore
    build_program_embedding,
    build_subgraph_embedding_from_program,
)
from modules.adapter import (  # type: ignore
    generate_lora_from_embedding,
    generate_lora_from_embedding_torch,


"""
A tiny, self-contained attention demo used for repo-grounding smoke tests.

The functions intentionally include clear docstrings and simple math so
retrieval by the runner can surface these lines in context windows.
"""

from __future__ import annotations

from typing import Tuple
import math

def attention_score(query: float, key: float) -> float:
    """
    Compute an unnormalized attention score between a scalar query and key.

    This mirrors the dot-product component in scaled dot-product attention for
    the 1-D case. In higher dimensions, the score is the dot product of the
    query and key vectors.
    """


modules/embedding.py

from typing import Any, Dict, List, Optional, Tuple
import os
import re
import numpy as np
from concurrent.futures import ThreadPoolExecutor, as_completed
try:
    import pathspec  # type: ignore
except Exception:  # pragma: no cover
    pathspec = None  # type: ignore
from modules.program_graph import ProgramGraph, Artifact  # type: ignore

EMBED_DIM_DEFAULT = 128
HASH_SEEDS = [1337, 1338, 1339, 1340, 1341, 1342, 1343, 1344]

def auto_model_dims(model_id: str, cache_dir: Optional[str]) -> Tuple[int, int]:
    try:
        from transformers import AutoConfig  # type: ignore

        cfg = AutoConfig.from_pretrained(model_id, cache_dir=cache_dir)
        n_layers = int(getattr(cfg, "num_hidden_layers", 0) or 0)
        d_model = int(getattr(cfg, "hidden_size", 0) or 0)


modules/runtime.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
import os
import subprocess

import numpy as np
import torch

# Reuse existing module utilities
# No example-specific imports in core runtime
from model.inspect import detect_target_names_from_model_full

@dataclass
class OTFFlags:
    of_sources: str = "question"  # "zoom" | "question"
    zoom_symbol: Optional[str] = None
    zoom_radius: int = 1
    include_text: bool = True
    text_max_bytes: int = 250_000


modules/caches.py

from __future__ import annotations

import os
import json
from typing import Dict, Iterable, Iterator, List, Optional, Tuple

def _ensure_path(p: Optional[str]) -> Optional[str]:
    if not p:
        return None
    pth = os.path.abspath(os.path.expanduser(os.path.expandvars(p)))
    return pth if os.path.exists(pth) else None

def load_manifest(adapters_dir: str) -> Dict:
    mf = os.path.join(adapters_dir, "manifest.json")
    try:
        return json.loads(open(mf, "r", encoding="utf-8").read())
    except Exception:
        return {}

def resolve_cache_path(manifest: Dict, key: str, default_path: str) -> str:
    try:


from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, Protocol, Optional, Tuple, Dict, List

# Core types
EntityId = str

@dataclass(frozen=True)
class Span:
    start_line: int
    end_line: int  # inclusive, 1-based

@dataclass(frozen=True)
class Entity:
    uri: str
    id: EntityId
    kind: str


modules/context.py

# function-first packer (lift from run_repo_adapter.py)

from typing import List, Optional, Tuple
import os
import torch  # type: ignore

def pack_context_heads(program_root: str, files: List[str], tok, budget_tokens: int) -> str:
    lines_out: List[str] = ["Program snippets:"]
    used = 0
    for rel in files:
        abs_fp = rel if os.path.isabs(rel) else os.path.abspath(os.path.join(program_root, rel))
        try:
            src_lines = open(abs_fp, "r", encoding="utf-8", errors="ignore").read().splitlines()
        except Exception:
            continue
        head_n = min(len(src_lines), 120)
        block = [f"[ctx] path: {os.path.relpath(abs_fp, program_root)}:1-{head_n}"] + src_lines[:head_n] + [""]
        text = "\n".join(block) + "\n"
        t = len(tok(text).input_ids)
        if used + t > budget_tokens:


modules/priors.py

# kbann_priors derivation, round_lora utilities (move from run_repo_adapter.py/build_repo_adapter.py)


modules/targets.py

from typing import Optional, Dict, Tuple

def parse_target_shapes(arg: Optional[str]) -> Optional[Dict[str, Tuple[int, int]]]:
    if not arg:
        return None
    result: Dict[str, Tuple[int, int]] = {}
    try:
        parts = [p.strip() for p in str(arg).split(",") if p.strip()]
        for p in parts:
            if "=" not in p or ":" not in p:
                continue
            name, dims = p.split("=", 1)
            a, b = dims.split(":", 1)
            result[name.strip()] = (int(a), int(b))
        return result or None
    except Exception:
        return None


modules/capacity.py

from __future__ import annotations

from typing import Dict, List, Tuple

import numpy as np

def _parse_weights(spec: str) -> Dict[str, float]:
    # accept both "program" and "repo" as primary component key
    out: Dict[str, float] = {"program": 0.4, "subgraph": 0.4, "question": 0.2}
    try:
        parts = [p.strip() for p in str(spec).split(",") if p.strip()]
        for p in parts:
            if "=" not in p:
                continue
            k, v = p.split("=", 1)
            out[str(k).strip()] = float(v)
    except Exception:
        pass
    s = float(sum(max(0.0, v) for v in out.values()))
    if s > 0:


from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Tuple, Iterable, Optional, List, Set
import math

from .program_graph import ProgramGraph, Entity, Edge, Span

@dataclass
class RetrievalMix:
    sim: float = 0.6
    struct: float = 0.4
    temp: float = 0.7

    @staticmethod
    def parse(spec: Optional[str], temp: Optional[float]) -> "RetrievalMix":
        if not spec:
            return RetrievalMix(temp=(temp if isinstance(temp, (int, float)) else 0.7))
        w_sim = 0.0
        w_struct = 0.0


modules/provenance.py

import subprocess
from typing import Optional

def git_commit_sha(repo_root: str) -> Optional[str]:
    try:
        sha = subprocess.check_output(
            ["git", "-C", repo_root, "rev-parse", "HEAD"], stderr=subprocess.DEVNULL, text=True
        ).strip()
        return sha if sha else None
    except Exception:
        return None

def git_tree_sha(repo_root: str) -> Optional[str]:
    """Return the HEAD tree SHA if available (pins exact tracked file set)."""
    try:
        sha = subprocess.check_output(
            ["git", "-C", repo_root, "rev-parse", "HEAD^{tree}"], stderr=subprocess.DEVNULL, text=True
        ).strip()
        return sha if sha else None
    except Exception:


modules/adapter.py

from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import math
import os
import json
import re
from .embedding import _stable_hash

def _make_random_matrix(shape: Tuple[int, int], *, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    # Xavier uniform
    limit = math.sqrt(6.0 / float(shape[0] + shape[1]))
    return rng.uniform(-limit, limit, size=shape).astype(np.float32)

def generate_lora_from_embedding(
    z: np.ndarray,
    *,
    d_model: int,
    num_layers: int,
    rank: int = 8,
    seed: int = 0,


modules/mixing.py

# delta-cap guard (lift from modular.py)
import os
import re
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch

from blocks.targets import targets_map
from blocks.utils import getattr_nested

def register_hook_mixed_adapters(
    model: Any,
    base_layers: List[Dict[str, Dict[str, np.ndarray]]],
    sub_layers: Optional[List[Dict[str, Dict[str, np.ndarray]]]],
    *,
    alpha_star: float,
    g_sub: float,
    rank: int,
    beta: float,
    target_weights: Optional[Dict[str, float]] = None,


modules/telemetry.py

# telemetry meta assembly and write (from run_repo_adapter.py)


from __future__ import annotations

from dataclasses import dataclass, asdict
from typing import Any, Dict, List, Tuple, Optional
import json
import hashlib
import os
import time

@dataclass
class ProgramState:
    root: str
    candidates_modules: List[str]
    candidates_files: List[str]
    citations: List[Tuple[str, int, int]]
    vec: Optional[List[float]] = None
    vec_weight: float = 1.0
    H: float = 0.0
    behavior_log: List[Dict[str, Any]] = None  # type: ignore[assignment]


from .adapter import (
    generate_lora_from_embedding,
    generate_lora_from_embedding_torch,
    save_npz,
    load_adapters_npz
)


from __future__ import annotations

from typing import Dict, List, Any, Tuple

from .program_graph import ProgramGraph, Entity
from .retrieval_policy import RetrievalPolicy
from .citations import CitationManager, CitationPolicy

def select_region(query: str, pg: ProgramGraph, policy: RetrievalPolicy, top_k: int = 16) -> List[str]:
    scores = policy.score_entities(query, pg)
    ids = sorted(scores.keys(), key=lambda k: scores.get(k, 0.0), reverse=True)[:max(1, int(top_k))]
    return ids

def prepare_citations(units: List[Dict[str, Any]], region_entity_ids: List[str], pg: ProgramGraph, citations_policy: Dict[str, Any], manifest: Dict[str, Any]) -> List[Dict[str, Any]]:
    cm = CitationManager(
        policy=CitationPolicy(
            enforce=bool(citations_policy.get("enforce", True)),
            per_paragraph=bool(citations_policy.get("per_paragraph", False)),
            repair=bool(citations_policy.get("repair", True)),


modules/interpret.py

# ActivationTracer hooks, capture pipeline, write JSON (lift from run_repo_adapter.py)

from typing import Dict, Tuple
import torch

def is_block(name: str, _m: torch.nn.Module) -> bool:
    if not name.startswith("model.layers."):
        return False
    rest = name[len("model.layers."):]
    return rest.isdigit() and "." not in rest

def block_out_hook(_key: str, _m: torch.nn.Module, _inputs: Tuple[torch.Tensor, ...], output: torch.Tensor) -> torch.Tensor | None:
    try:
        if isinstance(output, torch.Tensor):
            return output
        if isinstance(output, (tuple, list)) and output:
            v = output[0]
            return v if isinstance(v, torch.Tensor) else None


modules/peft.py

from typing import Dict, Any, List, Optional
import os
import re
from model.inspect import detect_target_names_from_model_full

def infer_target_names(model_id: str) -> Dict[str, str]:
    names = detect_target_names_from_model_full(model_id, target_regex=None) or {}
    # names maps short -> path within the first layer subtree
    # Example: {"q_proj": "self_attn.q_proj", "o_proj": "self_attn.o_proj", ...}
    return names

def save_peft_like(out_dir: str, adapters: Dict[str, Any], *, r: int, alpha: int, target_modules: List[str], bias: str = "none", int8: bool = False, target_paths: Optional[Dict[str, str]] = None) -> None:
    """Write a minimal PEFT LoRA config + tensors for quick benchmarking.

    Note: This is a best-effort exporter; users may still need to map names depending on the model arch.
    """
    try:
        import json as _json
        cfg = {
            "peft_type": "LORA",
            "r": int(r),


modules/runner.py

from __future__ import annotations

import os
import json
import time
from typing import Optional, Tuple, List, Dict, Any, Set

import numpy as np
import torch
from data.tokenizer import LocalLlamaTokenizer

from modules.adapter import (
    load_adapters_npz,
    generate_lora_from_embedding,
)
from model.llama_bootstrap import build_local_llama_from_snapshot
from model.hf_snapshot import ensure_snapshot
from blocks.targets import targets_map
from blocks.inspect import infer_target_shapes
from model.inspect import detect_target_names_from_model_full
from modules.mixing import (


modules/citations.py

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Any, Optional, Tuple
import hashlib
import time
import re

from .program_graph import ProgramGraph, Span, ResolvedAnchor

@dataclass
class CitationPolicy:
    enforce: bool = True
    per_paragraph: bool = False
    repair: bool = True

def _extract_identifier_tokens(text: str) -> List[str]:
    toks = re.findall(r"`([A-Za-z_][A-Za-z0-9_\.]+)`", text or "")
    toks2 = re.findall(r"\b([A-Za-z_][A-Za-z0-9_]*)\s*\(", text or "")


from __future__ import annotations

import os
import sys
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Any, List, Tuple

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader

@dataclass
class EvalPaths:
	adapters_dir: Path
	out_dir: Path
	state_path: Path


from __future__ import annotations

import os
import sys
import json
from dataclasses import dataclass
from pathlib import Path
from typing import List, Dict, Any

import torch
import torch.nn.functional as F

# Minimal dependencies: transformers + peft + datasets

@dataclass
class TrainPaths:
	adapters_dir: Path
	out_dir: Path


from __future__ import annotations

import sys
import os
import json
from pathlib import Path
from typing import List, Dict, Any
import math
import random
import numpy as np

import torch
from torch.utils.data import Dataset, DataLoader

ROOT = EX_DIR
PG_BACKEND = "examples.scripts.dataset_graph:DatasetProgramGraph"
try:
	from .program_config import load_program_config  # type: ignore  # noqa: E402
except Exception:
	from program_config import load_program_config  # type: ignore  # noqa: E402


from __future__ import annotations

import sys
import os
import json
from pathlib import Path
from typing import List, Dict

import torch
from torch.utils.data import Dataset, DataLoader

ROOT = EX_DIR
PG_BACKEND = "examples.scripts.dataset_graph:DatasetProgramGraph"
try:
	from .program_config import load_program_config  # type: ignore  # noqa: E402
except Exception:
	from program_config import load_program_config  # type: ignore  # noqa: E402

def _run(cmd: List[str]) -> int:
	print("[run]", " ".join(cmd))


from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Optional

@dataclass(frozen=True)
class ProgramContracts:
	require_citations: bool = False
	citations_per_paragraph: bool = False
	retrieval_policy: str = "sim:0.5,struct:0.5"
	retrieval_temp: float = 0.6

@dataclass(frozen=True)
class ProgramPaths:
	adapters_dir: Path
	out_dir: Path
	telemetry_path: Path
	program_state_path: Optional[Path] = None


from __future__ import annotations

import sys
import os
import json
from pathlib import Path
from typing import List, Dict

import torch
from torch.utils.data import Dataset

ROOT = EX_DIR
try:
	from .program_config import load_program_config  # type: ignore  # noqa: E402
except Exception:
	from program_config import load_program_config  # type: ignore  # noqa: E402

class SimpleLmDataset(Dataset):
	def __init__(self, tokenizer, texts: List[str], max_len: int = 128) -> None:
		self.tok = tokenizer


from __future__ import annotations

import sys
import subprocess
from pathlib import Path

# Resolve program_conditioned_adapter directory directly from this file
EX_DIR = Path(__file__).resolve().parents[2]
SMOKE_REPO = EX_DIR / "smoke_repo"
PG_BACKEND = "examples.scripts.python_repo_graph:PythonRepoGraph"
# Program configuration
from .program_config import load_program_config  # noqa: E402

def _run(cmd: list[str]) -> int:
    print("[run]", " ".join(cmd))
    return subprocess.call(cmd)

def main() -> None:
    cfg = load_program_config(str(SMOKE_REPO))


from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Optional

@dataclass(frozen=True)
class ProgramContracts:
    require_citations: bool = True
    citations_per_paragraph: bool = False
    retrieval_policy: str = "sim:0.6,struct:0.4,plan:0.2"
    retrieval_temp: float = 0.7

@dataclass(frozen=True)
class ProgramPaths:
    adapters_dir: Path
    knowledge_dir: Path
    program_state_path: Optional[Path] = None