mirror of https://github.com/razor-ai/soup.git
feat(rag): RAFT span-mask trainer + RA-DIT auto-link + live steering + eval citation (v0.71.10)
Lifts the v0.62.0 RAG-family schema-only stubs to live, validated on SmolLM2-135M: - #199 RAFT: data.format=raft trains answer-only (prompt span masked to -100, [doc-N] citation ids, deterministic doc shuffle by raft_shuffle_seed); rows whose prompt fills max_length are dropped with a warning. New utils/raft.py + trainer/raft.py (RaftDataCollator + weighted-CE _RaftTrainer). - #200 soup ra-dit: one-shot two-stage orchestrator (train retriever -> record it as the generator's paired retriever -> train generator); a generator-stage `soup train` with no retriever set auto-links the latest RA-DIT retriever from the Registry. New utils/ra_dit_run.py + commands/ra_dit.py. - #201 soup steer train/apply + soup serve --steer: live CAA/ITI/RepE fit from {positive, negative} pairs + decode-time forward hook (transformers backend). Lifts the steering.py apply_steering/build_steering_vector stubs. - #202 soup eval citation + citation-span per-token loss boost + 7th `citation` failure mode in soup diagnose. New commands/_eval_v07110.py + diagnose/citation.py. Review fixes (3 agents, all CRITICAL->LOW): markup-escaped autolink advisory; shared enforce_under_cwd_and_no_symlink + O_NOFOLLOW on every new file read; steering-artifact containment; honest RA-DIT docs (records pairing, no weight fusion); public validate_ra_dit_config_path + render_raft_prompt; repe/iti require >=2 pairs; eval citation --shuffle-seed. Full suite: 13034 passed, 106 skipped (13142 collected). ruff clean.
This commit is contained in:
Alpamys
parent
0b62b5d0b7
commit
a2287dd6a5
33
CHANGELOG.md
33
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 <name> --steer-strength <s>`).
|
||||
- **`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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
```
|
||||
|
||||
|
|
|
|||
27
README.md
27
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 <m> --after-model <m> --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 <name>`** — 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).
|
||||
|
||||
|
|
|
|||
|
|
@ -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 <name>` 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
|
||||
|
||||
|
|
|
|||
|
|
@ -171,6 +171,12 @@ soup eval unlearning <run-id> --benchmark tofu|muse|wmdp Forget Quality + Model
|
|||
soup edit set --base <m> --method rome|memit|alphaedit|grace --subject "..." --target "..." [--output <dir>] [--device cpu] [--governor/--no-governor] [--registry-id <id>] Live surgical knowledge edit (--plan-only available)
|
||||
soup edit diff <before-run> <after-run> --probes p.jsonl [--before-model <m> --after-model <m>] 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 <r.yaml> --generator-config <g.yaml> [--retriever-model <m>] [--plan-only] One-shot two-stage RA-DIT: train retriever → record pairing → train generator
|
||||
soup eval citation <data> [--style bracket|inline|footnote] [--shuffle-seed N] [--output o.json] Citation precision/recall/F1 over predictions or RAFT rows
|
||||
soup steer train --base <m> --method caa|iti|repe --name <id> --pairs <jsonl> Fit a CAA/ITI/RepE activation-steering vector from {positive, negative} pairs
|
||||
soup steer apply --name <id> --strength <s> Preview a stored steering vector; soup steer list lists them
|
||||
soup serve --steer <name> [--steer-strength <s>] Apply a steering vector at decode time via a forward hook (transformers backend)
|
||||
soup ingest --source langfuse|langsmith|helicone|openpipe|otel|openai-stored --logs <jsonl> Universal trace importer (6 SaaS adapters → normalised JSONL)
|
||||
soup prune-prompt --input <jsonl> --output <jsonl> --min-frequency 0.95 Detect + strip shared system-prompt prefix
|
||||
soup prune-prompt ... --tokenizer <id-or-path> Tokenizer-aware prefix detection (decodes remaining ids, boundary-safe)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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 <m>` to override.
|
||||
|
||||
|
||||
## Curriculum-Aware Training (BETA)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
"""Soup CLI — Fine-tune LLMs in one command."""
|
||||
|
||||
__version__ = "0.71.9"
|
||||
__version__ = "0.71.10"
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)}")
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 <yaml> --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",
|
||||
)
|
||||
)
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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/<name>). 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
|
||||
<name> --steer-strength <s>`` (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")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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<bracketed>[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<id>[A-Za-z0-9][A-Za-z0-9._\-]{0,127})\]"),
|
||||
"inline": re.compile(r"\((?P<id>[A-Za-z0-9][A-Za-z0-9._\-]{0,127})\)"),
|
||||
"footnote": re.compile(r"\[\^(?P<id>[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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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 <yaml> --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 <p> --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,
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
``<output_dir>/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/<name>`` directory if it
|
||||
holds a ``steering_config.json``; (2) the most recent Registry entry named
|
||||
``<name>`` 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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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 ----------
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
Loading…
Reference in New Issue