"""REFRACT v0.1 CLI. Usage: python3 -m refract.cli score \\ --model MODEL.gguf \\ --reference "ctk=f16,ctv=f16" \\ --candidate "ctk=q8_0,ctv=turbo4,attn_rot_v=0" \\ --prompts refract/prompts/v0.1.jsonl \\ --corpus ~/local_llms/llama.cpp/wikitext-2-raw/wiki.test.raw \\ --chunks 32 -c 512 -ngl 99 """ from __future__ import annotations import argparse import json import sys from pathlib import Path from .axes.gtm import run_gtm from .axes.kld import run_kld from .axes.plad import run_plad from .axes.rniah import run_rniah from .axes.trajectory import run_trajectory from .report import json_report, text_report, to_json_string from .runner import KVConfig from .score import MIN_FLOOR, composite_score _SCORE_DESCRIPTION = """\ REFRACT v0.2 — score a KV-cache quantization config against a fp16-KV reference on the same model. Returns a 0–100 composite + EXCELLENT/PASS/DEGRADED/FAIL band, with per-axis bands and one-line plain-English interpretation. Four axes (A and B are cheap and run by default; C and D are opt-in): A — GTM / Trajectory Greedy Trajectory Match. Greedy-decode N tokens from each prompt under both the reference and candidate KV configs; compare the emitted token IDs. Score = fraction of the candidate's decoded tokens that match the reference. Catches "the candidate generates different text" failures. B — KLD@D KL Divergence at the Decoder. Per-token KL between the candidate's next-token distribution and the reference's, averaged across a natural-text corpus. Bit-exact zero on Metal when candidate == reference, so any non-zero is real signal. Catches "distributions silently shift even though argmax stayed the same" failures. C — R-NIAH (--full) Retrieval Needle-In-A-Haystack. Insert a sentinel fact ("the password is X") into a long context at fractional positions; ask the model to retrieve it. Score = 100 * (1 - mean candidate degradation vs reference per (length, position) cell). Catches "scores 99 on KLD@D but fails at 32K context" failures that short-window axes miss. D — PLAD (--full) Perturbation-Locality Aware Drift. Generate anchor completions; perturb each prompt minimally (typo, casing, punctuation, paraphrase); compare candidate's drift vs reference's drift via token edit distance. Catches "works on the demo, breaks on real users with typos" brittleness. Cost on a 7B Q8 model (Apple Silicon ballpark): default Trajectory + KLD ~5-7 min --full adds R-NIAH + PLAD ~25-30 min Composite = harmonic mean of all axes scored. Any single axis being broken drops the composite hard — the framework is intentionally fail-loud. """ _SCORE_EPILOG = """\ Examples: Quick go/no-go (default; just trajectory + KLD): python3 -m refract.cli score \\ --model model.gguf \\ --candidate "ctk=q8_0,ctv=q8_0" \\ --prompts refract/prompts/v0.1.jsonl \\ --corpus ~/local_llms/llama.cpp/wikitext-2-raw/wiki.test.raw \\ --axis-a trajectory Full audit (all four axes): python3 -m refract.cli score \\ --model model.gguf \\ --candidate "ctk=q8_0,ctv=turbo4" \\ --prompts refract/prompts/v0.1.jsonl \\ --corpus ~/local_llms/llama.cpp/wikitext-2-raw/wiki.test.raw \\ --axis-a trajectory \\ --full \\ --rniah-haystack ~/local_llms/llama.cpp/wikitext-2-raw/wiki.train.raw \\ --rniah-ctx-max 16384 \\ --json-out report.json Borderline result triage (start cheap, add R-NIAH only if you need it): # First pass: composite=72 DEGRADED → see which axis dragged it down python3 -m refract.cli score ... (default) # If KLD is fine but trajectory is bad → don't run R-NIAH; the issue is # short-context distribution drift, not long-context retrieval. # If KLD is fine AND trajectory is fine → run --axis-rniah; long-context # retrieval is the surface most likely to be the hidden problem. KV config syntax: comma-separated key=value pairs, e.g. "ctk=q8_0,ctv=turbo4". Recognised keys: ctk, ctv, attn_rot_k, attn_rot_v, attn_rot_disable. """ _REFRACT_CACHE = Path.home() / ".cache" / "refract" _WIKITEXT_2_URL = "https://huggingface.co/datasets/ggml-org/ci/resolve/main/wikitext-2-raw-v1.zip" def _ensure_wikitext_2(cache_dir: Path = _REFRACT_CACHE, silent: bool = False) -> Path: """Make sure wikitext-2-raw is downloaded + extracted under ``cache_dir/wikitext-2-raw/``. Returns that directory. Idempotent: re-running is a no-op when files already exist. Network: ~10 MB single zip, ~30 s on a typical home connection. No third-party deps; uses stdlib urllib + zipfile. """ import urllib.request import zipfile target = cache_dir / "wikitext-2-raw" test_p = target / "wiki.test.raw" train_p = target / "wiki.train.raw" if test_p.exists() and train_p.exists(): return target cache_dir.mkdir(parents=True, exist_ok=True) zip_path = cache_dir / "wikitext-2-raw-v1.zip" if not silent: print(f"Downloading wikitext-2-raw (~10MB) → {target} ...") urllib.request.urlretrieve(_WIKITEXT_2_URL, zip_path) if not silent: print("Unzipping ...") with zipfile.ZipFile(zip_path) as zf: zf.extractall(cache_dir) try: zip_path.unlink() except OSError: pass if not test_p.exists() or not train_p.exists(): raise RuntimeError( f"Wikitext-2 unzip didn't produce expected files at {target}. " f"Inspect {cache_dir} or fetch manually." ) if not silent: print(f"✓ Cached at {target}") return target def _resolve_default_paths(args, *, need_corpus: bool, need_haystack: bool): """Auto-resolve --corpus and --rniah-haystack from the wikitext-2 cache. If the cache is missing AND ``--no-auto-fetch`` is NOT set, this triggers a one-time wikitext-2 download. With ``--no-auto-fetch``, raises with a clear remediation message. """ if not (need_corpus or need_haystack): return have_corpus = bool(getattr(args, "corpus", None)) have_haystack = bool(getattr(args, "rniah_haystack", None)) if (need_corpus and have_corpus) and (not need_haystack or have_haystack): return no_fetch = bool(getattr(args, "no_auto_fetch", False)) target = _REFRACT_CACHE / "wikitext-2-raw" have_cache = (target / "wiki.test.raw").exists() and (target / "wiki.train.raw").exists() if not have_cache: if no_fetch: raise SystemExit( "Missing --corpus (and/or --rniah-haystack) and " "--no-auto-fetch was passed. Either pass paths explicitly " "or drop --no-auto-fetch to let REFRACT download " "wikitext-2-raw to ~/.cache/refract/ (~10MB)." ) _ensure_wikitext_2() if need_corpus and not have_corpus: args.corpus = target / "wiki.test.raw" print(f" using cached corpus : {args.corpus}") if need_haystack and not have_haystack: args.rniah_haystack = target / "wiki.train.raw" print(f" using cached haystack: {args.rniah_haystack}") def _add_score_parser(sub): import argparse as _ap p = sub.add_parser( "score", help="Score a candidate KV config vs the fp16-KV reference (4 axes).", description=_SCORE_DESCRIPTION, epilog=_SCORE_EPILOG, formatter_class=_ap.RawDescriptionHelpFormatter, ) p.add_argument("--model", required=True, type=Path, help="Path to the GGUF model.") p.add_argument("--reference", default="ctk=f16,ctv=f16", help="Reference KV config (default: ctk=f16,ctv=f16).") p.add_argument("--candidate", required=True, help="Candidate KV config to score.") p.add_argument("--prompts", type=Path, default=None, help="Path to JSONL prompts file. If omitted, REFRACT uses " "the bundled refract/prompts/v0.1.jsonl shipped in the " "wheel (30 CC0 prompts; see prompts/README.md).") p.add_argument("--corpus", type=Path, default=None, help="Path to plain-text corpus for KLD axis. " "If omitted, REFRACT auto-downloads wikitext-2-raw " "(~10MB) to ~/.cache/refract/ and uses wiki.test.raw. " "Pass --no-auto-fetch to require an explicit path.") p.add_argument("--no-auto-fetch", action="store_true", help="Disable auto-download of wikitext-2-raw. Requires " "--corpus and --rniah-haystack to be passed " "explicitly.") p.add_argument("--chunks", type=int, default=32, help="--chunks for llama-perplexity (default: 32).") p.add_argument("-c", "--ctx", type=int, default=512, help="Context size (default: 512).") p.add_argument("-ngl", "--n-gpu-layers", type=int, default=99, help="-ngl flag (default: 99).") p.add_argument("--n-predict", type=int, default=128, help="Tokens to greedy-decode per GTM prompt (default: 128).") p.add_argument("--seed", type=int, default=42, help="Greedy seed (default: 42).") p.add_argument("--measure-floor", action="store_true", help="Measure REFRACT(ref, ref) and abort if < %.1f." % MIN_FLOOR) p.add_argument("--skip-gtm", action="store_true", help="Skip Axis A. Composite uses KLD only (debug).") p.add_argument("--skip-kld", action="store_true", help="Skip Axis B. Composite uses GTM only (debug).") p.add_argument("--axis-a", choices=["gtm", "trajectory"], default="trajectory", help="Axis A implementation. 'trajectory' (default, v0.1.4+) " "captures token IDs at decode time via the patched " "llama-completion binary — no detokenize round-trip. " "'gtm' (deprecated, v0.1.x) compares retokenized text " "and has known unit-mismatch artifacts.") p.add_argument("--backend", choices=["auto", "llamacpp", "mlx", "vllm"], default="auto", help="Inference backend. 'auto' (default) infers from " "model path: .gguf → llamacpp; directory → mlx. " "Override via REFRACT_BACKEND env var.") p.add_argument("--full", action="store_true", help="Enable all v0.2 axes (R-NIAH + PLAD). Equivalent to " "passing --axis-rniah --axis-plad. R-NIAH still " "requires --rniah-haystack and --rniah-ctx-max.") p.add_argument("--axis-rniah", action="store_true", help="Enable Axis C (R-NIAH, v0.2). Probes long-context " "retrieval degradation. Requires --rniah-haystack " "(auto-resolved from cache when omitted). Cost: " "~10–15 min on a 7B Q8 at 16K, scaling roughly " "linearly with max length.") p.add_argument("--rniah-haystack", type=Path, default=None, help="Path to long-text haystack corpus for R-NIAH. " "Auto-resolved from ~/.cache/refract/ when omitted " "(wiki.train.raw, fits cells up to ~16K cleanly).") p.add_argument("--rniah-up-to", type=int, default=16384, help="R-NIAH max context length to test. Lengths are " "auto-generated as a doubling step-up starting at " "4096 (4K, 8K, 16K, 32K, ... up to this value). " "Default: 16384 (4K, 8K, 16K). Pass 32768 / 65536 / " "131072 for deeper long-context audits. Cells are " "always run at positions 0.10/0.50/0.90 unless " "overridden via --rniah-positions.") p.add_argument("--rniah-ctx-max", type=int, default=None, help="(Power-user) hard ceiling for R-NIAH cells. Cells " "longer than this are skipped. Defaults to " "--rniah-up-to. Useful when the model cannot " "actually do its nominal max context.") p.add_argument("--rniah-lengths", type=str, default=None, help="(Power-user) comma-separated R-NIAH lengths, " "overrides --rniah-up-to. Default: synthesised " "from --rniah-up-to.") p.add_argument("--rniah-positions", type=str, default=None, help="Comma-separated R-NIAH needle positions as " "fractions of length. Default: 0.10,0.50,0.90.") p.add_argument("--rniah-trials", type=int, default=1, help="Trials per cell. Default: 1.") p.add_argument("--axis-plad", action="store_true", help="Enable Axis D (PLAD, v0.2). Probes brittleness " "to small prompt perturbations. Reuses --prompts. " "Cost: ~5–7 min on a 7B Q8 " "(prompts × (1 + n_perturbations) × 2 generations).") p.add_argument("--json-out", type=Path, default=None, help="Path to write the JSON report to.") p.add_argument("--html-out", type=Path, default=None, help="Path to write a self-contained HTML report (single " "file with inline CSS, no JS deps). Includes hardware " "info, model params, and the exact repro command.") p.add_argument("--no-progress", action="store_true", help="Suppress per-prompt progress output.") return p def _run_score(args) -> int: from . import __version__ from .backends import auto_backend, get_backend from .runner import set_active_backend # If --prompts wasn't passed, fall back to the bundled prompts file # shipped inside the wheel via package_data. if args.prompts is None: try: import importlib.resources args.prompts = Path( str(importlib.resources.files("refract").joinpath( "prompts/v0.1.jsonl")) ) except Exception as e: print( f"ERROR: --prompts not given and bundled prompts not found " f"({e}). Pass --prompts /path/to/prompts.jsonl explicitly." ) return 2 ref_kv = KVConfig.parse(args.reference) cand_kv = KVConfig.parse(args.candidate) # v0.3.1: select backend per --backend flag (or auto-detect from model path). if args.backend == "auto": backend = auto_backend(args.model) else: backend = get_backend(args.backend) set_active_backend(backend) # --full is sugar for the two opt-in v0.2 axes. if args.full: args.axis_rniah = True args.axis_plad = True # v0.3.2: auto-resolve corpus + haystack from ~/.cache/refract/ when # the user didn't pass them. Triggers a one-time wikitext-2-raw fetch # if the cache is empty and --no-auto-fetch isn't set. _resolve_default_paths( args, need_corpus=not args.skip_kld, need_haystack=args.axis_rniah, ) # Cost hint up front so the user knows what they're committing to. cost_axes = ["A (~2 min)", "B (~5 min)"] if args.axis_rniah: cost_axes.append("C R-NIAH (~10-15 min)") if args.axis_plad: cost_axes.append("D PLAD (~5-7 min)") cost_hint = " + ".join(cost_axes) print(f"REFRACT v{__version__}") print(f" model : {args.model}") print(f" reference : {ref_kv.label()}") print(f" candidate : {cand_kv.label()}") print(f" backend : {backend.name}") print(f" axes : {cost_hint} (estimates for a 7B Q8 model on Apple Silicon)") # v0.3.1: thinking-mode auto-detect at startup. Reasoning is auto-disabled # via -rea off in run_completion when backend supports it; we surface the # detection so the user knows what we did and why. thinking_detected, thinking_markers = False, [] try: thinking_detected, thinking_markers = backend.detect_thinking_mode( model=args.model, timeout=30, ) except Exception as e: print(f" thinking : probe failed ({e}); assuming non-thinking") if thinking_detected: print(f" thinking : DETECTED (markers: {thinking_markers}) — " f"reasoning disabled, n_predict-aware fallback active") else: print(f" thinking : not detected") print() # ---- Floor check ------------------------------------------------------ floor_score = None if args.measure_floor: print("Measuring noise floor: REFRACT(ref, ref) ...") floor_gtm = run_gtm( model=args.model, reference_kv=ref_kv, candidate_kv=ref_kv, prompts_path=args.prompts, n_predict=args.n_predict, ctx=args.ctx, n_gpu_layers=args.n_gpu_layers, seed=args.seed, progress=not args.no_progress, ) floor_kld = run_kld( model=args.model, corpus=args.corpus, reference_kv=ref_kv, candidate_kv=ref_kv, chunks=args.chunks, ctx=args.ctx, n_gpu_layers=args.n_gpu_layers, progress=not args.no_progress, ) floor_composite = composite_score(floor_gtm.score, floor_kld.score) floor_score = floor_composite.composite print(f" floor: {floor_score:.2f} (gtm={floor_gtm.score:.2f}, " f"kld={floor_kld.score:.2f}, kld nats={floor_kld.mean_kld:.6f})") if floor_score < MIN_FLOOR: print() print(f"ERROR: noise floor {floor_score:.2f} < {MIN_FLOOR}.") print("The reference itself is non-deterministic on this build.") print("KLD deltas vs this reference cannot be trusted. Aborting.") return 2 # v0.1.3: composite-level floor passes if KLD is ~100 (bit-exact zero # on Metal) even when GTM is in a "high-score broken" state. Add a # GTM-level byte-identity check: ref-vs-ref greedy must match for # the FULL retokenized candidate length on every prompt. If # mean_prefix != mean_cand, the tokenizer or runner is broken. if floor_gtm.mean_cand_length > 0: ratio = ( floor_gtm.mean_prefix_agreement_length / floor_gtm.mean_cand_length ) else: ratio = 0.0 if abs(ratio - 1.0) > 1e-9: print() print( f"ERROR: GTM ref-vs-ref ratio = {ratio:.6f} (expected 1.0). " f"mean_prefix={floor_gtm.mean_prefix_agreement_length:.2f}, " f"mean_cand_length={floor_gtm.mean_cand_length:.2f}." ) print( "GTM ref-vs-ref isn't byte-identical, your tokenizer or " "runner is broken. Aborting." ) return 2 print() # ---- Axis A (GTM or Trajectory) --------------------------------------- if args.skip_gtm: gtm = _stub_gtm() else: if args.axis_a == "trajectory": print("Running Axis A (Trajectory, v0.1.4)...") gtm = run_trajectory( model=args.model, reference_kv=ref_kv, candidate_kv=cand_kv, prompts_path=args.prompts, n_predict=args.n_predict, ctx=args.ctx, n_gpu_layers=args.n_gpu_layers, seed=args.seed, progress=not args.no_progress, ) print(f" Trajectory score: {gtm.score:.2f}") else: print("Running Axis A (GTM, v0.1.x)...") gtm = run_gtm( model=args.model, reference_kv=ref_kv, candidate_kv=cand_kv, prompts_path=args.prompts, n_predict=args.n_predict, ctx=args.ctx, n_gpu_layers=args.n_gpu_layers, seed=args.seed, progress=not args.no_progress, ) print(f" GTM score: {gtm.score:.2f}") # ---- KLD -------------------------------------------------------------- if args.skip_kld: kld = _stub_kld(args.chunks, args.ctx) else: print("Running Axis B (KLD@D, corpus proxy)...") kld = run_kld( model=args.model, corpus=args.corpus, reference_kv=ref_kv, candidate_kv=cand_kv, chunks=args.chunks, ctx=args.ctx, n_gpu_layers=args.n_gpu_layers, progress=not args.no_progress, ) print(f" KLD score: {kld.score:.2f} (mean KLD = {kld.mean_kld:.6f} nats)") # ---- R-NIAH (v0.2 opt-in) --------------------------------------------- rniah = None if args.axis_rniah: if args.rniah_haystack is None: print("ERROR: --axis-rniah requires --rniah-haystack " "(or run `refract fetch` first to populate the cache).") return 2 # v0.3.2: synthesize lengths from --rniah-up-to (doubling step-up # from 4K) when --rniah-lengths isn't set. Default --rniah-ctx-max # to --rniah-up-to so the user only has one knob to think about. if args.rniah_lengths: lengths = tuple(int(x) for x in args.rniah_lengths.split(",")) else: up_to = args.rniah_up_to n = 4096 synth: list[int] = [] while n <= up_to: synth.append(n) n *= 2 if not synth: synth = [up_to] lengths = tuple(synth) if args.rniah_ctx_max is None: args.rniah_ctx_max = max(lengths) positions = ( tuple(float(x) for x in args.rniah_positions.split(",")) if args.rniah_positions else None ) kwargs: dict = { "model": args.model, "haystack_corpus": args.rniah_haystack, "reference_kv": ref_kv, "candidate_kv": cand_kv, "ctx_max": args.rniah_ctx_max, "n_trials": args.rniah_trials, "n_gpu_layers": args.n_gpu_layers, "seed": args.seed, "progress": not args.no_progress, } if lengths is not None: kwargs["lengths"] = lengths if positions is not None: kwargs["positions"] = positions print("Running Axis C (R-NIAH, v0.2)...") rniah = run_rniah(**kwargs) print(f" R-NIAH score: {rniah.score:.2f} ({rniah.n_cells} cells)") # ---- PLAD (v0.2 opt-in) ----------------------------------------------- plad = None if args.axis_plad: print("Running Axis D (PLAD, v0.2)...") plad = run_plad( model=args.model, prompts_path=args.prompts, reference_kv=ref_kv, candidate_kv=cand_kv, n_predict=args.n_predict, ctx=args.ctx, n_gpu_layers=args.n_gpu_layers, seed=args.seed, progress=not args.no_progress, ) print(f" PLAD score: {plad.score:.2f}") # ---- Composite -------------------------------------------------------- # Skipped axes contribute None to the composite (dropped before the # harmonic mean) rather than a stub 100, which would inflate the # composite and read as EXCELLENT in the report. gtm_for_composite = None if args.skip_gtm else gtm.score kld_for_composite = None if args.skip_kld else kld.score composite = composite_score( gtm_for_composite, kld_for_composite, rniah_score=rniah.score if rniah else None, plad_score=plad.score if plad else None, floor_score=floor_score, ) print() print(text_report( model=str(args.model), reference_label=ref_kv.label(), candidate_label=cand_kv.label(), composite=composite, gtm=gtm, kld=kld, rniah=rniah, plad=plad, )) rep = None if args.json_out or args.html_out: rep = json_report( model=str(args.model), reference_label=ref_kv.label(), candidate_label=cand_kv.label(), composite=composite, gtm=gtm, kld=kld, rniah=rniah, plad=plad, ) if args.json_out: args.json_out.write_text(to_json_string(rep)) print(f"\nJSON report written to {args.json_out}") if args.html_out: from .report_html import html_report html = html_report( model=str(args.model), reference_label=ref_kv.label(), candidate_label=cand_kv.label(), composite=composite, gtm=gtm, kld=kld, rniah=rniah, plad=plad, raw_json=rep, ) args.html_out.write_text(html) print(f"HTML report written to {args.html_out}") return 0 # Stubs for --skip-gtm / --skip-kld dev modes (composite still computes). def _stub_gtm(): from .axes.gtm import GTMResult return GTMResult( score=100.0, full_match_rate=1.0, median_first_divergence=None, mean_prefix_agreement_length=0.0, mean_cand_length=0.0, mean_ref_length=0.0, n_prompts=0, n_tokens_each=0, per_prompt=[], ) def _stub_kld(chunks: int, ctx: int): from .axes.kld import KLDResult return KLDResult( score=100.0, mean_kld=0.0, ppl=None, rms_dp_pct=None, same_topp_pct=None, base_path="", chunks=chunks, ctx=ctx, is_self_reference=False, ) _TOP_DESCRIPTION = """\ REFRACT — REFerence-anchored Robust Acid-test for Compressed Transformers. A benchmaxx-resistant alternative to corpus PPL for evaluating KV-cache quantization quality on a model. Replaces "lower PPL = better" — a metric that can invert sign on instruct-tuned models — with a 4-axis composite anchored to the fp16-KV reference. Subcommands: score run REFRACT on a candidate KV config vs the reference Run `refract score --help` for full options + usage examples. """ def _add_selftest_parser(sub): p = sub.add_parser( "selftest", help="Preflight check before running score/repeatability.", description=( "Verifies the environment is ready for REFRACT: binaries / " "imports, required CLI flags (--jinja, REFRACT_TRAJECTORY env), " "KV cache types compiled in, model loadable, fp16-vs-fp16 " "sanity generation. Bails out with a useful message on any " "failure so you don't burn a long run finding out your setup " "is broken.\n\n" "Tip: pass --model for a real generation probe (~10-30s " "depending on model size). Without --model, only static " "binary/import checks run (~1s)." ), ) p.add_argument("--backend", choices=["auto", "llamacpp", "mlx", "vllm"], default="auto", help="Backend to test. 'auto' picks llamacpp for .gguf " "models and mlx for directories. Override via " "REFRACT_BACKEND env.") p.add_argument("--model", type=Path, default=None, help="Optional model to use for a real generation probe. " "Accepts both .gguf files (llamacpp) and MLX model " "directories. If omitted, only static checks run.") return p def _run_selftest(args) -> int: from . import __version__ from .backends import get_backend, auto_backend, BackendCapabilityError print(f"REFRACT v{__version__} selftest") print() failures: list[str] = [] warnings: list[str] = [] # --- Backend selection --- if args.backend == "auto": if args.model: backend = auto_backend(args.model) else: from .backends.llamacpp import LlamaCppBackend backend = LlamaCppBackend() else: backend = get_backend(args.backend) print(f"Backend: {backend.name}") # --- Backend-specific binaries --- if backend.name == "llamacpp": from .runner import DEFAULT_BIN_DIR print(f" bin dir : {DEFAULT_BIN_DIR}") for tool in ("llama-cli", "llama-completion", "llama-tokenize", "llama-perplexity"): p = DEFAULT_BIN_DIR / tool if p.exists(): print(f" ✓ {tool}") else: print(f" ✗ {tool} (set LLAMA_CPP_BIN_DIR or rebuild)") failures.append(f"missing binary: {tool}") # Check for --jinja flag support. Capture stderr too — when the # binary can't load (e.g. Linux missing libllama.so on PATH, Windows # missing the DLL next to the exe), --help fails *before* it can # print its help text. Without checking the launch result, an empty # stdout would cause a misleading "--jinja missing" diagnosis. try: import subprocess as _sp proc = _sp.run( [str(DEFAULT_BIN_DIR / "llama-completion"), "--help"], capture_output=True, text=True, timeout=10, ) help_text = proc.stdout launch_failed = ( proc.returncode != 0 and not help_text.strip() ) stderr_tail = (proc.stderr or "")[-400:] if launch_failed: print(f" ✗ llama-completion can't launch (rc={proc.returncode})") if stderr_tail: print(f" stderr: {stderr_tail.strip()}") # Linux + Windows shared-library hints hint = "" if "cannot open shared object" in stderr_tail or \ "libllama" in stderr_tail.lower(): hint = (" — Linux: libllama.so/libggml*.so not on the " "loader path. Add the bin dir to LD_LIBRARY_PATH " "or register it via ldconfig.") elif "dll" in stderr_tail.lower() or \ "0xc0000135" in stderr_tail.lower(): hint = (" — Windows: required DLLs (llama.dll, ggml-*.dll) " "not next to the .exe. Copy them from the build " "dir or add the bin dir to PATH.") failures.append( f"llama-completion failed to launch (rc={proc.returncode}){hint}" ) elif "--jinja" in help_text: print(" ✓ --jinja chat template flag supported") else: print(" ✗ --jinja not in llama-completion help") failures.append("--jinja missing — chat templates won't apply") if "REFRACT_TRAJECTORY" in help_text or "trajectory" in help_text.lower(): print(" ✓ REFRACT_TRAJECTORY likely supported") elif not launch_failed: # The patch is env-var triggered, no help-string evidence. # Run a probe instead. warnings.append( "Could not verify REFRACT_TRAJECTORY support from --help. " "Run a quick trajectory probe to confirm." ) except Exception as e: warnings.append(f"Could not probe llama-completion --help: {e}") elif backend.name == "mlx": try: from .backends.mlx import _require_mlx _require_mlx() print(" ✓ mlx + mlx_lm importable") except BackendCapabilityError as e: print(f" ✗ mlx not importable: {e}") failures.append(str(e)) elif backend.name == "vllm": try: import vllm # noqa print(f" ✓ vllm importable (version: {getattr(vllm, '__version__', 'unknown')})") except ImportError: print(" ✗ vllm not importable; pip install vllm") failures.append("vllm not importable") # --- Model probe (optional) --- if args.model: if not args.model.exists(): print(f" ✗ model not found: {args.model}") failures.append(f"model missing: {args.model}") else: print(f"\nProbing model: {args.model.name}") gen_ok = False try: result = backend.run_completion( model=args.model, prompt="What is 2+2?", kv_config_str="ctk=f16,ctv=f16", n_predict=16, ctx=128, seed=42, temperature=0.0, timeout=120, ) preview = (result.text or "").replace("\n", " ")[:80] print(f" ✓ generation works → {preview!r}") gen_ok = True except Exception as e: print(f" ✗ generation failed: {e}") failures.append(f"model generation: {e}") # Thinking probe — only run if generation succeeded. If gen # failed, the probe uses the same broken path and "no markers # detected" would be misleading. if gen_ok: try: detected, markers = backend.detect_thinking_mode(model=args.model) if detected: print(f" ℹ thinking-mode detected (markers: {markers}). " "REFRACT will handle this automatically.") else: print(" ✓ no thinking-mode markers (faster runs)") except Exception as e: warnings.append(f"thinking-mode probe failed: {e}") # --- Summary --- print() if failures: print(f"FAILED {len(failures)} issue(s):") for f in failures: print(f" - {f}") return 2 if warnings: print(f"OK with {len(warnings)} warning(s):") for w in warnings: print(f" - {w}") else: print("OK — selftest passed.") print() print("You're ready to run `refract score --model ... --candidate ... ...`.") return 0 def _add_compare_parser(sub): p = sub.add_parser( "compare", help="Side-by-side comparison of multiple report JSONs.", ) p.add_argument("reports", type=Path, nargs="+", help="Two or more report.json files to compare.") return p def _run_compare(args) -> int: import json as _json rows = [] for path in args.reports: try: d = _json.loads(path.read_text()) except Exception as e: print(f"skip {path}: {e}") continue rows.append({ "label": path.stem, "composite": d.get("composite"), "band": d.get("band"), "summary": d.get("summary"), "axes": d.get("axes", {}), "version": d.get("framework_version"), "backend": d.get("environment", {}).get("backend"), }) if not rows: print("no reports parsed") return 1 # Print a markdown-style comparison table print() print(f"{'Report':<32} {'Comp':>7} {'Band':<10} {'Traj':>7} {'KLD':>7} {'R-NIAH':>7} {'PLAD':>7}") print("-" * 80) for r in rows: a = r["axes"] def fmt(d, k): try: ax = a[k] if ax.get("skipped") or ax.get("score") is None: return "skip" return f"{ax['score']:.2f}" except Exception: return "—" comp_val = r["composite"] comp_str = f"{comp_val:>7.2f}" if isinstance(comp_val, (int, float)) else f"{'—':>7}" band_str = r["band"] if isinstance(r["band"], str) else "—" print(f"{r['label'][:32]:<32} {comp_str} {band_str:<10} " f"{fmt(a, 'gtm'):>7} {fmt(a, 'kld'):>7} " f"{fmt(a, 'rniah'):>7} {fmt(a, 'plad'):>7}") print() return 0 def _add_fetch_parser(sub): p = sub.add_parser( "fetch", help="Download wikitext-2-raw corpus + haystack to ~/.cache/refract/.", description=( "Pre-fetch the wikitext-2-raw corpus (10MB zip → wiki.test.raw " "+ wiki.train.raw) into ~/.cache/refract/. Subsequent score / " "repeatability invocations will auto-find these files when " "--corpus / --rniah-haystack are omitted.\n\n" "Idempotent: re-running with the cache already populated is a " "no-op. Skips any download if files already exist." ), ) p.add_argument("--cache-dir", type=Path, default=_REFRACT_CACHE, help=f"Override cache location (default: {_REFRACT_CACHE}).") return p def _run_fetch(args) -> int: target = _ensure_wikitext_2(cache_dir=args.cache_dir) print() print(f" test : {target / 'wiki.test.raw'}") print(f" train : {target / 'wiki.train.raw'}") print() print("`refract score` will auto-find these when --corpus / --rniah-haystack are omitted.") return 0 def _add_repeatability_parser(sub): p = sub.add_parser( "repeatability", help="Run the same scoring config N times and report the spread.", description=( "Verifies REFRACT scores are reproducible on the same model + " "candidate. Runs ``score`` N times back-to-back, then prints " "min/median/max + stdev for the composite and each axis. " "Healthy framework: composite stdev ≤ 1.0, per-axis stdev ≤ 2.0 " "(R-NIAH may be noisier on n_trials=1)." ), ) p.add_argument("--model", required=True, type=Path) p.add_argument("--candidate", required=True) p.add_argument("--reference", default="ctk=f16,ctv=f16") p.add_argument("--prompts", required=True, type=Path) p.add_argument("--corpus", required=True, type=Path) p.add_argument("--runs", type=int, default=4, help="Number of repeat runs (default: 4).") p.add_argument("--n-predict", type=int, default=50) p.add_argument("--ctx", "-c", type=int, default=512) p.add_argument("--chunks", type=int, default=32) p.add_argument("--n-gpu-layers", "-ngl", type=int, default=99) p.add_argument("--seed", type=int, default=42, help="Seed used for ALL runs. Identical seed gives upper " "bound on reproducibility — variance comes purely " "from non-determinism in the engine, not RNG.") p.add_argument("--axis-a", choices=["gtm", "trajectory"], default="trajectory") p.add_argument("--full", action="store_true", help="Include R-NIAH + PLAD in each repeat run.") p.add_argument("--rniah-haystack", type=Path, default=None) p.add_argument("--rniah-ctx-max", type=int, default=None) p.add_argument("--backend", choices=["auto", "llamacpp", "mlx", "vllm"], default="auto") p.add_argument("--out-dir", type=Path, default=None, help="Directory for per-run JSON outputs. Default: tmp.") return p def _run_repeatability(args) -> int: """Run ``score`` N times and report the spread per-axis.""" import json as _json import statistics as _stats import tempfile as _tf from . import __version__ from .backends import auto_backend, get_backend from .runner import set_active_backend # Set backend up front so all runs share it backend = (auto_backend(args.model) if args.backend == "auto" else get_backend(args.backend)) set_active_backend(backend) out_dir = args.out_dir or Path(_tf.mkdtemp(prefix="refract-repeatability-")) out_dir.mkdir(parents=True, exist_ok=True) print(f"REFRACT v{__version__} repeatability — {args.runs} runs") print(f" model : {args.model}") print(f" candidate : {args.candidate}") print(f" output : {out_dir}") print() composite_scores: list[float] = [] axis_scores: dict[str, list[float]] = { "trajectory": [], "kld": [], "rniah": [], "plad": [], } # Build a one-shot args namespace for each run import argparse as _ap base_args = _ap.Namespace( model=args.model, reference=args.reference, candidate=args.candidate, prompts=args.prompts, corpus=args.corpus, chunks=args.chunks, ctx=args.ctx, n_gpu_layers=args.n_gpu_layers, n_predict=args.n_predict, seed=args.seed, measure_floor=False, skip_gtm=False, skip_kld=False, axis_a=args.axis_a, full=args.full, axis_rniah=args.full, rniah_haystack=args.rniah_haystack, rniah_ctx_max=args.rniah_ctx_max, rniah_lengths=None, rniah_positions=None, rniah_trials=1, axis_plad=args.full, json_out=None, html_out=None, no_progress=True, backend=args.backend, ) for i in range(1, args.runs + 1): json_path = out_dir / f"run-{i:02d}.json" run_args = _ap.Namespace(**vars(base_args)) run_args.json_out = json_path print(f"=== run {i}/{args.runs} → {json_path.name} ===") rc = _run_score(run_args) if rc != 0: print(f" run {i} failed (rc={rc}); aborting") return rc try: d = _json.loads(json_path.read_text()) composite_scores.append(float(d.get("composite", 0))) for axis, key in (("trajectory", "gtm"), ("kld", "kld"), ("rniah", "rniah"), ("plad", "plad")): ax = d.get("axes", {}).get(key, {}) if "score" in ax and isinstance(ax["score"], (int, float)): axis_scores[axis].append(float(ax["score"])) except Exception as e: print(f" warning: could not parse {json_path}: {e}") # Spread report print() print("=" * 72) print(f"REPEATABILITY ({args.runs} runs)") print("=" * 72) def _spread(values: list[float]) -> str: if not values: return "no data" if len(values) == 1: return f"single run: {values[0]:.2f}" return ( f"min={min(values):.2f} med={_stats.median(values):.2f} " f"max={max(values):.2f} stdev={_stats.stdev(values):.2f} " f"range={max(values) - min(values):.2f}" ) print(f"composite : {_spread(composite_scores)}") for axis, vals in axis_scores.items(): print(f" {axis:<10}: {_spread(vals)}") # Health check print() if composite_scores: if len(composite_scores) > 1: cs_stdev = _stats.stdev(composite_scores) if cs_stdev <= 1.0: print(f"✓ HEALTHY composite stdev {cs_stdev:.2f} ≤ 1.0") elif cs_stdev <= 3.0: print(f"⚠ NOISY composite stdev {cs_stdev:.2f} (1.0-3.0): " "results stable but with some run-to-run jitter") else: print(f"✗ UNSTABLE composite stdev {cs_stdev:.2f} > 3.0: " "framework or engine is non-deterministic. Investigate.") return 0 def main(argv=None) -> int: import argparse as _ap parser = argparse.ArgumentParser( prog="refract", description=_TOP_DESCRIPTION, formatter_class=_ap.RawDescriptionHelpFormatter, ) sub = parser.add_subparsers(dest="cmd", required=True) _add_score_parser(sub) _add_selftest_parser(sub) _add_compare_parser(sub) _add_repeatability_parser(sub) _add_fetch_parser(sub) args = parser.parse_args(argv) if args.cmd == "score": return _run_score(args) if args.cmd == "selftest": return _run_selftest(args) if args.cmd == "compare": return _run_compare(args) if args.cmd == "repeatability": return _run_repeatability(args) if args.cmd == "fetch": return _run_fetch(args) parser.error(f"unknown subcommand: {args.cmd}") return 1 if __name__ == "__main__": sys.exit(main())