diff --git a/CHANGELOG.md b/CHANGELOG.md index ce977ab..0b95641 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,39 @@ reproducing 70+ versions of notes. ## [Unreleased] +## [0.71.10] - 2026-06-03 + +### Added +- **RAG family — live wiring** (closes #199, #200, #201, #202). The four + retrieval / steering surfaces that shipped schema-only in v0.62.0 are now + real, validated on SmolLM2-135M. +- **RAFT span-mask training is live** (#199). `data.format: raft` rows + (`{query, golden_doc, distractor_docs, answer}`) now train answer-only: the + prompt span is masked to `-100` and each document is labelled `[doc-N]` so + the model learns to cite the supporting document. Documents are shuffled + reproducibly (`data.raft_shuffle_seed`). Rows whose prompt fills + `max_length` (answer fully truncated) are dropped with a warning rather than + silently shrinking the effective dataset. +- **`soup ra-dit` — one-shot two-stage orchestrator** (#200). Trains the + retriever (stage 1, embedding/contrastive) then the generator (stage 2, + RAFT-SFT) in a single command, recording the trained retriever as the + generator's paired retriever. A `soup train` of a generator-stage config + with no retriever model set now auto-links the most-recent RA-DIT retriever + run from the Registry. `--plan-only` validates both configs without + training; `--retriever-model` overrides the auto-link. +- **`soup steer train` / `apply` + `soup serve --steer` are live** (#201). + Fit a CAA (contrastive activation addition), ITI (inference-time + intervention) or RepE (representation-engineering PCA) control vector from + `{positive, negative}` contrastive pairs, persist it as a safetensors + + config artifact, and apply it at decode time via a forward hook + (`soup serve --steer --steer-strength `). +- **`soup eval citation` + citation-span loss boost are live** (#202). Score + citation precision / recall / F1 over `{predicted, expected_ids}` or RAFT + rows (`--shuffle-seed` aligns the golden `[doc-N]` id with what the model + saw at train time). When `citation_faithful: true`, bracketed `[doc-id]` + spans in the answer get a boosted per-token loss weight. A new `citation` + failure mode is available in `soup diagnose`. + ## [0.71.9] - 2026-06-03 ### Added diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 044f681..b16033f 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -120,7 +120,7 @@ src/soup_cli/ templates/ - 17 built-in soup.yaml templates (YAML + manifest.json) with load_template loader (v0.39.0, +bco v0.40.0) ui/ - Web UI (FastAPI + HTML/JS SPA) -tests/ - Test suite (279 files, 12999 tests) +tests/ - Test suite (280 files, 13142 tests) examples/ - Real-world config examples and datasets ``` diff --git a/README.md b/README.md index 179aa3e..78ca733 100644 --- a/README.md +++ b/README.md @@ -49,20 +49,21 @@ infrastructure instead of improving models. Soup fixes that. ## What's New -**v0.71.9 — Knowledge edit + unlearn (live).** The surgical-edit and unlearning surfaces are now -real, validated on SmolLM2-135M: +**v0.71.10 — RAG family (live).** The retrieval-augmented fine-tuning, steering, and citation +surfaces are now real, validated on SmolLM2-135M: -- **`soup edit set --method rome|memit|alphaedit`** — surgical rank-1 weight edit that patches one - fact without re-training. On a tiny model a ROME edit moved P("Lyon" | "The capital of France is") - from 0.0016 → 0.96. `--output` saves the edited model; `--governor` refuses edits after norm blowup. -- **`soup edit diff --before-model --after-model --probes p.jsonl`** — generate completions - through both models and surface the probes whose output changed. -- **`soup edit set --method grace`** — GRACE codebook edit: stores the fact in a discrete - (key, value) sidecar applied at decode time via a forward hook (survives thousands of edits). -- **Sequential-edit governor persistence** — per-base-model edit count + verdict survive across - separate `soup edit set` runs (SQLite, cross-process-locked). -- **`soup train --task unlearn`** — NPO / SimNPO / RMU unlearning from a `forget_set` (+ optional - `retain_set`); NPO/SimNPO drive the forget-set loss down while the retain set anchors capability. +- **`data.format: raft`** — RAFT (retrieval-augmented fine-tuning): train on a query + golden + document mixed with distractors, answer-only loss, each doc labelled `[doc-N]` so the model + learns to cite the supporting source and ignore noise. +- **`soup ra-dit`** — one-shot two-stage orchestrator: trains the retriever then the generator + and records the trained retriever as the generator's paired retriever. A `soup train` of a + generator stage auto-links the latest RA-DIT retriever from the Registry. +- **`soup steer train --method caa|iti|repe` + `soup serve --steer `** — fit a + contrastive-activation / inference-time-intervention / representation-engineering control + vector from `{positive, negative}` pairs and apply it at decode time via a forward hook. +- **`soup eval citation`** — score citation precision / recall / F1 over predictions or RAFT + rows; with `citation_faithful: true`, `[doc-id]` spans get a boosted per-token loss weight. + A new `citation` failure mode joins `soup diagnose`. Full history: [CHANGELOG.md](CHANGELOG.md) · [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases). diff --git a/docs/adapters-and-governance.md b/docs/adapters-and-governance.md index 7511761..e463f06 100644 --- a/docs/adapters-and-governance.md +++ b/docs/adapters-and-governance.md @@ -184,6 +184,8 @@ soup steer list Steering names are validated against a strict regex (`^[A-Za-z0-9][A-Za-z0-9._\-]{0,127}$` — no path separators, no shell metacharacters); strength is bounded `|s| <= 10.0`. The trained vectors land in the Soup Registry under the `steering_vector` artifact kind so lineage is preserved. +As of v0.71.10 the fit and the decode hook are **live** (validated on SmolLM2-135M): `soup steer train` captures residual-stream activations (CAA / RepE) or per-head `o_proj`-input activations (ITI) on the contrastive pairs, computes the control vector, and persists `steering_vector.safetensors` + `steering_config.json`. `soup serve --steer ` installs a forward hook on the loaded model that adds `strength × vector` at decode time (transformers backend; `--steer` is rejected with a clear error on vLLM/SGLang). RepE / ITI need at least two contrastive pairs; CAA works from one. + ## GRACE Codebook — Lifelong Knowledge Edits diff --git a/docs/commands.md b/docs/commands.md index 8ddecf6..bc692b3 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -171,6 +171,12 @@ soup eval unlearning --benchmark tofu|muse|wmdp Forget Quality + Model soup edit set --base --method rome|memit|alphaedit|grace --subject "..." --target "..." [--output ] [--device cpu] [--governor/--no-governor] [--registry-id ] Live surgical knowledge edit (--plan-only available) soup edit diff --probes p.jsonl [--before-model --after-model ] Knowledge-injection diff (live before/after generation when both models given) soup train --task unlearn NPO/SimNPO/RMU unlearning from data.forget_set (+ optional data.retain_set) +soup train # data.format='raft' Answer-only span-mask RAFT training (golden+distractor docs, [doc-N] citations); generator-stage configs auto-link the latest RA-DIT retriever +soup ra-dit --retriever-config --generator-config [--retriever-model ] [--plan-only] One-shot two-stage RA-DIT: train retriever → record pairing → train generator +soup eval citation [--style bracket|inline|footnote] [--shuffle-seed N] [--output o.json] Citation precision/recall/F1 over predictions or RAFT rows +soup steer train --base --method caa|iti|repe --name --pairs Fit a CAA/ITI/RepE activation-steering vector from {positive, negative} pairs +soup steer apply --name --strength Preview a stored steering vector; soup steer list lists them +soup serve --steer [--steer-strength ] Apply a steering vector at decode time via a forward hook (transformers backend) soup ingest --source langfuse|langsmith|helicone|openpipe|otel|openai-stored --logs Universal trace importer (6 SaaS adapters → normalised JSONL) soup prune-prompt --input --output --min-frequency 0.95 Detect + strip shared system-prompt prefix soup prune-prompt ... --tokenizer Tokenizer-aware prefix detection (decodes remaining ids, boundary-safe) diff --git a/docs/evaluation.md b/docs/evaluation.md index 62d3142..8c68a8a 100644 --- a/docs/evaluation.md +++ b/docs/evaluation.md @@ -253,7 +253,7 @@ Default threshold 0.2 matches v0.43.0 KL-delta quant-check thresholds. Webhooks ## Diagnose (Post-Training Report Card) -`soup diagnose` scores six independent failure modes for a trained adapter and renders an OK / MINOR / MAJOR verdict per mode plus an overall headline — same taxonomy as Quant-Lobotomy. Useful for catching adapter regressions that a loss curve cannot distinguish from a healthy run. +`soup diagnose` scores seven independent failure modes for a trained adapter and renders an OK / MINOR / MAJOR verdict per mode plus an overall headline — same taxonomy as Quant-Lobotomy. Useful for catching adapter regressions that a loss curve cannot distinguish from a healthy run. ```bash # Neutral report (no model load — runs as a sanity check) @@ -279,7 +279,7 @@ soup diagnose my-run-id --output diag.json --attach-to-registry abc123 a built-in probe set; `format` only fires when the dataset's own targets look like JSON; `contamination` stays neutral unless a benchmark corpus is supplied. Validated on SmolLM2-135M. -**Six failure-mode probes:** +**Seven failure-mode probes:** | Mode | What it catches | Score range | |------|-----------------|-------------| @@ -289,6 +289,7 @@ a built-in probe set; `format` only fires when the dataset's own targets look li | `mode_collapse` | Diversity collapse at T=0 and T=1 | pairwise n-gram Jaccard distance | | `memorization` | Verbatim training-prefix echo on partial prompts | 1 − echo_rate | | `contamination` | Training data overlapping public benchmarks | 1 − contamination_rate | +| `citation` | RAFT model stopped citing the supporting `[doc-N]` (v0.71.10) | fraction of answers citing the golden doc | **Verdict pill colours:** OK (≥ 0.85) green / MINOR (≥ 0.60) amber / MAJOR (< 0.60) red. `soup diagnose` exits 2 when the overall verdict is MAJOR — wire into CI to fail the build on regression. diff --git a/docs/training.md b/docs/training.md index 292f935..747d7a4 100644 --- a/docs/training.md +++ b/docs/training.md @@ -886,6 +886,19 @@ score = score_citations( Citation-faithful FT is gated to `task in {sft, pretrain}` + `data.format='raft'` — misconfigured runs fail at config load with a named-field message. +Under the hood, a `format: raft` run trains **answer-only**: each row is composed into a prompt (golden + distractor docs, shuffled deterministically by `data.raft_shuffle_seed`, each labelled `[doc-N]`) followed by the answer; the prompt span is masked out of the loss and — when `citation_faithful: true` — the bracketed `[doc-id]` spans in the answer get a boosted per-token loss weight. Rows whose prompt fills `max_length` (answer fully truncated) are dropped with a warning rather than silently shrinking the dataset. + +Score a trained model's citations from the CLI: + +```bash +# {predicted, expected_ids} rows, OR RAFT rows scored against their own golden [doc-N] +soup eval citation preds.jsonl --style bracket +# RAFT rows: pass the train-time shuffle seed so the golden id lines up +soup eval citation raft.jsonl --shuffle-seed 0 --output citation.json +``` + +`soup diagnose` also gains a `citation` failure mode that flags a model that stopped citing the supporting document. + ## RA-DIT — Retrieval-Augmented Dual Instruction Tuning @@ -903,6 +916,16 @@ soup train The schema enforces stage-task pairing — `ra_dit_stage: retriever` requires `task: embedding`; `ra_dit_stage: generator` requires `task: sft`. A misconfigured recipe fails at config load with a named-field message. +Run both stages in one command with `soup ra-dit`: + +```bash +soup ra-dit --retriever-config retriever.yaml --generator-config generator.yaml +# preview the plan + the resolved retriever link without training: +soup ra-dit -r retriever.yaml -g generator.yaml --plan-only +``` + +It trains the retriever, then **records** that trained retriever as the generator's paired retriever (writing its output dir into the generator's `training.ra_dit_retriever_model`) and trains the generator RAFT-style. The recorded retriever is the one used at deploy/serve time — stage-2 does not fuse the retriever weights. A plain `soup train` of a generator-stage config with no retriever model set **auto-links** the most-recent RA-DIT retriever run from the Registry; pass `--retriever-model ` to override. + ## Curriculum-Aware Training (BETA) diff --git a/pyproject.toml b/pyproject.toml index 4d49fda..c81df18 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "soup-cli" -version = "0.71.9" +version = "0.71.10" description = "Fine-tune LLMs in one command. No SSH, no config hell." readme = "README.md" license = "Apache-2.0" diff --git a/src/soup_cli/__init__.py b/src/soup_cli/__init__.py index 4e96d19..631ab7f 100644 --- a/src/soup_cli/__init__.py +++ b/src/soup_cli/__init__.py @@ -1,3 +1,3 @@ """Soup CLI — Fine-tune LLMs in one command.""" -__version__ = "0.71.9" +__version__ = "0.71.10" diff --git a/src/soup_cli/cli.py b/src/soup_cli/cli.py index 3b9a93d..6f61207 100644 --- a/src/soup_cli/cli.py +++ b/src/soup_cli/cli.py @@ -484,6 +484,15 @@ app.add_typer( help="Iterative DPO loop driver (v0.70.0 Part E).", ) +# v0.71.10 #200 — `soup ra-dit` (two-stage RA-DIT orchestrator). +from soup_cli.commands import ra_dit as _ra_dit_cmd # noqa: E402 + +app.add_typer( + _ra_dit_cmd.app, + name="ra-dit", + help="RA-DIT two-stage orchestrator: retriever -> generator (v0.71.10).", +) + def _rewrite_advise_argv(argv: list) -> list: """Inject `run` between `advise` and a non-subcommand first argument. diff --git a/src/soup_cli/commands/_eval_v07110.py b/src/soup_cli/commands/_eval_v07110.py new file mode 100644 index 0000000..1a8a267 --- /dev/null +++ b/src/soup_cli/commands/_eval_v07110.py @@ -0,0 +1,259 @@ +"""v0.71.10 #202 — `soup eval citation`. + +Scores citation precision / recall / F1 per row + an aggregate, routed through +``citation_faithful.score_citations``. Attached to the existing ``soup eval`` +Typer app via :func:`register` (mirrors the v0.55.0 / v0.61.0 / v0.65.0 +registration pattern so ``commands/eval.py`` stays under length cap). + +Input JSONL accepts two row shapes: + +* ``{"predicted": str, "expected_ids": [str, ...]}`` — scored directly. +* RAFT rows ``{"query", "golden_doc", "distractor_docs", "answer"}`` — the + ``answer`` is treated as the prediction and the ground-truth citation is the + golden document's deterministic ``[doc-N]`` id (the same id the RAFT trainer + assigns), so a model trained with ``citation_faithful`` can be scored against + its own training data. +""" + +from __future__ import annotations + +import json +import os +import stat +from typing import List, Mapping, Optional, Tuple + +import typer +from rich.console import Console +from rich.markup import escape +from rich.panel import Panel +from rich.table import Table + +_MAX_FILE_BYTES = 256 * 1024 * 1024 # 256 MiB +_MAX_ROWS = 1_000_000 + + +def _load_jsonl_rows(path: str, console: Console) -> List[dict]: + """Load JSONL rows with cwd containment + O_NOFOLLOW + size/row caps. + + Mirrors the v0.65.0 ``irt.load_response_rows`` TOCTOU-safe reader: the + cwd-containment helper rejects out-of-tree / symlink paths, then the open + uses ``O_NOFOLLOW`` (POSIX) so a symlink swapped in after the check still + fails. Malformed JSON lines are skipped (counted toward the row cap). + """ + from soup_cli.utils.paths import enforce_under_cwd_and_no_symlink + + enforce_under_cwd_and_no_symlink(path, "citation data") + no_follow = getattr(os, "O_NOFOLLOW", 0) + try: + fd = os.open(path, os.O_RDONLY | no_follow) + except FileNotFoundError as exc: + raise typer.BadParameter( + f"citation data not found: {os.path.basename(path)}" + ) from exc + except OSError as exc: + raise typer.BadParameter( + f"citation data cannot be opened (symlink?): {type(exc).__name__}" + ) from exc + try: + st = os.fstat(fd) + if stat.S_ISLNK(st.st_mode) or not stat.S_ISREG(st.st_mode): + raise typer.BadParameter("citation data must be a regular file") + if st.st_size > _MAX_FILE_BYTES: + raise typer.BadParameter( + f"citation data exceeds {_MAX_FILE_BYTES} bytes" + ) + rows: List[dict] = [] + skipped = 0 + with os.fdopen(fd, "r", encoding="utf-8-sig") as handle: + for count, line in enumerate(handle): + if count >= _MAX_ROWS: + break + line = line.strip() + if not line: + continue + try: + obj = json.loads(line) + except json.JSONDecodeError: + skipped += 1 + continue + if isinstance(obj, dict): + rows.append(obj) + else: + skipped += 1 + except Exception: + # fdopen may not have taken ownership of fd on early raise. + try: + os.close(fd) + except OSError: + pass + raise + if skipped: + console.print(f"[dim]Skipped {skipped} malformed / non-object row(s).[/]") + return rows + + +def _row_predicted_expected( + row: Mapping[str, object], *, row_index: int, shuffle_seed: Optional[int] = None +) -> Optional[Tuple[str, Tuple[str, ...]]]: + """Resolve a row into ``(predicted_text, expected_ids)`` or None to skip. + + For a RAFT-shaped row the golden ``[doc-N]`` id is derived from + ``build_raft_prompt`` — which shuffles document order by + ``(shuffle_seed, row_index)``. To score a model against its OWN training + data, ``shuffle_seed`` MUST match the ``data.raft_shuffle_seed`` used at + train time (code-review L4); otherwise the golden id won't line up. + """ + predicted = row.get("predicted") + expected = row.get("expected_ids") + if isinstance(predicted, str) and isinstance(expected, list): + ids = tuple(e for e in expected if isinstance(e, str) and e) + return predicted, ids + # RAFT-shaped row → answer is the prediction, golden id the expectation. + if ( + isinstance(row.get("query"), str) + and isinstance(row.get("golden_doc"), str) + and isinstance(row.get("answer"), str) + ): + from soup_cli.utils.raft import build_raft_prompt + + composed = build_raft_prompt( + row, shuffle_seed=shuffle_seed, row_index=row_index + ) + return composed.answer, (composed.golden_doc_id,) + return None + + +def register(app: typer.Typer, console: Console) -> None: + """Attach ``soup eval citation`` to the eval Typer app.""" + + @app.command(name="citation") + def citation( + data: str = typer.Argument( + ..., + help=( + "JSONL of {predicted, expected_ids} rows OR RAFT rows " + "{query, golden_doc, distractor_docs, answer}." + ), + ), + style: str = typer.Option( + "bracket", "--style", + help="Citation style: bracket / inline / footnote.", + ), + shuffle_seed: Optional[int] = typer.Option( + None, "--shuffle-seed", + help=( + "For RAFT-shaped rows: the data.raft_shuffle_seed used at " + "train time, so the golden [doc-N] id matches what the model " + "saw. Ignored for {predicted, expected_ids} rows." + ), + ), + output: Optional[str] = typer.Option( + None, "--output", "-o", + help="Write the per-row + aggregate CitationScore JSON here.", + ), + ) -> None: + """Score citation precision / recall / F1 over a JSONL (#202).""" + from soup_cli.utils.citation_faithful import ( + score_citations, + validate_citation_style, + ) + + try: + canonical_style = validate_citation_style(style) + except (TypeError, ValueError) as exc: + console.print(f"[red]Invalid --style:[/] {escape(str(exc))}") + raise typer.Exit(2) from exc + + rows = _load_jsonl_rows(data, console) + if not rows: + console.print("[red]No usable rows in citation data.[/]") + raise typer.Exit(2) + + table = Table(title="Citation scores", border_style="cyan") + table.add_column("Row", justify="right") + table.add_column("Precision", justify="right") + table.add_column("Recall", justify="right") + table.add_column("F1", justify="right") + + scored: List[dict] = [] + prec_sum = rec_sum = f1_sum = 0.0 + for index, row in enumerate(rows): + resolved = _row_predicted_expected( + row, row_index=index, shuffle_seed=shuffle_seed + ) + if resolved is None: + continue + predicted, expected_ids = resolved + try: + cs = score_citations( + predicted=predicted, + expected_ids=list(expected_ids), + style=canonical_style, + ) + except (TypeError, ValueError) as exc: + console.print( + f"[yellow]Row {index} skipped:[/] {escape(str(exc))}" + ) + continue + prec_sum += cs.precision + rec_sum += cs.recall + f1_sum += cs.f1 + scored.append({ + "row": index, + "precision": cs.precision, + "recall": cs.recall, + "f1": cs.f1, + "predicted_count": cs.predicted_count, + "expected_count": cs.expected_count, + }) + if index < 50: # keep the table bounded + table.add_row( + str(index), + f"{cs.precision:.3f}", + f"{cs.recall:.3f}", + f"{cs.f1:.3f}", + ) + + if not scored: + console.print( + "[red]No scorable rows[/] — need {predicted, expected_ids} " + "or RAFT {query, golden_doc, answer}." + ) + raise typer.Exit(2) + + n = len(scored) + aggregate = { + "precision": prec_sum / n, + "recall": rec_sum / n, + "f1": f1_sum / n, + } + console.print(table) + console.print( + Panel( + f"Rows scored: [bold]{n}[/]\n" + f"Style: [bold]{escape(canonical_style)}[/]\n" + f"Mean precision: [bold]{aggregate['precision']:.3f}[/]\n" + f"Mean recall: [bold]{aggregate['recall']:.3f}[/]\n" + f"Mean F1: [bold]{aggregate['f1']:.3f}[/]", + title="Citation aggregate", + border_style="green", + ) + ) + + if output is not None: + from soup_cli.utils.paths import enforce_under_cwd_and_no_symlink + + try: + enforce_under_cwd_and_no_symlink(output, "citation output") + except (TypeError, ValueError) as exc: + console.print(f"[red]Invalid --output:[/] {escape(str(exc))}") + raise typer.Exit(2) from exc + payload = { + "style": canonical_style, + "n_rows": n, + "aggregate": aggregate, + "rows": scored, + } + with open(output, "w", encoding="utf-8") as fh: + json.dump(payload, fh, indent=2, allow_nan=False) + console.print(f"Wrote citation report -> {escape(output)}") diff --git a/src/soup_cli/commands/eval.py b/src/soup_cli/commands/eval.py index 73c5042..1cc3df3 100644 --- a/src/soup_cli/commands/eval.py +++ b/src/soup_cli/commands/eval.py @@ -1177,3 +1177,8 @@ _register_v0610(app, console) from soup_cli.commands._eval_v0650 import register as _register_v0650 # noqa: E402 _register_v0650(app, console) + +# Register v0.71.10 subcommand (eval citation — #202) +from soup_cli.commands._eval_v07110 import register as _register_v07110 # noqa: E402 + +_register_v07110(app, console) diff --git a/src/soup_cli/commands/ra_dit.py b/src/soup_cli/commands/ra_dit.py new file mode 100644 index 0000000..775ddfb --- /dev/null +++ b/src/soup_cli/commands/ra_dit.py @@ -0,0 +1,126 @@ +"""v0.71.10 #200 — ``soup ra-dit`` one-shot two-stage RA-DIT orchestrator. + +Chains the RA-DIT retriever (stage 1, embedding/contrastive) and generator +(stage 2, RAFT-style SFT) in a single invocation, **recording** the trained +retriever as the generator's paired retriever (it writes the retriever's +output dir into the generator config's ``training.ra_dit_retriever_model``). +The generator is trained RAFT-style to be robust to retrieved distractors; +the recorded retriever is the one used at deploy/serve time for the actual +retrieval step (per the Meta RA-DIT recipe) — stage-2 training does not fuse +the retriever weights. Mirrors the v0.62.0 Part B schema; this command lifts +the deferred "live orchestration" note. + +``--plan-only`` validates both config paths + renders the resolved plan +without training. The live path runs each stage as a subprocess +(``soup train --config --yes``) via :mod:`soup_cli.utils.ra_dit_run`. +""" + +from __future__ import annotations + +from typing import Optional + +import typer +from rich.console import Console +from rich.markup import escape +from rich.panel import Panel + +console = Console() + +app = typer.Typer( + name="ra-dit", + help=( + "RA-DIT (Retrieval-Augmented Dual Instruction Tuning) one-shot " + "orchestrator: train the retriever then the generator, recording the " + "trained retriever as the generator's paired retriever (v0.71.10)." + ), + no_args_is_help=True, + rich_markup_mode="rich", +) + + +@app.callback(invoke_without_command=True) +def main( + retriever_config: str = typer.Option( + ..., "--retriever-config", "-r", + help="Stage-1 embedding/contrastive recipe (cwd-contained).", + ), + generator_config: str = typer.Option( + ..., "--generator-config", "-g", + help="Stage-2 RAFT-SFT recipe (cwd-contained).", + ), + retriever_model: Optional[str] = typer.Option( + None, "--retriever-model", + help="Manual retriever-model override; skips Registry auto-link.", + ), + timeout: int = typer.Option( + 6 * 60 * 60, "--timeout", + help="Per-stage hard timeout in seconds (60..21600).", + ), + plan_only: bool = typer.Option( + False, "--plan-only", + help="Validate config paths + render the plan; skip training.", + ), +) -> None: + """Run (or plan) a two-stage RA-DIT pipeline.""" + from soup_cli.utils.ra_dit_run import ( + resolve_retriever_for_generator, + run_ra_dit, + validate_ra_dit_config_path, + ) + + # Validate both config paths up front (containment + existence) so a typo + # fails fast with a clear message — for plan-only AND live. + try: + retr_path = validate_ra_dit_config_path("--retriever-config", retriever_config) + gen_path = validate_ra_dit_config_path("--generator-config", generator_config) + except (TypeError, ValueError, FileNotFoundError) as exc: + console.print(f"[red]Invalid ra-dit config:[/] {escape(str(exc))}") + raise typer.Exit(2) from exc + + if plan_only: + # Preview the retriever link without touching the Registry write path. + if retriever_model is not None: + link_preview = ( + f"manual override: {escape(retriever_model)}" + ) + else: + resolved, advisory = resolve_retriever_for_generator(None) + link_preview = escape(advisory) + console.print( + Panel( + f"[bold]Stage 1 (retriever):[/] {escape(retr_path)}\n" + f"[bold]Stage 2 (generator):[/] {escape(gen_path)}\n" + f"[bold]Retriever link:[/] {link_preview}\n" + f"[bold]Per-stage timeout:[/] {timeout}s", + title="soup ra-dit (plan-only)", + border_style="cyan", + ) + ) + console.print( + "[green]Plan-only mode — config paths validated. " + "Drop --plan-only to run both stages.[/]" + ) + return + + try: + result = run_ra_dit( + retriever_config, + generator_config, + retriever_model=retriever_model, + timeout_seconds=timeout, + ) + except (TypeError, ValueError, FileNotFoundError, RuntimeError) as exc: + console.print(f"[red]RA-DIT run failed:[/] {escape(str(exc))}") + raise typer.Exit(2) from exc + + link_note = "auto-linked" if result.autolinked else "manual override" + console.print( + Panel( + f"[bold]Retriever output:[/] {escape(result.retriever_output)}\n" + f"[bold]Generator output:[/] {escape(result.generator_output)}\n" + f"[bold]Retriever model used:[/] " + f"{escape(result.retriever_model_used)} ({link_note})", + title="RA-DIT pipeline complete", + border_style="green", + ) + ) diff --git a/src/soup_cli/commands/serve.py b/src/soup_cli/commands/serve.py index 5b7a975..034c0a0 100644 --- a/src/soup_cli/commands/serve.py +++ b/src/soup_cli/commands/serve.py @@ -236,9 +236,9 @@ def serve( ), ): """Start a local inference server with OpenAI-compatible API.""" - # v0.62.0 Part C — validate `--steer` name + strength up front so a - # typo surfaces before backend init. Live decode-hook wiring lands - # in v0.62.1; this release just locks the flag surface. + # v0.62.0 Part C / v0.71.10 #201 — validate `--steer` name + strength up + # front so a typo surfaces before backend init. The live decode hook is + # installed in the transformers branch after model load. if steer is not None: from rich.markup import escape as _rich_escape @@ -258,10 +258,13 @@ def serve( f"[red]Invalid --steer:[/] {_rich_escape(str(exc))}" ) raise typer.Exit(code=2) from exc - console.print( - f"[yellow]--steer={_rich_escape(steer)!r} accepted; live " - "decode hook ships in v0.62.1.[/]" - ) + if backend.lower() != "transformers": + console.print( + "[red]--steer requires --backend transformers[/] " + "(activation steering installs a forward hook on the loaded " + "model; vLLM / SGLang / MII are not supported)." + ) + raise typer.Exit(code=2) # v0.53.10 #152 — pre-fetch base from a non-HF hub before serve starts. if hub and hub != "hf": @@ -605,6 +608,32 @@ def serve( ) console.print("[bold green]Model loaded![/]") + # v0.71.10 #201 — install the activation-steering decode hook. The + # handle persists for the server's lifetime (process-global model). + if steer is not None: + from rich.markup import escape as _esc + + from soup_cli.utils.steering import ( + install_steering_hook, + load_steering_artifact, + resolve_steering_dir, + ) + + try: + steer_dir = resolve_steering_dir(steer) + loaded_steer = load_steering_artifact(steer_dir) + install_steering_hook( + model_obj, loaded_steer, strength=steer_strength + ) + except (TypeError, ValueError, OSError) as exc: + console.print(f"[red]--steer:[/] {_esc(str(exc))}") + raise typer.Exit(2) from exc + console.print( + f"[green]Steering active:[/] {_esc(loaded_steer.name)} " + f"({_esc(loaded_steer.method)}, layer {loaded_steer.layer}, " + f"strength {steer_strength})" + ) + # Load draft model for speculative decoding (transformers backend) draft_model = None if speculative_model: diff --git a/src/soup_cli/commands/steer.py b/src/soup_cli/commands/steer.py index 8f5ca1b..2e20149 100644 --- a/src/soup_cli/commands/steer.py +++ b/src/soup_cli/commands/steer.py @@ -73,14 +73,25 @@ def train_steer( ), layer: Optional[int] = typer.Option( None, "--layer", "-l", - help="MLP layer index to extract the residual-stream vector from.", + help="Decoder layer index to extract the residual-stream vector from " + "(default: the middle layer).", + ), + output: Optional[str] = typer.Option( + None, "--output", "-o", + help="Directory to write the steering vector to (default: " + "./steering/). Cwd-contained.", + ), + device: Optional[str] = typer.Option( + None, "--device", + help="torch device (cpu / cuda). Defaults to CUDA when available.", + ), + top_k: int = typer.Option( + 8, "--top-k", "-k", + help="ITI only: number of attention heads to intervene on (1-256).", ), plan_only: bool = typer.Option( False, "--plan-only", - help=( - "Validate inputs + print the resolved plan; skip the " - "deferred-live training (v0.62.1)." - ), + help="Validate inputs + print the resolved plan; skip the live fit.", ), registry_id: Optional[str] = typer.Option( None, "--registry-id", @@ -141,6 +152,11 @@ def train_steer( f"[red]Invalid --layer:[/] must satisfy 0 <= layer <= 2048, got {layer}" ) raise typer.Exit(2) + if top_k < 1 or top_k > 256: + console.print( + f"[red]Invalid --top-k:[/] must satisfy 1 <= top_k <= 256, got {top_k}" + ) + raise typer.Exit(2) spec = get_steering_method_spec(canonical_method) panel_body = ( @@ -148,7 +164,7 @@ def train_steer( f"Method: {escape(canonical_method)}\n" f"Name: {escape(canonical_name)}\n" f"Pairs: {escape(pairs_path)}\n" - f"Layer: {layer if layer is not None else 'auto'}\n" + f"Layer: {layer if layer is not None else 'auto (middle)'}\n" f"Description: {escape(spec.description)}" ) console.print( @@ -161,30 +177,55 @@ def train_steer( if plan_only: console.print( - "[yellow]Plan-only mode:[/] live training deferred to v0.62.1. " - "Inputs validated and plan rendered." + "[green]Plan-only mode — inputs validated, plan rendered. " + "Drop --plan-only to fit the vector.[/]" ) return try: - build_steering_vector( + artifact = build_steering_vector( method=canonical_method, name=canonical_name, pairs_path=pairs_path, + base=base, layer=layer, + device=device, + output_dir=output, + top_k=top_k, ) - except NotImplementedError as exc: - console.print( - Panel( - f"[yellow]Live steer-train deferred to v0.62.1.[/]\n\n{escape(str(exc))}", - title="Deferred", - border_style="yellow", + except (TypeError, ValueError, FileNotFoundError) as exc: + console.print(f"[red]Steer-train failed:[/] {escape(str(exc))}") + raise typer.Exit(2) from exc + + console.print( + Panel( + f"[bold]Method:[/] {escape(artifact.method)}\n" + f"[bold]Name:[/] {escape(artifact.name)}\n" + f"[bold]Layer:[/] {artifact.layer}\n" + f"[bold]Intervention:[/] {escape(artifact.intervention_point)}\n" + f"[bold]Hidden dim:[/] {artifact.hidden_dim}\n" + f"[bold]Pairs:[/] {artifact.num_pairs}\n" + f"[bold]Saved:[/] {escape(artifact.output_dir)}", + title="Steering vector trained", + border_style="green", + ) + ) + + if registry_id is not None: + from soup_cli.registry.attach import attach_artifact + + try: + attach_artifact( + registry_id, path=artifact.output_dir, kind="steering_vector" + ) + console.print( + f"[green]Attached steering_vector to Registry entry " + f"{escape(registry_id)}.[/]" + ) + except (ValueError, FileNotFoundError) as exc: + console.print( + f"[yellow]Could not attach to Registry:[/] {escape(str(exc))}" ) - ) - # Match v0.61.0 Part C policy: exit code 3 distinguishes - # "deferred / not yet shipped" from exit code 2 = "validation - # rejection". - raise typer.Exit(3) from exc @app.command(name="apply") @@ -198,8 +239,16 @@ def apply_steer( help="Steering strength multiplier (|s| <= 10.0).", ), ) -> None: - """Apply a stored steering vector at decode time (preview-only in v0.62.0).""" + """Resolve + load a stored steering vector and print its metadata. + + The live decode-time intervention is applied by ``soup serve --steer + --steer-strength `` (which installs the forward hook on the + running model). This subcommand resolves the vector by name and confirms it + loads cleanly so operators can verify an artifact before serving. + """ from soup_cli.utils.steering import ( + load_steering_artifact, + resolve_steering_dir, validate_steering_name, validate_steering_strength, ) @@ -211,21 +260,31 @@ def apply_steer( console.print(f"[red]Invalid steer-apply input:[/] {escape(str(exc))}") raise typer.Exit(2) from exc + try: + steer_dir = resolve_steering_dir(canonical_name) + loaded = load_steering_artifact(steer_dir) + except (TypeError, ValueError, OSError) as exc: + console.print(f"[red]Cannot load steering vector:[/] {escape(str(exc))}") + raise typer.Exit(2) from exc + console.print( Panel( - ( - f"Vector: {escape(canonical_name)}\n" - f"Strength: {canonical_strength}" - ), + f"[bold]Vector:[/] {escape(canonical_name)}\n" + f"[bold]Method:[/] {escape(loaded.method)}\n" + f"[bold]Layer:[/] {loaded.layer}\n" + f"[bold]Intervention:[/] {escape(loaded.intervention_point)}\n" + f"[bold]Dim:[/] {len(loaded.vector)}\n" + f"[bold]Strength:[/] {canonical_strength}\n" + f"[bold]Dir:[/] {escape(steer_dir)}", title="soup steer apply", border_style="cyan", ) ) console.print( - "[yellow]Apply path deferred to v0.62.1.[/] Use `soup serve --steer " - f"{escape(canonical_name)}` once the live decode hook ships." + "[green]Vector loaded.[/] Apply it at decode time with " + f"`soup serve --steer {escape(canonical_name)} " + f"--steer-strength {canonical_strength}`." ) - raise typer.Exit(3) @app.command(name="list") diff --git a/src/soup_cli/commands/train.py b/src/soup_cli/commands/train.py index 0c2d5a2..bf75df4 100644 --- a/src/soup_cli/commands/train.py +++ b/src/soup_cli/commands/train.py @@ -273,6 +273,18 @@ def train( console.print(f"[dim]Loading config from {config_path}...[/]") cfg = load_config(config_path) + # --- RA-DIT generator-stage auto-link (v0.71.10 #200) --- + # When a generator stage has no retriever model set, splice in the latest + # RA-DIT retriever output from the Registry. A manual value always wins. + if getattr(cfg.training, "ra_dit_stage", None) == "generator": + from soup_cli.utils.ra_dit_run import autolink_generator_retriever + + advisory = autolink_generator_retriever(cfg) + if advisory: + # `advisory` embeds a Registry-derived `output` path — escape it + # before printing into the Rich-markup console (security MEDIUM). + console.print(f"[yellow]RA-DIT:[/] {markup_escape(advisory)}") + # --- Echo-trap tokenizer-aware shortcut --- if echo_trap_tokenizer_aware: if not cfg.training.echo_trap_enabled: diff --git a/src/soup_cli/config/schema.py b/src/soup_cli/config/schema.py index a9cd46e..321dc4f 100644 --- a/src/soup_cli/config/schema.py +++ b/src/soup_cli/config/schema.py @@ -244,6 +244,26 @@ class DataConfig(BaseModel): "shipped template (errors loudly if absent). (v0.36.0)" ), ) + raft_shuffle_seed: Optional[int] = Field( + default=None, + ge=0, + le=2_147_483_647, + description=( + "Seed for the RAFT golden/distractor document shuffle " + "(data.format='raft'). Documents are always shuffled for " + "distractor robustness; this knob fixes which reproducible " + "permutation. None = seed 0. (v0.71.10 #199)" + ), + ) + + @field_validator("raft_shuffle_seed", mode="before") + @classmethod + def _validate_raft_shuffle_seed(cls, v): + # Bool is a subclass of int — reject before Pydantic coerces True->1 + # (project bool-as-int policy). + if isinstance(v, bool): + raise ValueError("raft_shuffle_seed must not be a bool") + return v # --- v0.42.0 Data Pipeline Pro ----------------------------------------- video_dir: Optional[str] = Field( diff --git a/src/soup_cli/trainer/raft.py b/src/soup_cli/trainer/raft.py new file mode 100644 index 0000000..1cf1c44 --- /dev/null +++ b/src/soup_cli/trainer/raft.py @@ -0,0 +1,121 @@ +"""v0.71.10 #199 — live RAFT trainer pieces. + +``make_raft_trainer_class(base_cls)`` builds an HF ``Trainer`` subclass whose +``compute_loss`` does a per-token WEIGHTED cross-entropy: the answer-only loss +mask is expressed via ``loss_weights`` (0.0 on the prompt span, 1.0 on the +answer, boosted on bracketed citation spans when ``citation_faithful`` is set +— #202). With all-1.0 answer weights this reduces exactly to answer-only CE. + +``RaftDataCollator`` pads the pre-tokenised ``{input_ids, attention_mask, +labels, loss_weights}`` rows produced by ``utils.raft.tokenize_raft_example`` +(``DataCollatorForSeq2Seq`` cannot pad the custom ``loss_weights`` column). + +Mirrors the v0.53.11 ``make_prm_trainer_class`` / v0.53.2 ``_DistillTrainer`` +factory pattern. Heavy imports (torch) are local. +""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Any, List + + +class RaftDataCollator: + """Pad pre-tokenised RAFT rows (incl. the custom ``loss_weights`` column).""" + + def __init__(self, tokenizer: Any) -> None: + pad_id = getattr(tokenizer, "pad_token_id", None) + if pad_id is None: + pad_id = getattr(tokenizer, "eos_token_id", None) + if pad_id is None: + pad_id = 0 + self.pad_id = int(pad_id) + + def __call__(self, features: List[dict]) -> dict: + import torch + + if not features: + raise ValueError("RaftDataCollator received an empty batch") + max_len = max(len(f["input_ids"]) for f in features) + input_ids: List[List[int]] = [] + attention_mask: List[List[int]] = [] + labels: List[List[int]] = [] + loss_weights: List[List[float]] = [] + for f in features: + ids = list(f["input_ids"]) + n = len(ids) + pad = max_len - n + input_ids.append(ids + [self.pad_id] * pad) + attention_mask.append(list(f.get("attention_mask", [1] * n)) + [0] * pad) + labels.append(list(f["labels"]) + [-100] * pad) + loss_weights.append( + list(f.get("loss_weights", [1.0] * n)) + [0.0] * pad + ) + return { + "input_ids": torch.tensor(input_ids, dtype=torch.long), + "attention_mask": torch.tensor(attention_mask, dtype=torch.long), + "labels": torch.tensor(labels, dtype=torch.long), + "loss_weights": torch.tensor(loss_weights, dtype=torch.float32), + } + + +@lru_cache(maxsize=None) +def make_raft_trainer_class(base_cls: type) -> type: + """Build a ``Trainer`` subclass with per-token weighted-CE ``compute_loss``. + + Cached so repeated ``raft`` runs against the same base class share one + subclass (mirrors ``make_prm_trainer_class`` / ``make_multipack_trainer_class``). + """ + + class _RaftTrainer(base_cls): # type: ignore[valid-type, misc] + def compute_loss( + self, + model, + inputs, + return_outputs: bool = False, + num_items_in_batch=None, + ): + import torch + from torch.nn.functional import cross_entropy + + labels = inputs.get("labels") + loss_weights = inputs.get("loss_weights") + model_inputs = { + k: v + for k, v in inputs.items() + if k not in ("labels", "loss_weights") + } + outputs = model(**model_inputs) + logits = outputs.logits + shift_logits = logits[:, :-1, :].contiguous() + shift_labels = labels[:, 1:].contiguous() + per_token = cross_entropy( + shift_logits.view(-1, shift_logits.size(-1)), + shift_labels.view(-1), + ignore_index=-100, + reduction="none", + ) # [B*S'] — 0.0 at ignored positions. + if loss_weights is not None: + w = loss_weights[:, 1:].contiguous().reshape(-1).to(per_token.dtype) + # Zero the weight wherever the label is ignored so padded / + # prompt tokens never enter the weighted mean even if a + # non-zero weight leaked in. + valid = (shift_labels.reshape(-1) != -100).to(per_token.dtype) + w = w * valid + else: + w = (shift_labels.reshape(-1) != -100).to(per_token.dtype) + denom = w.sum().clamp(min=1.0) + loss = (per_token * w).sum() / denom + if not torch.isfinite(loss): + # Degenerate batch (all-masked) OR a non-finite forward — + # return a STRUCTURAL zero (never NaN even if per_token has a + # NaN: `nan * 0.0 == nan`, so anchor on a fresh zeros tensor + # with grad rather than `per_token.mean() * 0.0`). + loss = torch.zeros( + (), device=per_token.device, dtype=per_token.dtype, + requires_grad=True, + ) + return (loss, outputs) if return_outputs else loss + + _RaftTrainer.__name__ = f"_RaftTrainer_{base_cls.__name__}" + return _RaftTrainer diff --git a/src/soup_cli/trainer/sft.py b/src/soup_cli/trainer/sft.py index 0e66a09..ed27f9b 100644 --- a/src/soup_cli/trainer/sft.py +++ b/src/soup_cli/trainer/sft.py @@ -109,6 +109,7 @@ class SFTTrainerWrapper: self.model = None self.tokenizer = None self.trainer = None + self._is_raft = False # set in setup() when data.format == 'raft' # Resolve once — raises ValueError if model needs custom code but # the user did not opt in. Result is cached on the wrapper for use # by every from_pretrained() call below. @@ -222,9 +223,14 @@ class SFTTrainerWrapper: # v0.53.7 #86 — short-circuit tokenization when caller pre-tokenized # via `soup data preprocess`. Skips the format_row + tokenizer pass # entirely; rows already carry input_ids/labels/attention_mask. + # v0.71.10 #199 — RAFT format: golden/distractor-doc rows are NOT + # {messages}; build a pre-tokenised answer-only-mask dataset instead. + self._is_raft = cfg.data.format == "raft" pretok = _maybe_load_pretokenized(cfg.data, cfg.base, console) if pretok is not None: train_ds, eval_ds = pretok + elif self._is_raft: + train_ds, eval_ds = self._prepare_raft_dataset(dataset, cfg, tcfg) elif use_vision: train_ds, eval_ds = self._prepare_vision_dataset(dataset) elif use_audio: @@ -393,7 +399,36 @@ class SFTTrainerWrapper: # so two ``multipack: true`` runs against the same base class share # the same subclass. use_multipack = bool(getattr(tcfg, "multipack", False)) - if use_multipack: + if self._is_raft: + # v0.71.10 #199 — RAFT uses a plain Trainer + weighted-CE loss + # (answer-only mask via loss_weights; citation-span boost when + # training.citation_faithful is set — #202). The pre-tokenised + # rows + custom collator skip SFTTrainer's text-column processing. + from transformers import Trainer + + from soup_cli.trainer.raft import ( + RaftDataCollator, + make_raft_trainer_class, + ) + + raft_cls = make_raft_trainer_class(Trainer) + self.trainer = raft_cls( + model=self.model, + args=training_args, + train_dataset=train_ds, + eval_dataset=eval_ds, + data_collator=RaftDataCollator(self.tokenizer), + processing_class=self.tokenizer, + ) + if tcfg.citation_faithful: + console.print( + "[green]RAFT + citation-faithful:[/] answer-only mask " + f"with boosted [{tcfg.citation_style or 'bracket'}] " + "citation spans" + ) + else: + console.print("[green]RAFT trainer enabled:[/] answer-only loss mask") + elif use_multipack: from soup_cli.utils.multipack_sampler import ( validate_multipack_architecture, ) @@ -423,6 +458,64 @@ class SFTTrainerWrapper: self._output_dir = str(output_dir) self._batch_size = batch_size + def _prepare_raft_dataset(self, dataset: dict, cfg, tcfg): + """v0.71.10 #199 — build pre-tokenised RAFT rows (answer-only mask). + + Each ``{query, golden_doc, distractor_docs, answer}`` row is composed + into a prompt + answer with deterministic ``[doc-N]`` ids + (:func:`soup_cli.utils.raft.build_raft_prompt`) then tokenised with the + prompt span masked. When ``citation_faithful`` is set, citation spans + in the answer get a boosted ``loss_weights`` entry. + """ + from datasets import Dataset + + from soup_cli.utils.raft import build_raft_prompt, tokenize_raft_example + + shuffle_seed = cfg.data.raft_shuffle_seed + citation = bool(tcfg.citation_faithful) + style = tcfg.citation_style or "bracket" + max_length = cfg.data.max_length + tokenizer = self.tokenizer + + def _fmt(example: dict, idx: int) -> dict: + composed = build_raft_prompt( + example, shuffle_seed=shuffle_seed, row_index=idx + ) + return tokenize_raft_example( + tokenizer, + composed, + max_length=max_length, + citation_faithful=citation, + citation_style=style, + ) + + def _has_trainable_tokens(example: dict) -> bool: + # A row whose prompt fills `max_length` truncates the answer away, + # leaving an all-masked (loss_weights all 0.0) row that contributes + # a zero gradient. Drop such rows so they don't silently shrink the + # effective dataset (code-review M4). + return any(w > 0.0 for w in example["loss_weights"]) + + def _map_and_filter(raw, split: str): + mapped = raw.map( + _fmt, with_indices=True, remove_columns=raw.column_names + ) + kept = mapped.filter(_has_trainable_tokens) + dropped = len(mapped) - len(kept) + if dropped: + console.print( + f"[yellow]RAFT:[/] dropped {dropped} {split} row(s) whose " + f"prompt filled max_length={max_length} (answer fully " + "truncated -> all-masked). Raise max_length to keep them." + ) + return kept + + train_ds = _map_and_filter(Dataset.from_list(dataset["train"]), "train") + eval_ds = None + if "val" in dataset and dataset["val"]: + eval_ds = _map_and_filter(Dataset.from_list(dataset["val"]), "val") + return train_ds, eval_ds + def _resolve_mixed_precision(self, tcfg, base_model: str) -> tuple[bool, bool]: """Return ``(bf16, fp16)`` flags for TrainingArguments. diff --git a/src/soup_cli/utils/citation_faithful.py b/src/soup_cli/utils/citation_faithful.py index 185e650..8bc4c64 100644 --- a/src/soup_cli/utils/citation_faithful.py +++ b/src/soup_cli/utils/citation_faithful.py @@ -32,13 +32,31 @@ _MAX_STYLE_LEN: int = 32 _MAX_PREDICTED_LEN: int = 2_000_000 # 2 MB cap on per-row predicted text. _MAX_EXPECTED_IDS: int = 10_000 # Per-row expected-citation cap. -# Default extraction regex — matches both ``[doc-id]`` brackets and bare -# ``doc-id`` tokens. The bracket form is the canonical RAFT default; -# ``inline`` and ``footnote`` use the same characters today (live -# per-style extractors ship in v0.62.1 once we benchmark variations). -_CITATION_RE: re.Pattern[str] = re.compile( - r"\[(?P[A-Za-z0-9][A-Za-z0-9._\-]{0,127})\]" -) +# Per-style citation regexes (v0.71.10 #202). Each captures the document id in +# group ``id``; the FULL match span (incl. delimiters) is used by the RAFT +# span-mask so the boost covers the brackets too. +# +# * ``bracket`` — ``[doc-id]`` (the canonical Stanford RAFT default). +# * ``inline`` — ``(doc-id)`` parenthetical references. +# * ``footnote`` — ``[^id]`` markdown-footnote markers. +# +# The id char class (alnum-leading, ``._-`` body, <=128 chars) is shared. The +# bracket and footnote forms are disjoint: ``[^1]`` starts with ``^`` after the +# bracket, which fails the alnum-leading bracket pattern. +_STYLE_RE: dict[str, re.Pattern[str]] = { + "bracket": re.compile(r"\[(?P[A-Za-z0-9][A-Za-z0-9._\-]{0,127})\]"), + "inline": re.compile(r"\((?P[A-Za-z0-9][A-Za-z0-9._\-]{0,127})\)"), + "footnote": re.compile(r"\[\^(?P[A-Za-z0-9][A-Za-z0-9._\-]{0,127})\]"), +} + +# Back-compat alias — the bracket form was the v0.62.0 default. +_CITATION_RE: re.Pattern[str] = _STYLE_RE["bracket"] + + +def _resolve_style_re(style: str) -> re.Pattern[str]: + """Return the compiled regex for ``style`` (validated) or raise.""" + canonical = validate_citation_style(style) + return _STYLE_RE[canonical] @dataclass(frozen=True) @@ -108,12 +126,14 @@ def validate_citation_threshold(value: object) -> float: return fval -def extract_citation_ids(text: str) -> tuple[str, ...]: - """Extract every ``[doc-id]`` citation from ``text``. +def extract_citation_ids(text: str, *, style: str = "bracket") -> tuple[str, ...]: + """Extract every citation id from ``text`` for the given ``style``. - Returns a tuple of IDs in encounter order. Duplicates are preserved - so the caller can compute precision honestly (a model that cites - the same doc three times should not silently dedupe). + Returns a tuple of IDs in encounter order. Duplicates are preserved so the + caller can compute precision honestly (a model that cites the same doc + three times should not silently dedupe). ``style`` selects the per-style + extractor: ``bracket`` (``[doc-id]``) / ``inline`` (``(doc-id)``) / + ``footnote`` (``[^id]``). Defaults to ``bracket`` for back-compat. """ if not isinstance(text, str): raise TypeError( @@ -123,13 +143,32 @@ def extract_citation_ids(text: str) -> tuple[str, ...]: raise ValueError( f"text must be <= {_MAX_PREDICTED_LEN} chars for citation extract" ) - return tuple(m.group("bracketed") for m in _CITATION_RE.finditer(text)) + pattern = _resolve_style_re(style) + return tuple(m.group("id") for m in pattern.finditer(text)) + + +def citation_spans(text: str, *, style: str = "bracket") -> tuple[tuple[int, int], ...]: + """Return the char ``(start, end)`` span of every citation in ``text``. + + The span covers the FULL match (delimiters included) so a token-level + span mask boosts the brackets as well as the id. Used by the RAFT + span-mask (#199) + citation-faithful loss boost (#202). + """ + if not isinstance(text, str): + raise TypeError(f"text must be str, got {type(text).__name__}") + if len(text) > _MAX_PREDICTED_LEN: + raise ValueError( + f"text must be <= {_MAX_PREDICTED_LEN} chars for citation extract" + ) + pattern = _resolve_style_re(style) + return tuple((m.start(), m.end()) for m in pattern.finditer(text)) def score_citations( *, predicted: object, expected_ids: object, + style: str = "bracket", ) -> CitationScore: """Compute citation precision / recall / F1. @@ -138,6 +177,7 @@ def score_citations( Undefined denominators (empty predicted / expected) return 0.0 by convention — same policy as v0.43.0 BLEU on zero-precision n-grams. + ``style`` selects the per-style extractor (bracket / inline / footnote). """ if isinstance(predicted, bool): raise TypeError( @@ -174,7 +214,7 @@ def score_citations( ) expected_set.add(eid) - predicted_ids = extract_citation_ids(predicted) + predicted_ids = extract_citation_ids(predicted, style=style) predicted_count = len(predicted_ids) expected_count = len(expected_set) diff --git a/src/soup_cli/utils/diagnose/citation.py b/src/soup_cli/utils/diagnose/citation.py new file mode 100644 index 0000000..b2d0268 --- /dev/null +++ b/src/soup_cli/utils/diagnose/citation.py @@ -0,0 +1,97 @@ +"""Citation-recall regression probe (v0.71.10 #202). + +A RAFT-trained model is supposed to cite the supporting ``[doc-id]`` in its +answer. This probe runs each RAFT-shaped row (``{query, golden_doc, +distractor_docs, answer}``) through the adapter generator, extracts the cited +document ids, and measures recall against the ground-truth golden-doc id. + +The ``FailureScore`` follows the same OK/MINOR/MAJOR taxonomy as every other +diagnose mode (``classify_score`` on the mean recall). The training-config +``citation_recall_threshold`` is a separate eval-gate knob — the diagnose +badge uses the shared 0.85/0.60 bands so it reads consistently next to the +other six probes. + +Heavy imports are avoided here: ``utils.raft`` + ``citation_faithful`` are +pure-Python; only the caller's generator closure touches a model. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence + +from soup_cli.utils.diagnose._common import ( + GeneratorFn, + call_generator, + merge_evidence, +) +from soup_cli.utils.diagnose.report import FailureScore, classify_score + +_MAX_CITATION_ROWS = 256 + + +def is_raft_row(row: object) -> bool: + """True when ``row`` carries the RAFT fields the citation probe needs.""" + if not isinstance(row, Mapping): + return False + return ( + isinstance(row.get("query"), str) + and isinstance(row.get("golden_doc"), str) + and isinstance(row.get("answer"), str) + ) + + +def score_citation( + rows: Sequence[Mapping[str, object]], + generator: GeneratorFn, + *, + citation_style: str = "bracket", + shuffle_seed: int | None = None, + max_rows: int = _MAX_CITATION_ROWS, +) -> FailureScore: + """Measure citation recall of ``generator`` over RAFT rows. + + For each row the prompt is composed (with deterministic ``[doc-N]`` ids), + the generator produces an answer, and recall = ``1.0`` when the golden + doc id is among the cited ids (else ``0.0``). The mean recall becomes the + score; the verdict follows ``classify_score``. + + Raises ``ValueError`` when no usable RAFT row is supplied (so the caller + falls back to a neutral score rather than reporting a misleading 0.0). + """ + from soup_cli.utils.citation_faithful import extract_citation_ids + from soup_cli.utils.raft import build_raft_prompt + + if not isinstance(rows, Sequence) or isinstance(rows, (str, bytes)): + raise TypeError("rows must be a sequence of mappings") + if isinstance(max_rows, bool) or not isinstance(max_rows, int) or max_rows < 1: + raise ValueError("max_rows must be a positive int") + + raft_rows = [r for r in rows if is_raft_row(r)][:max_rows] + if not raft_rows: + raise ValueError("no RAFT-shaped rows for the citation probe") + + recalls: list[float] = [] + cited_total = 0 + for index, row in enumerate(raft_rows): + composed = build_raft_prompt(row, shuffle_seed=shuffle_seed, row_index=index) + output = call_generator(generator, composed.prompt) + cited = set(extract_citation_ids(output, style=citation_style)) + cited_total += len(cited) + recalls.append(1.0 if composed.golden_doc_id in cited else 0.0) + + mean_recall = sum(recalls) / len(recalls) + score = max(0.0, min(1.0, mean_recall)) + evidence = merge_evidence( + { + "rows": len(raft_rows), + "mean_recall": mean_recall, + "cited_total": cited_total, + "style": citation_style, + } + ) + return FailureScore( + mode="citation", + score=score, + verdict=classify_score(score), + evidence=evidence, + ) diff --git a/src/soup_cli/utils/diagnose/live.py b/src/soup_cli/utils/diagnose/live.py index 7e0061d..62cf283 100644 --- a/src/soup_cli/utils/diagnose/live.py +++ b/src/soup_cli/utils/diagnose/live.py @@ -304,6 +304,16 @@ def run_live_diagnose( neutral_score("contamination", "no benchmark corpus supplied"), ) + # --- citation (v0.71.10 #202 — RAFT-shaped rows only) --- + if rows: + from soup_cli.utils.diagnose.citation import is_raft_row, score_citation + + if any(is_raft_row(r) for r in rows): + try: + scores["citation"] = score_citation(rows, adapter_gen) + except (ValueError, TypeError): + scores["citation"] = neutral_score("citation", "probe failed") + # Fill any still-missing modes (no dataset → forgetting/format/etc neutral). for mode in FAILURE_MODES: scores.setdefault(mode, neutral_score(mode, "probe inputs unavailable")) diff --git a/src/soup_cli/utils/diagnose/report.py b/src/soup_cli/utils/diagnose/report.py index 387f853..b28408e 100644 --- a/src/soup_cli/utils/diagnose/report.py +++ b/src/soup_cli/utils/diagnose/report.py @@ -19,6 +19,8 @@ FAILURE_MODES: Tuple[str, ...] = ( "mode_collapse", "memorization", "contamination", + # v0.71.10 #202 — citation-recall regression on RAFT-shaped data. + "citation", ) VERDICTS: Tuple[str, ...] = ("OK", "MINOR", "MAJOR") diff --git a/src/soup_cli/utils/ra_dit_run.py b/src/soup_cli/utils/ra_dit_run.py new file mode 100644 index 0000000..0f867c8 --- /dev/null +++ b/src/soup_cli/utils/ra_dit_run.py @@ -0,0 +1,373 @@ +"""v0.71.10 #200 — live RA-DIT two-stage orchestration + retriever auto-link. + +Lifts the v0.62.0 Part B deferred note ("Live orchestration that chains the +two stages in a single call is deferred"). Two surfaces: + +* ``autolink_generator_retriever(cfg)`` — when a ``soup train`` of a generator + stage (``training.ra_dit_stage='generator'``) has no + ``training.ra_dit_retriever_model`` set, discover the most-recent + RA-DIT retriever run from the Registry and splice its output in. A manual + value always wins (no overwrite); a yellow advisory string is returned so + the caller can surface what happened. +* ``run_ra_dit(retriever_config, generator_config, ...)`` — the one-shot + orchestrator behind ``soup ra-dit``: run the retriever stage, link its + output into the generator config, run the generator stage. Subprocess + invocation mirrors ``utils.mix_proxy`` (list argv, no shell, per-stage + ``soup train --config --yes``). + +Security: +- Config paths are containment-checked via the shared + :func:`enforce_under_cwd_and_no_symlink` helper (TOCTOU symlink rejection). +- The rewritten generator YAML is staged in a ``tempfile.mkdtemp`` dir and + cleaned up in ``finally`` (matches v0.53.5 ``mix_proxy`` policy). +- ``timeout_seconds`` capped to ``[60, 6*3600]``; ``TimeoutExpired`` raises + ``RuntimeError`` so the CLI can report a clean per-stage failure. +- No top-level torch / transformers import (CLI stays fast). +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from typing import Any, Callable, Optional, Tuple + +_MIN_TIMEOUT_S = 60 +_MAX_TIMEOUT_S = 6 * 60 * 60 # retriever + generator can each be long. +_MAX_PATH_LEN = 4096 +_MAX_YAML_BYTES = 256 * 1024 +_RETRIEVER_STAGE = "retriever" +_GENERATOR_STAGE = "generator" + +__all__ = [ + "RaDitRunResult", + "autolink_generator_retriever", + "discover_latest_retriever", + "resolve_retriever_for_generator", + "run_ra_dit", + "validate_ra_dit_config_path", +] + + +@dataclass(frozen=True) +class RaDitRunResult: + """Outcome of a two-stage :func:`run_ra_dit`. Frozen post-construction.""" + + retriever_output: str + generator_output: str + retriever_model_used: str + autolinked: bool + + +def discover_latest_retriever(*, store: Any = None) -> Optional[str]: + """Return the output dir of the most-recent RA-DIT retriever Registry run. + + Returns the ``output`` field of the first ``task='embedding'`` entry whose + stored config declares ``training.ra_dit_stage == 'retriever'``. + ``RegistryStore.list`` orders ``created_at DESC, id DESC`` (see + ``registry/store.py``), so the first match IS the most-recent retriever + run. Returns ``None`` when none exist. Registry/IO errors degrade to + ``None`` (auto-link is advisory; a missing Registry must never crash a + training run). The discovered ``output`` is run through + ``validate_ra_dit_retriever_model`` so a corrupt Registry row (null byte / + oversize) is skipped rather than flowing into the training config. + """ + from soup_cli.utils.ra_dit import ( # noqa: PLC0415 + validate_ra_dit_retriever_model, + ) + + own_store = False + if store is None: + try: + from soup_cli.registry.store import RegistryStore # noqa: PLC0415 + + store = RegistryStore() + own_store = True + except Exception: # noqa: BLE001 — Registry is optional. + return None + try: + entries = store.list(task="embedding", limit=100) + except Exception: # noqa: BLE001 — degrade to "no retriever". + return None + finally: + if own_store: + _close_quietly(store) + + for entry in entries: + raw = entry.get("config_json") + if not isinstance(raw, str): + continue + try: + cfg = json.loads(raw) + except (ValueError, TypeError): + continue + if not isinstance(cfg, dict): + continue + training = cfg.get("training") + if not isinstance(training, dict): + continue + if training.get("ra_dit_stage") != _RETRIEVER_STAGE: + continue + output = cfg.get("output") + if not (isinstance(output, str) and output): + continue + try: + return validate_ra_dit_retriever_model(output) + except (TypeError, ValueError): + # Corrupt Registry row — skip it, keep looking for a clean one. + continue + return None + + +def _close_quietly(store: Any) -> None: + closer = getattr(store, "close", None) + if callable(closer): + try: + closer() + except Exception: # noqa: BLE001 — best-effort cleanup. + pass + + +def resolve_retriever_for_generator( + retriever_model: Optional[str], *, store: Any = None +) -> Tuple[Optional[str], str]: + """Resolve the retriever model for a generator stage. + + A manual ``retriever_model`` always wins. Otherwise the latest RA-DIT + retriever run is discovered from the Registry. Returns + ``(resolved_or_None, advisory_message)``. + """ + if retriever_model is not None: + return ( + retriever_model, + f"Using manual --retriever-model override: {retriever_model}", + ) + discovered = discover_latest_retriever(store=store) + if discovered is not None: + return ( + discovered, + f"Auto-linked latest RA-DIT retriever from Registry: {discovered}", + ) + return ( + None, + "No RA-DIT retriever found in Registry to auto-link; train the " + "retriever stage first or pass --retriever-model.", + ) + + +def autolink_generator_retriever(cfg: Any) -> Optional[str]: + """Auto-link a generator-stage config to the latest retriever run. + + Mutates ``cfg.training.ra_dit_retriever_model`` in place when (a) the + stage is ``generator`` and (b) no retriever model is already set. Returns + an advisory string describing what happened, or ``None`` when no action + was taken (non-generator stage, or a manual value already present). + """ + training = getattr(cfg, "training", None) + if training is None: + return None + if getattr(training, "ra_dit_stage", None) != _GENERATOR_STAGE: + return None + if getattr(training, "ra_dit_retriever_model", None) is not None: + # Manual value present — never overwrite. + return None + resolved, advisory = resolve_retriever_for_generator(None) + if resolved is not None: + training.ra_dit_retriever_model = resolved + return advisory + + +def validate_ra_dit_config_path(name: str, raw: str) -> str: + """Containment-check an operator-supplied RA-DIT config path. + + Public so the ``soup ra-dit`` CLI can reuse it without importing a private + symbol across modules (code-review M5). + """ + from soup_cli.utils.paths import enforce_under_cwd_and_no_symlink # noqa: PLC0415 + + if not isinstance(raw, str): + raise TypeError(f"{name} must be str, got {type(raw).__name__}") + if len(raw) > _MAX_PATH_LEN: + raise ValueError(f"{name} length {len(raw)} exceeds {_MAX_PATH_LEN}") + enforce_under_cwd_and_no_symlink(raw, name) + if not os.path.isfile(os.path.realpath(raw)): + raise FileNotFoundError( + f"{name} not found: {os.path.basename(raw)!r}" + ) + return raw + + +# Back-compat private alias (pre-v0.71.10-review callers). +_validate_config_path = validate_ra_dit_config_path + + +def _validate_timeout(timeout_seconds: object) -> int: + if isinstance(timeout_seconds, bool): + raise ValueError("timeout_seconds must be int, not bool") + if not isinstance(timeout_seconds, int): + raise TypeError( + f"timeout_seconds must be int, got {type(timeout_seconds).__name__}" + ) + if timeout_seconds < _MIN_TIMEOUT_S or timeout_seconds > _MAX_TIMEOUT_S: + raise ValueError( + f"timeout_seconds must be in [{_MIN_TIMEOUT_S}, " + f"{_MAX_TIMEOUT_S}], got {timeout_seconds}" + ) + return timeout_seconds + + +def _load_yaml_config(path: str) -> dict: + """Read + parse a config YAML, re-rejecting a symlink at read time. + + ``_validate_config_path`` already ran ``enforce_under_cwd_and_no_symlink`` + on this path, but the file is opened again here — so use ``O_NOFOLLOW`` + (POSIX) + ``fstat`` to close the TOCTOU window a symlink-swap-after-check + would otherwise open (mirrors ``_eval_v07110._load_jsonl_rows`` policy). + """ + import stat # noqa: PLC0415 + + import yaml # noqa: PLC0415 + + no_follow = getattr(os, "O_NOFOLLOW", 0) + fd = os.open(path, os.O_RDONLY | no_follow) + try: + st = os.fstat(fd) + if stat.S_ISLNK(st.st_mode) or not stat.S_ISREG(st.st_mode): + raise ValueError("config YAML must be a regular file") + if st.st_size > _MAX_YAML_BYTES: + raise ValueError( + f"config YAML exceeds {_MAX_YAML_BYTES // 1024}KB cap " + f"(got {st.st_size} bytes)" + ) + with os.fdopen(fd, "r", encoding="utf-8") as fh: + text = fh.read() + except Exception: + try: + os.close(fd) + except OSError: + pass + raise + raw = yaml.safe_load(text) + if not isinstance(raw, dict): + raise ValueError("config YAML must be a top-level mapping") + return raw + + +def _run_train_subprocess(config_path: str, *, timeout_seconds: int) -> None: + """Default stage runner: ``python -m soup_cli.cli train --config

