run

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:
    p = argparse.ArgumentParser()
    # Model & sources
    p.add_argument("--model", default="meta-llama/Llama-3.1-8B-Instruct")
    p.add_argument("--sources", default=str(_root()), help="Program sources root (path or URI)")
    p.add_argument("--program", default=None, help="Alias for --sources (program root path or URI)")
    p.add_argument("--prompt", default=(
        "Explain how generation works for this program. Cite path:line for each claim."
    ))

    # PCA program options (program-agnostic; backend not enforced here)
    p.add_argument("--backend", default=None, help="Program backend id (informational; use --pg-backend to enable ProgramGraph features)")
    p.add_argument("--retrieval-policy", dest="retrieval_policy", default=None, help='Retrieval mix like "sim:0.6,struct:0.4"')
    p.add_argument("--retrieval-temp", dest="retrieval_temp", type=float, default=None)
    p.add_argument("--use-cache", action="store_true", help="Prefer caches when available (symbol/windows/facts)")
    p.add_argument("--citations-enforce", action="store_true", help="Require citations policy in PCA layer")
    p.add_argument("--citations-repair", action="store_true")
    p.add_argument("--pg-backend", default=None, help="Dotted path to ProgramGraph factory, e.g. 'examples.python_repo_grounded_qa.python_repo_graph:PythonRepoGraph'")
    p.add_argument("--program-state", default=None, help="Optional path to .program_state.json (defaults to <adapters_dir>/.program_state.json if present)")

    # Base adapters generation
    p.add_argument("--adapters-dir", default=None, help="Directory to save/find base adapters (defaults under examples/.../artifacts/base_adapters)")
    p.add_argument("--base-rank", type=int, default=8)
    p.add_argument("--embed-dim", type=int, default=1536)
    p.add_argument("--include-text", action="store_true")
    p.add_argument("--text-max-bytes", type=int, default=0)
    p.add_argument("--max-text-tokens", type=int, default=0)
    p.add_argument("--seed", type=int, default=0)

    # Runner mixing and capacity
    p.add_argument("--alpha", type=float, default=20.0)
    p.add_argument("--rank", type=int, default=12)
    p.add_argument("--gsub", type=float, default=0.75)
    p.add_argument("--mix-beta", type=float, default=0.1)
    p.add_argument("--target-weights", default=None, help="CSV like q_proj=1,o_proj=1.1,up_proj=1.1,down_proj=1.05")
    p.add_argument("--knowledge-preset", action="store_true", help="Use a preset of target weights tuned for knowledge recall (boost o/up/down; modest q/k/v)")
    p.add_argument("--code-recall-preset", action="store_true", help="Opt-in preset tuned for code recall (o,v,up,down,gate emphasized; light q,k)")
    p.add_argument("--delta-cap", type=float, default=0.05, help="AB-norm clipping cap relative to base weight norm per layer (0 disables)")
    # Entropy-aware capacity (forward to enhanced runner)
    p.add_argument("--entropy-aware", action="store_true")
    p.add_argument("--rank-min", type=int, default=8)
    p.add_argument("--rank-max", type=int, default=32)
    p.add_argument("--gsub-min", type=float, default=0.6)
    p.add_argument("--gsub-max", type=float, default=0.9)
    p.add_argument("--entropy-weights", default="program=0.4,subgraph=0.4,question=0.2")

    # Selection & context packing
    p.add_argument("--of-sources", choices=["question", "zoom"], default="question")
    p.add_argument("--zoom-symbol", default=None)
    p.add_argument("--zoom-radius", type=int, default=0)
    p.add_argument("--pack-context", action="store_true")
    p.add_argument("--pack-mode", choices=["heads", "windows"], default="heads")
    p.add_argument("--context-tokens", type=int, default=3000)
    p.add_argument("--require-citations", action="store_true")
    p.add_argument("--citations-per-paragraph", action="store_true")
    p.add_argument("--function-first", action="store_true")
    p.add_argument("--ff-max-candidates", type=int, default=24)
    p.add_argument("--ff-window-lines", type=int, default=80)
    p.add_argument("--ff-threshold", type=float, default=0.55)
    p.add_argument("--ff-noise-penalty", type=float, default=0.30)
    # Layer schedule and q-aware weights (opt-in)
    p.add_argument("--layer-schedule", action="store_true", help="Enable a light per-layer multiplier rising toward top third (additive; default off)")
    p.add_argument("--q-aware-weights", action="store_true", help="Heuristic reweighting of targets by question intent (additive; default off)")
    p.add_argument("--per-target-rank-schedule", action="store_true", help="Trim per-target effective rank at run-time (additive; default off)")
    p.add_argument("--rank-budget", type=int, default=0, help="Optional per-layer rank budget across targets; rescale keeps to meet budget (0 disables)")
    p.add_argument("--ablate-attn", action="store_true", help="Ablate attention targets (q/k/v/o) by zeroing their weights")
    p.add_argument("--ablate-mlp", action="store_true", help="Ablate MLP targets (up/down/gate) by zeroing their weights")
    p.add_argument("--layer-rank-tiers", action="store_true", help="Use top/mid/low thirds per-layer rank keeps by target group (opt-in)")
    # Mixture bank (opt-in)
    p.add_argument("--mixture-m", type=int, default=0, help="Top-m subgraph adapters from bank to mix (0 disables)")
    p.add_argument("--adapters-bank", default=None, help="Path to a bank of sub_adapters (from build --per-module)")
    # Adapter mapping is always enhanced; no CLI toggle
    # Alpha warmup and decoding hooks (opt-in)
    p.add_argument("--alpha-warmup", action="store_true", help="Use a lighter alpha on first attempt (or first structured pass), then full alpha on retry/subsequent passes")
    p.add_argument("--adapter-aware-decoding", action="store_true", help="Slightly relax sampling and prompt pointer-first when citations are required")
    # Telemetry verification (opt-in)
    p.add_argument("--telemetry-tests", action="store_true", help="Attempt simple verify_with_tests() on a few selected modules and record results (structured path only)")
    # Reranking
    p.add_argument("--rerank", action="store_true")
    p.add_argument("--self-queries", default=None, help="Path to self_queries.jsonl for retrieval boosts")

    # Sampling controls
    p.add_argument("--do-sample", action="store_true")
    p.add_argument("--temperature", type=float, default=0.7)
    p.add_argument("--top-p", type=float, default=0.9)
    p.add_argument("--repetition-penalty", type=float, default=1.1)
    p.add_argument("--min-new-tokens", type=int, default=64)
    p.add_argument("--max-new-tokens", type=int, default=512)
    # Local generation/perf controls
    p.add_argument("--kv-window", type=int, default=0, help="Optional sliding KV window length; 0 disables")
    p.add_argument("--head-device", choices=["same", "cpu", "auto"], default="same", help="Place lm_head on cpu/auto to save VRAM")

    # Device & telemetry
    p.add_argument("--device-map", default="auto", choices=["auto", "none"])
    p.add_argument("--gpu-ids", default=None)
    p.add_argument("--max-memory", default=None)
    p.add_argument("--telemetry-out", default=None)
    # DCPO/Structured controls
    p.add_argument("--structured", action="store_true")
    p.add_argument("--lfp-iters", type=int, default=1)
    p.add_argument("--budget-H", type=float, default=0.0)
    p.add_argument("--monotone-selection", action="store_true")
    p.add_argument("--samples", type=int, default=1)
    p.add_argument("--cone-join", choices=["concat", "weighted"], default="concat")
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints", help="Cache directory for HF models/tokenizers")
    p.add_argument("--verbose", action="store_true")
    p.add_argument("--no-adapters", action="store_true", help="Disable applying adapters in the enhanced runner")
    p.add_argument("--commit-footer", action="store_true", help="Append 'answer valid for commit X' footer")
    args = p.parse_args()

    # Normalize program root (prefer --sources; allow --program alias)
    src_arg = args.sources
    if getattr(args, "program", None):
        src_arg = args.program if args.program else src_arg
    program_root = Path(src_arg)
    # Example directory for artifacts colocated with this run.py
    example_dir = Path(__file__).resolve().parent
    artifacts = example_dir / "artifacts"
    base_dir = Path(args.adapters_dir) if args.adapters_dir else (artifacts / "base_adapters")
    base_dir.mkdir(parents=True, exist_ok=True)
    cache_dir = Path(args.cache_dir)
    cache_dir.mkdir(parents=True, exist_ok=True)
    adapters_npz = base_dir / "adapters.npz"
    manifest_path = base_dir / "manifest.json"
    # Resolve program state path
    ps_path = None
    if args.program_state:
        ps_path = Path(args.program_state)
    else:
        default_ps = base_dir / ".program_state.json"
        if default_ps.exists():
            ps_path = default_ps

    # Generate base adapters if not present by delegating to build.py (program-agnostic)
    if not adapters_npz.exists():
        cmd = [
            sys.executable,
            str(Path(__file__).resolve().parent / "build.py"),
            "--sources", str(program_root),
            "--model", str(args.model),
            "--adapters-dir", str(base_dir),
            "--embed-dim", str(int(args.embed_dim)),
            "--seed", str(int(args.seed)),
        ]
        if bool(args.include_text):
            cmd.append("--include-text")
        if int(args.text_max_bytes) > 0:
            cmd += ["--text-max-bytes", str(int(args.text_max_bytes))]
        if int(args.max_text_tokens) > 0:
            cmd += ["--max-text-tokens", str(int(args.max_text_tokens))]
        if bool(args.knowledge_preset):
            cmd.append("--knowledge-preset")
        if bool(args.code_recall_preset):
            cmd.append("--code-recall-preset")
        subprocess.check_call(cmd)

    # Run directly via orchestrator
    from modules.runner import generate_answer
    from modules.runner import generate_answer_structured

    if args.gpu_ids and str(args.gpu_ids).strip():
        os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_ids).strip()
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

    # ProgramGraph + RetrievalPolicy selection (seed zoom) if requested
    zoom_seed = None
    try:
        if args.pg_backend:
            # Load PG factory dynamically only when explicitly provided
            def _load_symbol(path: str):
                mod, _, attr = path.partition(":")
                m = importlib.import_module(mod)
                return getattr(m, attr)
            pg_ctor = _load_symbol(args.pg_backend)
            pg = pg_ctor(str(program_root), ignore=None)
            pol = RetrievalPolicy.from_spec(args.retrieval_policy, temp=args.retrieval_temp)
            region_ids = select_region(args.prompt, pg, pol, top_k=8)
            ents_by_id = {e.id: e for e in pg.entities()}
            names = []
            for eid in region_ids[:4]:
                e = ents_by_id.get(eid)
                if e and e.name:
                    names.append(e.name)
            if names:
                zoom_seed = ",".join(names)
    except Exception:
        zoom_seed = None

    if bool(args.structured):
        res = generate_answer_structured(
            model_id=args.model,
            adapters_npz=str(adapters_npz),
            program_root=str(program_root),
            delta_cap=float(max(0.0, args.delta_cap)),
            prompt=args.prompt,
            cache_dir=str(cache_dir),
            device_map=str(args.device_map),
            alpha=float(args.alpha),
            rank=int(args.rank),
            gsub=float(args.gsub),
            beta=float(args.mix_beta),
            of_sources=("zoom" if (zoom_seed and not args.zoom_symbol) else args.of_sources),
            zoom_symbol=(args.zoom_symbol or (zoom_seed or None)),
            zoom_radius=int(args.zoom_radius),
            pack_context=bool(args.pack_context),
            pack_mode=args.pack_mode,
            context_tokens=int(args.context_tokens),
            require_citations=bool(args.require_citations),
            citations_per_paragraph=bool(args.citations_per_paragraph),
            function_first=bool(args.function_first),
            ff_max_candidates=int(args.ff_max_candidates),
            ff_window_lines=int(args.ff_window_lines),
            ff_threshold=float(args.ff_threshold),
            ff_noise_penalty=float(args.ff_noise_penalty),
            do_sample=bool(args.do_sample),
            temperature=float(args.temperature),
            top_p=float(args.top_p),
            repetition_penalty=float(args.repetition_penalty),
            min_new_tokens=int(args.min_new_tokens),
            max_new_tokens=int(args.max_new_tokens),
            kv_window=int(args.kv_window),
            head_device=str(args.head_device),
            seed=int(args.seed),
            entropy_aware=bool(args.entropy_aware),
            rank_min=int(args.rank_min),
            rank_max=int(args.rank_max),
            gsub_min=float(args.gsub_min),
            gsub_max=float(args.gsub_max),
            entropy_weights=str(args.entropy_weights),
            target_weights=(str(args.target_weights) if args.target_weights else (
                "o_proj=1.15,v_proj=1.10,up_proj=1.10,down_proj=1.05,gate_proj=1.00,q_proj=0.95,k_proj=0.90" if args.code_recall_preset else (
                    "q_proj=0.95,k_proj=0.95,v_proj=0.95,o_proj=1.10,up_proj=1.10,down_proj=1.05" if args.knowledge_preset else None
                )
            )),
            rerank=bool(args.rerank),
            self_queries_path=(str(args.self_queries) if args.self_queries else None),
            commit_footer=bool(args.commit_footer),
            verbose=bool(args.verbose),
            lfp_iters=int(args.lfp_iters),
            budget_H=float(args.budget_H),
            monotone_selection=bool(args.monotone_selection),
            program_state_path=(str(ps_path) if ps_path else None),
            samples=int(args.samples),
            cone_join=str(args.cone_join),
            telemetry_out=(str(args.telemetry_out) if args.telemetry_out else None),
            layer_schedule=bool(args.layer_schedule),
            q_aware_weights=bool(args.q_aware_weights),
            mixture_m=int(args.mixture_m),
            adapters_bank=(str(args.adapters_bank) if args.adapters_bank else None),
            # Forward additional knobs to structured path
            per_target_rank_schedule=bool(args.per_target_rank_schedule),
            rank_budget=int(args.rank_budget),
            ablate_attn=bool(args.ablate_attn),
            ablate_mlp=bool(args.ablate_mlp),
            alpha_warmup=bool(args.alpha_warmup),
            adapter_aware_decoding=bool(args.adapter_aware_decoding),
            layer_rank_tiers=bool(args.layer_rank_tiers),
            telemetry_verify_tests=bool(args.telemetry_tests),
        )
        # PCA evidence/provenance stamping when enabled
        try:
            if (args.citations_enforce or args.citations_per_paragraph or args.citations_repair) and args.pg_backend:
                # Load manifest (for provenance)
                mf_obj: Dict[str, Any] = {}
                try:
                    if os.path.exists(manifest_path):
                        mf_obj = json.loads(open(manifest_path, "r", encoding="utf-8").read())
                except Exception:
                    mf_obj = {}
                # Minimal unit
                unit = {"text": res.get("text", ""), "evidence": []}
                citations_policy = {
                    "enforce": bool(args.citations_enforce),
                    "per_paragraph": bool(args.citations_per_paragraph),
                    "repair": bool(args.citations_repair),
                }
                # Build ProgramGraph for stamping when provided
                def _load_symbol2(path: str):
                    mod, _, attr = path.partition(":")
                    m = importlib.import_module(mod)
                    return getattr(m, attr)
                pg_ctor2 = _load_symbol2(args.pg_backend)
                pg2 = pg_ctor2(str(program_root), ignore=None)
                pol2 = RetrievalPolicy.from_spec(args.retrieval_policy, temp=args.retrieval_temp)
                region_ids2 = select_region(args.prompt, pg2, pol2, top_k=8)
                stamped = prepare_citations([unit], region_ids2, pg2, citations_policy=citations_policy, manifest=mf_obj)
                if stamped:
                    res["evidence"] = stamped[0].get("evidence", [])
                    res["provenance"] = stamped[0].get("provenance", res.get("provenance"))
        except Exception:
            pass
        print(res.get("text", ""))
        if args.verbose:
            try:
                import json as _json
                print(_json.dumps({
                    "must": len(res.get("must", [])),
                    "may": len(res.get("may", [])),
                    "lfp_passes": res.get("lfp_passes"),
                    "converged": res.get("converged"),
                    "confidence": res.get("confidence"),
                }, indent=2))
            except Exception:
                pass
    else:
        text = generate_answer(
            model_id=args.model,
            adapters_npz=str(adapters_npz),
            program_root=str(program_root),
            delta_cap=float(max(0.0, args.delta_cap)),
            prompt=args.prompt,
            cache_dir=str(cache_dir),
            device_map=str(args.device_map),
            alpha=float(args.alpha),
            rank=int(args.rank),
            gsub=float(args.gsub),
            beta=float(args.mix_beta),
            of_sources=("zoom" if (zoom_seed and not args.zoom_symbol) else args.of_sources),
            zoom_symbol=(args.zoom_symbol or (zoom_seed or None)),
            zoom_radius=int(args.zoom_radius),
            pack_context=bool(args.pack_context),
            pack_mode=args.pack_mode,
            context_tokens=int(args.context_tokens),
            require_citations=bool(args.require_citations),
            citations_per_paragraph=bool(args.citations_per_paragraph),
            function_first=bool(args.function_first),
            ff_max_candidates=int(args.ff_max_candidates),
            ff_window_lines=int(args.ff_window_lines),
            ff_threshold=float(args.ff_threshold),
            ff_noise_penalty=float(args.ff_noise_penalty),
            do_sample=bool(args.do_sample),
            temperature=float(args.temperature),
            top_p=float(args.top_p),
            repetition_penalty=float(args.repetition_penalty),
            min_new_tokens=int(args.min_new_tokens),
            max_new_tokens=int(args.max_new_tokens),
            kv_window=int(args.kv_window),
            head_device=str(args.head_device),
            seed=int(args.seed),
            entropy_aware=bool(args.entropy_aware),
            rank_min=int(args.rank_min),
            rank_max=int(args.rank_max),
            gsub_min=float(args.gsub_min),
            gsub_max=float(args.gsub_max),
            entropy_weights=str(args.entropy_weights),
            target_weights=(str(args.target_weights) if args.target_weights else (
                "o_proj=1.15,v_proj=1.10,up_proj=1.10,down_proj=1.05,gate_proj=1.00,q_proj=0.95,k_proj=0.90" if args.code_recall_preset else (
                    "q_proj=0.95,k_proj=0.95,v_proj=0.95,o_proj=1.10,up_proj=1.10,down_proj=1.05" if args.know
# ... truncated ...


            cmd += ["--max-text-tokens", str(int(args.max_text_tokens))]
        if bool(args.knowledge_preset):
            cmd.append("--knowledge-preset")
        if bool(args.code_recall_preset):
            cmd.append("--code-recall-preset")
        subprocess.check_call(cmd)

    # Run directly via orchestrator
    from modules.runner import generate_answer
    from modules.runner import generate_answer_structured

    if args.gpu_ids and str(args.gpu_ids).strip():
        os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_ids).strip()
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

    # ProgramGraph + RetrievalPolicy selection (seed zoom) if requested
    zoom_seed = None
    try:
        if args.pg_backend:
            # Load PG factory dynamically only when explicitly provided
            def _load_symbol(path: str):
                mod, _, attr = path.partition(":")
                m = importlib.import_module(mod)
                return getattr(m, attr)
            pg_ctor = _load_symbol(args.pg_backend)
            pg = pg_ctor(str(program_root), ignore=None)
            pol = RetrievalPolicy.from_spec(args.retrieval_policy, temp=args.retrieval_temp)
            region_ids = select_region(args.prompt, pg, pol, top_k=8)
            ents_by_id = {e.id: e for e in pg.entities()}
            names = []
            for eid in region_ids[:4]:
                e = ents_by_id.get(eid)
                if e and e.name:
                    names.append(e.name)
            if names:
                zoom_seed = ",".join(names)
    except Exception:
        zoom_seed = None

    if bool(args.structured):
        res = generate_answer_structured(
            model_id=args.model,
            adapters_npz=str(adapters_npz),
            program_root=str(program_root),


            telemetry_verify_tests=bool(args.telemetry_tests),
        )
        # PCA evidence/provenance stamping when enabled
        try:
            if (args.citations_enforce or args.citations_per_paragraph or args.citations_repair) and args.pg_backend:
                # Load manifest (for provenance)
                mf_obj: Dict[str, Any] = {}
                try:
                    if os.path.exists(manifest_path):
                        mf_obj = json.loads(open(manifest_path, "r", encoding="utf-8").read())
                except Exception:
                    mf_obj = {}
                # Minimal unit
                unit = {"text": res.get("text", ""), "evidence": []}
                citations_policy = {
                    "enforce": bool(args.citations_enforce),
                    "per_paragraph": bool(args.citations_per_paragraph),
                    "repair": bool(args.citations_repair),
                }
                # Build ProgramGraph for stamping when provided
                def _load_symbol2(path: str):
                    mod, _, attr = path.partition(":")
                    m = importlib.import_module(mod)
                    return getattr(m, attr)
                pg_ctor2 = _load_symbol2(args.pg_backend)
                pg2 = pg_ctor2(str(program_root), ignore=None)
                pol2 = RetrievalPolicy.from_spec(args.retrieval_policy, temp=args.retrieval_temp)
                region_ids2 = select_region(args.prompt, pg2, pol2, top_k=8)
                stamped = prepare_citations([unit], region_ids2, pg2, citations_policy=citations_policy, manifest=mf_obj)
                if stamped:
                    res["evidence"] = stamped[0].get("evidence", [])
                    res["provenance"] = stamped[0].get("provenance", res.get("provenance"))
        except Exception:
            pass
        print(res.get("text", ""))
        if args.verbose:
            try:
                import json as _json
                print(_json.dumps({
                    "must": len(res.get("must", [])),
                    "may": len(res.get("may", [])),
                    "lfp_passes": res.get("lfp_passes"),
                    "converged": res.get("converged"),
                    "confidence": res.get("confidence"),


build

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,
    save_npz,
)
from model.hf_snapshot import ensure_snapshot
from model.inspect import detect_target_shapes_from_model
from modules.capacity import entropy_score  # type: ignore
#
# NOTE: Build is program-agnostic; no backend-specific imports here.
from modules.program_graph import ProgramGraph  # type: ignore

def _load_symbol(path: Optional[str]):
    if not path:
        return None
    mod, _, attr = path.partition(":")
    if not mod or not attr:
        raise ValueError(f"Invalid symbol path '{path}', expected 'module:attr'")
    m = importlib.import_module(mod)
    return getattr(m, attr)

def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--sources", required=False, help="Program sources root (path or URI)")
    p.add_argument("--program", required=False, help="Alias for --sources (program root path or URI)")
    p.add_argument("--model", required=True)
    p.add_argument("--adapters-dir", required=True)
    # Embedding
    p.add_argument("--embed-dim", type=int, default=1536)
    p.add_argument("--include-text", action="store_true")
    p.add_argument("--text-max-bytes", type=int, default=0)
    p.add_argument("--max-text-tokens", type=int, default=0)
    p.add_argument("--text-weight", type=float, default=0.25)
    p.add_argument("--graph-prop-hops", type=int, default=0)
    p.add_argument("--graph-prop-damp", type=float, default=0.85)
    p.add_argument("--ignore", action="append", default=None)
    # Base adapter
    p.add_argument("--base-rank", type=int, default=8)
    p.add_argument("--target-weights", default=None)
    p.add_argument("--knowledge-preset", action="store_true")
    # Adapter mapping is always enhanced; no CLI toggle
    p.add_argument("--code-recall-preset", action="store_true", help="Opt-in preset tuned for code recall (o,v,up,down,gate emphasized; light q,k)")
    # Priors & rounding
    p.add_argument("--kbann-priors", action="store_true")
    p.add_argument("--kbann-strong", action="store_true", help="Stronger inhibitory priors: boost o/up/down; damp v,k (slightly q)")
    p.add_argument("--round-lora", action="store_true")
    p.add_argument("--round-threshold", type=float, default=0.5)
    p.add_argument("--round-mode", choices=["none", "hard", "soft"], default=None, help="Rounding/sparsification mode for adapter matrices")
    p.add_argument("--round-soft-kp", type=float, default=10.0, help="Soft mode: percent of entries to keep (top-|.|) per axis")
    p.add_argument("--round-axis", choices=["row", "col", "global"], default="row", help="Soft mode axis for top-k sparsification")
    p.add_argument("--zero-b", action="store_true", help="After generation, set all B matrices to zero (official LoRA init)")
    p.add_argument("--learn-bias", action="store_true", help="Export zero bias vectors so downstream can fine-tune bias only")
    # Contracts (optional snapshot into manifest and embedding contracts channel)
    p.add_argument("--contracts-require-citations", action="store_true")
    p.add_argument("--contracts-retrieval-policy", default=None)
    p.add_argument("--contracts-retrieval-temp", type=float, default=None)
    p.add_argument("--contracts-weight", type=float, default=0.10)
    # Per-module export (built by default; use --no-per-module to skip)
    p.add_argument("--per-module", action="store_true", help="DEPRECATED: no effect; per-module bank is built by default. Use --no-per-module to disable.")
    p.add_argument("--no-per-module", action="store_true", help="Disable building per-module sub_adapters bank (built by default)")
    p.add_argument("--include-deps", action="store_true")
    p.add_argument("--max-deps", type=int, default=4)
    p.add_argument("--sub-rank", type=int, default=8)
    p.add_argument("--files-only", action="store_true", help="Export a files-only sub-adapter using an explicit file allowlist")
    p.add_argument(
        "--files-allowlist",
        action="append",
        default=None,
        help="Files-only mode: relative file paths to include (repeatable)",
    )
    # HF cache/model
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints")
    p.add_argument("--probe-full", action="store_true")
    p.add_argument("--gen-backend", choices=["numpy", "torch"], default="numpy")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--verbose", action="store_true")
    p.add_argument("--fallback-topology", action="store_true", help="If embedding is empty, fall back to topology-only component z_top")
    # ProgramGraph/Embedder plugins (backend-agnostic)
    p.add_argument("--pg-backend", default=None, help="Dotted path to ProgramGraph factory, e.g. 'examples.python_repo_grounded_qa.python_repo_graph:PythonRepoGraph'")
    p.add_argument("--embedder-fn", default=None, help="Dotted path to embedder function '(pg, sources_root, **opts)->{z,...}'. Defaults to PCA program embedder.")
    # Capacity scheduling
    p.add_argument("--auto-rank", action="store_true", help="Auto-schedule rank based on program complexity")
    p.add_argument("--rank-min", type=int, default=None)
    p.add_argument("--rank-max", type=int, default=None)
    p.add_argument("--mdl-budget-params", type=int, default=0, help="Optional MDL-style global parameter budget across all layers/targets (approximate)")
    # Program seeding
    p.add_argument("--init-program-state", action="store_true", help="Initialize or update baseline ProgramState (.program_state.json) after build")
    p.add_argument("--program-state-path", default=None, help="Optional explicit path to write ProgramState JSON; defaults to <sources_root>/.program_state.json")
    # (self-tune is no longer part of build; use modules/tune.py externally if desired)
    args = p.parse_args()

    # Warn on deprecated/ignored flags for per-module export
    try:
        if bool(args.per_module) and bool(args.no_per_module):
            print("[pca.build] Warning: Both --per-module and --no-per-module were provided; proceeding with --no-per-module (disables per-module bank).", file=sys.stderr)
        elif bool(args.per_module):
            print("[pca.build] Warning: --per-module is deprecated and has no effect; per-module bank is built by default. Use --no-per-module to disable.", file=sys.stderr)
    except Exception:
        pass
    # Determinism: set seeds and hash seed (best-effort)
    try:
        os.environ["PYTHONHASHSEED"] = str(int(args.seed))
    except Exception:
        pass
    try:
        random.seed(int(args.seed))
    except Exception:
        pass
    try:
        np.random.seed(int(args.seed))
    except Exception:
        pass
    torch_seeded = False
    try:
        import torch as _torch  # type: ignore
        _torch.manual_seed(int(args.seed))
        torch_seeded = True
    except Exception:
        torch_seeded = False

    out_dir = os.path.abspath(os.path.expanduser(os.path.expandvars(args.adapters_dir)))
    os.makedirs(out_dir, exist_ok=True)

    # Normalize sources root (prefer --program, else --sources)
    src_arg = args.program if getattr(args, "program", None) else args.sources
    if not src_arg:
        raise SystemExit("--sources or --program is required")
    sources_root = os.path.abspath(os.path.expanduser(os.path.expandvars(src_arg)))

    # Construct ProgramGraph via plugin (if provided)
    pg: Optional[ProgramGraph] = None
    try:
        pg_ctor = _load_symbol(args.pg_backend)
        if pg_ctor:
            pg = pg_ctor(sources_root, ignore=[s for s in (args.ignore or []) if s])
    except Exception:
        pg = None

    # Build embedding (program-agnostic) via plugin or default PCA program embedder
    embedder = _load_symbol(args.embedder_fn) if args.embedder_fn else None
    if embedder is not None:
        try:
            # Prefer new program-graph signature
            emb = embedder(
                pg,
                sources_root=sources_root,
                dim=int(args.embed_dim),
                seed=int(args.seed),
                include_text=bool(args.include_text),
                text_max_bytes=int(args.text_max_bytes),
                max_text_tokens=int(args.max_text_tokens),
                text_weight=float(args.text_weight),
                calls_weight=0.25,
                types_weight=0.20,
                tests_weight=0.15,
                graph_prop_hops=int(args.graph_prop_hops),
                graph_prop_damp=float(args.graph_prop_damp),
                ignore=[s for s in (args.ignore or []) if s],
            )
        except Exception:
            # Backward-compat: allow sources-root signature
            emb = embedder(
                sources_root,
                dim=int(args.embed_dim),
                seed=int(args.seed),
                include_text=bool(args.include_text),
                text_max_bytes=int(args.text_max_bytes),
                max_text_tokens=int(args.max_text_tokens),
                text_weight=float(args.text_weight),
                calls_weight=0.25,
                types_weight=0.20,
                tests_weight=0.15,
                graph_prop_hops=int(args.graph_prop_hops),
                graph_prop_damp=float(args.graph_prop_damp),
                ignore=[s for s in (args.ignore or []) if s],
            )
    else:
        emb = build_program_embedding(
            pg if pg is not None else None,  # type: ignore[arg-type]
            sources_root=sources_root,
            dim=int(args.embed_dim),
            seed=int(args.seed),
            include_text=bool(args.include_text),
            text_max_bytes=int(args.text_max_bytes),
            max_text_tokens=int(args.max_text_tokens),
            text_weight=float(args.text_weight),
            calls_weight=0.25,
            types_weight=0.20,
            tests_weight=0.15,
            contracts_kv={
                "require_citations": (True if bool(args.contracts_require_citations) else False),
                "retrieval_policy": (str(args.contracts_retrieval_policy) if args.contracts_retrieval_policy else ""),
                "retrieval_temp": (str(args.contracts_retrieval_temp) if args.contracts_retrieval_temp is not None else ""),
            },
            contracts_weight=float(max(0.0, args.contracts_weight)),
            graph_prop_hops=int(args.graph_prop_hops),
            graph_prop_damp=float(args.graph_prop_damp),
            ignore=[s for s in (args.ignore or []) if s],
        )

    # Crash early on empty embeddings; optionally fall back to topology-only
    try:
        z = emb.get("z")
        z_norm = float(np.linalg.norm(z)) if isinstance(z, np.ndarray) else 0.0
        fams = ["z_sym", "z_doc", "z_mod", "z_top", "z_text"]
        fam_norms = {k: (float(np.linalg.norm(emb[k])) if isinstance(emb.get(k), np.ndarray) else 0.0) for k in fams if k in emb}
        embed_empty = (z_norm == 0.0)
        fallback_used = False
        if embed_empty:
            if bool(args.fallback_topology) and (fam_norms.get("z_top", 0.0) > 0.0):
                emb["z"] = emb["z_top"].astype(np.float32)
                fallback_used = True
            else:
                detail = {
                    "z_norm": z_norm,
                    "family_norms": fam_norms,
                    "hint": "Relax --ignore, enable --include-text, increase --text-max-bytes/--max-text-tokens, or pass --fallback-topology",
                }
                raise RuntimeError(f"Empty program embedding (||z||=0). Details: {json.dumps(detail)}")
    except Exception as _e:
        if not isinstance(_e, RuntimeError):
            raise

    # Shapes & dims
    # Prefer explicit cache_dir; else env; else project root (/..../checkpoints)
    if args.cache_dir:
        cache_dir = args.cache_dir
    else:
        env_cache = os.environ.get("TRANSFORMER_CACHE_DIR") or os.environ.get("HF_HOME")
        if env_cache:
            cache_dir = env_cache
        else:
            mod_dir = os.path.dirname(__file__)
            proj_root = os.path.abspath(os.path.join(mod_dir, "..", "..", ".."))
            cache_dir = os.path.join(proj_root, "checkpoints")
    try:
        os.makedirs(cache_dir, exist_ok=True)
    except Exception:
        pass
    # Infer shapes/dims from local snapshot config.json (no transformers)
    snap_dir = ensure_snapshot(args.model, cache_dir)
    cfg_path = os.path.join(snap_dir, "config.json")
    cfg_obj = json.load(open(cfg_path, "r", encoding="utf-8"))
    d_model = int(cfg_obj.get("hidden_size", 4096))
    num_layers = int(cfg_obj.get("num_hidden_layers", 32))
    n_heads = int(cfg_obj.get("num_attention_heads", 32))
    n_kv_heads = int(cfg_obj.get("num_key_value_heads", n_heads))
    d_ff = int(cfg_obj.get("intermediate_size", 11008))
    head_dim = int(cfg_obj.get("head_dim", d_model // max(1, n_heads)))
    target_shapes: Optional[Dict[str, Tuple[int, int]]] = {
        "q_proj": (n_heads * head_dim, d_model),
        "k_proj": (n_kv_heads * head_dim, d_model),
        "v_proj": (n_kv_heads * head_dim, d_model),
        "o_proj": (d_model, n_heads * head_dim),
        "up_proj": (d_ff, d_model),
        "gate_proj": (d_ff, d_model),
        "down_proj": (d_model, d_ff),
    }

    # Target weights: parse or preset
    def _parse_tw(spec: Optional[str]) -> Optional[Dict[str, float]]:
        if not spec:
            return None
        out: Dict[str, float] = {}
        try:
            parts = [p.strip() for p in str(spec).split(",") if p.strip()]
            for p in parts:
                if "=" not in p:
                    continue
                n, v = p.split("=", 1)
                out[n.strip()] = float(v)
            return out or None
        except Exception:
            return None
    tw = _parse_tw(args.target_weights)
    if tw is None:
        if bool(args.code_recall_preset):
            tw = {"o_proj": 1.15, "v_proj": 1.10, "up_proj": 1.10, "down_proj": 1.05, "gate_proj": 1.00, "q_proj": 0.95, "k_proj": 0.90}
        elif bool(args.knowledge_preset):
            tw = {"q_proj": 0.95, "k_proj": 0.95, "v_proj": 0.95, "o_proj": 1.10, "up_proj": 1.10, "down_proj": 1.05}

    # Optional KBANN priors (domain-derived boosts) - program-agnostic
    kbann_mode = None
    if bool(args.kbann_priors):
        try:
            # Estimate structural complexity and density from ProgramGraph (if available)
            ents_cnt = 0
            edges_cnt = 0
            try:
                if 'pg' in locals() and (pg is not None):
                    try:
                        ents_cnt = sum(1 for _ in pg.entities())
                    except Exception:
                        ents_cnt = 0
                    try:
                        edges_cnt = sum(1 for _ in pg.edges())  # type: ignore[attr-defined]
                    except Exception:
                        edges_cnt = 0
            except Exception:
                ents_cnt = 0
                edges_cnt = 0
            scomp = float(max(0.0, min(1.0, ents_cnt / 10000.0)))  # normalize entity count
            density = float(edges_cnt) / float(max(1, ents_cnt))
            dens_term = 0.05 * float(min(1.0, density / 8.0))
            if bool(args.kbann_strong):
                kbann_mode = "strong"
                kb_tw = {
                    "o_proj": 1.0 + 0.22 * scomp + 1.2 * dens_term,
                    "up_proj": 1.0 + 0.22 * scomp + 1.2 * dens_term,
                    "down_proj": 1.0 + 0.15 * scomp + dens_term,
                    "v_proj": 0.85,   # stronger inhibitory prior
                    "k_proj": 0.95,
                    "q_proj": 0.98,
                }
            else:
                kbann_mode = "standard"
                # Prior emphasis: o/up/down (composition/usage), mild downweight v, neutral k/q
                kb_tw = {
                    "o_proj": 1.0 + 0.15 * scomp + dens_term,
                    "up_proj": 1.0 + 0.15 * scomp + dens_term,
                    "down_proj": 1.0 + 0.10 * scomp + (dens_term * 0.8),
                    "v_proj": 0.95,
                    "k_proj": 1.00,
                    "q_proj": 1.00,
                }
            if tw is None:
                tw = kb_tw
            else:
                for k, v in kb_tw.items():
                    tw[k] = float(tw.get(k, 1.0)) * float(v)
        except Exception:
            pass

    # Targets & shape safety
    NON_SQUARE_TARGETS = {"k_proj", "v_proj", "up_proj", "down_proj", "gate_proj"}
    if target_shapes:
        selected_targets: List[str] = list(target_shapes.keys())
        # shape sanity
        for t, shp in target_shapes.items():
            if not isinstance(shp, tuple) or len(shp) != 2:
                raise RuntimeError(f"Invalid target shape for {t}: {shp}")
            a, b = int(shp[0]), int(shp[1])
            if a <= 0 or b <= 0:
                raise RuntimeError(f"Non-positive dims for {t}: {shp}")
    else:
        # Without shapes, only allow square-safe targets
        selected_targets = ["q_proj", "o_proj"]
        # If user expects non-square behavior (e.g., via presets), hard-fail
        if any(t in NON_SQUARE_TARGETS for t in selected_targets):
            raise RuntimeError("Non-square targets requested without inferred shapes. Use --probe-full or restrict targets to q_proj,o_proj.")

    # Per-layer/dims consistency
    if int(num_layers) <= 0 or int(d_model) <= 0:
        raise RuntimeError("Could not infer model dims (layers/d_model). Pass a valid --model or use --probe-full.")

    # Final normalization/clipping of target weights (single place)
    if tw:
        try:
            vals = [float(v) for v in tw.values() if v is not None]
            if vals:
                mean_v = float(sum(vals) / float(len(vals)))
                if mean_v > 0:
                    for k in list(tw.keys()):
                        tw[k] = float(tw[k]) / mean_v
                # clip to [0.7, 1.3]
                for k in list(tw.keys()):
                    v = float(tw[k])
                    if v < 0.7:
                        tw[k] = 0.7
                    elif v > 1.3:
                        tw[k] = 1.3
        except Exception:
            pass

    # Auto rank scheduling
    rank_min = int(args.rank_min) if args.rank_min is not None else int(args.base_rank)
    rank_max = int(args.rank_max) if args.rank_max is not None else int(max(rank_min, args.base_rank))
    rank_global = int(args.base_rank)
    scomp_ar = 0.5
    density_ar = 0.0
    if bool(args.auto_rank):
        try:
            ents_cnt_ar = 0
            edges_cnt_ar = 0
            if 'pg' in locals() and (pg is not None):
                try:
                    ents_cnt_ar = sum(1 for _ in pg.entities())
                except Exception:
                    ents_cnt_ar = 0
                try:
                    edges_cnt_ar = sum(1 for _ in pg.edges())  # type: ignore[attr-defined]
                except Exception:
                    edges_cnt_ar = 0
            density_ar = float(edges_cnt_ar) / float(max(1, ents_cnt_ar))
            scomp_ar = float(max(0.0, min(1.0, ents_cnt_ar / 10000.0)))
            comp = max(0.0, min(1.0, 0.5 * scomp_ar + 0.5 * min(1.0, density_ar / 8.0)))
            rank_global = int(max(1, round(rank_min + comp * float(max(0, rank_max - rank_min)))))
        except Exception:
            rank_global = int(args.base_rank)

    # Generate base adapters (backend selectable)
    backend_used = str(args.gen_backend)
    try:
        if backend_used == "torch":
            adapters = generate_lora_from_embedding_torch(
# ... truncated ...


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,
    save_npz,
)
from model.hf_snapshot import ensure_snapshot
from model.inspect import detect_target_shapes_from_model
from modules.capacity import entropy_score  # type: ignore
#
# NOTE: Build is program-agnostic; no backend-specific imports here.
from modules.program_graph import ProgramGraph  # type: ignore

def _load_symbol(path: Optional[str]):
    if not path:
        return None
    mod, _, attr = path.partition(":")
    if not mod or not attr:
        raise ValueError(f"Invalid symbol path '{path}', expected 'module:attr'")
    m = importlib.import_module(mod)
    return getattr(m, attr)

def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--sources", required=False, help="Program sources root (path or URI)")
    p.add_argument("--program", required=False, help="Alias for --sources (program root path or URI)")
    p.add_argument("--model", required=True)
    p.add_argument("--adapters-dir", required=True)
    # Embedding
    p.add_argument("--embed-dim", type=int, default=1536)
    p.add_argument("--include-text", action="store_true")
    p.add_argument("--text-max-bytes", type=int, default=0)
    p.add_argument("--max-text-tokens", type=int, default=0)
    p.add_argument("--text-weight", type=float, default=0.25)
    p.add_argument("--graph-prop-hops", type=int, default=0)
    p.add_argument("--graph-prop-damp", type=float, default=0.85)
    p.add_argument("--ignore", action="append", default=None)
    # Base adapter
    p.add_argument("--base-rank", type=int, default=8)
    p.add_argument("--target-weights", default=None)
    p.add_argument("--knowledge-preset", action="store_true")


main

    generate_lora_from_embedding,
    generate_lora_from_embedding_torch,
    save_npz,
)
from model.hf_snapshot import ensure_snapshot
from model.inspect import detect_target_shapes_from_model
from modules.capacity import entropy_score  # type: ignore
#
# NOTE: Build is program-agnostic; no backend-specific imports here.
from modules.program_graph import ProgramGraph  # type: ignore

def _load_symbol(path: Optional[str]):
    if not path:
        return None
    mod, _, attr = path.partition(":")
    if not mod or not attr:
        raise ValueError(f"Invalid symbol path '{path}', expected 'module:attr'")
    m = importlib.import_module(mod)
    return getattr(m, attr)

def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--sources", required=False, help="Program sources root (path or URI)")
    p.add_argument("--program", required=False, help="Alias for --sources (program root path or URI)")
    p.add_argument("--model", required=True)
    p.add_argument("--adapters-dir", required=True)
    # Embedding
    p.add_argument("--embed-dim", type=int, default=1536)
    p.add_argument("--include-text", action="store_true")
    p.add_argument("--text-max-bytes", type=int, default=0)
    p.add_argument("--max-text-tokens", type=int, default=0)
    p.add_argument("--text-weight", type=float, default=0.25)
    p.add_argument("--graph-prop-hops", type=int, default=0)
    p.add_argument("--graph-prop-damp", type=float, default=0.85)
    p.add_argument("--ignore", action="append", default=None)
    # Base adapter
    p.add_argument("--base-rank", type=int, default=8)
    p.add_argument("--target-weights", default=None)
    p.add_argument("--knowledge-preset", action="store_true")
    # Adapter mapping is always enhanced; no CLI toggle
    p.add_argument("--code-recall-preset", action="store_true", help="Opt-in preset tuned for code recall (o,v,up,down,gate emphasized; light q,k)")
    # Priors & rounding
    p.add_argument("--kbann-priors", action="store_true")
    p.add_argument("--kbann-strong", action="store_true", help="Stronger inhibitory priors: boost o/up/down; damp v,k (slightly q)")
    p.add_argument("--round-lora", action="store_true")
    p.add_argument("--round-threshold", type=float, default=0.5)
    p.add_argument("--round-mode", choices=["none", "hard", "soft"], default=None, help="Rounding/sparsification mode for adapter matrices")
    p.add_argument("--round-soft-kp", type=float, default=10.0, help="Soft mode: percent of entries to keep (top-|.|) per axis")
    p.add_argument("--round-axis", choices=["row", "col", "global"], default="row", help="Soft mode axis for top-k sparsification")
    p.add_argument("--zero-b", action="store_true", help="After generation, set all B matrices to zero (official LoRA init)")
    p.add_argument("--learn-bias", action="store_true", help="Export zero bias vectors so downstream can fine-tune bias only")
    # Contracts (optional snapshot into manifest and embedding contracts channel)
    p.add_argument("--contracts-require-citations", action="store_true")
    p.add_argument("--contracts-retrieval-policy", default=None)
    p.add_argument("--contracts-retrieval-temp", type=float, default=None)
    p.add_argument("--contracts-weight", type=float, default=0.10)
    # Per-module export (built by default; use --no-per-module to skip)
    p.add_argument("--per-module", action="store_true", help="DEPRECATED: no effect; per-module bank is built by default. Use --no-per-module to disable.")
    p.add_argument("--no-per-module", action="store_true", help="Disable building per-module sub_adapters bank (built by default)")
    p.add_argument("--include-deps", action="store_true")
    p.add_argument("--max-deps", type=int, default=4)
    p.add_argument("--sub-rank", type=int, default=8)
    p.add_argument("--files-only", action="store_true", help="Export a files-only sub-adapter using an explicit file allowlist")
    p.add_argument(
        "--files-allowlist",
        action="append",
        default=None,
        help="Files-only mode: relative file paths to include (repeatable)",
    )
    # HF cache/model
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints")
    p.add_argument("--probe-full", action="store_true")
    p.add_argument("--gen-backend", choices=["numpy", "torch"], default="numpy")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--verbose", action="store_true")
    p.add_argument("--fallback-topology", action="store_true", help="If embedding is empty, fall back to topology-only component z_top")
    # ProgramGraph/Embedder plugins (backend-agnostic)
    p.add_argument("--pg-backend", default=None, help="Dotted path to ProgramGraph factory, e.g. 'examples.python_repo_grounded_qa.python_repo_graph:PythonRepoGraph'")
    p.add_argument("--embedder-fn", default=None, help="Dotted path to embedder function '(pg, sources_root, **opts)->{z,...}'. Defaults to PCA program embedder.")
    # Capacity scheduling
    p.add_argument("--auto-rank", action="store_true", help="Auto-schedule rank based on program complexity")
    p.add_argument("--rank-min", type=int, default=None)
    p.add_argument("--rank-max", type=int, default=None)
    p.add_argument("--mdl-budget-params", type=int, default=0, help="Optional MDL-style global parameter budget across all layers/targets (approximate)")
    # Program seeding
    p.add_argument("--init-program-state", action="store_true", help="Initialize or update baseline ProgramState (.program_state.json) after build")
    p.add_argument("--program-state-path", default=None, help="Optional explicit path to write ProgramState JSON; defaults to <sources_root>/.program_state.json")
    # (self-tune is no longer part of build; use modules/tune.py externally if desired)
    args = p.parse_args()

    # Warn on deprecated/ignored flags for per-module export
    try:
        if bool(args.per_module) and bool(args.no_per_module):
            print("[pca.build] Warning: Both --per-module and --no-per-module were provided; proceeding with --no-per-module (disables per-module bank).", file=sys.stderr)
        elif bool(args.per_module):
            print("[pca.build] Warning: --per-module is deprecated and has no effect; per-module bank is built by default. Use --no-per-module to disable.", file=sys.stderr)
    except Exception:
        pass
    # Determinism: set seeds and hash seed (best-effort)
    try:
        os.environ["PYTHONHASHSEED"] = str(int(args.seed))
    except Exception:
        pass
    try:
        random.seed(int(args.seed))
    except Exception:
        pass
    try:
        np.random.seed(int(args.seed))
    except Exception:
        pass
    torch_seeded = False
    try:
        import torch as _torch  # type: ignore
        _torch.manual_seed(int(args.seed))
        torch_seeded = True
    except Exception:
        torch_seeded = False

    out_dir = os.path.abspath(os.path.expanduser(os.path.expandvars(args.adapters_dir)))
    os.makedirs(out_dir, exist_ok=True)

    # Normalize sources root (prefer --program, else --sources)
    src_arg = args.program if getattr(args, "program", None) else args.sources
    if not src_arg:
        raise SystemExit("--sources or --program is required")
    sources_root = os.path.abspath(os.path.expanduser(os.path.expandvars(src_arg)))

    # Construct ProgramGraph via plugin (if provided)
    pg: Optional[ProgramGraph] = None
    try:
        pg_ctor = _load_symbol(args.pg_backend)
        if pg_ctor:
            pg = pg_ctor(sources_root, ignore=[s for s in (args.ignore or []) if s])
    except Exception:
        pg = None

    # Build embedding (program-agnostic) via plugin or default PCA program embedder
    embedder = _load_symbol(args.embedder_fn) if args.embedder_fn else None
    if embedder is not None:
        try:
            # Prefer new program-graph signature
            emb = embedder(
                pg,
                sources_root=sources_root,
                dim=int(args.embed_dim),
                seed=int(args.seed),
                include_text=bool(args.include_text),
                text_max_bytes=int(args.text_max_bytes),
                max_text_tokens=int(args.max_text_tokens),
                text_weight=float(args.text_weight),
                calls_weight=0.25,
                types_weight=0.20,
                tests_weight=0.15,
                graph_prop_hops=int(args.graph_prop_hops),
                graph_prop_damp=float(args.graph_prop_damp),
                ignore=[s for s in (args.ignore or []) if s],
            )
        except Exception:
            # Backward-compat: allow sources-root signature
            emb = embedder(
                sources_root,
                dim=int(args.embed_dim),
                seed=int(args.seed),
                include_text=bool(args.include_text),
                text_max_bytes=int(args.text_max_bytes),
                max_text_tokens=int(args.max_text_tokens),
                text_weight=float(args.text_weight),
                calls_weight=0.25,
                types_weight=0.20,
                tests_weight=0.15,
                graph_prop_hops=int(args.graph_prop_hops),
                graph_prop_damp=float(args.graph_prop_damp),
                ignore=[s for s in (args.ignore or []) if s],
            )
    else:
        emb = build_program_embedding(
            pg if pg is not None else None,  # type: ignore[arg-type]
            sources_root=sources_root,
            dim=int(args.embed_dim),
            seed=int(args.seed),
            include_text=bool(args.include_text),
            text_max_bytes=int(args.text_max_bytes),
            max_text_tokens=int(args.max_text_tokens),
            text_weight=float(args.text_weight),
            calls_weight=0.25,
            types_weight=0.20,
            tests_weight=0.15,
            contracts_kv={
                "require_citations": (True if bool(args.contracts_require_citations) else False),
                "retrieval_policy": (str(args.contracts_retrieval_policy) if args.contracts_retrieval_policy else ""),
                "retrieval_temp": (str(args.contracts_retrieval_temp) if args.contracts_retrieval_temp is not None else ""),
            },
            contracts_weight=float(max(0.0, args.contracts_weight)),
            graph_prop_hops=int(args.graph_prop_hops),
            graph_prop_damp=float(args.graph_prop_damp),
            ignore=[s for s in (args.ignore or []) if s],
        )

    # Crash early on empty embeddings; optionally fall back to topology-only
    try:
        z = emb.get("z")
        z_norm = float(np.linalg.norm(z)) if isinstance(z, np.ndarray) else 0.0
        fams = ["z_sym", "z_doc", "z_mod", "z_top", "z_text"]
        fam_norms = {k: (float(np.linalg.norm(emb[k])) if isinstance(emb.get(k), np.ndarray) else 0.0) for k in fams if k in emb}
        embed_empty = (z_norm == 0.0)
        fallback_used = False
        if embed_empty:
            if bool(args.fallback_topology) and (fam_norms.get("z_top", 0.0) > 0.0):
                emb["z"] = emb["z_top"].astype(np.float32)
                fallback_used = True
            else:
                detail = {
                    "z_norm": z_norm,
                    "family_norms": fam_norms,
                    "hint": "Relax --ignore, enable --include-text, increase --text-max-bytes/--max-text-tokens, or pass --fallback-topology",
                }
                raise RuntimeError(f"Empty program embedding (||z||=0). Details: {json.dumps(detail)}")
    except Exception as _e:
        if not isinstance(_e, RuntimeError):
            raise

    # Shapes & dims
    # Prefer explicit cache_dir; else env; else project root (/..../checkpoints)
    if args.cache_dir:
        cache_dir = args.cache_dir
    else:
        env_cache = os.environ.get("TRANSFORMER_CACHE_DIR") or os.environ.get("HF_HOME")
        if env_cache:
            cache_dir = env_cache
        else:
            mod_dir = os.path.dirname(__file__)
            proj_root = os.path.abspath(os.path.join(mod_dir, "..", "..", ".."))
            cache_dir = os.path.join(proj_root, "checkpoints")
    try:
        os.makedirs(cache_dir, exist_ok=True)
    except Exception:
        pass
    # Infer shapes/dims from local snapshot config.json (no transformers)
    snap_dir = ensure_snapshot(args.model, cache_dir)
    cfg_path = os.path.join(snap_dir, "config.json")
    cfg_obj = json.load(open(cfg_path, "r", encoding="utf-8"))
    d_model = int(cfg_obj.get("hidden_size", 4096))
    num_layers = int(cfg_obj.get("num_hidden_layers", 32))
    n_heads = int(cfg_obj.get("num_attention_heads", 32))
    n_kv_heads = int(cfg_obj.get("num_key_value_heads", n_heads))
    d_ff = int(cfg_obj.get("intermediate_size", 11008))
    head_dim = int(cfg_obj.get("head_dim", d_model // max(1, n_heads)))
    target_shapes: Optional[Dict[str, Tuple[int, int]]] = {
        "q_proj": (n_heads * head_dim, d_model),
        "k_proj": (n_kv_heads * head_dim, d_model),
        "v_proj": (n_kv_heads * head_dim, d_model),
        "o_proj": (d_model, n_heads * head_dim),
        "up_proj": (d_ff, d_model),
        "gate_proj": (d_ff, d_model),
        "down_proj": (d_model, d_ff),
    }

    # Target weights: parse or preset
    def _parse_tw(spec: Optional[str]) -> Optional[Dict[str, float]]:
        if not spec:
            return None
        out: Dict[str, float] = {}
        try:
            parts = [p.strip() for p in str(spec).split(",") if p.strip()]
            for p in parts:
                if "=" not in p:
                    continue
                n, v = p.split("=", 1)
                out[n.strip()] = float(v)
            return out or None
        except Exception:
            return None
    tw = _parse_tw(args.target_weights)
    if tw is None:
        if bool(args.code_recall_preset):
            tw = {"o_proj": 1.15, "v_proj": 1.10, "up_proj": 1.10, "down_proj": 1.05, "gate_proj": 1.00, "q_proj": 0.95, "k_proj": 0.90}
        elif bool(args.knowledge_preset):
            tw = {"q_proj": 0.95, "k_proj": 0.95, "v_proj": 0.95, "o_proj": 1.10, "up_proj": 1.10, "down_proj": 1.05}

    # Optional KBANN priors (domain-derived boosts) - program-agnostic
    kbann_mode = None
    if bool(args.kbann_priors):
        try:
            # Estimate structural complexity and density from ProgramGraph (if available)
            ents_cnt = 0
            edges_cnt = 0
            try:
                if 'pg' in locals() and (pg is not None):
                    try:
                        ents_cnt = sum(1 for _ in pg.entities())
                    except Exception:
                        ents_cnt = 0
                    try:
                        edges_cnt = sum(1 for _ in pg.edges())  # type: ignore[attr-defined]
                    except Exception:
                        edges_cnt = 0
            except Exception:
                ents_cnt = 0
                edges_cnt = 0
            scomp = float(max(0.0, min(1.0, ents_cnt / 10000.0)))  # normalize entity count
            density = float(edges_cnt) / float(max(1, ents_cnt))
            dens_term = 0.05 * float(min(1.0, density / 8.0))
            if bool(args.kbann_strong):
                kbann_mode = "strong"
                kb_tw = {
                    "o_proj": 1.0 + 0.22 * scomp + 1.2 * dens_term,
                    "up_proj": 1.0 + 0.22 * scomp + 1.2 * dens_term,
                    "down_proj": 1.0 + 0.15 * scomp + dens_term,
                    "v_proj": 0.85,   # stronger inhibitory prior
                    "k_proj": 0.95,
                    "q_proj": 0.98,
                }
            else:
                kbann_mode = "standard"
                # Prior emphasis: o/up/down (composition/usage), mild downweight v, neutral k/q
                kb_tw = {
                    "o_proj": 1.0 + 0.15 * scomp + dens_term,
                    "up_proj": 1.0 + 0.15 * scomp + dens_term,
                    "down_proj": 1.0 + 0.10 * scomp + (dens_term * 0.8),
                    "v_proj": 0.95,
                    "k_proj": 1.00,
                    "q_proj": 1.00,
                }
            if tw is None:
                tw = kb_tw
            else:
                for k, v in kb_tw.items():
                    tw[k] = float(tw.get(k, 1.0)) * float(v)
        except Exception:
            pass

    # Targets & shape safety
    NON_SQUARE_TARGETS = {"k_proj", "v_proj", "up_proj", "down_proj", "gate_proj"}
    if target_shapes:
        selected_targets: List[str] = list(target_shapes.keys())
        # shape sanity
        for t, shp in target_shapes.items():
            if not isinstance(shp, tuple) or len(shp) != 2:
                raise RuntimeError(f"Invalid target shape for {t}: {shp}")
            a, b = int(shp[0]), int(shp[1])
            if a <= 0 or b <= 0:
                raise RuntimeError(f"Non-positive dims for {t}: {shp}")
    else:
        # Without shapes, only allow square-safe targets
        selected_targets = ["q_proj", "o_proj"]
        # If user expects non-square behavior (e.g., via presets), hard-fail
        if any(t in NON_SQUARE_TARGETS for t in selected_targets):
            raise RuntimeError("Non-square targets requested without inferred shapes. Use --probe-full or restrict targets to q_proj,o_proj.")

    # Per-layer/dims consistency
    if int(num_layers) <= 0 or int(d_model) <= 0:
        raise RuntimeError("Could not infer model dims (layers/d_model). Pass a valid --model or use --probe-full.")

    # Final normalization/clipping of target weights (single place)
    if tw:
        try:
            vals = [float(v) for v in tw.values() if v is not None]
            if vals:
                mean_v = float(sum(vals) / float(len(vals)))
                if mean_v > 0:
                    for k in list(tw.keys()):
                        tw[k] = float(tw[k]) / mean_v
                # clip to [0.7, 1.3]
                for k in list(tw.keys()):
                    v = float(tw[k])
                    if v < 0.7:
                        tw[k] = 0.7
                    elif v > 1.3:
                        tw[k] = 1.3
        except Exception:
            pass

    # Auto rank scheduling
    rank_min = int(args.rank_min) if args.rank_min is not None else int(args.base_rank)
    rank_max = int(args.rank_max) if args.rank_max is not None else int(max(rank_min, args.base_rank))
    rank_global = int(args.base_rank)
    scomp_ar = 0.5
    density_ar = 0.0
    if bool(args.auto_rank):
        try:
            ents_cnt_ar = 0
            edges_cnt_ar = 0
            if 'pg' in locals() and (pg is not None):
                try:
                    ents_cnt_ar = sum(1 for _ in pg.entities())
                except Exception:
                    ents_cnt_ar = 0
                try:
                    edges_cnt_ar = sum(1 for _ in pg.edges())  # type: ignore[attr-defined]
                except Exception:
                    edges_cnt_ar = 0
            density_ar = float(edges_cnt_ar) / float(max(1, ents_cnt_ar))
            scomp_ar = float(max(0.0, min(1.0, ents_cnt_ar / 10000.0)))
            comp = max(0.0, min(1.0, 0.5 * scomp_ar + 0.5 * min(1.0, density_ar / 8.0)))
            rank_global = int(max(1, round(rank_min + comp * float(max(0, rank_max - rank_min)))))
        except Exception:
            rank_global = int(args.base_rank)

    # Generate base adapters (backend selectable)
    backend_used = str(args.gen_backend)
    try:
        if backend_used == "torch":
            adapters = generate_lora_from_embedding_torch(
                emb["z"],
                d_model=int(d_model),
                num_layers=int(num_layers),
                rank=int(rank_global),
                seed=int(args.seed),
                targets=selected_targets,
                target_shapes=target_shapes,
                target_weights=tw,
            )
            if bool(args.learn_bias):
                # Add zero bias vectors for downstream fine-t
# ... truncated ...


    snap_dir = ensure_snapshot(args.model, cache_dir)
    cfg_path = os.path.join(snap_dir, "config.json")
    cfg_obj = json.load(open(cfg_path, "r", encoding="utf-8"))
    d_model = int(cfg_obj.get("hidden_size", 4096))
    num_layers = int(cfg_obj.get("num_hidden_layers", 32))
    n_heads = int(cfg_obj.get("num_attention_heads", 32))
    n_kv_heads = int(cfg_obj.get("num_key_value_heads", n_heads))
    d_ff = int(cfg_obj.get("intermediate_size", 11008))
    head_dim = int(cfg_obj.get("head_dim", d_model // max(1, n_heads)))
    target_shapes: Optional[Dict[str, Tuple[int, int]]] = {
        "q_proj": (n_heads * head_dim, d_model),
        "k_proj": (n_kv_heads * head_dim, d_model),
        "v_proj": (n_kv_heads * head_dim, d_model),
        "o_proj": (d_model, n_heads * head_dim),
        "up_proj": (d_ff, d_model),
        "gate_proj": (d_ff, d_model),
        "down_proj": (d_model, d_ff),
    }

    # Target weights: parse or preset
    def _parse_tw(spec: Optional[str]) -> Optional[Dict[str, float]]:
        if not spec:
            return None
        out: Dict[str, float] = {}
        try:
            parts = [p.strip() for p in str(spec).split(",") if p.strip()]
            for p in parts:
                if "=" not in p:
                    continue
                n, v = p.split("=", 1)
                out[n.strip()] = float(v)
            return out or None
        except Exception:
            return None
    tw = _parse_tw(args.target_weights)
    if tw is None:
        if bool(args.code_recall_preset):
            tw = {"o_proj": 1.15, "v_proj": 1.10, "up_proj": 1.10, "down_proj": 1.05, "gate_proj": 1.00, "q_proj": 0.95, "k_proj": 0.90}
        elif bool(args.knowledge_preset):
            tw = {"q_proj": 0.95, "k_proj": 0.95, "v_proj": 0.95, "o_proj": 1.10, "up_proj": 1.10, "down_proj": 1.05}

    # Optional KBANN priors (domain-derived boosts) - program-agnostic
    kbann_mode = None
    if bool(args.kbann_priors):
        try:
            # Estimate structural complexity and density from ProgramGraph (if available)
            ents_cnt = 0
            edges_cnt = 0
            try:
                if 'pg' in locals() and (pg is not None):
                    try:
                        ents_cnt = sum(1 for _ in pg.entities())
                    except Exception:
                        ents_cnt = 0


    # Optional capacity schedule per target (effective rank via zeroing)
    per_target_keep: Dict[str, int] = {}
    if bool(args.auto_rank) and rank_global > 0:
        base_frac: Dict[str, float] = {
            "o_proj": 1.00, "up_proj": 1.00, "down_proj": 0.90, "gate_proj": 0.80,
            "q_proj": 0.70, "k_proj": 0.65, "v_proj": 0.60,
        }
        comp_adj = max(0.85, min(1.15, 1.0 + 0.15 * (scomp_ar - 0.5)))
        for t in selected_targets:
            frac = base_frac.get(t, 0.8) * comp_adj
            keep = int(max(1, min(rank_global, round(rank_global * frac))))
            per_target_keep[t] = keep
        # MDL-style global budget: cap total params across layers/targets
        try:
            budget = int(max(0, int(args.mdl_budget_params)))
        except Exception:
            budget = 0
        if budget > 0:
            # estimate params per target per layer: keep * (d_out + d_in)
            def _shape_for(t: str) -> Tuple[int, int]:
                if target_shapes and t in target_shapes:
                    a, b = target_shapes[t]
                    return int(a), int(b)
                # fallback square
                return int(d_model), int(d_model)
            total = 0
            per_t_cost: Dict[str, int] = {}
            for t, k in per_target_keep.items():
                a, b = _shape_for(t)
                cost = int(max(1, k)) * int(max(1, a + b))
                per_t_cost[t] = cost
                total += cost
            total *= int(num_layers)
            if total > budget:
                scale = float(budget) / float(total)
                # rescale keeps proportionally, ensure at least 1
                for t in list(per_target_keep.keys()):
                    k = int(per_target_keep[t])
                    k2 = int(max(1, round(float(k) * scale)))
                    per_target_keep[t] = min(int(rank_global), k2)
        try:
            for i in range(len(adapters.get("layers", []))):
                for name, tensors in adapters["layers"][i].items():
                    keep = int(per_target_keep.get(name, rank_global))
                    A = tensors.get("A"); B = tensors.get("B")


    # Optional rounding/sparsification
    try:
        mode = (args.round_mode or ("hard" if bool(args.round_lora) else "none"))
        if mode == "hard":
            thr = float(max(0.0, args.round_threshold))
            for i in range(len(adapters.get("layers", []))):
                for name, tensors in adapters["layers"][i].items():
                    for key in ("A", "B"):
                        arr = tensors.get(key)
                        if not isinstance(arr, np.ndarray):
                            continue
                        q = float(np.median(np.abs(arr))) if arr.size > 0 else 0.0
                        if q <= 0:
                            continue
                        out = np.where(np.abs(arr) < (thr * q), 0.0, np.sign(arr) * q).astype(np.float32)
                        tensors[key] = out
        elif mode == "soft":
            kp = float(max(0.0, args.round_soft_kp))
            axis = str(args.round_axis)
            keep_frac = max(0.0, min(100.0, kp)) / 100.0
            def _sparsify_topk(a: np.ndarray) -> np.ndarray:
                if a.size == 0 or keep_frac <= 0.0:
                    return np.zeros_like(a)
                if keep_frac >= 1.0:
                    return a
                if axis == "global":
                    k = int(np.ceil(keep_frac * a.size))
                    if k <= 0:
                        return np.zeros_like(a)
                    flat_idx = np.argpartition(np.abs(a).ravel(), -k)[-k:]
                    mask = np.zeros(a.size, dtype=bool)
                    mask[flat_idx] = True
                    return (a.ravel() * mask).reshape(a.shape)
                elif axis == "row":
                    rows, cols = a.shape
                    k = int(np.ceil(keep_frac * cols))
                    if k <= 0:
                        return np.zeros_like(a)
                    out = np.zeros_like(a)
                    for r in range(rows):
                        idx = np.argpartition(np.abs(a[r]), -k)[-k:]
                        out[r, idx] = a[r, idx]
                    return out
                else:  # col
                    rows, cols = a.shape
                    k = int(np.ceil(keep_frac * rows))
                    if k <= 0:
                        return np.zeros_like(a)
                    out = np.zeros_like(a)
                    for c in range(cols):
                        idx = np.argpartition(np.abs(a[:, c]), -k)[-k:]
                        out[idx, c] = a[idx, c]
                    return out
            for i in range(len(adapters.get("layers", []))):
                for name, tensors in adapters["layers"][i].items():
                    for key in ("A", "B"):
                        arr = tensors.get(key)
                        if not isinstance(arr, np.ndarray) or arr.ndim != 2:
                            continue
                        tensors[key] = _sparsify_topk(arr).astype(np.float32)
        # else: none -> no rounding
    except Exception:
        pass

    # Git metadata (best-effort)
    commit_sha: Optional[str] = None
    tree_sha: Optional[str] = None
    try:
        commit_sha = subprocess.check_output(["git", "-C", sources_root, "rev-parse", "HEAD"], text=True).strip()
    except Exception:
        commit_sha = None
    try:
        tree_sha = subprocess.check_output(["git", "-C", sources_root, "rev-parse", "HEAD^{tree}"], text=True).strip()


"""
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.
    """
    return float(query * key)

def scaled_dot_product_attention(
    query: Tuple[float, ...], key: Tuple[float, ...], value: Tuple[float, ...]
) -> float:
    """
    Minimal scaled dot-product attention for 1-token, 1-head toy inputs.

    Steps:
    1) score = sum_i query[i] * key[i]
    2) scale by sqrt(d) where d is the dimensionality of query/key
    3) softmax over a single score is 1.0 (degenerate 1-token case)
    4) output = weight * sum_i value[i] (here weight==1.0)
    """
    if len(query) != len(key):
        raise ValueError("query and key must have same length")
    d = max(1, len(query))
    score = sum(float(q) * float(k) for q, k in zip(query, key))
    scaled = score / math.sqrt(float(d))
    # With a single score, softmax(scaled) == 1.0; keep explicit for clarity
    weight = 1.0 if math.isfinite(scaled) else 0.0
    return float(weight * sum(float(v) for v in value))

def explain_attention_brief() -> str:
    """
    Return a brief, human-readable explanation of attention suitable for tests.
    """
    return (
        "Attention computes similarity between a query and keys, normalizes the "
        "scores (softmax), and uses them to weight values. The scaled "
        "dot-product form divides by sqrt(d) to keep gradients stable."
    )


"""
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.
    """
    return float(query * key)

def scaled_dot_product_attention(
    query: Tuple[float, ...], key: Tuple[float, ...], value: Tuple[float, ...]
) -> float:
    """
    Minimal scaled dot-product attention for 1-token, 1-head toy inputs.

    Steps:
    1) score = sum_i query[i] * key[i]
    2) scale by sqrt(d) where d is the dimensionality of query/key
    3) softmax over a single score is 1.0 (degenerate 1-token case)
    4) output = weight * sum_i value[i] (here weight==1.0)
    """
    if len(query) != len(key):
        raise ValueError("query and key must have same length")
    d = max(1, len(query))
    score = sum(float(q) * float(k) for q, k in zip(query, key))
    scaled = score / math.sqrt(float(d))
    # With a single score, softmax(scaled) == 1.0; keep explicit for clarity


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.
    """
    return float(query * key)

def scaled_dot_product_attention(
    query: Tuple[float, ...], key: Tuple[float, ...], value: Tuple[float, ...]
) -> float:
    """
    Minimal scaled dot-product attention for 1-token, 1-head toy inputs.

    Steps:
    1) score = sum_i query[i] * key[i]
    2) scale by sqrt(d) where d is the dimensionality of query/key
    3) softmax over a single score is 1.0 (degenerate 1-token case)
    4) output = weight * sum_i value[i] (here weight==1.0)
    """
    if len(query) != len(key):
        raise ValueError("query and key must have same length")
    d = max(1, len(query))
    score = sum(float(q) * float(k) for q, k in zip(query, key))
    scaled = score / math.sqrt(float(d))
    # With a single score, softmax(scaled) == 1.0; keep explicit for clarity
    weight = 1.0 if math.isfinite(scaled) else 0.0
    return float(weight * sum(float(v) for v in value))

def explain_attention_brief() -> str:
    """
    Return a brief, human-readable explanation of attention suitable for tests.
    """
    return (
        "Attention computes similarity between a query and keys, normalizes the "
        "scores (softmax), and uses them to weight values. The scaled "
        "dot-product form divides by sqrt(d) to keep gradients stable."
    )


) -> float:
    """
    Minimal scaled dot-product attention for 1-token, 1-head toy inputs.

    Steps:
    1) score = sum_i query[i] * key[i]
    2) scale by sqrt(d) where d is the dimensionality of query/key
    3) softmax over a single score is 1.0 (degenerate 1-token case)
    4) output = weight * sum_i value[i] (here weight==1.0)
    """
    if len(query) != len(key):
        raise ValueError("query and key must have same length")
    d = max(1, len(query))
    score = sum(float(q) * float(k) for q, k in zip(query, key))
    scaled = score / math.sqrt(float(d))
    # With a single score, softmax(scaled) == 1.0; keep explicit for clarity
    weight = 1.0 if math.isfinite(scaled) else 0.0
    return float(weight * sum(float(v) for v in value))

def explain_attention_brief() -> str:
    """
    Return a brief, human-readable explanation of attention suitable for tests.
    """
    return (
        "Attention computes similarity between a query and keys, normalizes the "
        "scores (softmax), and uses them to weight values. The scaled "
        "dot-product form divides by sqrt(d) to keep gradients stable."
    )