--yes``. + + List argv, no shell (matches v0.53.5 ``mix_proxy`` / v0.40.4 ``execvp`` + policy). Raises ``RuntimeError`` on non-zero rc or timeout. + """ + argv = [ + sys.executable, + "-m", + "soup_cli.cli", + "train", + "--config", + config_path, + "--yes", + ] + try: + result = subprocess.run( # noqa: S603 — argv list, no shell. + argv, + capture_output=True, + check=False, + timeout=timeout_seconds, + ) + except subprocess.TimeoutExpired as exc: + raise RuntimeError( + f"RA-DIT stage exceeded {timeout_seconds}s timeout" + ) from exc + if result.returncode != 0: + raise RuntimeError( + f"RA-DIT stage failed (rc={result.returncode})" + ) + + +def run_ra_dit( + retriever_config: str, + generator_config: str, + *, + retriever_model: Optional[str] = None, + timeout_seconds: int = 6 * 60 * 60, + _runner: Optional[Callable[[str], None]] = None, +) -> RaDitRunResult: + """Chain the two RA-DIT stages: retriever → auto-link → generator. + + Args: + retriever_config: Path (under cwd) to the stage-1 embedding recipe. + generator_config: Path (under cwd) to the stage-2 RAFT-SFT recipe. + retriever_model: Optional manual override; when set, skips Registry + auto-link and uses this value as the generator's retriever model. + timeout_seconds: Per-stage hard timeout (60s..6h). + _runner: Test seam — a ``(config_path) -> None`` callable that stands + in for the subprocess train call. + + Returns: + :class:`RaDitRunResult` describing both stages' outputs + the link. + + Raises: + TypeError / ValueError / FileNotFoundError: input validation. + RuntimeError: a training stage failed. + """ + retr_path = _validate_config_path("retriever_config", retriever_config) + gen_path = _validate_config_path("generator_config", generator_config) + timeout = _validate_timeout(timeout_seconds) + runner = _runner if _runner is not None else ( + lambda p: _run_train_subprocess(p, timeout_seconds=timeout) + ) + + retr_cfg = _load_yaml_config(retr_path) + gen_cfg = _load_yaml_config(gen_path) + retriever_output = retr_cfg.get("output") + if not isinstance(retriever_output, str) or not retriever_output: + raise ValueError( + "retriever config must declare a non-empty 'output' directory" + ) + generator_output = gen_cfg.get("output") + if not isinstance(generator_output, str) or not generator_output: + raise ValueError( + "generator config must declare a non-empty 'output' directory" + ) + + # Stage 1 — train the retriever. + runner(retr_path) + + # Link: manual override wins; else the retriever's own output dir. + if retriever_model is not None: + model_used = retriever_model + autolinked = False + else: + model_used = retriever_output + autolinked = True + + # Rewrite the generator config to carry the resolved retriever model. + training = gen_cfg.get("training") + if not isinstance(training, dict): + training = {} + gen_cfg["training"] = training + training["ra_dit_retriever_model"] = model_used + + tmp_dir = tempfile.mkdtemp(prefix=".soup_ra_dit.") + try: + import yaml # noqa: PLC0415 + + tmp_gen = os.path.join(tmp_dir, "generator.yaml") + with open(tmp_gen, "w", encoding="utf-8") as fh: + yaml.safe_dump(gen_cfg, fh, sort_keys=False) + # Stage 2 — train the generator with the linked retriever. + runner(tmp_gen) + finally: + shutil.rmtree(tmp_dir, ignore_errors=True) + + return RaDitRunResult( + retriever_output=retriever_output, + generator_output=generator_output, + retriever_model_used=model_used, + autolinked=autolinked, + ) diff --git a/src/soup_cli/utils/raft.py b/src/soup_cli/utils/raft.py new file mode 100644 index 0000000..77fdfe2 --- /dev/null +++ b/src/soup_cli/utils/raft.py @@ -0,0 +1,282 @@ +"""v0.71.10 #199 — live RAFT (Retrieval-Augmented Fine-Tuning) data builders. + +RAFT (Stanford 2024) trains a model to answer a query given a *golden* document +mixed with *distractor* documents, so the model learns to attend to the relevant +context and ignore noise. This module ships the pure-Python composition + answer +-only loss-mask tokenisation consumed by the live RAFT trainer +(:mod:`soup_cli.trainer.raft`). + +Design: + +* ``build_raft_prompt`` combines ``golden_doc`` + ``distractor_docs`` (shuffled + for distractor robustness — reproducibly when ``shuffle_seed`` is set), + labels every document with a deterministic ``[doc-N]`` id, and returns the + prompt + answer + the golden document's id. The ``[doc-N]`` ids let the + answer cite the supporting document verbatim (composes with the v0.62.0 + citation-faithful kernel — #202). +* ``tokenize_raft_example`` builds ``{input_ids, attention_mask, labels, + loss_weights}`` with the prompt span masked to ``-100``. When + ``citation_faithful`` is set it boosts ``loss_weights`` on bracketed citation + spans inside the answer (the span-mask from #199 / #202). + +No heavy imports — everything operates on python lists / strings so it runs +inside a ``datasets.Dataset.map`` without torch. ``return_offsets_mapping`` +needs a fast tokenizer; when unavailable the citation boost degrades to a flat +answer mask (logged at DEBUG, never raises). +""" + +from __future__ import annotations + +import logging +import random +from dataclasses import dataclass +from typing import Any, List, Mapping, Optional, Tuple + +logger = logging.getLogger(__name__) + +# Caps mirror the v0.62.0 ``_convert_raft`` limits. +_MAX_DOCS = 65 # golden + up to 64 distractors +_MAX_FIELD_LEN = 65_536 +_DEFAULT_CITATION_BOOST = 5.0 +_MAX_CITATION_BOOST = 100.0 + +#: Instruction prepended to every RAFT prompt. The bracket-citation hint pairs +#: with the v0.62.0 citation-faithful span mask (#202). +RAFT_INSTRUCTION = ( + "Answer the question using only the relevant document(s) below. " + "Cite the supporting document id in [brackets]." +) + + +@dataclass(frozen=True) +class RaftComposed: + """One composed RAFT example: prompt + answer + the golden doc id.""" + + prompt: str + answer: str + golden_doc_id: str + doc_ids: Tuple[str, ...] + + +def _doc_id(index: int) -> str: + return f"doc-{index}" + + +def _check_str(name: str, value: object) -> str: + if not isinstance(value, str): + raise ValueError(f"RAFT {name!r} must be a string, got {type(value).__name__}") + if not value: + raise ValueError(f"RAFT {name!r} must be non-empty") + if "\x00" in value: + raise ValueError(f"RAFT {name!r} must not contain null bytes") + if len(value) > _MAX_FIELD_LEN: + raise ValueError(f"RAFT {name!r} must be <= {_MAX_FIELD_LEN} chars") + return value + + +def build_raft_prompt( + row: Mapping[str, Any], + *, + shuffle_seed: Optional[int] = None, + row_index: int = 0, +) -> RaftComposed: + """Compose a RAFT row into ``(prompt, answer, golden_doc_id, doc_ids)``. + + ``row`` is a ``{query, golden_doc, distractor_docs, answer}`` mapping (the + output of ``data.formats._convert_raft``). The golden doc and distractors + are concatenated and shuffled deterministically — when ``shuffle_seed`` is + ``None`` a fixed seed (0) is used so the order is still reproducible (the + documents ARE shuffled for distractor robustness; the knob only controls + *which* reproducible permutation). Each document is labelled ``[doc-N]`` in + its shuffled position, and the golden document's id is returned so callers + (citation scorer / span-mask) know the ground-truth citation. + + Raises ``ValueError`` on malformed rows (mirrors ``_convert_raft`` policy). + """ + if not isinstance(row, Mapping): + raise ValueError(f"RAFT row must be a mapping, got {type(row).__name__}") + if isinstance(shuffle_seed, bool) or ( + shuffle_seed is not None and not isinstance(shuffle_seed, int) + ): + raise TypeError("shuffle_seed must be int or None") + if isinstance(row_index, bool) or not isinstance(row_index, int) or row_index < 0: + raise ValueError("row_index must be a non-negative int") + + query = _check_str("query", row.get("query")) + golden_doc = _check_str("golden_doc", row.get("golden_doc")) + answer = _check_str("answer", row.get("answer")) + raw_distractors = row.get("distractor_docs", []) or [] + if not isinstance(raw_distractors, list): + raise ValueError("RAFT 'distractor_docs' must be a list") + distractors = [ + _check_str(f"distractor_docs[{i}]", d) for i, d in enumerate(raw_distractors) + ] + if 1 + len(distractors) > _MAX_DOCS: + raise ValueError(f"RAFT example has > {_MAX_DOCS} documents") + + # (text, is_golden) pairs; shuffle deterministically. + docs: List[Tuple[str, bool]] = [(golden_doc, True)] + docs.extend((d, False) for d in distractors) + seed_base = 0 if shuffle_seed is None else int(shuffle_seed) + rng = random.Random(seed_base * 1_000_003 + row_index) + rng.shuffle(docs) + + doc_lines: List[str] = [] + doc_ids: List[str] = [] + golden_id = _doc_id(0) + for index, (text, is_golden) in enumerate(docs): + did = _doc_id(index) + doc_ids.append(did) + if is_golden: + golden_id = did + doc_lines.append(f"[{did}] {text}") + + prompt = ( + f"{RAFT_INSTRUCTION}\n\n" + f"Question: {query}\n\n" + f"Documents:\n" + "\n".join(doc_lines) + "\n\nAnswer:" + ) + return RaftComposed( + prompt=prompt, + answer=answer, + golden_doc_id=golden_id, + doc_ids=tuple(doc_ids), + ) + + +def citation_span_token_weights( + answer: str, + offsets: List[Tuple[int, int]], + *, + style: str = "bracket", + boost: float = _DEFAULT_CITATION_BOOST, +) -> List[float]: + """Per-token loss weights for ``answer`` tokens, boosting citation spans. + + ``offsets`` is the per-token ``(start_char, end_char)`` mapping for the + answer (from a fast tokenizer's ``return_offsets_mapping``). Tokens whose + char span overlaps any ``[doc-id]`` (per ``style``) citation get ``boost``; + every other token gets ``1.0``. Returns one weight per offset entry. + """ + if isinstance(boost, bool) or not isinstance(boost, (int, float)): + raise TypeError("boost must be a number") + boost_f = float(boost) + if boost_f < 1.0 or boost_f > _MAX_CITATION_BOOST: + raise ValueError(f"boost must be in [1.0, {_MAX_CITATION_BOOST}]") + + # Local import keeps utils/raft import-light + avoids a circular import. + from soup_cli.utils.citation_faithful import citation_spans + + spans = citation_spans(answer, style=style) + weights: List[float] = [] + for start, end in offsets: + # A zero-width offset (special token) never overlaps a citation span. + overlaps = any(start < s_end and end > s_start for (s_start, s_end) in spans) + weights.append(boost_f if overlaps else 1.0) + return weights + + +def tokenize_raft_example( + tokenizer: Any, + composed: RaftComposed, + *, + max_length: int, + citation_faithful: bool = False, + citation_style: str = "bracket", + citation_boost: float = _DEFAULT_CITATION_BOOST, +) -> dict: + """Tokenise a composed RAFT example into a pre-tokenised training row. + + Returns ``{input_ids, attention_mask, labels, loss_weights}`` (python + lists). The prompt span is masked with ``-100`` in ``labels`` and ``0.0`` + in ``loss_weights`` so only the answer contributes to the loss. When + ``citation_faithful`` is set, ``loss_weights`` on bracketed citation spans + inside the answer are boosted to ``citation_boost``. + """ + if isinstance(max_length, bool) or not isinstance(max_length, int) or max_length < 8: + raise ValueError("max_length must be an int >= 8") + if not isinstance(composed, RaftComposed): + raise TypeError("composed must be a RaftComposed") + + prompt_text = _render_prompt(tokenizer, composed.prompt) + prompt_ids = tokenizer(prompt_text, add_special_tokens=False)["input_ids"] + + answer_ids: List[int] + answer_weights: List[float] + if citation_faithful: + answer_ids, answer_weights = _tokenize_answer_with_citation( + tokenizer, composed.answer, style=citation_style, boost=citation_boost + ) + else: + answer_ids = tokenizer(composed.answer, add_special_tokens=False)["input_ids"] + answer_weights = [1.0] * len(answer_ids) + + eos = getattr(tokenizer, "eos_token_id", None) + if eos is not None: + answer_ids = answer_ids + [eos] + answer_weights = answer_weights + [1.0] + + input_ids = (prompt_ids + answer_ids)[:max_length] + labels = ([-100] * len(prompt_ids) + answer_ids)[:max_length] + loss_weights = ([0.0] * len(prompt_ids) + answer_weights)[:max_length] + attention_mask = [1] * len(input_ids) + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels": labels, + "loss_weights": loss_weights, + } + + +def _tokenize_answer_with_citation( + tokenizer: Any, answer: str, *, style: str, boost: float +) -> Tuple[List[int], List[float]]: + """Tokenise the answer and compute per-token citation-boost weights. + + Uses ``return_offsets_mapping`` (fast tokenizer). When offsets are + unavailable (slow tokenizer), degrades to a flat ``1.0`` answer mask and + logs at DEBUG — never raises. + """ + try: + enc = tokenizer( + answer, add_special_tokens=False, return_offsets_mapping=True + ) + answer_ids = list(enc["input_ids"]) + offsets = enc.get("offset_mapping") + except (NotImplementedError, ValueError, TypeError): + offsets = None + answer_ids = list(tokenizer(answer, add_special_tokens=False)["input_ids"]) + if not offsets or len(offsets) != len(answer_ids): + logger.debug( + "citation span boost skipped: tokenizer offsets unavailable " + "(slow tokenizer?); using flat answer mask" + ) + return answer_ids, [1.0] * len(answer_ids) + weights = citation_span_token_weights( + answer, [tuple(o) for o in offsets], style=style, boost=boost + ) + return answer_ids, weights + + +def render_raft_prompt(tokenizer: Any, prompt: str) -> str: + """Render a single user turn through the tokenizer's chat template. + + Falls back to the raw prompt when the tokenizer has no chat template or the + template raises. Public so other modules (e.g. ``steering._capture_attn_heads``) + can reuse it without importing a private symbol across modules (code-review L3). + """ + chat_template = getattr(tokenizer, "chat_template", None) + if chat_template: + try: + return tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=False, + add_generation_prompt=True, + ) + except Exception: # noqa: BLE001 — malformed template → raw prompt + return prompt + return prompt + + +# Back-compat private alias (internal callers + pre-review imports). +_render_prompt = render_raft_prompt diff --git a/src/soup_cli/utils/steering.py b/src/soup_cli/utils/steering.py index 9275a3c..bf02332 100644 --- a/src/soup_cli/utils/steering.py +++ b/src/soup_cli/utils/steering.py @@ -20,7 +20,7 @@ import math import re from dataclasses import dataclass from types import MappingProxyType -from typing import Mapping, Optional +from typing import Any, Mapping, Optional, Tuple SUPPORTED_STEERING_METHODS: frozenset[str] = frozenset({"caa", "iti", "repe"}) @@ -177,21 +177,225 @@ def get_steering_method_spec(name: str) -> SteeringMethodSpec: return _STEERING_METHOD_METADATA[canonical] -def apply_steering(method: str) -> None: - """Apply a steering vector during decoding — deferred to v0.62.1. +# --- v0.71.10 #201 — live CAA / ITI / RepE fitting + decode hook ---------- - Validates the method name first so the deferred-live error - distinguishes between "unknown method" and "method is on the - allowlist but not yet wired". Mirrors v0.50.0 ``apply_variant_loss`` / - v0.61.0 ``apply_unlearn_loss`` policy. +_MAX_PAIRS = 2000 # cap on contrastive-pair count (DoS guard on the JSONL). +_DEFAULT_TOP_K_HEADS = 8 # ITI: top-K attention heads to intervene on. +_MAX_FILE_BYTES = 256 * 1024 * 1024 # 256 MiB cap on the pairs JSONL. +_MAX_PAIR_FIELD_LEN = 65_536 # per-field cap (parity with raft._check_str). +_MIN_PAIRS_FOR_PCA = 2 # repe / iti need >= 2 pairs (SVD degenerate at N=1). +_CONFIG_NAME = "steering_config.json" +_VECTOR_NAME = "steering_vector.safetensors" + + +@dataclass(frozen=True) +class SteeringArtifact: + """On-disk result of ``build_steering_vector`` (frozen).""" + + method: str + name: str + layer: int + hidden_dim: int + intervention_point: str # "residual" | "attn_o_proj_input" + output_dir: str + base: str + num_pairs: int + + +@dataclass(frozen=True) +class LoadedSteering: + """A steering vector loaded from disk, ready for the decode hook.""" + + method: str + name: str + layer: int + intervention_point: str + vector: Any # numpy float32 [D] + default_strength: float + + +def compute_caa_vector(positive: Any, negative: Any) -> Any: + """CAA control vector = ``mean(positive) - mean(negative)`` (Panickssery 2023). + + ``positive`` / ``negative`` are ``[N, D]`` activation matrices (one mean- + pooled residual-stream vector per prompt). The raw mean-difference is + returned (NOT unit-normalised) so its magnitude reflects one "contrast + unit" — the decode-time ``strength`` multiplier scales it. Returns a + float32 ``[D]`` numpy vector. """ - canonical = validate_steering_method(method) - raise NotImplementedError( - f"apply_steering({canonical!r}) is deferred to v0.62.1. " - "Schema accepts the method now so callers can write soup.yaml " - "today, but the live forward-hook + decode-time intervention " - "land in v0.62.1." - ) + import numpy as np + + pos = _steer_as_2d(positive, "positive") + neg = _steer_as_2d(negative, "negative") + if pos.shape[1] != neg.shape[1]: + raise ValueError( + f"hidden-dim mismatch: positive[{pos.shape[1]}] vs negative[{neg.shape[1]}]" + ) + vec = (pos.mean(axis=0) - neg.mean(axis=0)).astype(np.float32) + if not np.all(np.isfinite(vec)): + raise ValueError("CAA vector is not finite") + return vec + + +def compute_repe_direction(diffs: Any) -> Any: + """RepE behavioural direction (Zou 2023) — top PCA component of the diffs. + + ``diffs`` is the ``[N, D]`` per-example ``positive - negative`` matrix. The + top principal component is extracted via SVD, sign-aligned so the mean diff + projects positively, and scaled by that mean projection so the magnitude is + comparable to the CAA vector. Returns float32 ``[D]``. + """ + import numpy as np + + d = _steer_as_2d(diffs, "diffs") + centered = d - d.mean(axis=0, keepdims=True) + # full_matrices=False keeps the SVD cheap; Vt rows are the right-singular + # vectors (principal directions in feature space). + _, _, vt = np.linalg.svd(centered, full_matrices=False) + pc = vt[0] + mean_diff = d.mean(axis=0) + scale = float(mean_diff @ pc) + if scale < 0: + pc = -pc + scale = -scale + vec = (pc * scale).astype(np.float32) + if not np.all(np.isfinite(vec)): + raise ValueError("RepE direction is not finite") + return vec + + +def compute_iti_directions( + pos_heads: Any, neg_heads: Any, *, top_k: int = _DEFAULT_TOP_K_HEADS +) -> Tuple[Any, Tuple[int, ...]]: + """ITI per-head intervention directions (Li 2023). + + ``pos_heads`` / ``neg_heads`` are ``[N, H, Dh]`` per-example, per-head + activations (the input to ``o_proj`` reshaped into heads). The per-head + mean difference is computed, heads are ranked by the L2 norm of their diff, + and the top ``top_k`` heads keep their direction (others zeroed). Returns + ``(directions[H, Dh] float32, selected_heads tuple)``. + """ + import numpy as np + + pos = _steer_as_3d(pos_heads, "pos_heads") + neg = _steer_as_3d(neg_heads, "neg_heads") + if pos.shape[1:] != neg.shape[1:]: + raise ValueError( + f"head-shape mismatch: pos{pos.shape[1:]} vs neg{neg.shape[1:]}" + ) + if isinstance(top_k, bool) or not isinstance(top_k, int) or top_k < 1: + raise ValueError("top_k must be a positive int") + num_heads = pos.shape[1] + diff = pos.mean(axis=0) - neg.mean(axis=0) # [H, Dh] + head_scores = np.linalg.norm(diff, axis=1) # [H] + k = min(top_k, num_heads) + selected = tuple(sorted(int(i) for i in np.argsort(head_scores)[-k:])) + directions = np.zeros_like(diff, dtype=np.float32) + for head in selected: + directions[head] = diff[head] + if not np.all(np.isfinite(directions)): + raise ValueError("ITI directions are not finite") + return directions, selected + + +def _steer_as_2d(value: Any, field: str) -> Any: + import numpy as np + + arr = np.asarray(value, dtype=np.float64) + if arr.ndim != 2 or arr.shape[0] == 0 or arr.shape[1] == 0: + raise ValueError(f"{field} must be a non-empty 2D array") + if not np.all(np.isfinite(arr)): + raise ValueError(f"{field} must be finite") + return arr + + +def _steer_as_3d(value: Any, field: str) -> Any: + import numpy as np + + arr = np.asarray(value, dtype=np.float64) + if arr.ndim != 3 or arr.shape[0] == 0: + raise ValueError(f"{field} must be a non-empty 3D [N, H, Dh] array") + if not np.all(np.isfinite(arr)): + raise ValueError(f"{field} must be finite") + return arr + + +def load_contrastive_pairs(path: str) -> list: + """Load ``[(positive, negative), ...]`` from a JSONL of ``{positive, negative}``. + + cwd-contained + symlink-rejected (via the shared + ``enforce_under_cwd_and_no_symlink`` helper) + ``O_NOFOLLOW`` open (closes + the TOCTOU window the inline ``lexists``-gated lstat left) + size-capped. + Skips malformed / incomplete rows; per-field cap ``_MAX_PAIR_FIELD_LEN``. + Capped at ``_MAX_PAIRS``. + """ + import json + import os + import stat + + from soup_cli.utils.paths import enforce_under_cwd_and_no_symlink + + enforce_under_cwd_and_no_symlink(path, "pairs_path") + no_follow = getattr(os, "O_NOFOLLOW", 0) + try: + fd = os.open(path, os.O_RDONLY | no_follow) + except FileNotFoundError as exc: + raise FileNotFoundError(f"pairs file not found: {path!r}") from exc + except OSError as exc: + raise ValueError( + f"pairs file cannot be opened (symlink?): {type(exc).__name__}" + ) from exc + pairs: list = [] + try: + st = os.fstat(fd) + if stat.S_ISLNK(st.st_mode) or not stat.S_ISREG(st.st_mode): + raise ValueError("pairs_path must be a regular file") + if st.st_size > _MAX_FILE_BYTES: + raise ValueError(f"pairs file exceeds {_MAX_FILE_BYTES} bytes") + with os.fdopen(fd, "r", encoding="utf-8-sig") as fh: + for line in fh: + if len(pairs) >= _MAX_PAIRS: + break + line = line.strip() + if not line: + continue + try: + row = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(row, dict): + continue + pos = row.get("positive") + neg = row.get("negative") + if ( + isinstance(pos, str) and isinstance(neg, str) + and pos and neg + and len(pos) <= _MAX_PAIR_FIELD_LEN + and len(neg) <= _MAX_PAIR_FIELD_LEN + ): + pairs.append((pos, neg)) + except Exception: + try: + os.close(fd) + except OSError: + pass + raise + if not pairs: + raise ValueError(f"pairs file {path!r} yielded no usable (positive, negative) rows") + return pairs + + +def _validated_steer_output_dir(output_dir: str) -> str: + """Reject an output dir outside cwd or pointing at a symlink (TOCTOU). + + Delegates to the shared ``enforce_under_cwd_and_no_symlink`` helper for the + containment + leaf-symlink rejection (SEC LOW-1 — was an inline + ``lexists``-gated lstat). + """ + from soup_cli.utils.paths import enforce_under_cwd_and_no_symlink + + enforce_under_cwd_and_no_symlink(output_dir, "output dir") + return output_dir def build_steering_vector( @@ -199,36 +403,318 @@ def build_steering_vector( method: str, name: str, pairs_path: Optional[str] = None, + base: Optional[str] = None, layer: Optional[int] = None, -) -> None: - """Train a steering vector from contrastive pairs — deferred to v0.62.1. + device: Optional[str] = None, + output_dir: Optional[str] = None, + top_k: int = _DEFAULT_TOP_K_HEADS, + trust_remote_code: bool = False, +) -> SteeringArtifact: + """Fit a steering vector from contrastive pairs and persist it. - Validates inputs first; the deferred-live error fires only after the - method + name + (optional) layer all pass shape/range checks. + Validates method + name FIRST (so a typo surfaces before the model load), + then requires ``base`` + ``pairs_path``. Loads the model, captures residual + -stream (CAA / RepE) or per-head (ITI) activations on the contrastive + prompts, computes the control vector, and writes + ``/steering_vector.safetensors`` + ``steering_config.json``. + + Returns a frozen :class:`SteeringArtifact`. Heavy imports are local. """ + import json + import os + + import numpy as np + canonical = validate_steering_method(method) canonical_name = validate_steering_name(name) - if pairs_path is not None: - if not isinstance(pairs_path, str): - raise TypeError( - f"pairs_path must be str, got {type(pairs_path).__name__}" - ) - if not pairs_path: - raise ValueError("pairs_path must be non-empty") - if "\x00" in pairs_path: - raise ValueError("pairs_path must not contain null bytes") + if not base or not isinstance(base, str): + raise ValueError("build_steering_vector requires a non-empty base model") + if not pairs_path or not isinstance(pairs_path, str): + raise ValueError("build_steering_vector requires a pairs_path (JSONL)") if layer is not None: - if isinstance(layer, bool): - raise TypeError("layer must not be bool") - if not isinstance(layer, int): - raise TypeError( - f"layer must be int, got {type(layer).__name__}" - ) + if isinstance(layer, bool) or not isinstance(layer, int): + raise TypeError("layer must be int or None") if layer < 0 or layer > 2048: - raise ValueError( - f"layer must satisfy 0 <= layer <= 2048, got {layer}" - ) - raise NotImplementedError( - f"build_steering_vector(method={canonical!r}, " - f"name={canonical_name!r}) is deferred to v0.62.1." + raise ValueError(f"layer must satisfy 0 <= layer <= 2048, got {layer}") + if isinstance(top_k, bool) or not isinstance(top_k, int) or top_k < 1: + raise ValueError("top_k must be a positive int") + + pairs = load_contrastive_pairs(pairs_path) + # repe (PCA) and iti (per-head ranking) need >= 2 pairs — a single pair + # makes the centered SVD degenerate (code-review L5). CAA (mean-diff) is + # well-defined at N=1. + if canonical in ("repe", "iti") and len(pairs) < _MIN_PAIRS_FOR_PCA: + raise ValueError( + f"method {canonical!r} requires >= {_MIN_PAIRS_FOR_PCA} contrastive " + f"pairs (got {len(pairs)}); use 'caa' for a single pair" + ) + pos_prompts = [p for p, _ in pairs] + neg_prompts = [n for _, n in pairs] + + from soup_cli.utils.edit_kernels import _locate_decoder_layers + from soup_cli.utils.live_eval import ( + extract_layer_activations, + load_model_and_tokenizer, ) + + model, tokenizer, dev = load_model_and_tokenizer( + base, device=device, trust_remote_code=trust_remote_code + ) + layers = _locate_decoder_layers(model) + n_layers = len(layers) + resolved_layer = n_layers // 2 if layer is None else layer + if resolved_layer < 0 or resolved_layer >= n_layers: + raise ValueError( + f"layer {resolved_layer} out of range for a {n_layers}-layer model" + ) + + extra: dict = {} + if canonical in ("caa", "repe"): + layer_path = f"model.layers.{resolved_layer}" + pos = extract_layer_activations( + model, tokenizer, pos_prompts, layer=layer_path, device=dev, pool="mean" + ) + neg = extract_layer_activations( + model, tokenizer, neg_prompts, layer=layer_path, device=dev, pool="mean" + ) + if canonical == "caa": + vector = compute_caa_vector(pos, neg) + else: + vector = compute_repe_direction(np.asarray(pos) - np.asarray(neg)) + intervention = "residual" + else: # iti + num_heads = int(getattr(model.config, "num_attention_heads", 0)) or 1 + pos_h = _capture_attn_heads( + model, tokenizer, pos_prompts, resolved_layer, num_heads, dev + ) + neg_h = _capture_attn_heads( + model, tokenizer, neg_prompts, resolved_layer, num_heads, dev + ) + dirs, selected = compute_iti_directions(pos_h, neg_h, top_k=top_k) + vector = dirs.reshape(-1).astype(np.float32) + intervention = "attn_o_proj_input" + extra = { + "num_heads": num_heads, + "head_dim": int(dirs.shape[1]), + "selected_heads": list(selected), + } + + hidden_dim = int(vector.shape[0]) + out_dir = output_dir or os.path.join("steering", canonical_name) + out_dir = _validated_steer_output_dir(out_dir) + os.makedirs(out_dir, exist_ok=True) + + from safetensors.numpy import save_file + + save_file( + {"vector": np.ascontiguousarray(vector, dtype=np.float32)}, + os.path.join(out_dir, _VECTOR_NAME), + ) + config = { + "method": canonical, + "name": canonical_name, + "layer": resolved_layer, + "hidden_dim": hidden_dim, + "intervention_point": intervention, + "base": base, + "num_pairs": len(pairs), + "default_strength": 1.0, + **extra, + } + with open(os.path.join(out_dir, _CONFIG_NAME), "w", encoding="utf-8") as fh: + json.dump(config, fh, indent=2, allow_nan=False) + + return SteeringArtifact( + method=canonical, + name=canonical_name, + layer=resolved_layer, + hidden_dim=hidden_dim, + intervention_point=intervention, + output_dir=out_dir, + base=base, + num_pairs=len(pairs), + ) + + +def _capture_attn_heads(model, tokenizer, prompts, layer_idx, num_heads, device): + """Capture per-head ``o_proj``-input activations: ``[N, H, Dh]`` (ITI).""" + import numpy as np + import torch + + from soup_cli.utils.edit_kernels import _locate_decoder_layers + from soup_cli.utils.raft import render_raft_prompt + + layers = _locate_decoder_layers(model) + attn = getattr(layers[layer_idx], "self_attn", None) + o_proj = getattr(attn, "o_proj", None) if attn is not None else None + if o_proj is None or not hasattr(o_proj, "weight"): + raise ValueError( + f"layer {layer_idx} has no self_attn.o_proj (unsupported arch for ITI)" + ) + captured: list = [] + + def _pre_hook(_mod, args): + # args[0]: [batch, seq, in_features] + captured.append(args[0][0].detach().to(torch.float32).mean(dim=0).cpu()) + + handle = o_proj.register_forward_pre_hook(_pre_hook) + rows: list = [] + try: + model.eval() + with torch.no_grad(): + for prompt in prompts: + text = render_raft_prompt(tokenizer, prompt) + inputs = tokenizer( + text, return_tensors="pt", truncation=True, max_length=1024 + ).to(device) + captured.clear() + model(**inputs) + if captured: + rows.append(captured[-1].numpy()) + finally: + handle.remove() + if not rows: + raise ValueError("no attention activations captured for ITI") + flat = np.stack(rows).astype(np.float32) # [N, in_features] + in_features = flat.shape[1] + head_dim = in_features // num_heads + if head_dim * num_heads != in_features: + raise ValueError( + f"o_proj in_features {in_features} not divisible by num_heads {num_heads}" + ) + return flat.reshape(flat.shape[0], num_heads, head_dim) + + +def resolve_steering_dir(name: str) -> str: + """Resolve a steering-vector NAME to its on-disk directory. + + Resolution order: (1) the default ``./steering/`` directory if it + holds a ``steering_config.json``; (2) the most recent Registry entry named + ```` carrying a ``steering_vector`` artifact. Raises ``ValueError`` + when neither resolves. + """ + import os + + from soup_cli.utils.paths import is_under_cwd + + canonical = validate_steering_name(name) + local = os.path.join("steering", canonical) + if os.path.isfile(os.path.join(local, _CONFIG_NAME)): + return local + try: + from soup_cli.registry.store import RegistryStore + + with RegistryStore() as store: + for entry in store.list(): + if str(entry.get("name", "")) != canonical: + continue + for art in store.get_artifacts(entry["id"]): + if art.get("kind") == "steering_vector": + art_path = str(art.get("path", "")) + # Only trust a Registry-supplied path that is still + # under cwd (a shared/copied Registry DB could hold an + # absolute out-of-tree path) — SEC LOW-2. + if ( + art_path + and is_under_cwd(art_path) + and os.path.isfile(os.path.join(art_path, _CONFIG_NAME)) + ): + return art_path + except (ImportError, OSError, ValueError): + pass + raise ValueError( + f"no steering vector named {canonical!r} (looked under ./steering/ " + "and the Registry). Train one with `soup steer train`." + ) + + +def load_steering_artifact(dir_path: str) -> LoadedSteering: + """Load a steering vector + config from ``dir_path`` into a :class:`LoadedSteering`.""" + import json + import os + + import numpy as np + from safetensors import safe_open + + from soup_cli.utils.paths import enforce_under_cwd_and_no_symlink + + # The dir (and its two leaf files) may come from a Registry row — contain + # them under cwd + reject symlinks before reading (SEC LOW-2). + enforce_under_cwd_and_no_symlink(dir_path, "steering dir") + config_path = os.path.join(dir_path, _CONFIG_NAME) + vector_path = os.path.join(dir_path, _VECTOR_NAME) + if not os.path.isfile(config_path) or not os.path.isfile(vector_path): + raise ValueError(f"steering dir {dir_path!r} missing config / vector file") + enforce_under_cwd_and_no_symlink(config_path, "steering config") + enforce_under_cwd_and_no_symlink(vector_path, "steering vector") + with open(config_path, encoding="utf-8") as fh: + config = json.load(fh) + method = validate_steering_method(config.get("method", "caa")) + name = validate_steering_name(config.get("name", "steer")) + layer = int(config.get("layer", 0)) + intervention = config.get("intervention_point", "residual") + if intervention not in ("residual", "attn_o_proj_input"): + raise ValueError(f"unknown intervention_point {intervention!r}") + with safe_open(vector_path, framework="numpy") as handle: + if "vector" not in handle.keys(): + raise ValueError("steering safetensors missing 'vector'") + vector = np.asarray(handle.get_tensor("vector"), dtype=np.float32) + if vector.ndim != 1 or not np.all(np.isfinite(vector)): + raise ValueError("steering vector must be a finite 1D array") + default_strength = float(config.get("default_strength", 1.0)) + return LoadedSteering( + method=method, + name=name, + layer=layer, + intervention_point=intervention, + vector=vector, + default_strength=default_strength, + ) + + +def install_steering_hook(model: Any, loaded: LoadedSteering, *, strength: float): + """Install a decode-time forward hook that adds ``strength * vector``. + + * ``residual`` (CAA / RepE) — a forward hook on the decoder layer adds the + vector to the residual-stream output. + * ``attn_o_proj_input`` (ITI) — a forward pre-hook on ``self_attn.o_proj`` + shifts the per-head activations before the output projection. + + Returns the hook handle so the caller can ``.remove()`` it on shutdown. + """ + import torch + + if not isinstance(loaded, LoadedSteering): + raise TypeError("loaded must be a LoadedSteering") + strength_f = validate_steering_strength(strength) + + from soup_cli.utils.edit_kernels import _locate_decoder_layers + + layers = _locate_decoder_layers(model) + if loaded.layer < 0 or loaded.layer >= len(layers): + raise ValueError( + f"steering layer {loaded.layer} out of range for " + f"{len(layers)}-layer model" + ) + block = layers[loaded.layer] + param = next(model.parameters()) + vec = torch.tensor(loaded.vector, dtype=param.dtype, device=param.device) + + if loaded.intervention_point == "residual": + def _hook(_mod, _args, output): + if isinstance(output, (tuple, list)): + hidden = output[0] + strength_f * vec + return (hidden,) + tuple(output[1:]) + return output + strength_f * vec + + return block.register_forward_hook(_hook) + + o_proj = getattr(getattr(block, "self_attn", None), "o_proj", None) + if o_proj is None: + raise ValueError("ITI steering requires self_attn.o_proj on the layer") + + def _pre_hook(_mod, args): + shifted = args[0] + strength_f * vec + return (shifted,) + tuple(args[1:]) + + return o_proj.register_forward_pre_hook(_pre_hook) diff --git a/tests/test_v0560.py b/tests/test_v0560.py index bd5bd81..9c88c5c 100644 --- a/tests/test_v0560.py +++ b/tests/test_v0560.py @@ -752,7 +752,8 @@ class TestSourceWiring: assert "_diagnose_cmd" in source def test_no_top_level_heavy_imports(self) -> None: - # All 6 probe + report modules must be torch/transformers-free. + # All probe + report modules must be torch/transformers-free + # (citation.py added in v0.71.10 #202). probe_modules = [ "src/soup_cli/utils/diagnose/__init__.py", "src/soup_cli/utils/diagnose/report.py", @@ -763,6 +764,7 @@ class TestSourceWiring: "src/soup_cli/utils/diagnose/mode_collapse.py", "src/soup_cli/utils/diagnose/memorization.py", "src/soup_cli/utils/diagnose/contamination.py", + "src/soup_cli/utils/diagnose/citation.py", "src/soup_cli/utils/diagnose/badge.py", "src/soup_cli/utils/diagnose/runner.py", ] diff --git a/tests/test_v0620_followups.py b/tests/test_v0620_followups.py index 4ae55d2..1e34be1 100644 --- a/tests/test_v0620_followups.py +++ b/tests/test_v0620_followups.py @@ -195,11 +195,14 @@ class TestRaftNullByteRejection: }) -# ---------- H5 — `soup steer train --plan-only` deferred-marker assertion ---------- +# ---------- `soup steer train --plan-only` (live as of v0.71.10 #201) ---------- class TestSteerPlanOnlyMarker: - def test_plan_only_emits_v0621_marker(self, tmp_path, monkeypatch): + def test_plan_only_validates_and_renders(self, tmp_path, monkeypatch): + # v0.71.10 #201 — live fitting shipped; --plan-only now just validates + # inputs + renders the plan (no deferred-version marker) and exits 0 + # WITHOUT loading a model. from soup_cli.commands.steer import app monkeypatch.chdir(tmp_path) @@ -219,9 +222,9 @@ class TestSteerPlanOnlyMarker: assert result.exit_code == 0, ( result.output, repr(result.exception) ) - # Defence: the panel MUST surface the deferred-version marker so - # operators know live training ships in v0.62.1. - assert "v0.62.1" in result.output + assert "Plan-only" in result.output + # Live fitting shipped — the deferred-version marker is gone. + assert "v0.62.1" not in result.output # ---------- H6 — apply_edit v0.61.1 marker not regressed for legacy methods ---------- diff --git a/tests/test_v0620_part_c.py b/tests/test_v0620_part_c.py index 3c63212..8369d58 100644 --- a/tests/test_v0620_part_c.py +++ b/tests/test_v0620_part_c.py @@ -18,14 +18,14 @@ class TestModuleSurface: from soup_cli.utils.steering import ( SUPPORTED_STEERING_METHODS, SteeringMethodSpec, - apply_steering, build_steering_vector, get_steering_method_spec, + install_steering_hook, # v0.71.10 #201 — replaces apply_steering stub validate_steering_method, ) assert callable(validate_steering_method) assert callable(get_steering_method_spec) - assert callable(apply_steering) + assert callable(install_steering_hook) assert callable(build_steering_vector) assert dataclasses.is_dataclass(SteeringMethodSpec) assert isinstance(SUPPORTED_STEERING_METHODS, frozenset) @@ -230,23 +230,19 @@ class TestSpec: class TestDeferredStubs: - def test_apply_steering_deferred(self): - from soup_cli.utils.steering import apply_steering + def test_apply_steering_removed(self): + # v0.71.10 #201 — the apply_steering stub is replaced by the live + # install_steering_hook runtime. + import soup_cli.utils.steering as steering - with pytest.raises(NotImplementedError, match="v0.62.1"): - apply_steering("caa") + assert not hasattr(steering, "apply_steering") - def test_apply_steering_validates_first(self): - from soup_cli.utils.steering import apply_steering - - # Unknown method rejected BEFORE the deferred-live raise. - with pytest.raises(ValueError): - apply_steering("nonsense") - - def test_build_steering_vector_deferred(self): + def test_build_steering_vector_requires_base(self): + # v0.71.10 #201 — no longer NotImplementedError; missing base/pairs is + # a loud ValueError (validation happens after method+name checks). from soup_cli.utils.steering import build_steering_vector - with pytest.raises(NotImplementedError, match="v0.62.1"): + with pytest.raises(ValueError, match="base"): build_steering_vector(method="caa", name="safety-v1") def test_build_steering_vector_validates_method_first(self): diff --git a/tests/test_v07110.py b/tests/test_v07110.py new file mode 100644 index 0000000..e4766c9 --- /dev/null +++ b/tests/test_v07110.py @@ -0,0 +1,1914 @@ +"""v0.71.10 "RAG family" — RAFT / RA-DIT / steering / citation live wiring. + +Closes #199 (RAFT span-mask trainer), #200 (RA-DIT auto-link), #201 (live +CAA/ITI/RepE steering + serve --steer decode hook), #202 (citation-span +loss-mask + soup eval citation + diagnose citation mode). + +Pure-Python / CPU tests + tiny-tensor torch tests; the live model paths are +step-6 smoked on SmolLM2-135M (RTX 3050). +""" + +from __future__ import annotations + +import re +from typing import List, Tuple + +import pytest + +# --------------------------------------------------------------------------- +# Shared fakes +# --------------------------------------------------------------------------- + + +class _FakeTokenizer: + """Word-level fake tokenizer with offset-mapping support. + + Tokenises on non-whitespace runs; each token gets a deterministic id and + its ``(start, end)`` char offset. Supports the subset of the HF tokenizer + API that ``utils.raft`` uses. + """ + + eos_token_id = 99 + pad_token_id = 0 + chat_template = None + + _WORD = re.compile(r"\S+") + + def __init__(self, fast: bool = True): + self._fast = fast + self._vocab: dict[str, int] = {} + + def _id(self, tok: str) -> int: + if tok not in self._vocab: + self._vocab[tok] = len(self._vocab) + 1 # 1-based; 0 = pad + return self._vocab[tok] + + def __call__(self, text, add_special_tokens=False, return_offsets_mapping=False): + toks = list(self._WORD.finditer(text)) + ids = [self._id(m.group(0)) for m in toks] + out = {"input_ids": ids} + if return_offsets_mapping: + if not self._fast: + raise NotImplementedError("slow tokenizer has no offsets") + out["offset_mapping"] = [(m.start(), m.end()) for m in toks] + return out + + +class _FakeBaseTrainer: + """Minimal stand-in for transformers.Trainer (compute_loss test).""" + + def __init__(self, *args, **kwargs): + pass + + +class _FakeOutputs: + def __init__(self, logits): + self.logits = logits + + +class _FakeModel: + """Returns fixed logits regardless of input (compute_loss test).""" + + def __init__(self, logits): + self._logits = logits + + def __call__(self, **kwargs): + return _FakeOutputs(self._logits) + + +def _raft_row(distractors: int = 2) -> dict: + return { + "query": "What is the capital of France?", + "golden_doc": "Paris has been the capital of France since 987.", + "distractor_docs": [f"Distractor number {i} about geese." for i in range(distractors)], + "answer": "The capital is Paris [doc-0].", + } + + +# --------------------------------------------------------------------------- +# #199 — build_raft_prompt +# --------------------------------------------------------------------------- + + +class TestBuildRaftPrompt: + def test_basic_prompt_shape(self): + from soup_cli.utils.raft import RAFT_INSTRUCTION, build_raft_prompt + + composed = build_raft_prompt(_raft_row(2)) + assert RAFT_INSTRUCTION in composed.prompt + assert "Question: What is the capital of France?" in composed.prompt + assert "Documents:" in composed.prompt + assert composed.prompt.rstrip().endswith("Answer:") + assert composed.answer == "The capital is Paris [doc-0]." + + def test_doc_ids_assigned_and_golden_tracked(self): + from soup_cli.utils.raft import build_raft_prompt + + composed = build_raft_prompt(_raft_row(2)) + # 1 golden + 2 distractors → doc-0, doc-1, doc-2. + assert composed.doc_ids == ("doc-0", "doc-1", "doc-2") + assert composed.golden_doc_id in composed.doc_ids + # The golden doc text must appear next to its assigned id. + assert f"[{composed.golden_doc_id}] Paris has been the capital" in composed.prompt + + def test_shuffle_reproducible_for_same_seed(self): + from soup_cli.utils.raft import build_raft_prompt + + a = build_raft_prompt(_raft_row(4), shuffle_seed=42, row_index=3) + b = build_raft_prompt(_raft_row(4), shuffle_seed=42, row_index=3) + assert a.prompt == b.prompt + assert a.golden_doc_id == b.golden_doc_id + + def test_different_row_index_differs(self): + from soup_cli.utils.raft import build_raft_prompt + + # Across a handful of indices the golden position should vary at least + # once (deterministic shuffle keyed on index). + positions = { + build_raft_prompt(_raft_row(6), shuffle_seed=1, row_index=i).golden_doc_id + for i in range(8) + } + assert len(positions) > 1 + + def test_no_distractors_ok(self): + from soup_cli.utils.raft import build_raft_prompt + + composed = build_raft_prompt(_raft_row(0)) + assert composed.doc_ids == ("doc-0",) + assert composed.golden_doc_id == "doc-0" + + def test_missing_field_rejected(self): + from soup_cli.utils.raft import build_raft_prompt + + with pytest.raises(ValueError, match="query"): + build_raft_prompt({"golden_doc": "x", "answer": "y"}) + + def test_shuffle_seed_bool_rejected(self): + from soup_cli.utils.raft import build_raft_prompt + + with pytest.raises(TypeError, match="shuffle_seed"): + build_raft_prompt(_raft_row(), shuffle_seed=True) + + def test_row_index_negative_rejected(self): + from soup_cli.utils.raft import build_raft_prompt + + with pytest.raises(ValueError, match="row_index"): + build_raft_prompt(_raft_row(), row_index=-1) + + def test_too_many_docs_rejected(self): + from soup_cli.utils.raft import build_raft_prompt + + row = _raft_row(0) + row["distractor_docs"] = [f"d{i}" for i in range(65)] + with pytest.raises(ValueError, match="documents"): + build_raft_prompt(row) + + def test_non_mapping_rejected(self): + from soup_cli.utils.raft import build_raft_prompt + + with pytest.raises(ValueError, match="mapping"): + build_raft_prompt(["not", "a", "dict"]) + + +# --------------------------------------------------------------------------- +# #199 — tokenize_raft_example +# --------------------------------------------------------------------------- + + +class TestTokenizeRaftExample: + def _composed(self): + from soup_cli.utils.raft import build_raft_prompt + + return build_raft_prompt(_raft_row(1)) + + def test_answer_only_mask(self): + from soup_cli.utils.raft import tokenize_raft_example + + tok = _FakeTokenizer() + row = tokenize_raft_example(tok, self._composed(), max_length=512) + labels = row["labels"] + weights = row["loss_weights"] + # The leading run (prompt) is masked -100 / weight 0; the answer tail + # is unmasked / weight 1. + assert any(x == -100 for x in labels) + assert any(x != -100 for x in labels) + # Prompt positions: label -100 ⇔ weight 0.0. + for label, w in zip(labels, weights): + if label == -100: + assert w == 0.0 + else: + assert w >= 1.0 + assert len(row["input_ids"]) == len(labels) == len(weights) + assert row["attention_mask"] == [1] * len(row["input_ids"]) + + def test_eos_appended(self): + from soup_cli.utils.raft import tokenize_raft_example + + tok = _FakeTokenizer() + row = tokenize_raft_example(tok, self._composed(), max_length=512) + assert row["input_ids"][-1] == tok.eos_token_id + + def test_truncation_respects_max_length(self): + from soup_cli.utils.raft import tokenize_raft_example + + tok = _FakeTokenizer() + row = tokenize_raft_example(tok, self._composed(), max_length=8) + assert len(row["input_ids"]) == 8 + assert len(row["labels"]) == 8 + assert len(row["loss_weights"]) == 8 + + def test_citation_boost_applied(self): + from soup_cli.utils.raft import tokenize_raft_example + + tok = _FakeTokenizer() + # Answer with a [doc-0] citation token → that token's weight boosted. + composed = self._composed() + row = tokenize_raft_example( + tok, composed, max_length=512, citation_faithful=True + ) + # At least one answer token gets the boost (> 1.0). + assert any(w > 1.0 for w in row["loss_weights"]) + + def test_no_citation_boost_when_disabled(self): + from soup_cli.utils.raft import tokenize_raft_example + + tok = _FakeTokenizer() + row = tokenize_raft_example( + tok, self._composed(), max_length=512, citation_faithful=False + ) + assert all(w in (0.0, 1.0) for w in row["loss_weights"]) + + def test_slow_tokenizer_degrades_to_flat_mask(self): + from soup_cli.utils.raft import tokenize_raft_example + + tok = _FakeTokenizer(fast=False) # no offset mapping + row = tokenize_raft_example( + tok, self._composed(), max_length=512, citation_faithful=True + ) + # No offsets → flat answer weights (no boost), never raises. + assert all(w in (0.0, 1.0) for w in row["loss_weights"]) + + def test_bad_max_length_rejected(self): + from soup_cli.utils.raft import tokenize_raft_example + + tok = _FakeTokenizer() + with pytest.raises(ValueError, match="max_length"): + tokenize_raft_example(tok, self._composed(), max_length=4) + with pytest.raises(ValueError, match="max_length"): + tokenize_raft_example(tok, self._composed(), max_length=True) + + +# --------------------------------------------------------------------------- +# #199 — citation_span_token_weights +# --------------------------------------------------------------------------- + + +class TestCitationSpanTokenWeights: + def test_overlapping_tokens_boosted(self): + from soup_cli.utils.raft import citation_span_token_weights + + answer = "Paris [doc-0] is." + # offsets for: "Paris"(0-5) "[doc-0]"(6-13) "is."(14-17) + offsets: List[Tuple[int, int]] = [(0, 5), (6, 13), (14, 17)] + weights = citation_span_token_weights(answer, offsets, boost=5.0) + assert weights == [1.0, 5.0, 1.0] + + def test_no_citation_all_one(self): + from soup_cli.utils.raft import citation_span_token_weights + + weights = citation_span_token_weights("plain text", [(0, 5), (6, 10)]) + assert weights == [1.0, 1.0] + + def test_boost_below_one_rejected(self): + from soup_cli.utils.raft import citation_span_token_weights + + with pytest.raises(ValueError, match="boost"): + citation_span_token_weights("x", [(0, 1)], boost=0.5) + + def test_boost_bool_rejected(self): + from soup_cli.utils.raft import citation_span_token_weights + + with pytest.raises(TypeError, match="boost"): + citation_span_token_weights("x", [(0, 1)], boost=True) + + +# --------------------------------------------------------------------------- +# #199 — RaftDataCollator +# --------------------------------------------------------------------------- + + +class TestRaftDataCollator: + def test_pads_ragged_batch(self): + import torch + + from soup_cli.trainer.raft import RaftDataCollator + + collate = RaftDataCollator(_FakeTokenizer()) + batch = collate([ + {"input_ids": [1, 2, 3], "attention_mask": [1, 1, 1], + "labels": [-100, 2, 3], "loss_weights": [0.0, 1.0, 1.0]}, + {"input_ids": [4, 5], "attention_mask": [1, 1], + "labels": [-100, 5], "loss_weights": [0.0, 1.0]}, + ]) + assert batch["input_ids"].shape == (2, 3) + # Row 2 padded with pad_id=0 in input, -100 in labels, 0.0 in weights. + assert batch["input_ids"][1].tolist() == [4, 5, 0] + assert batch["labels"][1].tolist() == [-100, 5, -100] + assert batch["loss_weights"][1].tolist() == [0.0, 1.0, 0.0] + assert batch["attention_mask"][1].tolist() == [1, 1, 0] + assert batch["loss_weights"].dtype == torch.float32 + + def test_empty_batch_rejected(self): + from soup_cli.trainer.raft import RaftDataCollator + + with pytest.raises(ValueError, match="empty batch"): + RaftDataCollator(_FakeTokenizer())([]) + + +# --------------------------------------------------------------------------- +# #199 — make_raft_trainer_class compute_loss +# --------------------------------------------------------------------------- + + +class TestRaftTrainerComputeLoss: + def test_factory_caches(self): + from soup_cli.trainer.raft import make_raft_trainer_class + + a = make_raft_trainer_class(_FakeBaseTrainer) + b = make_raft_trainer_class(_FakeBaseTrainer) + assert a is b + assert "_RaftTrainer" in a.__name__ + + def test_weighted_loss_finite(self): + import torch + + from soup_cli.trainer.raft import make_raft_trainer_class + + cls = make_raft_trainer_class(_FakeBaseTrainer) + trainer = cls() + vocab = 10 + logits = torch.randn(1, 4, vocab) + model = _FakeModel(logits) + inputs = { + "input_ids": torch.tensor([[1, 2, 3, 4]]), + "attention_mask": torch.tensor([[1, 1, 1, 1]]), + "labels": torch.tensor([[-100, -100, 5, 6]]), + "loss_weights": torch.tensor([[0.0, 0.0, 1.0, 5.0]]), + } + loss = trainer.compute_loss(model, inputs) + assert torch.isfinite(loss) + assert loss.item() >= 0.0 + + def test_all_one_weights_equals_answer_only_ce(self): + import torch + from torch.nn.functional import cross_entropy + + from soup_cli.trainer.raft import make_raft_trainer_class + + cls = make_raft_trainer_class(_FakeBaseTrainer) + trainer = cls() + torch.manual_seed(0) + logits = torch.randn(1, 4, 10) + model = _FakeModel(logits) + labels = torch.tensor([[-100, -100, 5, 6]]) + inputs = { + "input_ids": torch.tensor([[1, 2, 3, 4]]), + "labels": labels, + "loss_weights": torch.tensor([[0.0, 0.0, 1.0, 1.0]]), + } + weighted = trainer.compute_loss(model, inputs).item() + # Reference answer-only CE. + ref = cross_entropy( + logits[:, :-1, :].reshape(-1, 10), + labels[:, 1:].reshape(-1), + ignore_index=-100, + ).item() + assert weighted == pytest.approx(ref, abs=1e-5) + + +# --------------------------------------------------------------------------- +# #199 — schema raft_shuffle_seed +# --------------------------------------------------------------------------- + + +class TestRaftShuffleSeedSchema: + def _yaml(self, seed_line: str) -> str: + return ( + "base: hf-internal-testing/tiny-random-gpt2\n" + "task: sft\n" + "data:\n" + " train: ./data/raft.jsonl\n" + " format: raft\n" + f"{seed_line}" + "training:\n" + " epochs: 1\n" + "output: ./output\n" + ) + + def test_accepts_int(self): + from soup_cli.config.loader import load_config_from_string + + cfg = load_config_from_string(self._yaml(" raft_shuffle_seed: 42\n")) + assert cfg.data.raft_shuffle_seed == 42 + + def test_default_none(self): + from soup_cli.config.loader import load_config_from_string + + cfg = load_config_from_string(self._yaml("")) + assert cfg.data.raft_shuffle_seed is None + + def test_bool_rejected(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(Exception, match="raft_shuffle_seed"): + load_config_from_string(self._yaml(" raft_shuffle_seed: true\n")) + + def test_negative_rejected(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(Exception, match="raft_shuffle_seed"): + load_config_from_string(self._yaml(" raft_shuffle_seed: -1\n")) + + +# --------------------------------------------------------------------------- +# #199 — SFT wiring source-grep regression guards +# --------------------------------------------------------------------------- + + +class TestSftRaftWiring: + def _sft_src(self) -> str: + import soup_cli.trainer.sft as sft + + with open(sft.__file__, encoding="utf-8") as fh: + return fh.read() + + def test_sft_routes_raft(self): + src = self._sft_src() + assert "_prepare_raft_dataset" in src + assert "make_raft_trainer_class" in src + assert "RaftDataCollator" in src + assert 'cfg.data.format == "raft"' in src + + def test_raft_modules_have_no_top_level_torch(self): + import soup_cli.utils.raft as raft + + with open(raft.__file__, encoding="utf-8") as fh: + src = fh.read() + assert "\nimport torch" not in src + assert "\nfrom torch" not in src + + +# =========================================================================== +# #202 — per-style citation extractors + eval citation + diagnose mode +# =========================================================================== + + +class TestPerStyleExtractors: + def test_bracket_default(self): + from soup_cli.utils.citation_faithful import extract_citation_ids + + assert extract_citation_ids("see [doc-1] and [doc-2]") == ("doc-1", "doc-2") + + def test_inline_style(self): + from soup_cli.utils.citation_faithful import extract_citation_ids + + assert extract_citation_ids("see (doc-1) here", style="inline") == ("doc-1",) + # bracket extractor must NOT match the parenthetical form. + assert extract_citation_ids("see (doc-1) here") == () + + def test_footnote_style(self): + from soup_cli.utils.citation_faithful import extract_citation_ids + + assert extract_citation_ids("fact[^3] more", style="footnote") == ("3",) + # footnote `[^3]` must not be picked up by the bracket extractor + # (starts with ^ after the bracket → fails alnum-leading). + assert extract_citation_ids("fact[^3] more") == () + + def test_citation_spans_cover_delimiters(self): + from soup_cli.utils.citation_faithful import citation_spans + + spans = citation_spans("ab [doc-0] cd") + assert spans == ((3, 10),) # covers the full "[doc-0]" + + def test_score_citations_with_style(self): + from soup_cli.utils.citation_faithful import score_citations + + cs = score_citations( + predicted="answer (doc-0)", expected_ids=["doc-0"], style="inline" + ) + assert cs.precision == 1.0 + assert cs.recall == 1.0 + + def test_unknown_style_rejected(self): + from soup_cli.utils.citation_faithful import extract_citation_ids + + with pytest.raises(ValueError, match="citation_style"): + extract_citation_ids("x", style="bogus") + + +class TestEvalCitationCli: + def _runner(self): + from typer.testing import CliRunner + + return CliRunner() + + def test_citation_in_eval_help(self): + from soup_cli.cli import app + + result = self._runner().invoke(app, ["eval", "--help"]) + assert result.exit_code == 0 + assert "citation" in result.output + + def test_citation_predicted_expected(self): + import json + + from soup_cli.cli import app + + runner = self._runner() + with runner.isolated_filesystem(): + with open("c.jsonl", "w", encoding="utf-8") as fh: + fh.write( + json.dumps({"predicted": "Paris [doc-0].", "expected_ids": ["doc-0"]}) + + "\n" + ) + fh.write( + json.dumps({"predicted": "Berlin [doc-2].", "expected_ids": ["doc-0"]}) + + "\n" + ) + result = runner.invoke( + app, ["eval", "citation", "c.jsonl", "--output", "out.json"] + ) + assert result.exit_code == 0, (result.output, result.exception) + with open("out.json", encoding="utf-8") as fh: + payload = json.load(fh) + assert payload["n_rows"] == 2 + # First row recall 1.0, second 0.0 → mean 0.5. + assert payload["aggregate"]["recall"] == pytest.approx(0.5) + + def test_citation_raft_rows(self): + import json + + from soup_cli.cli import app + + runner = self._runner() + with runner.isolated_filesystem(): + with open("raft.jsonl", "w", encoding="utf-8") as fh: + # answer cites doc-0; with no shuffle (default seed 0) golden + # id is deterministic — the row scores its own answer. + fh.write(json.dumps({ + "query": "q", "golden_doc": "g", "distractor_docs": [], + "answer": "see [doc-0]", + }) + "\n") + result = runner.invoke(app, ["eval", "citation", "raft.jsonl"]) + assert result.exit_code == 0, (result.output, result.exception) + assert "Citation aggregate" in result.output + + def test_invalid_style_exit_2(self): + from soup_cli.cli import app + + runner = self._runner() + with runner.isolated_filesystem(): + with open("c.jsonl", "w", encoding="utf-8") as fh: + fh.write('{"predicted": "x", "expected_ids": ["a"]}\n') + result = runner.invoke(app, ["eval", "citation", "c.jsonl", "--style", "bogus"]) + assert result.exit_code == 2 + + def test_missing_file_exit(self): + from soup_cli.cli import app + + runner = self._runner() + with runner.isolated_filesystem(): + result = runner.invoke(app, ["eval", "citation", "nope.jsonl"]) + assert result.exit_code != 0 + + def test_no_scorable_rows_exit_2(self): + from soup_cli.cli import app + + runner = self._runner() + with runner.isolated_filesystem(): + with open("c.jsonl", "w", encoding="utf-8") as fh: + fh.write('{"unrelated": "row"}\n') + result = runner.invoke(app, ["eval", "citation", "c.jsonl"]) + assert result.exit_code == 2 + + +class TestDiagnoseCitationMode: + def test_citation_in_failure_modes(self): + from soup_cli.utils.diagnose.report import FAILURE_MODES + + assert "citation" in FAILURE_MODES + + def test_is_raft_row(self): + from soup_cli.utils.diagnose.citation import is_raft_row + + assert is_raft_row({"query": "q", "golden_doc": "g", "answer": "a"}) + assert not is_raft_row({"prompt": "p"}) + assert not is_raft_row("not a dict") + + def test_score_citation_recall_ok(self): + from soup_cli.utils.diagnose.citation import score_citation + + rows = [_raft_row(2)] + + def gen(prompt: str) -> str: + # Always cite the golden doc id present in the prompt's [doc-N]. + from soup_cli.utils.raft import build_raft_prompt + + golden = build_raft_prompt(rows[0]).golden_doc_id + return f"the answer is correct [{golden}]" + + result = score_citation(rows, gen) + assert result.mode == "citation" + assert result.score == pytest.approx(1.0) + assert result.verdict == "OK" + + def test_score_citation_no_citation_major(self): + from soup_cli.utils.diagnose.citation import score_citation + + rows = [_raft_row(2), _raft_row(2)] + result = score_citation(rows, lambda p: "no citation at all") + assert result.score == 0.0 + assert result.verdict == "MAJOR" + + def test_score_citation_no_raft_rows_raises(self): + from soup_cli.utils.diagnose.citation import score_citation + + with pytest.raises(ValueError, match="RAFT"): + score_citation([{"prompt": "p"}], lambda p: "x") + + def test_build_report_fills_citation_neutral(self): + from soup_cli.utils.diagnose.runner import build_report + + report = build_report(run_id="r", base="b", adapter="a", scores={}) + assert "citation" in report.scores + assert report.scores["citation"].verdict == "OK" + + def test_citation_probe_torch_free(self): + import soup_cli.utils.diagnose.citation as cit + + with open(cit.__file__, encoding="utf-8") as fh: + src = fh.read() + assert "\nimport torch" not in src + + +# =========================================================================== +# #201 — live CAA / ITI / RepE steering + serve --steer decode hook +# =========================================================================== + + +class TestSteeringMath: + def test_caa_mean_difference(self): + import numpy as np + + from soup_cli.utils.steering import compute_caa_vector + + pos = np.array([[2.0, 0.0], [4.0, 0.0]]) # mean [3,0] + neg = np.array([[0.0, 1.0], [0.0, 3.0]]) # mean [0,2] + vec = compute_caa_vector(pos, neg) + assert np.allclose(vec, [3.0, -2.0]) + assert vec.dtype == np.float32 + + def test_caa_dim_mismatch_rejected(self): + import numpy as np + + from soup_cli.utils.steering import compute_caa_vector + + with pytest.raises(ValueError, match="mismatch"): + compute_caa_vector(np.zeros((2, 3)), np.zeros((2, 4))) + + def test_caa_empty_rejected(self): + import numpy as np + + from soup_cli.utils.steering import compute_caa_vector + + with pytest.raises(ValueError): + compute_caa_vector(np.zeros((0, 3)), np.zeros((2, 3))) + + def test_repe_direction_aligned(self): + import numpy as np + + from soup_cli.utils.steering import compute_repe_direction + + # Diffs scattered along the +x axis → top PC ≈ x, sign-aligned positive. + diffs = np.array([[2.0, 0.1], [3.0, -0.1], [4.0, 0.05]]) + vec = compute_repe_direction(diffs) + assert vec.shape == (2,) + assert vec[0] > 0 # points along the dominant +x diff direction + + def test_iti_selects_top_heads(self): + import numpy as np + + from soup_cli.utils.steering import compute_iti_directions + + # 3 heads; head 1 has the biggest pos/neg separation. + pos = np.zeros((2, 3, 2)) + neg = np.zeros((2, 3, 2)) + pos[:, 1, :] = 5.0 # head 1 strongly separated + pos[:, 0, :] = 0.5 + dirs, selected = compute_iti_directions(pos, neg, top_k=1) + assert selected == (1,) + assert np.allclose(dirs[1], [5.0, 5.0]) + assert np.allclose(dirs[0], [0.0, 0.0]) # non-selected zeroed + + def test_iti_top_k_bool_rejected(self): + import numpy as np + + from soup_cli.utils.steering import compute_iti_directions + + with pytest.raises(ValueError, match="top_k"): + compute_iti_directions(np.zeros((1, 2, 2)), np.zeros((1, 2, 2)), top_k=True) + + +class TestLoadContrastivePairs: + def test_loads_pairs(self, tmp_path, monkeypatch): + import json + + from soup_cli.utils.steering import load_contrastive_pairs + + monkeypatch.chdir(tmp_path) + p = tmp_path / "pairs.jsonl" + with open(p, "w", encoding="utf-8") as fh: + fh.write(json.dumps({"positive": "be kind", "negative": "be rude"}) + "\n") + fh.write(json.dumps({"positive": "help", "negative": "refuse"}) + "\n") + fh.write("not json\n") # skipped + fh.write(json.dumps({"positive": "x"}) + "\n") # incomplete, skipped + pairs = load_contrastive_pairs("pairs.jsonl") + assert pairs == [("be kind", "be rude"), ("help", "refuse")] + + def test_outside_cwd_rejected(self, tmp_path): + from soup_cli.utils.steering import load_contrastive_pairs + + outside = tmp_path / "pairs.jsonl" + outside.write_text('{"positive":"a","negative":"b"}\n', encoding="utf-8") + with pytest.raises(ValueError, match="cwd"): + load_contrastive_pairs(str(outside)) + + def test_empty_file_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.steering import load_contrastive_pairs + + monkeypatch.chdir(tmp_path) + (tmp_path / "empty.jsonl").write_text("\n\n", encoding="utf-8") + with pytest.raises(ValueError, match="no usable"): + load_contrastive_pairs("empty.jsonl") + + +class TestBuildSteeringVectorValidation: + def test_method_validated_first(self): + from soup_cli.utils.steering import build_steering_vector + + with pytest.raises(ValueError, match="steering method"): + build_steering_vector(method="nonsense", name="x", base="m", pairs_path="p") + + def test_name_validated(self): + from soup_cli.utils.steering import build_steering_vector + + with pytest.raises(ValueError, match="steering name"): + build_steering_vector(method="caa", name="bad/path", base="m", pairs_path="p") + + def test_base_required(self): + from soup_cli.utils.steering import build_steering_vector + + with pytest.raises(ValueError, match="base"): + build_steering_vector(method="caa", name="ok", pairs_path="p") + + def test_pairs_required(self): + from soup_cli.utils.steering import build_steering_vector + + with pytest.raises(ValueError, match="pairs_path"): + build_steering_vector(method="caa", name="ok", base="m") + + +class TestSteeringArtifactRoundtrip: + def _write_artifact(self, dir_path, *, method="caa", intervention="residual", vec=None): + import json + import os + + import numpy as np + from safetensors.numpy import save_file + + os.makedirs(dir_path, exist_ok=True) + if vec is None: + vec = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32) + save_file({"vector": vec}, os.path.join(dir_path, "steering_vector.safetensors")) + config = { + "method": method, + "name": "safety-v1", + "layer": 1, + "hidden_dim": int(vec.shape[0]), + "intervention_point": intervention, + "base": "tiny", + "default_strength": 1.0, + } + with open(os.path.join(dir_path, "steering_config.json"), "w", encoding="utf-8") as fh: + json.dump(config, fh) + + def test_load_roundtrip(self, tmp_path, monkeypatch): + import numpy as np + + from soup_cli.utils.steering import load_steering_artifact + + monkeypatch.chdir(tmp_path) + self._write_artifact("steering/safety-v1") + loaded = load_steering_artifact("steering/safety-v1") + assert loaded.method == "caa" + assert loaded.layer == 1 + assert loaded.intervention_point == "residual" + assert np.allclose(loaded.vector, [0.1, 0.2, 0.3, 0.4]) + + def test_resolve_steering_dir_local_fallback(self, tmp_path, monkeypatch): + from soup_cli.utils.steering import resolve_steering_dir + + monkeypatch.chdir(tmp_path) + self._write_artifact("steering/safety-v1") + resolved = resolve_steering_dir("safety-v1") + assert resolved.replace("\\", "/").endswith("steering/safety-v1") + + def test_resolve_unknown_raises(self, tmp_path, monkeypatch): + from soup_cli.utils.steering import resolve_steering_dir + + monkeypatch.chdir(tmp_path) + with pytest.raises(ValueError, match="no steering vector"): + resolve_steering_dir("does-not-exist") + + def test_load_missing_files_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.steering import load_steering_artifact + + monkeypatch.chdir(tmp_path) + (tmp_path / "empty").mkdir() + with pytest.raises(ValueError, match="missing"): + load_steering_artifact("empty") + + +class TestInstallSteeringHook: + def _fake_model(self, d=4, n=2): + import torch + import torch.nn as nn + + class FakeLayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Module() + self.self_attn.o_proj = nn.Linear(d, d, bias=False) + with torch.no_grad(): + self.self_attn.o_proj.weight.copy_(torch.eye(d)) + self.dummy = nn.Parameter(torch.zeros(1)) + + def forward(self, x): + return (x,) + + class FakeInner(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([FakeLayer() for _ in range(n)]) + + class FakeModel(nn.Module): + def __init__(self): + super().__init__() + self.model = FakeInner() + + return FakeModel() + + def _loaded(self, intervention, vec): + from soup_cli.utils.steering import LoadedSteering + + return LoadedSteering( + method="caa" if intervention == "residual" else "iti", + name="t", + layer=0, + intervention_point=intervention, + vector=vec, + default_strength=1.0, + ) + + def test_residual_hook_adds_vector(self): + import numpy as np + import torch + + from soup_cli.utils.steering import install_steering_hook + + model = self._fake_model(d=4) + vec = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32) + handle = install_steering_hook(model, self._loaded("residual", vec), strength=2.0) + x = torch.zeros(1, 3, 4) + out = model.model.layers[0](x)[0] + # residual hook adds strength*vec = [2,0,0,0] to every position. + assert torch.allclose(out[0, 0], torch.tensor([2.0, 0.0, 0.0, 0.0])) + handle.remove() + + def test_iti_pre_hook_shifts_o_proj_input(self): + import numpy as np + import torch + + from soup_cli.utils.steering import install_steering_hook + + model = self._fake_model(d=4) + vec = np.array([0.0, 3.0, 0.0, 0.0], dtype=np.float32) + handle = install_steering_hook(model, self._loaded("attn_o_proj_input", vec), strength=1.0) + o_proj = model.model.layers[0].self_attn.o_proj + x = torch.zeros(1, 2, 4) + # o_proj is identity → output == shifted input == x + vec. + out = o_proj(x) + assert torch.allclose(out[0, 0], torch.tensor([0.0, 3.0, 0.0, 0.0])) + handle.remove() + + def test_non_loaded_rejected(self): + from soup_cli.utils.steering import install_steering_hook + + with pytest.raises(TypeError, match="LoadedSteering"): + install_steering_hook(self._fake_model(), {"not": "loaded"}, strength=1.0) + + +class TestServeSteerCli: + def _runner(self): + from typer.testing import CliRunner + + return CliRunner() + + def test_steer_flag_in_serve_help(self): + from soup_cli.cli import app + + result = self._runner().invoke(app, ["serve", "--help"]) + assert result.exit_code == 0 + assert "--steer" in result.output + + def test_steer_requires_transformers_backend(self): + from soup_cli.cli import app + + result = self._runner().invoke( + app, ["serve", "-m", "model", "--steer", "safety-v1", "--backend", "vllm"] + ) + assert result.exit_code == 2 + assert "transformers" in result.output + + def test_steer_bad_name_rejected(self): + from soup_cli.cli import app + + result = self._runner().invoke( + app, ["serve", "-m", "model", "--steer", "bad/name"] + ) + assert result.exit_code == 2 + + +class TestSteerCommandPlumbing: + def test_steer_apply_loads_artifact(self, tmp_path, monkeypatch): + import json + import os + + import numpy as np + from safetensors.numpy import save_file + from typer.testing import CliRunner + + from soup_cli.commands.steer import app + + monkeypatch.chdir(tmp_path) + d = "steering/safety-v1" + os.makedirs(d, exist_ok=True) + save_file( + {"vector": np.array([0.1, 0.2], dtype=np.float32)}, + os.path.join(d, "steering_vector.safetensors"), + ) + with open(os.path.join(d, "steering_config.json"), "w", encoding="utf-8") as fh: + json.dump({ + "method": "caa", "name": "safety-v1", "layer": 1, + "hidden_dim": 2, "intervention_point": "residual", + "base": "tiny", "default_strength": 1.0, + }, fh) + result = CliRunner().invoke(app, ["apply", "--name", "safety-v1"]) + assert result.exit_code == 0, (result.output, result.exception) + assert "Vector loaded" in result.output + + def test_steer_train_help_has_output(self): + from typer.testing import CliRunner + + from soup_cli.commands.steer import app + + result = CliRunner().invoke(app, ["train", "--help"]) + assert result.exit_code == 0 + assert "--output" in result.output + assert "--top-k" in result.output + + +# =========================================================================== +# #200 — live RA-DIT two-stage orchestrator + retriever auto-link +# =========================================================================== + + +def _seed_registry(db_path, *, embedding_output, with_generator=False): + """Push a retriever (embedding) entry whose config marks ra_dit_stage.""" + import os + + os.environ["SOUP_REGISTRY_DB_PATH"] = str(db_path) + from soup_cli.registry.store import RegistryStore + + with RegistryStore() as store: + # A non-RA-DIT embedding run (should NOT be picked). + store.push( + name="plain-embed", + tag="v1", + base_model="sentence-transformers/all-MiniLM-L6-v2", + task="embedding", + run_id=None, + config={"task": "embedding", "output": "./other-embed"}, + ) + # The RA-DIT retriever stage (SHOULD be picked). + rid = store.push( + name="ra-dit-retriever", + tag="v1", + base_model="sentence-transformers/all-MiniLM-L6-v2", + task="embedding", + run_id=None, + config={ + "task": "embedding", + "output": embedding_output, + "training": {"ra_dit_stage": "retriever"}, + }, + ) + if with_generator: + store.push( + name="ra-dit-gen", + tag="v1", + base_model="meta-llama/Llama-3.1-8B", + task="sft", + run_id=None, + config={ + "task": "sft", + "output": "./gen-out", + "training": {"ra_dit_stage": "generator"}, + }, + ) + return rid + + +class TestDiscoverLatestRetriever: + def test_finds_ra_dit_retriever_output(self, tmp_path, monkeypatch): + from soup_cli.utils.ra_dit_run import discover_latest_retriever + + db = tmp_path / "reg.db" + _seed_registry(db, embedding_output="./retriever-out") + monkeypatch.setenv("SOUP_REGISTRY_DB_PATH", str(db)) + found = discover_latest_retriever() + assert found == "./retriever-out" + + def test_returns_none_when_no_retriever(self, tmp_path, monkeypatch): + import os + + from soup_cli.registry.store import RegistryStore + from soup_cli.utils.ra_dit_run import discover_latest_retriever + + db = tmp_path / "reg.db" + os.environ["SOUP_REGISTRY_DB_PATH"] = str(db) + monkeypatch.setenv("SOUP_REGISTRY_DB_PATH", str(db)) + with RegistryStore() as store: + store.push( + name="plain", tag="v1", + base_model="m", task="embedding", run_id=None, + config={"task": "embedding", "output": "./x"}, + ) + assert discover_latest_retriever() is None + + def test_empty_registry_returns_none(self, tmp_path, monkeypatch): + from soup_cli.utils.ra_dit_run import discover_latest_retriever + + db = tmp_path / "reg.db" + monkeypatch.setenv("SOUP_REGISTRY_DB_PATH", str(db)) + assert discover_latest_retriever() is None + + +class TestResolveRetrieverForGenerator: + def test_manual_override_wins(self, tmp_path, monkeypatch): + from soup_cli.utils.ra_dit_run import resolve_retriever_for_generator + + db = tmp_path / "reg.db" + _seed_registry(db, embedding_output="./auto-retriever") + monkeypatch.setenv("SOUP_REGISTRY_DB_PATH", str(db)) + resolved, advisory = resolve_retriever_for_generator("my/manual-retriever") + assert resolved == "my/manual-retriever" + assert "override" in advisory.lower() or "manual" in advisory.lower() + + def test_autolinks_when_none(self, tmp_path, monkeypatch): + from soup_cli.utils.ra_dit_run import resolve_retriever_for_generator + + db = tmp_path / "reg.db" + _seed_registry(db, embedding_output="./auto-retriever") + monkeypatch.setenv("SOUP_REGISTRY_DB_PATH", str(db)) + resolved, advisory = resolve_retriever_for_generator(None) + assert resolved == "./auto-retriever" + assert "auto" in advisory.lower() + + def test_not_found_advisory(self, tmp_path, monkeypatch): + from soup_cli.utils.ra_dit_run import resolve_retriever_for_generator + + db = tmp_path / "reg.db" + monkeypatch.setenv("SOUP_REGISTRY_DB_PATH", str(db)) + resolved, advisory = resolve_retriever_for_generator(None) + assert resolved is None + assert "no" in advisory.lower() and "retriever" in advisory.lower() + + +class TestRunRaDit: + def _write_configs(self, tmp_path): + retr = tmp_path / "retriever.yaml" + retr.write_text( + "base: sentence-transformers/all-MiniLM-L6-v2\n" + "task: embedding\n" + "output: ./ra-out/retriever\n" + "training:\n ra_dit_stage: retriever\n" + "data:\n train: ./triples.jsonl\n format: embedding\n", + encoding="utf-8", + ) + gen = tmp_path / "generator.yaml" + gen.write_text( + "base: meta-llama/Llama-3.1-8B\n" + "task: sft\n" + "output: ./ra-out/generator\n" + "training:\n ra_dit_stage: generator\n" + "data:\n train: ./raft.jsonl\n format: raft\n", + encoding="utf-8", + ) + return retr, gen + + def test_chains_two_stages_with_autolink(self, tmp_path, monkeypatch): + from soup_cli.utils.ra_dit_run import run_ra_dit + + monkeypatch.chdir(tmp_path) + retr, gen = self._write_configs(tmp_path) + calls = [] + + def fake_runner(config_path): + calls.append(config_path) + + result = run_ra_dit( + "retriever.yaml", "generator.yaml", _runner=fake_runner + ) + # Two subprocess stages ran, retriever first. + assert len(calls) == 2 + assert "retriever" in calls[0].replace("\\", "/") + # Generator stage ran via a rewritten temp yaml carrying the link. + assert result.retriever_output.replace("\\", "/").endswith("ra-out/retriever") + # Auto-linked retriever model == retriever output dir. + assert result.retriever_model_used.replace("\\", "/").endswith( + "ra-out/retriever" + ) + assert result.autolinked is True + + def test_manual_override_skips_autolink(self, tmp_path, monkeypatch): + from soup_cli.utils.ra_dit_run import run_ra_dit + + monkeypatch.chdir(tmp_path) + self._write_configs(tmp_path) + result = run_ra_dit( + "retriever.yaml", + "generator.yaml", + retriever_model="my/explicit-retriever", + _runner=lambda p: None, + ) + assert result.retriever_model_used == "my/explicit-retriever" + assert result.autolinked is False + + def test_rewrites_generator_with_retriever_model(self, tmp_path, monkeypatch): + import yaml + + from soup_cli.utils.ra_dit_run import run_ra_dit + + monkeypatch.chdir(tmp_path) + self._write_configs(tmp_path) + seen_yaml = {} + + def fake_runner(config_path): + # On the 2nd call (generator), capture the rewritten YAML. + with open(config_path, encoding="utf-8") as fh: + seen_yaml[config_path] = yaml.safe_load(fh) + + run_ra_dit("retriever.yaml", "generator.yaml", _runner=fake_runner) + # The generator temp yaml must carry the retriever model. + gen_cfgs = [ + c for c in seen_yaml.values() + if c.get("task") == "sft" + ] + assert gen_cfgs, "generator config not captured" + training = gen_cfgs[0].get("training", {}) + assert "ra_dit_retriever_model" in training + + def test_outside_cwd_config_rejected(self, tmp_path): + from soup_cli.utils.ra_dit_run import run_ra_dit + + retr, gen = self._write_configs(tmp_path) + with pytest.raises(ValueError, match="cwd"): + run_ra_dit(str(retr), str(gen), _runner=lambda p: None) + + +class TestTrainAutolinkHook: + def test_generator_stage_autolinks_in_train(self, tmp_path, monkeypatch): + """`soup train` of a generator stage with no retriever model auto-links.""" + import os + + from soup_cli.config.loader import load_config_from_string + + db = tmp_path / "reg.db" + _seed_registry(db, embedding_output="./linked-retriever") + monkeypatch.setenv("SOUP_REGISTRY_DB_PATH", str(db)) + os.environ["SOUP_REGISTRY_DB_PATH"] = str(db) + + cfg = load_config_from_string( + "base: meta-llama/Llama-3.1-8B\n" + "task: sft\n" + "output: ./out\n" + "training:\n ra_dit_stage: generator\n" + "data:\n train: ./raft.jsonl\n format: raft\n" + ) + assert cfg.training.ra_dit_retriever_model is None + from soup_cli.utils.ra_dit_run import autolink_generator_retriever + + advisory = autolink_generator_retriever(cfg) + assert cfg.training.ra_dit_retriever_model == "./linked-retriever" + assert advisory is not None and "auto" in advisory.lower() + + def test_no_autolink_when_not_generator(self, tmp_path, monkeypatch): + from soup_cli.config.loader import load_config_from_string + from soup_cli.utils.ra_dit_run import autolink_generator_retriever + + db = tmp_path / "reg.db" + monkeypatch.setenv("SOUP_REGISTRY_DB_PATH", str(db)) + cfg = load_config_from_string( + "base: m\ntask: sft\noutput: ./out\n" + "data:\n train: ./x.jsonl\n format: chatml\n" + ) + advisory = autolink_generator_retriever(cfg) + assert advisory is None + + def test_manual_retriever_model_not_overwritten(self, tmp_path, monkeypatch): + from soup_cli.config.loader import load_config_from_string + from soup_cli.utils.ra_dit_run import autolink_generator_retriever + + db = tmp_path / "reg.db" + _seed_registry(db, embedding_output="./auto") + monkeypatch.setenv("SOUP_REGISTRY_DB_PATH", str(db)) + cfg = load_config_from_string( + "base: m\ntask: sft\noutput: ./out\n" + "training:\n ra_dit_stage: generator\n" + " ra_dit_retriever_model: my/explicit\n" + "data:\n train: ./x.jsonl\n format: raft\n" + ) + advisory = autolink_generator_retriever(cfg) + assert cfg.training.ra_dit_retriever_model == "my/explicit" + assert advisory is None or "explicit" in advisory or "manual" in advisory.lower() + + +class TestRaDitCli: + def _runner(self): + from typer.testing import CliRunner + + return CliRunner() + + def test_help(self): + from soup_cli.commands.ra_dit import app + + result = self._runner().invoke(app, ["--help"]) + assert result.exit_code == 0 + assert "retriever" in result.output.lower() + + def test_cli_registered(self): + from soup_cli.cli import app + + result = self._runner().invoke(app, ["--help"]) + assert result.exit_code == 0 + assert "ra-dit" in result.output.lower() + + def test_plan_only(self, tmp_path, monkeypatch): + from soup_cli.commands.ra_dit import app + + monkeypatch.chdir(tmp_path) + (tmp_path / "retriever.yaml").write_text( + "base: st/mini\ntask: embedding\noutput: ./r\n" + "training:\n ra_dit_stage: retriever\n" + "data:\n train: ./t.jsonl\n format: embedding\n", + encoding="utf-8", + ) + (tmp_path / "generator.yaml").write_text( + "base: meta-llama/Llama-3.1-8B\ntask: sft\noutput: ./g\n" + "training:\n ra_dit_stage: generator\n" + "data:\n train: ./raft.jsonl\n format: raft\n", + encoding="utf-8", + ) + result = self._runner().invoke(app, [ + "--retriever-config", "retriever.yaml", + "--generator-config", "generator.yaml", + "--plan-only", + ]) + assert result.exit_code == 0, (result.output, result.exception) + assert "retriever" in result.output.lower() + assert "generator" in result.output.lower() + + def test_missing_config_rejected(self, tmp_path, monkeypatch): + from soup_cli.commands.ra_dit import app + + monkeypatch.chdir(tmp_path) + result = self._runner().invoke(app, [ + "--retriever-config", "nope.yaml", + "--generator-config", "alsonope.yaml", + "--plan-only", + ]) + assert result.exit_code == 2 + + +class TestRaDitRunTorchFree: + def test_module_torch_free(self): + import soup_cli.utils.ra_dit_run as rr + + with open(rr.__file__, encoding="utf-8") as fh: + src = fh.read() + assert "\nimport torch" not in src + + +# =========================================================================== +# v0.71.10 review-fix coverage (TDD wave) + regression guards for code fixes +# =========================================================================== + + +class TestValidateSteeringStrengthBounds: + """#201 — validate_steering_strength boundary suite (HIGH).""" + + def test_zero_ok(self): + from soup_cli.utils.steering import validate_steering_strength + + assert validate_steering_strength(0.0) == 0.0 + + def test_upper_bound_ok(self): + from soup_cli.utils.steering import validate_steering_strength + + assert validate_steering_strength(10.0) == 10.0 + + def test_lower_bound_ok(self): + from soup_cli.utils.steering import validate_steering_strength + + assert validate_steering_strength(-10.0) == -10.0 + + def test_just_over_upper_rejected(self): + from soup_cli.utils.steering import validate_steering_strength + + with pytest.raises(ValueError, match="<="): + validate_steering_strength(10.0001) + + def test_just_under_lower_rejected(self): + from soup_cli.utils.steering import validate_steering_strength + + with pytest.raises(ValueError, match="<="): + validate_steering_strength(-10.0001) + + def test_nan_rejected(self): + from soup_cli.utils.steering import validate_steering_strength + + with pytest.raises(ValueError, match="finite"): + validate_steering_strength(float("nan")) + + def test_inf_rejected(self): + from soup_cli.utils.steering import validate_steering_strength + + with pytest.raises(ValueError, match="finite"): + validate_steering_strength(float("inf")) + + def test_bool_rejected(self): + from soup_cli.utils.steering import validate_steering_strength + + with pytest.raises(TypeError, match="bool"): + validate_steering_strength(True) + + def test_non_number_rejected(self): + from soup_cli.utils.steering import validate_steering_strength + + with pytest.raises(TypeError, match="number"): + validate_steering_strength("2.0") + + +class TestRepeDirectionSign: + """#201 — RepE sign-alignment + degenerate-input rejection (HIGH).""" + + def test_sign_aligned_with_mean(self): + import numpy as np + + from soup_cli.utils.steering import compute_repe_direction + + # Variance + mean both along -x → returned vector points -x AND its + # projection onto the mean diff is non-negative (sign-aligned). This + # invariant holds regardless of SVD's arbitrary sign choice. + diffs = np.array([[-2.0, 0.0], [-4.0, 0.0], [-3.0, 0.0]]) + vec = compute_repe_direction(diffs) + assert vec[0] < 0.0 + assert float(vec @ diffs.mean(axis=0)) >= 0.0 + + def test_empty_rejected(self): + import numpy as np + + from soup_cli.utils.steering import compute_repe_direction + + with pytest.raises(ValueError): + compute_repe_direction(np.zeros((0, 3))) + + def test_non_2d_rejected(self): + import numpy as np + + from soup_cli.utils.steering import compute_repe_direction + + with pytest.raises(ValueError): + compute_repe_direction(np.zeros((3,))) + + +class TestBuildSteeringVectorPreLoadValidation: + """#201 — layer / top_k validated BEFORE the model load (HIGH).""" + + def test_layer_out_of_range_rejected_pre_load(self): + from soup_cli.utils.steering import build_steering_vector + + # layer=99999 fails the bounds check before load_contrastive_pairs / + # the model load (base + pairs_path are never touched). + with pytest.raises(ValueError, match="layer"): + build_steering_vector( + method="caa", name="ok", base="m", pairs_path="p", layer=99999 + ) + + def test_layer_bool_rejected(self): + from soup_cli.utils.steering import build_steering_vector + + with pytest.raises(TypeError, match="layer"): + build_steering_vector( + method="caa", name="ok", base="m", pairs_path="p", layer=True + ) + + def test_top_k_zero_rejected_pre_load(self): + from soup_cli.utils.steering import build_steering_vector + + with pytest.raises(ValueError, match="top_k"): + build_steering_vector( + method="caa", name="ok", base="m", pairs_path="p", top_k=0 + ) + + def test_top_k_bool_rejected_pre_load(self): + from soup_cli.utils.steering import build_steering_vector + + with pytest.raises(ValueError, match="top_k"): + build_steering_vector( + method="caa", name="ok", base="m", pairs_path="p", top_k=True + ) + + def test_repe_requires_two_pairs(self, tmp_path, monkeypatch): + import json + + from soup_cli.utils.steering import build_steering_vector + + monkeypatch.chdir(tmp_path) + (tmp_path / "one.jsonl").write_text( + json.dumps({"positive": "a", "negative": "b"}) + "\n", encoding="utf-8" + ) + # repe needs >= 2 pairs — fails AFTER load_contrastive_pairs but BEFORE + # the model load (regression for code-review L5). + with pytest.raises(ValueError, match=">="): + build_steering_vector( + method="repe", name="ok", base="m", pairs_path="one.jsonl" + ) + + +class TestCitationBoostCap: + """#199 — citation boost upper bound (MEDIUM).""" + + def test_boost_100_ok(self): + from soup_cli.utils.raft import citation_span_token_weights + + weights = citation_span_token_weights( + "[doc-0] x", [(0, 7), (8, 9)], boost=100.0 + ) + assert weights[0] == 100.0 + + def test_boost_over_100_rejected(self): + from soup_cli.utils.raft import citation_span_token_weights + + with pytest.raises(ValueError, match="boost"): + citation_span_token_weights("x", [(0, 1)], boost=100.1) + + def test_tokenize_citation_boost_applied(self): + from soup_cli.utils.raft import RaftComposed, tokenize_raft_example + + composed = RaftComposed( + prompt="Q", answer="Paris [doc-0]", golden_doc_id="doc-0", + doc_ids=("doc-0",), + ) + row = tokenize_raft_example( + _FakeTokenizer(fast=True), composed, max_length=64, + citation_faithful=True, citation_boost=7.0, + ) + # Some answer token carries the boost weight. + assert 7.0 in row["loss_weights"] + + def test_tokenize_citation_boost_below_one_rejected(self): + from soup_cli.utils.raft import RaftComposed, tokenize_raft_example + + composed = RaftComposed( + prompt="Q", answer="Paris [doc-0]", golden_doc_id="doc-0", + doc_ids=("doc-0",), + ) + with pytest.raises(ValueError, match="boost"): + tokenize_raft_example( + _FakeTokenizer(fast=True), composed, max_length=64, + citation_faithful=True, citation_boost=0.5, + ) + + +class TestRaftDocCap: + """#199 — _MAX_DOCS + _MAX_FIELD_LEN boundaries (MEDIUM).""" + + def test_64_distractors_ok(self): + from soup_cli.utils.raft import build_raft_prompt + + composed = build_raft_prompt(_raft_row(64)) # 1 golden + 64 = 65 = cap + assert len(composed.doc_ids) == 65 + + def test_65_distractors_rejected(self): + from soup_cli.utils.raft import build_raft_prompt + + with pytest.raises(ValueError, match="documents"): + build_raft_prompt(_raft_row(65)) # 66 > 65 cap + + def test_oversize_field_rejected(self): + from soup_cli.utils.raft import build_raft_prompt + + row = _raft_row(1) + row["golden_doc"] = "x" * 70_000 # > _MAX_FIELD_LEN + with pytest.raises(ValueError, match="chars"): + build_raft_prompt(row) + + +class TestV07110FrozenDataclasses: + """LOW — frozen invariant on the new v0.71.10 dataclasses.""" + + def test_raft_composed_frozen(self): + import dataclasses + + from soup_cli.utils.raft import RaftComposed + + c = RaftComposed( + prompt="p", answer="a", golden_doc_id="doc-0", doc_ids=("doc-0",) + ) + with pytest.raises(dataclasses.FrozenInstanceError): + c.answer = "x" + + def test_steering_artifact_frozen(self): + import dataclasses + + from soup_cli.utils.steering import SteeringArtifact + + a = SteeringArtifact( + method="caa", name="n", layer=1, hidden_dim=4, + intervention_point="residual", output_dir="d", base="b", num_pairs=2, + ) + with pytest.raises(dataclasses.FrozenInstanceError): + a.layer = 9 + + def test_loaded_steering_frozen(self): + import dataclasses + + import numpy as np + + from soup_cli.utils.steering import LoadedSteering + + loaded = LoadedSteering( + method="caa", name="n", layer=0, intervention_point="residual", + vector=np.zeros(4, dtype=np.float32), default_strength=1.0, + ) + with pytest.raises(dataclasses.FrozenInstanceError): + loaded.layer = 1 + + def test_radit_run_result_frozen(self): + import dataclasses + + from soup_cli.utils.ra_dit_run import RaDitRunResult + + r = RaDitRunResult( + retriever_output="r", generator_output="g", + retriever_model_used="r", autolinked=True, + ) + with pytest.raises(dataclasses.FrozenInstanceError): + r.autolinked = False + + +class TestRunRaDitValidation: + """#200 — timeout + oversize-yaml validation (LOW).""" + + def _configs(self, tmp_path): + retr = tmp_path / "retriever.yaml" + retr.write_text( + "base: st/mini\ntask: embedding\noutput: ./r\n" + "training:\n ra_dit_stage: retriever\n" + "data:\n train: ./t.jsonl\n format: embedding\n", + encoding="utf-8", + ) + gen = tmp_path / "generator.yaml" + gen.write_text( + "base: m\ntask: sft\noutput: ./g\n" + "training:\n ra_dit_stage: generator\n" + "data:\n train: ./raft.jsonl\n format: raft\n", + encoding="utf-8", + ) + return retr, gen + + def test_bad_timeout_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.ra_dit_run import run_ra_dit + + monkeypatch.chdir(tmp_path) + self._configs(tmp_path) + with pytest.raises(ValueError, match="timeout"): + run_ra_dit( + "retriever.yaml", "generator.yaml", + timeout_seconds=5, _runner=lambda p: None, + ) + + def test_timeout_bool_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.ra_dit_run import run_ra_dit + + monkeypatch.chdir(tmp_path) + self._configs(tmp_path) + with pytest.raises(ValueError, match="bool"): + run_ra_dit( + "retriever.yaml", "generator.yaml", + timeout_seconds=True, _runner=lambda p: None, + ) + + def test_oversize_yaml_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.ra_dit_run import run_ra_dit + + monkeypatch.chdir(tmp_path) + retr, _gen = self._configs(tmp_path) + # Pad the retriever config past the 256KB cap with a comment line. + retr.write_text( + retr.read_text(encoding="utf-8") + "\n# " + "x" * (256 * 1024 + 10), + encoding="utf-8", + ) + with pytest.raises(ValueError, match="cap|exceeds|KB"): + run_ra_dit("retriever.yaml", "generator.yaml", _runner=lambda p: None) + + +class TestInstallSteeringHookExtras: + """#201 — strength cap / bad layer / post-remove revert (LOW).""" + + def _fake_model(self, d=4, n=2): + import torch + import torch.nn as nn + + class FakeLayer(nn.Module): + def __init__(self): + super().__init__() + self.dummy = nn.Parameter(torch.zeros(1)) + + def forward(self, x): + return (x,) + + class FakeInner(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([FakeLayer() for _ in range(n)]) + + class FakeModel(nn.Module): + def __init__(self): + super().__init__() + self.model = FakeInner() + + return FakeModel() + + def _loaded(self, layer, vec): + from soup_cli.utils.steering import LoadedSteering + + return LoadedSteering( + method="caa", name="t", layer=layer, + intervention_point="residual", vector=vec, default_strength=1.0, + ) + + def test_strength_cap_rejected(self): + import numpy as np + + from soup_cli.utils.steering import install_steering_hook + + model = self._fake_model() + vec = np.zeros(4, dtype=np.float32) + with pytest.raises(ValueError, match="<="): + install_steering_hook(model, self._loaded(0, vec), strength=11.0) + + def test_bad_layer_rejected(self): + import numpy as np + + from soup_cli.utils.steering import install_steering_hook + + model = self._fake_model(n=2) + vec = np.zeros(4, dtype=np.float32) + with pytest.raises(ValueError, match="out of range"): + install_steering_hook(model, self._loaded(99, vec), strength=1.0) + + def test_remove_reverts(self): + import numpy as np + import torch + + from soup_cli.utils.steering import install_steering_hook + + model = self._fake_model(d=4) + vec = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32) + handle = install_steering_hook(model, self._loaded(0, vec), strength=2.0) + x = torch.zeros(1, 1, 4) + shifted = model.model.layers[0](x)[0] + assert torch.allclose(shifted[0, 0], torch.tensor([2.0, 0.0, 0.0, 0.0])) + handle.remove() + reverted = model.model.layers[0](x)[0] + assert torch.allclose(reverted[0, 0], torch.zeros(4)) + + +class TestSteeringNoTopLevelTorch: + def test_steering_no_top_level_torch(self): + import soup_cli.utils.steering as steering + + with open(steering.__file__, encoding="utf-8") as fh: + src = fh.read() + assert "\nimport torch" not in src + assert "\nfrom torch" not in src + + +class TestPrepareRaftDatasetExecution: + """#199 — _prepare_raft_dataset map+filter real execution (MEDIUM).""" + + def _call(self, rows, *, max_length, citation=False): + import types + + from soup_cli.trainer.sft import SFTTrainerWrapper + + stub = types.SimpleNamespace(tokenizer=_FakeTokenizer(fast=True)) + cfg = types.SimpleNamespace( + data=types.SimpleNamespace( + raft_shuffle_seed=None, max_length=max_length + ) + ) + tcfg = types.SimpleNamespace( + citation_faithful=citation, citation_style="bracket" + ) + return SFTTrainerWrapper._prepare_raft_dataset( + stub, {"train": rows}, cfg, tcfg + ) + + def test_keeps_trainable_rows(self): + train, eval_ds = self._call([_raft_row(1)], max_length=512) + assert len(train) == 1 + assert eval_ds is None + assert any(w > 0.0 for w in train[0]["loss_weights"]) + + def test_drops_all_masked_rows(self): + # max_length=8 truncates the answer away → all-masked → dropped (M4). + train, _eval = self._call([_raft_row(1), _raft_row(1)], max_length=8) + assert len(train) == 0 + + +class TestRaftComputeLossEdges: + """#199 — compute_loss degenerate-zero + citation-boost shift (MEDIUM).""" + + def test_all_masked_returns_finite_zero(self): + import torch + + from soup_cli.trainer.raft import make_raft_trainer_class + + cls = make_raft_trainer_class(_FakeBaseTrainer) + trainer = cls() + model = _FakeModel(torch.randn(1, 4, 10)) + inputs = { + "input_ids": torch.tensor([[1, 2, 3, 4]]), + "labels": torch.tensor([[-100, -100, -100, -100]]), + "loss_weights": torch.tensor([[0.0, 0.0, 0.0, 0.0]]), + } + loss = trainer.compute_loss(model, inputs) + assert torch.isfinite(loss) + assert loss.item() == pytest.approx(0.0) + + def test_nan_logits_returns_structural_zero_with_grad(self): + import torch + + from soup_cli.trainer.raft import make_raft_trainer_class + + cls = make_raft_trainer_class(_FakeBaseTrainer) + trainer = cls() + logits = torch.full((1, 4, 10), float("nan"), requires_grad=True) + model = _FakeModel(logits) + inputs = { + "input_ids": torch.tensor([[1, 2, 3, 4]]), + "labels": torch.tensor([[-100, -100, 5, 6]]), + "loss_weights": torch.tensor([[0.0, 0.0, 1.0, 1.0]]), + } + loss = trainer.compute_loss(model, inputs) + # NaN forward → weighted mean NaN → structural zero (L2). + assert torch.isfinite(loss) + assert loss.item() == pytest.approx(0.0) + assert loss.requires_grad + + def test_citation_boost_shifts_loss(self): + import torch + + from soup_cli.trainer.raft import make_raft_trainer_class + + cls = make_raft_trainer_class(_FakeBaseTrainer) + trainer = cls() + torch.manual_seed(0) + model = _FakeModel(torch.randn(1, 4, 10)) + labels = torch.tensor([[-100, -100, 5, 6]]) + flat = { + "input_ids": torch.tensor([[1, 2, 3, 4]]), + "labels": labels, + "loss_weights": torch.tensor([[0.0, 0.0, 1.0, 1.0]]), + } + boosted = { + "input_ids": torch.tensor([[1, 2, 3, 4]]), + "labels": labels, + "loss_weights": torch.tensor([[0.0, 0.0, 1.0, 5.0]]), + } + l_flat = trainer.compute_loss(model, flat).item() + l_boost = trainer.compute_loss(model, boosted).item() + # Boosting one answer token's weight shifts the weighted mean. + assert l_flat != pytest.approx(l_boost) + + +class TestV07110ReviewFixRegressions: + """Regression guards for the v0.71.10 review-fix code changes.""" + + def test_validate_ra_dit_config_path_public_and_alias(self): + from soup_cli.utils import ra_dit_run + + assert hasattr(ra_dit_run, "validate_ra_dit_config_path") + assert "validate_ra_dit_config_path" in ra_dit_run.__all__ + # back-compat private alias still points at the public function (M5). + assert ( + ra_dit_run._validate_config_path + is ra_dit_run.validate_ra_dit_config_path + ) + + def test_render_raft_prompt_public_and_alias(self): + from soup_cli.utils import raft + + assert hasattr(raft, "render_raft_prompt") + # back-compat private alias (L3). + assert raft._render_prompt is raft.render_raft_prompt + + def test_discover_skips_corrupt_registry_output(self, tmp_path, monkeypatch): + import os + + from soup_cli.registry.store import RegistryStore + from soup_cli.utils.ra_dit_run import discover_latest_retriever + + db = tmp_path / "reg.db" + os.environ["SOUP_REGISTRY_DB_PATH"] = str(db) + monkeypatch.setenv("SOUP_REGISTRY_DB_PATH", str(db)) + with RegistryStore() as store: + # A corrupt retriever row (oversize output > 512 chars) — the + # discovered output is run through validate_ra_dit_retriever_model + # (M1), so this row is skipped rather than flowing into a config. + store.push( + name="ra-dit-retriever-bad", tag="v1", + base_model="st/mini", task="embedding", run_id=None, + config={ + "task": "embedding", "output": "x" * 600, + "training": {"ra_dit_stage": "retriever"}, + }, + ) + # No clean retriever row → discovery returns None, not the corrupt one. + assert discover_latest_retriever() is None + + def test_load_yaml_config_o_nofollow(self): + import soup_cli.utils.ra_dit_run as rr + + with open(rr.__file__, encoding="utf-8") as fh: + src = fh.read() + assert "O_NOFOLLOW" in src + assert "os.fstat" in src + + def test_load_steering_artifact_outside_cwd_rejected(self, tmp_path): + from soup_cli.utils.steering import load_steering_artifact + + # An absolute out-of-cwd dir is rejected by the containment helper. + with pytest.raises(ValueError, match="cwd"): + load_steering_artifact(str(tmp_path / "elsewhere")) + + def test_train_autolink_source_grep(self): + import soup_cli.commands.train as train_mod + + with open(train_mod.__file__, encoding="utf-8") as fh: + src = fh.read() + assert "autolink_generator_retriever" in src + # advisory is markup-escaped (SEC MED-1). + assert "escape" in src + + def test_serve_steer_strength_source_grep(self): + import soup_cli.commands.serve as serve_mod + + with open(serve_mod.__file__, encoding="utf-8") as fh: + src = fh.read() + assert "validate_steering_strength" in src + + def test_serve_bad_steer_name_message(self): + from typer.testing import CliRunner + + from soup_cli.cli import app + + result = CliRunner().invoke( + app, ["serve", "-m", "model", "--steer", "bad/name"] + ) + assert result.exit_code == 2 + assert "Invalid --steer" in result.output