From b1ee9e1bb9616320249010f5aee4c21564f70fea Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Mon, 14 Sep 2026 15:50:02 -0700 Subject: [PATCH] feat(distill): [US-006] add `tablassert distill-weigh` to emit weighted training rows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The captured corpus knew each run's outcome and could score it, but nothing composed the two into a trainable dataset — the owner had no command to turn `--distill` output into weighted rows for LoRA/QLoRA SFT. `tablassert distill-weigh --distill-dir --out train.ndjson`: - Partitions the corpus by CONTENT (`is_outcome_file` reads the first line's `record_type`, so a renamed file still partitions correctly), joins records to outcomes on `run_id` with a consistent last-wins index for both flattening and reward, and refuses derived or mixed files already sitting in the corpus (a previous `train.ndjson` re-ingested would double every example). - Computes `reward` with the resolved RewardConfig and `edge_ref` (corpus median over the same deduped outcome pool unless `--edge-ref`/config overrides, so a superseded duplicate can no longer shift weights), applies the chosen policy (`--policy threshold|best-of-n|replication`), and writes exactly one row per input record with the canonical `TRAIN_ROW_KEYS` schema — `replicas` stays a count; rows are never physically duplicated. - `--purpose agent|all|` filters which recorded calls train (judge and reflexion turns are error-recovery trajectories, not teaching signal); `--final-call-only` keeps each run's most complete conversation and RETAINS rows it cannot rank, warning and counting `unrankable_final_call` instead of silently dropping them. - Writes a deterministic manifest (atomic): resolved reward config, resolved `edge_ref` + source, policy knobs, join stats, `selected_count`, `selected_zero_weight` (best-of-n has no weight floor — an operator must see gated-failed trajectories it would select), `unmatched_count`, nested `distinct` diversity counters before/after selection (fixed to read the flattened `outcome_config_yaml_sha256` column), and absolute source paths. Two identical runs produce byte-identical output and manifest. - Fails loud (exit 2, actionable stderr) on every bad input: missing/empty corpus, malformed NDJSON naming file+line, `matched==0` with the pre-v2 hint, bad policy/knobs/purpose/config, `--out`/`--manifest` naming a directory or colliding or landing inside `--distill-dir`. - Zero-dependency (no `[distill]` extra): stdlib + pyyaml only. This is data selection for supervised fine-tuning, explicitly not RLHF — no reward model, no PPO/GRPO/DPO/KTO, no trainer import. Docs updated in the same commit per the repo's source-of-truth gates: `docs/cli.md` (seven commands, command-table row, full section with every flag/default and the manifest contents), `llms.txt` (token-count and module description), and the two sanctioned `test_docs_cli_coverage.py` registries. Testing: `uv run pytest tests/test_distill_weigh.py -q -n0 --no-cov` -> 36 passed, 0 skipped (base env, no extras); docs SSOT gate -> 167 passed; CLI regression -> 100 passed; `uv run pytest -q` -> 1491 passed, 15 skipped, 0 failed; `uv run pyright` -> 0 errors; ruff check + format --check clean; `mkdocs build --strict` clean. Independent Tier-2 review re-verified all ten repaired findings against source and reproduced the artifact exercises (exit-2 paths, retention counts, diversity counters, key-set invariants). --- docs/cli.md | 41 ++- llms.txt | 6 +- src/tablassert/cli.py | 267 +++++++++++++++- tests/test_distill_weigh.py | 524 ++++++++++++++++++++++++++++++++ tests/test_docs_cli_coverage.py | 2 + 5 files changed, 833 insertions(+), 7 deletions(-) create mode 100644 tests/test_distill_weigh.py diff --git a/docs/cli.md b/docs/cli.md index dddc34eb..e35ac736 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -1,7 +1,7 @@ # CLI Reference Tablassert extracts knowledge assertions from tabular data into KGX NDJSON. The `tablassert` app -exposes **six subcommands**: `agent`, `build-fullmap`, `build-kg`, `distill-export`, `validate`, +exposes **seven subcommands**: `agent`, `build-fullmap`, `build-kg`, `distill-export`, `distill-weigh`, `validate`, and `validate-kgx`, plus an app-level `--version` flag. Run `tablassert --help` (or ` --help`) for the live surface. @@ -13,6 +13,7 @@ for the live surface. | [`build-fullmap`](#build-fullmap) | Build the embedded fullmap redb used for entity resolution | | [`build-kg`](#build-kg) | Build a KGX NDJSON knowledge graph from a YAML configuration | | [`distill-export`](#distill-export) | Export a recorded distillation NDJSON dataset to an on-disk Hugging Face dataset | +| [`distill-weigh`](#distill-weigh) | Join distillation records to outcomes and prepare LoRA-SFT training rows | | [`validate`](#validate) | Validate a graph or table configuration without executing it | | [`validate-kgx`](#validate-kgx) | Validate built KGX NDJSON against the Biolink Model | @@ -95,7 +96,8 @@ Use this to convert a distillation dataset recorded with [`agent --distill`](#agent) into an on-disk Hugging Face dataset (`save_to_disk`). Requires the `[distill]` extra (`pip install "tablassert[distill]"`, pulls `datasets`). The raw NDJSON already loads directly in Unsloth Studio and via `datasets.load_dataset("json", ...)` — this export is -only needed for `datasets`-native workflows. +only needed for `datasets`-native workflows. Keep derived training output in a separate directory; +`distill-export` loads every `*.ndjson` under its input directory. ```bash tablassert distill-export --distill-dir .tablassert/agent/distill --out ./hf-dataset @@ -106,6 +108,41 @@ tablassert distill-export --distill-dir .tablassert/agent/distill --out ./hf-dat | `--distill-dir`, `-dd` | Path | Yes | n/a | Directory holding the recorded `*.ndjson` files (exit 2 when empty) | | `--out`, `-o` | Path | Yes | n/a | Destination directory for the `save_to_disk` dataset | +## distill-weigh + +Join `agent --distill` records to their sibling outcomes, compute deterministic reward weights, and +write one flat training row per input record for LoRA/QLoRA supervised fine-tuning. This is data +selection, not RLHF: no reward model or online trainer is involved. Keep the output outside the +input directory because `distill-export` loads every `*.ndjson` in its directory. + +```bash +tablassert distill-weigh --distill-dir .tablassert/agent/distill --out ./training/train.ndjson +# Then optionally convert the weighed rows to a Hugging Face dataset: +tablassert distill-export --distill-dir ./training --out ./hf-dataset +``` + +| Option | Type | Required | Default | Description | +| --- | --- | --- | --- | --- | +| `--distill-dir`, `-dd` | Path | Yes | n/a | Input directory containing record and outcome `*.ndjson` files | +| `--out`, `-o` | Path | Yes | n/a | Training NDJSON destination; must be outside `--distill-dir` | +| `--policy`, `-p` | string (`threshold`, `best-of-n`, or `replication`) | No | `threshold` | Selection policy; validated at runtime | +| `--threshold`, `-t` | float | No | `0.75` | Minimum weight for the `threshold` policy | +| `--top-n`, `-tn` | int | No | `2` | Number retained per `pmc_id` group for `best-of-n` | +| `--replication-k`, `-rk` | int | No | `2` | Replication slope for `replication`, bounded to 0–3 | +| `--reward-config`, `-rc` | Path | No | `None` | YAML/JSON reward configuration override | +| `--edge-ref` | float | No | `None` | Breadth reference override; otherwise the corpus median | +| `--purpose` | str | No | `agent` | Keep this purpose, or use literal `all` to disable filtering | +| `--final-call-only` | bool flag | No | `False` | Keep only the highest `call_index` per run | +| `--manifest` | Path | No | `.manifest.json` | Reproducibility manifest destination | + +The manifest JSON records the resolved reward configuration, resolved `edge_ref` and its source, join +statistics, selected/unmatched counts, and nested `distinct` diversity counters for `pmc_id` and +`config_yaml_sha256` before and after selection. Source paths are resolved absolute paths for portable +provenance. Malformed input, missing records/outcomes, invalid policies or knobs, unmatched records, +and an output inside the input directory fail with exit 2 and an actionable message. A corpus without a +comparable build emits a warning and records a null `edge_ref`; its breadth contribution is 0.0. Replicas +are counts on rows, not physical row duplication. + --- ## build-fullmap diff --git a/llms.txt b/llms.txt index 7179b484..47346fd2 100644 --- a/llms.txt +++ b/llms.txt @@ -9,7 +9,7 @@ Current CLI behavior queries an embedded redb database whose primary file is `fu ## Quickstart - [README](README.md): high-level overview, install snippets, and one-command graph build. - [Installation Guide](docs/installation.md): local/dev/tool installation paths. Optional extras: `rt` (runtime-compatible Polars), `aria2` (bundled aria2c downloader), `qc` (quality-control runtime), `agent` (autonomous agent), `optimize` (GEPA prompt optimization), `distill` (distillation dataset export), and `log` (loguru-backed logging). -- [CLI Reference](docs/cli.md): six commands (`agent`, `build-fullmap`, `build-kg`, `distill-export`, `validate`, and `validate-kgx`) plus the app-level `--version` flag. +- [CLI Reference](docs/cli.md): seven commands (`agent`, `build-fullmap`, `build-kg`, `distill-export`, `distill-weigh`, `validate`, and `validate-kgx`) plus the app-level `--version` flag. - [Tutorial](docs/tutorial.md): first end-to-end run from CSV input to KGX NDJSON output. ## YAML Authoring @@ -23,7 +23,7 @@ Current CLI behavior queries an embedded redb database whose primary file is `fu - [Agent Example Artifacts](examples/agent/README.md): prompt-optimization and distillation example artifacts. ## CLI and Runtime -- [CLI Entry Point](src/tablassert/cli.py): six command implementations, the app-level `--version` flag, and build pipeline stages. +- [CLI Entry Point](src/tablassert/cli.py): seven command implementations, the app-level `--version` flag, and build pipeline stages. - [Pydantic Models](src/tablassert/models.py): authoritative schema for `Section` and `Graph`. - [YAML Ingestion](src/tablassert/ingests.py): `from_yaml()`, `to_sections()`, and template/section merge behavior. - [Fullmap Guide](docs/fullmap.md): entity-resolution database build pipeline and redb schema. @@ -42,7 +42,7 @@ Current CLI behavior queries an embedded redb database whose primary file is `fu - [CLI Module](src/tablassert/cli.py): command implementations and build pipeline stages. - [Coercion Helpers](src/tablassert/coerce.py): annotation and value coercion helpers. - [Distillation Module](src/tablassert/distill.py): ChatML distillation record serialization. -- [Distill Reward Module](src/tablassert/distill_reward.py): pure per-run outcome assembly (`build_outcome`) and config-provenance checks (`provenance_ok`) over the distillation capture schema. +- [Distill Reward Module](src/tablassert/distill_reward.py): reward calculation, selection policies, record/outcome joins, outcome flattening, and training-row schema normalization alongside outcome assembly and config-provenance checks. - [Enums Catalog](src/tablassert/enums.py): Tablassert-owned configuration vocabularies (`Tokens`, `Repositories`, `InformationResources`, `Contributions`, `Comparisons`, `Functions`, `Files`, `EncodingMethods`, `FillMethods`, `SourceStatuses`, `ProvisionMechanisms`, `DataFormats`, `IngestCategories`, `ContentCategories`, `ModelingCategories`). - [Error Types](src/tablassert/errors.py): coded exceptions and validation warnings. - [Extras Registry](src/tablassert/extras.py): optional-dependency checks and install hints. diff --git a/src/tablassert/cli.py b/src/tablassert/cli.py index 6fbf8e6e..68139259 100644 --- a/src/tablassert/cli.py +++ b/src/tablassert/cli.py @@ -1,17 +1,21 @@ from __future__ import annotations import hashlib +import json +import math +import os import re import subprocess import sys import time -from collections.abc import Callable +from collections.abc import Callable, Mapping +from dataclasses import asdict from importlib import import_module from importlib.metadata import version as get_version from itertools import chain, pairwise from multiprocessing import Pool from pathlib import Path -from typing import TYPE_CHECKING, Annotated, Any, BinaryIO, Literal +from typing import TYPE_CHECKING, Annotated, Any, BinaryIO, Literal, NoReturn from urllib.error import HTTPError, URLError from urllib.request import Request, urlopen @@ -1088,6 +1092,265 @@ def metric(key: str, default: float) -> float: ) +@APP.command(name="distill-weigh") +def distill_weigh( + *, + distill_dir: Annotated[Path, cyclopts.Parameter(name=["--distill-dir", "-dd"])], + out: Annotated[Path, cyclopts.Parameter(name=["--out", "-o"])], + policy: Annotated[str, cyclopts.Parameter(name=["--policy", "-p"])] = "threshold", + threshold: Annotated[float, cyclopts.Parameter(name=["--threshold", "-t"])] = 0.75, + top_n: Annotated[int, cyclopts.Parameter(name=["--top-n", "-tn"])] = 2, + replication_k: Annotated[int, cyclopts.Parameter(name=["--replication-k", "-rk"])] = 2, + reward_config: Annotated[Path | None, cyclopts.Parameter(name=["--reward-config", "-rc"])] = None, + edge_ref: Annotated[float | None, cyclopts.Parameter(name="--edge-ref")] = None, + purpose: Annotated[str, cyclopts.Parameter(name="--purpose")] = "agent", + final_call_only: Annotated[bool, cyclopts.Parameter(name="--final-call-only", negative="")] = False, + manifest: Annotated[Path | None, cyclopts.Parameter(name="--manifest")] = None, +) -> None: + """Join distillation records to outcomes and write deterministic training rows.""" + from tablassert.distill_reward import ( + OUTCOME_COLUMN_PREFIX, + POLICIES, + RECORD_TYPE_OUTCOME, + SELECTION_KEYS, + RewardConfig, + RewardConfigError, + is_outcome_file, + iter_record_files, + join_records_outcomes, + load_reward_config, + median_edge_ref, + read_ndjson, + reward, + select, + write_ndjson, + ) + + def fail(message: str) -> NoReturn: + """Print one actionable user error and use the CLI's documented exit status.""" + print(f"tablassert distill-weigh: {message}", file=sys.stderr) + raise SystemExit(2) + + if not distill_dir.exists() or not distill_dir.is_dir(): + fail(f"--distill-dir must be an existing directory: {distill_dir}") + if not out.name or out.is_dir(): + fail(f"--out must name a file, not a directory: {out}") + if manifest is not None and (not manifest.name or manifest.is_dir()): + fail(f"--manifest must name a file, not a directory: {manifest}") + input_dir: Path = distill_dir.resolve() + output_path: Path = out.resolve() + try: + output_path.relative_to(input_dir) + except ValueError: + pass + else: + fail( + f"--out {out} is inside --distill-dir {distill_dir}; distill-export loads every *.ndjson there, " + "so this would duplicate raw records and mix schemas" + ) + manifest_path: Path = manifest.resolve() if manifest is not None else output_path.with_name(f"{output_path.name}.manifest.json") + if manifest_path == output_path: + fail("--manifest and --out must be different paths") + try: + manifest_path.relative_to(input_dir) + except ValueError: + pass + else: + fail(f"--manifest {manifest} is inside --distill-dir {distill_dir}; keep derived artifacts outside the input corpus") + + files: list[Path] = iter_record_files(input_dir) + if not files: + fail(f"no .ndjson files under {distill_dir}; provide a distillation corpus with records and outcomes") + + try: + rows_by_file: dict[Path, list[dict[str, Any]]] = {path: read_ndjson(path) for path in files} + except (OSError, ValueError) as exc: + fail(str(exc)) + + # Content inspection BEFORE the join. Every derived training row carries the selection keys and + # ``outcome_matched``; a raw capture line carries neither, so this separates a previous weigh + # output (or any foreign derived corpus) from raw input without a filename heuristic. + derived_markers: frozenset[str] = frozenset(SELECTION_KEYS) | {"outcome_matched"} + for path, file_rows in rows_by_file.items(): + marker: str | None = next((key for row in file_rows for key in sorted(derived_markers) if key in row), None) + if marker is not None: + fail( + f"derived training output detected in {path} (it carries {marker!r}); --distill-dir must contain " + "the raw records.ndjson/outcomes.ndjson written by agent --distill, not derived training output" + ) + + outcome_files: list[Path] = [path for path in files if is_outcome_file(path)] + record_files: list[Path] = [path for path in files if path not in outcome_files] + # is_outcome_file reads only the FIRST non-blank line, so a concatenated sink is classified as + # records and would otherwise be joined as if every line were a record. + for path in record_files: + if any(row.get("record_type") == RECORD_TYPE_OUTCOME for row in rows_by_file[path]): + fail( + f"mixed record/outcome content detected in {path}; --distill-dir must contain raw " + "records.ndjson/outcomes.ndjson in separate files, not one concatenated file" + ) + if not record_files: + fail(f"no record files under {distill_dir}; this directory contains only outcome files") + if not outcome_files: + fail(f"no outcomes file under {distill_dir}; outcomes may never have been written because a run crashed") + + records: list[dict[str, Any]] = [record for path in record_files for record in rows_by_file[path]] + outcomes: list[dict[str, Any]] = [outcome for path in outcome_files for outcome in rows_by_file[path]] + # A record-classified file may be empty; an outcome-classified file cannot be, because + # is_outcome_file only classifies a file after successfully parsing one non-blank outcome line. + if not records: + fail(f"record files under {distill_dir} contain no records") + + if policy not in POLICIES: + fail(f"unknown --policy {policy!r}; expected one of {', '.join(POLICIES)}") + if not 0.0 <= threshold <= 1.0: + fail(f"--threshold must be in [0, 1], got {threshold!r}") + if top_n < 1: + fail(f"--top-n must be at least 1, got {top_n!r}") + if not 0 <= replication_k <= 3: + fail(f"--replication-k must be in [0, 3], got {replication_k!r}") + + live_purposes: list[str] = sorted({value for value in (record.get("purpose") for record in records) if isinstance(value, str)}) + if purpose != "all" and purpose not in live_purposes: + present: str = ", ".join(live_purposes) if live_purposes else "none" + fail(f"unrecognized --purpose {purpose!r}; live purposes are {present}; also accepted: all") + + config: RewardConfig + try: + config = RewardConfig() if reward_config is None else load_reward_config(reward_config) + except (OSError, RewardConfigError, TypeError, ValueError) as exc: + fail(str(exc)) + + try: + rows, join_stats = join_records_outcomes(records, outcomes) + except (TypeError, ValueError) as exc: + fail(str(exc)) + if join_stats["matched"] == 0 and join_stats["records"] > 0: + fail("no records joined an outcome; likely a pre-v2 corpus with no run_id, or outcomes were never written because the run crashed") + + # Rebuild the same append-order index used by join_records_outcomes. The last line wins for + # both reward and flattening, and non-string run ids are unjoinable by definition. + outcome_by_run: dict[str, Mapping[str, Any]] = {} + for outcome in outcomes: + run_id: object = outcome.get("run_id") + if isinstance(run_id, str): + outcome_by_run[run_id] = outcome + comparable_outcomes: list[Mapping[str, Any]] = list(outcome_by_run.values()) + if edge_ref is not None and ( + isinstance(edge_ref, bool) or not isinstance(edge_ref, (int, float)) or not math.isfinite(edge_ref) or edge_ref <= 0 + ): + fail(f"--edge-ref must be a finite positive number, got {edge_ref!r}") + resolved_edge_ref: float | None + edge_ref_source: str + if edge_ref is not None: + resolved_edge_ref, edge_ref_source = float(edge_ref), "overridden" + elif config.edge_ref is not None: + resolved_edge_ref, edge_ref_source = config.edge_ref, "overridden" + else: + resolved_edge_ref = median_edge_ref(comparable_outcomes) + edge_ref_source = "derived" if resolved_edge_ref is not None else "null" + if resolved_edge_ref is None: + print("tablassert distill-weigh: warning: edge_ref is null; breadth contributes 0.0 for every row", file=sys.stderr) + + def final_call_rankable(record: Mapping[str, Any]) -> bool: + """True when a record can be ranked for --final-call-only: a string run id plus an int call index.""" + rank_run_id: object = record.get("run_id") + rank_call_index: object = record.get("call_index") + return isinstance(rank_run_id, str) and isinstance(rank_call_index, int) and not isinstance(rank_call_index, bool) + + unrankable_final_call: int = 0 + filtered_indices: list[int] = list(range(len(rows))) + if purpose != "all": + filtered_indices = [index for index in filtered_indices if records[index].get("purpose") == purpose] + if final_call_only: + final_by_run: dict[str, tuple[int, int]] = {} + rankable: set[int] = set() + for index in filtered_indices: + record: Mapping[str, Any] = records[index] + if not final_call_rankable(record): + continue # an unrankable record cannot be proven non-final, so it is retained, never dropped + rankable.add(index) + rank_run_id: str = record["run_id"] + rank_call_index: int = record["call_index"] + previous = final_by_run.get(rank_run_id) + if previous is None or (rank_call_index, index) > previous: + final_by_run[rank_run_id] = (rank_call_index, index) + retained: set[int] = {value[1] for value in final_by_run.values()} + unrankable_final_call = len(filtered_indices) - len(rankable) + if unrankable_final_call: + print( + f"tablassert distill-weigh: warning: --final-call-only retained {unrankable_final_call} record(s) " + "with no rankable run_id/call_index pair; they cannot be proven non-final", + file=sys.stderr, + ) + filtered_indices = [index for index in filtered_indices if index not in rankable or index in retained] + filtered_rows: list[dict[str, Any]] = [rows[index] for index in filtered_indices] + if not filtered_rows: + # Defensive: --purpose is validated against the live purposes above and an unrankable record + # is retained rather than dropped, so only a future narrowing filter can empty the selection. + narrowing: str = "--final-call-only" if final_call_only else f"purpose filter {purpose!r}" + fail(f"{narrowing} matched no records; live purposes are {', '.join(live_purposes) or 'none'}") + try: + for row in filtered_rows: + run_id = row.get("run_id") + nested = outcome_by_run.get(run_id) if isinstance(run_id, str) else None + row["weight"] = 0.0 if nested is None else reward(nested, config, edge_ref=resolved_edge_ref) + except (TypeError, ValueError) as exc: + fail(str(exc)) + try: + selected_rows = select(filtered_rows, policy=policy, threshold=threshold, top_n=top_n, replication_k=replication_k) + except (TypeError, ValueError) as exc: + fail(str(exc)) + + try: + rows_written: int = write_ndjson(output_path, selected_rows) + # The record schema carries no config hash of its own: the config identity rides on the + # flattened ``outcome_config_yaml_sha256`` column, so the diversity counter must read that. + config_column: str = f"{OUTCOME_COLUMN_PREFIX}config_yaml_sha256" + before_pmc = {row.get("pmc_id") for row in selected_rows if row.get("pmc_id") is not None} + before_config = {row.get(config_column) for row in selected_rows if row.get(config_column) is not None} + chosen = [row for row in selected_rows if row.get("selected") is True] + after_pmc = {row.get("pmc_id") for row in chosen if row.get("pmc_id") is not None} + after_config = {row.get(config_column) for row in chosen if row.get(config_column) is not None} + resolved_config: dict[str, Any] = asdict(config) + resolved_config["edge_ref"] = resolved_edge_ref + manifest_data: dict[str, Any] = { + "schema_version": 2, + "tablassert_version": get_version("tablassert"), + "reward_config": resolved_config, + "edge_ref": resolved_edge_ref, + "edge_ref_source": edge_ref_source, + "policy": policy, + "threshold": threshold, + "top_n": top_n, + "replication_k": replication_k, + "purpose": purpose, + "final_call_only": final_call_only, + "unrankable_final_call": unrankable_final_call, + "join_stats": join_stats, + "rows_written": rows_written, + "selected_count": len(chosen), + "selected_zero_weight": sum(1 for row in chosen if row.get("weight") == 0.0), + "unmatched_count": join_stats["unmatched"], + "distinct": { + "pmc_id": {"before_selection": len(before_pmc), "after_selection": len(after_pmc)}, + "config_yaml_sha256": {"before_selection": len(before_config), "after_selection": len(after_config)}, + }, + "source_files": [str(path) for path in files], + "source_files_by_kind": {"records": [str(path) for path in record_files], "outcomes": [str(path) for path in outcome_files]}, + } + manifest_path.parent.mkdir(parents=True, exist_ok=True) + temporary = manifest_path.with_name(f".{manifest_path.name}.tmp") + try: + temporary.write_text(json.dumps(manifest_data, ensure_ascii=False, sort_keys=True, indent=2) + "\n", encoding="utf-8") + os.replace(temporary, manifest_path) + finally: + temporary.unlink(missing_ok=True) + except (OSError, TypeError, ValueError) as exc: + fail(f"could not write output or manifest: {exc}") + print(f"tablassert distill-weigh: {rows_written} row(s), {len(chosen)} selected, policy {policy}, edge_ref {resolved_edge_ref}, output {out}") + + @APP.command(name="distill-export") def distill_export( *, distill_dir: Annotated[Path, cyclopts.Parameter(name=["--distill-dir", "-dd"])], out: Annotated[Path, cyclopts.Parameter(name=["--out", "-o"])] diff --git a/tests/test_distill_weigh.py b/tests/test_distill_weigh.py new file mode 100644 index 00000000..6e4121cd --- /dev/null +++ b/tests/test_distill_weigh.py @@ -0,0 +1,524 @@ +"""Base-environment tests for the distill-weigh training-row command.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +import pytest + +from tablassert.cli import distill_weigh +from tablassert.distill import RECORD_KEYS +from tablassert.distill_reward import SELECTION_KEYS, TRAIN_ROW_KEYS + + +def _record(run_id: str | None, *, purpose: str = "agent", call_index: int = 0, pmc_id: str = "PMC1") -> dict[str, object]: + """Return a minimal schema-v2 record.""" + return { + **dict.fromkeys(RECORD_KEYS), + "record_type": "record", + "schema_version": 2, + "run_id": run_id, + "purpose": purpose, + "call_index": call_index, + "messages": [{"role": "user", "content": "hello"}], + "pmc_id": pmc_id, + } + + +def _outcome(run_id: str | None, *, edge_count: int = 100, ok: bool = True, coverage: float = 1.0) -> dict[str, object]: + """Return a complete rewardable outcome.""" + return { + "record_type": "outcome", + "schema_version": 2, + "run_id": run_id, + "timestamp": "fixed", + "pmc_id": "PMC1", + "model_id": "model", + "run_status": "MAPPED", + "ok": ok, + "measured": True, + "head": False, + "coverage_pct": coverage, + "best_coverage": coverage, + "coverage_history_len": 1, + "section_coverages_len": 1, + "biolink_valid_pct": 1.0, + "biolink_valid_pct_strict": 1.0, + "demoted_edge_pct": 0.0, + "node_count": 10, + "edge_count": edge_count, + "unresolved_count": 0, + "predicate_advice_count": 0, + "multivalued_suspect_count": 0, + "error_codes": [], + "attempts": 1, + "config_chars": 10, + "config_yaml_sha256": "hash", + "provenance_ok": True, + "qc_pass_rate": 1.0, + "tool_calls": {"total": 1, "failed": 0, "wrong": 0, "redundant": 0}, + "tokens_total": 1, + "steps": 1, + "judge_score": None, + "judge_dimensions": None, + "gate": {"map_threshold": 0.25, "biolink_threshold": 0.0, "judge_threshold": None}, + "versions": {"tablassert": "18.0.0", "biolink_model": "4.4.4"}, + } + + +def _write(path: Path, rows: list[dict[str, object]]) -> None: + """Write deterministic JSON lines.""" + path.write_text("\n".join(json.dumps(row, sort_keys=True) for row in rows) + "\n", encoding="utf-8") + + +def _corpus(tmp_path: Path, *, records: list[dict[str, object]] | None = None, outcomes: list[dict[str, object]] | None = None) -> Path: + """Create a small records/outcomes corpus.""" + directory = tmp_path / "distill" + directory.mkdir() + _write(directory / "records.ndjson", records or [_record("run:PMC1"), _record("missing", pmc_id="PMC2")]) + _write(directory / "renamed.ndjson", outcomes or [_outcome("run:PMC1")]) + return directory + + +def test_distill_weigh_writes_a_weighted_dataset_and_manifest(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + """The command writes flat rows, last-wins outcomes, and reproducibility metadata. + + Why: this is the user-facing boundary where append-only capture becomes training data, so row + cardinality, key shape, duplicate accounting, and manifest provenance must be checked together. + """ + directory = _corpus(tmp_path, outcomes=[_outcome("run:PMC1", edge_count=10), _outcome("run:PMC1", edge_count=100)]) + output = tmp_path / "train" / "rows.ndjson" + distill_weigh(distill_dir=directory, out=output) + rows = [json.loads(line) for line in output.read_text(encoding="utf-8").splitlines()] + manifest = json.loads(output.with_name("rows.ndjson.manifest.json").read_text(encoding="utf-8")) + assert len(rows) == 2 + assert all(tuple(row) == TRAIN_ROW_KEYS for row in rows) + assert rows[0]["outcome_edge_count"] == 100 + assert rows[1]["outcome_matched"] is False + assert manifest["join_stats"] == {"records": 2, "outcomes": 2, "matched": 1, "unmatched": 1, "duplicate_run_ids": 1} + assert manifest["edge_ref"] == 100.0 + assert manifest["edge_ref_source"] == "derived" + assert manifest["rows_written"] == 2 + assert "selected" in capsys.readouterr().out + + +@pytest.mark.parametrize("policy", ["threshold", "best-of-n", "replication"]) +def test_distill_weigh_supports_every_selection_policy(tmp_path: Path, policy: str) -> None: + """Each documented selection policy writes the canonical annotations. + + Why: policy dispatch is a CLI boundary, so pure selector coverage alone would not catch a + misspelled option, an omitted manifest policy, or a policy-specific write failure. + """ + directory = _corpus(tmp_path) + output = tmp_path / f"{policy}.ndjson" + distill_weigh(distill_dir=directory, out=output, policy=policy) + rows = [json.loads(line) for line in output.read_text(encoding="utf-8").splitlines()] + manifest = json.loads(output.with_name(f"{policy}.ndjson.manifest.json").read_text(encoding="utf-8")) + assert len(rows) == 2 + assert manifest["policy"] == policy + assert all(set(SELECTION_KEYS) <= set(row) for row in rows) + if policy == "best-of-n": + assert sum(row["selected"] for row in rows) == 2 + if policy == "replication": + assert sum(row["selected"] for row in rows) == 1 + + +def test_distill_weigh_rejects_directory_outputs_and_manifest_collisions(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + """Directory destinations and manifest collisions fail before corpus processing. + + Why: deriving a manifest from an empty output name previously raised an uncaught ValueError, + while collisions could overwrite the training artifact or feed derived files back as input. + """ + directory = _corpus(tmp_path) + with pytest.raises(SystemExit) as out_exc: + distill_weigh(distill_dir=directory, out=Path(".")) + assert out_exc.value.code == 2 + assert "must name a file" in capsys.readouterr().err + + with pytest.raises(SystemExit) as dir_exc: + distill_weigh(distill_dir=directory, out=tmp_path) + assert dir_exc.value.code == 2 + assert "must name a file" in capsys.readouterr().err + + output = tmp_path / "rows.ndjson" + with pytest.raises(SystemExit) as collision_exc: + distill_weigh(distill_dir=directory, out=output, manifest=output) + assert collision_exc.value.code == 2 + assert "different paths" in capsys.readouterr().err + + with pytest.raises(SystemExit) as manifest_exc: + distill_weigh(distill_dir=directory, out=output, manifest=directory) + assert manifest_exc.value.code == 2 + assert "must name a file" in capsys.readouterr().err + + +def test_distill_weigh_rejects_manifest_inside_input_dir(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + """A manifest nested under the raw corpus is rejected. + + Why: derived JSON beside append-only capture would be rediscovered by later corpus scans and + make the same input directory change meaning after one successful weigh. + """ + directory = _corpus(tmp_path) + with pytest.raises(SystemExit) as exc_info: + distill_weigh(distill_dir=directory, out=tmp_path / "rows.ndjson", manifest=directory / "manifest.json") + assert exc_info.value.code == 2 + assert "inside --distill-dir" in capsys.readouterr().err + + +def test_distill_weigh_rejects_empty_record_and_outcome_files(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + """Empty raw files fail with a specific missing-content error. + + Why: an empty sink can result from a crashed or interrupted capture and must not produce a + plausible training artifact with silently missing records or outcomes. + """ + for empty_name, valid_name, expected in [ + ("records.ndjson", "outcomes.ndjson", "contain no records"), + ("outcomes.ndjson", "records.ndjson", "no outcomes"), + ]: + directory = tmp_path / empty_name.replace(".ndjson", "") + directory.mkdir() + (directory / empty_name).write_text("", encoding="utf-8") + _write(directory / valid_name, [_outcome("r")] if valid_name.startswith("outcome") else [_record("r")]) + with pytest.raises(SystemExit) as exc_info: + distill_weigh(distill_dir=directory, out=tmp_path / f"{empty_name}.out.ndjson") + assert exc_info.value.code == 2 + assert expected in capsys.readouterr().err + + +def test_distill_weigh_rejects_bad_reward_config(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + """An invalid reward configuration exits 2 instead of falling back to defaults. + + Why: reward coefficients decide which rows train, so a typo or malformed file must be visible + to the operator rather than silently changing selection semantics. + """ + directory = _corpus(tmp_path) + config = tmp_path / "bad.yaml" + config.write_text("w_coverage: not-a-number\n", encoding="utf-8") + with pytest.raises(SystemExit) as exc_info: + distill_weigh(distill_dir=directory, out=tmp_path / "rows.ndjson", reward_config=config) + assert exc_info.value.code == 2 + assert "Reward config invalid" in capsys.readouterr().err + + +def test_distill_weigh_records_edge_ref_overrides(tmp_path: Path) -> None: + """Explicit and reward-config edge references resolve as overrides in the manifest. + + Why: the breadth denominator changes reward values and therefore selection; the manifest must + make whether it was derived or overridden auditable for both supported override surfaces. + """ + directory = _corpus(tmp_path) + explicit = tmp_path / "explicit.ndjson" + distill_weigh(distill_dir=directory, out=explicit, edge_ref=25.0) + explicit_manifest = json.loads(explicit.with_name("explicit.ndjson.manifest.json").read_text(encoding="utf-8")) + assert explicit_manifest["edge_ref"] == 25.0 + assert explicit_manifest["edge_ref_source"] == "overridden" + + config = tmp_path / "reward.yaml" + config.write_text("edge_ref: 25\n", encoding="utf-8") + configured = tmp_path / "configured.ndjson" + distill_weigh(distill_dir=directory, out=configured, reward_config=config) + configured_manifest = json.loads(configured.with_name("configured.ndjson.manifest.json").read_text(encoding="utf-8")) + assert configured_manifest["edge_ref"] == 25.0 + assert configured_manifest["edge_ref_source"] == "overridden" + + +def test_distill_weigh_retains_unrankable_final_call_rows(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + """Final-call filtering retains rows lacking a provable ranking pair. + + Why: null/non-integer call metadata and non-string run IDs cannot prove that a row is not + final, so silently dropping them would violate the retained-unmatched invariant. + """ + unrankable_run = _record(None) + non_string_run = _record("r2") + non_string_run["run_id"] = 42 + directory = _corpus(tmp_path, records=[_record("run:PMC1", call_index=0), _record("run:PMC1", call_index=1), unrankable_run, non_string_run]) + output = tmp_path / "rows.ndjson" + distill_weigh(distill_dir=directory, out=output, final_call_only=True) + rows = [json.loads(line) for line in output.read_text(encoding="utf-8").splitlines()] + manifest = json.loads(output.with_name("rows.ndjson.manifest.json").read_text(encoding="utf-8")) + assert [row["call_index"] for row in rows] == [1, 0, 0] + assert manifest["unrankable_final_call"] == 2 + assert "cannot be proven non-final" in capsys.readouterr().err + + +def test_distill_weigh_rejects_derived_and_mixed_input_files(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + """Derived markers and mixed record/outcome files are rejected by content inspection. + + Why: filename-based partitioning cannot distinguish a prior weighted artifact or a concatenated + sink, and silently joining either can duplicate rows or misclassify outcome data as records. + """ + directory = _corpus(tmp_path) + _write(directory / "derived.ndjson", [{"record_type": "record", "weight": 0.5, "selected": True, "policy": "threshold", "outcome_matched": True}]) + with pytest.raises(SystemExit) as derived_exc: + distill_weigh(distill_dir=directory, out=tmp_path / "derived-out.ndjson") + assert derived_exc.value.code == 2 + assert "derived training output" in capsys.readouterr().err + + mixed = tmp_path / "mixed" + mixed.mkdir() + _write(mixed / "records.ndjson", [_record("r"), _outcome("r")]) + with pytest.raises(SystemExit) as mixed_exc: + distill_weigh(distill_dir=mixed, out=tmp_path / "mixed-out.ndjson") + assert mixed_exc.value.code == 2 + assert "mixed record/outcome" in capsys.readouterr().err + + +def test_distill_weigh_manifest_uses_canonical_counts_and_absolute_sources(tmp_path: Path) -> None: + """The manifest has one canonical count spelling and absolute source provenance. + + Why: stable names prevent downstream readers from choosing between duplicate aliases, while + resolved paths make manifests comparable when the command is invoked from different directories. + """ + directory = _corpus(tmp_path) + output = tmp_path / "rows.ndjson" + distill_weigh(distill_dir=directory, out=output) + manifest = json.loads(output.with_name("rows.ndjson.manifest.json").read_text(encoding="utf-8")) + assert "selected_count" in manifest + assert "selected" not in manifest + assert "unmatched_count" in manifest + assert "unmatched" not in manifest + assert set(manifest["distinct"]) == {"pmc_id", "config_yaml_sha256"} + assert all(Path(path).is_absolute() for path in manifest["source_files"]) + assert all(Path(path).is_absolute() for paths in manifest["source_files_by_kind"].values() for path in paths) + + +def test_distill_weigh_manifest_contains_diversity_and_join_counts(tmp_path: Path) -> None: + """The manifest exposes the resolved diversity and join accounting needed for audit. + + Why: selection can reduce prompt/config diversity even when row counts look healthy, so both + before/after nested counters and canonical selected/unmatched totals must be machine-readable. + """ + directory = _corpus(tmp_path) + output = tmp_path / "rows.ndjson" + distill_weigh(distill_dir=directory, out=output) + manifest = json.loads(output.with_name("rows.ndjson.manifest.json").read_text(encoding="utf-8")) + assert manifest["selected_count"] == 1 + assert manifest["unmatched_count"] == 1 + assert manifest["distinct"]["pmc_id"] == {"before_selection": 2, "after_selection": 1} + assert manifest["distinct"]["config_yaml_sha256"] == {"before_selection": 1, "after_selection": 1} + + +def test_distill_weigh_writes_unmatched_rows_with_zero_weight(tmp_path: Path) -> None: + """Unmatched records remain visible as unselected zero-weight rows. + + Why: unmatched rows are audit data rather than training candidates; dropping them or marking + them selectable would hide corpus-join problems and corrupt downstream selection counts. + """ + directory = _corpus(tmp_path) + output = tmp_path / "rows.ndjson" + distill_weigh(distill_dir=directory, out=output) + rows = [json.loads(line) for line in output.read_text(encoding="utf-8").splitlines()] + unmatched = [row for row in rows if row["outcome_matched"] is False] + assert len(rows) == 2 + assert len(unmatched) == 1 + assert unmatched[0]["weight"] == 0.0 + assert unmatched[0]["selected"] is False + assert unmatched[0]["replicas"] == 0 + + +def test_distill_weigh_manifest_source_paths_survive_a_relative_invocation(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """A relative ``--distill-dir`` still yields resolved absolute manifest source paths. + + Why: the manifest is the reproducibility record for a training artifact, so two identical runs + launched from different working directories must produce byte-comparable provenance. + """ + directory = _corpus(tmp_path) + other_cwd = tmp_path / "elsewhere" + other_cwd.mkdir() + monkeypatch.chdir(other_cwd) + output = tmp_path / "rows.ndjson" + distill_weigh(distill_dir=Path("../distill"), out=output) + manifest = json.loads(output.with_name("rows.ndjson.manifest.json").read_text(encoding="utf-8")) + assert manifest["source_files"] == [str(directory / "records.ndjson"), str(directory / "renamed.ndjson")] + assert manifest["source_files_by_kind"]["records"] == [str(directory / "records.ndjson")] + assert manifest["source_files_by_kind"]["outcomes"] == [str(directory / "renamed.ndjson")] + + +def test_distill_weigh_requires_no_optional_extra() -> None: + """The command source has no optional-extra preflight or gated dependency. + + Why: weighing is deliberately available in the base installation; requiring the distill extra + here would prevent the core path from producing the NDJSON consumed by later export steps. + """ + source = Path(__file__).parents[1] / "src/tablassert/cli.py" + text = source.read_text(encoding="utf-8") + command = text[text.index('@APP.command(name="distill-weigh")') : text.index('@APP.command(name="distill-export")')] + assert "extras.require" not in command + assert "extras.is_installed" not in command + + +@pytest.mark.parametrize( + ("case", "expected"), + [ + ("missing", "existing directory"), + ("empty", "no .ndjson"), + ("records_only", "no outcomes"), + ("outcomes_only", "only outcome"), + ("malformed", "malformed NDJSON"), + ("unmatched", "pre-v2 corpus"), + ("bad_policy", "unknown --policy"), + ("bad_threshold", "--threshold"), + ("bad_top_n", "--top-n"), + ("bad_replication", "--replication-k"), + ("bad_purpose", "unrecognized --purpose"), + ("bad_edge_ref", "--edge-ref"), + ("inside", "loads every *.ndjson"), + ], +) +def test_distill_weigh_fails_loud_on_every_bad_input(tmp_path: Path, capsys: pytest.CaptureFixture[str], case: str, expected: str) -> None: + """Every specified malformed corpus or knob exits 2 with actionable stderr. + + Why: silently producing a plausible-looking zero-weight file is worse than rejecting a torn, + stale, or misconfigured corpus because it can contaminate a training run undetected. + """ + directory = tmp_path / "distill" + output = tmp_path / "out.ndjson" + kwargs: dict[str, object] = {"distill_dir": directory, "out": output} + if case == "missing": + pass + elif case == "empty": + directory.mkdir() + elif case == "records_only": + directory.mkdir() + _write(directory / "records.ndjson", [_record("r")]) + elif case == "outcomes_only": + directory.mkdir() + _write(directory / "outcomes.ndjson", [_outcome("r")]) + elif case == "malformed": + directory.mkdir() + (directory / "records.ndjson").write_text('{"record_type":"record"}\nnot-json\n', encoding="utf-8") + _write(directory / "outcomes.ndjson", [_outcome("r")]) + elif case == "unmatched": + directory = _corpus(tmp_path, records=[_record("missing")], outcomes=[_outcome("other")]) + kwargs["distill_dir"] = directory + elif case == "bad_policy": + directory = _corpus(tmp_path) + kwargs.update(distill_dir=directory, policy="wat") + elif case == "bad_threshold": + directory = _corpus(tmp_path) + kwargs.update(distill_dir=directory, threshold=2.0) + elif case == "bad_top_n": + directory = _corpus(tmp_path) + kwargs.update(distill_dir=directory, top_n=0) + elif case == "bad_replication": + directory = _corpus(tmp_path) + kwargs.update(distill_dir=directory, replication_k=4) + elif case == "bad_purpose": + directory = _corpus(tmp_path) + kwargs.update(distill_dir=directory, purpose="other") + elif case == "bad_edge_ref": + directory = _corpus(tmp_path) + kwargs.update(distill_dir=directory, edge_ref=0.0) + else: + directory = _corpus(tmp_path) + kwargs.update(distill_dir=directory, out=directory / "train.ndjson") + with pytest.raises(SystemExit) as exc_info: + distill_weigh(**kwargs) # type: ignore[arg-type] + assert exc_info.value.code == 2 + assert expected.lower() in capsys.readouterr().err.lower() + + +def test_distill_weigh_rejects_out_inside_distill_dir(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + """An output nested under the corpus is rejected before any write. + + Why: distill-export globs every NDJSON in its input directory, so accepting this path would + make a later export load both raw records and derived rows as if they were one corpus. + """ + directory = _corpus(tmp_path) + with pytest.raises(SystemExit): + distill_weigh(distill_dir=directory, out=directory / "train.ndjson") + assert "loads every *.ndjson" in capsys.readouterr().err + + +def test_distill_weigh_fails_loud_when_no_record_joins_an_outcome(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + """A nonempty but wholly unmatched corpus is rejected. + + Why: an all-zero dataset can look successful while actually indicating a pre-v2 corpus or a + crashed run whose outcome sink was never written. + """ + directory = _corpus(tmp_path, records=[_record("missing")], outcomes=[_outcome("other")]) + with pytest.raises(SystemExit): + distill_weigh(distill_dir=directory, out=tmp_path / "rows.ndjson") + assert "pre-v2" in capsys.readouterr().err + + +def test_distill_weigh_default_purpose_is_agent_and_all_disables_the_filter(tmp_path: Path) -> None: + """The default excludes judge rows while purpose=all includes them. + + Why: judge and reflexion conversations are distinct trajectories and should not silently enter + the default supervised corpus, while an explicit all-purpose request must remain available. + """ + records = [_record("run:PMC1", purpose="agent"), _record("run:PMC1", purpose="judge", call_index=1)] + directory = _corpus(tmp_path, records=records) + default_out = tmp_path / "default.ndjson" + all_out = tmp_path / "all.ndjson" + distill_weigh(distill_dir=directory, out=default_out) + distill_weigh(distill_dir=directory, out=all_out, purpose="all") + assert len(default_out.read_text(encoding="utf-8").splitlines()) == 1 + assert len(all_out.read_text(encoding="utf-8").splitlines()) == 2 + + +def test_distill_weigh_final_call_only_keeps_the_highest_call_index(tmp_path: Path) -> None: + """Final-call filtering retains only the highest call_index per run. + + Why: the final conversation is the most complete trajectory and intermediate calls can contain + recovery noise that should not be mixed into the default training artifact. + """ + directory = _corpus(tmp_path, records=[_record("run:PMC1", call_index=0), _record("run:PMC1", call_index=2)]) + output = tmp_path / "rows.ndjson" + distill_weigh(distill_dir=directory, out=output, final_call_only=True) + rows = [json.loads(line) for line in output.read_text(encoding="utf-8").splitlines()] + assert len(rows) == 1 + assert rows[0]["call_index"] == 2 + + +def test_distill_weigh_is_byte_reproducible_across_runs(tmp_path: Path) -> None: + """Identical inputs produce byte-identical output and manifest files. + + Why: reproducibility lets an operator audit selection changes as policy/config changes rather + than confusing timestamp or filesystem noise with a training-data change. + """ + directory = _corpus(tmp_path) + first = tmp_path / "first.ndjson" + second = tmp_path / "second.ndjson" + distill_weigh(distill_dir=directory, out=first) + distill_weigh(distill_dir=directory, out=second) + assert hashlib.sha256(first.read_bytes()).digest() == hashlib.sha256(second.read_bytes()).digest() + assert ( + hashlib.sha256(first.with_name("first.ndjson.manifest.json").read_bytes()).digest() + == hashlib.sha256(second.with_name("second.ndjson.manifest.json").read_bytes()).digest() + ) + + +def test_distill_weigh_warns_and_records_null_edge_ref(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + """Head-only or otherwise incomparable outcomes warn and retain a null edge reference. + + Why: a corpus with no comparable full build is still useful, but breadth must not acquire a + fabricated denominator that changes its reward silently. + """ + outcome = _outcome("run:PMC1", edge_count=0) + outcome["head"] = True + directory = _corpus(tmp_path, outcomes=[outcome]) + output = tmp_path / "rows.ndjson" + distill_weigh(distill_dir=directory, out=output) + manifest = json.loads(output.with_name("rows.ndjson.manifest.json").read_text(encoding="utf-8")) + assert manifest["edge_ref"] is None + assert "breadth contributes 0.0" in capsys.readouterr().err + + +def test_distill_weigh_applies_a_reward_config_file(tmp_path: Path) -> None: + """A YAML reward configuration changes the resolved reward policy in the manifest. + + Why: operators must be able to retune data selection without re-recording the append-only + corpus, while still retaining the exact resolved configuration beside the output. + """ + directory = _corpus(tmp_path) + config = tmp_path / "reward.yaml" + config.write_text("w_coverage: 0.0\nw_biolink: 0.28\nw_specificity: 0.17\nw_cleanliness: 0.07\nw_breadth: 0.48\n", encoding="utf-8") + output = tmp_path / "rows.ndjson" + distill_weigh(distill_dir=directory, out=output, reward_config=config) + manifest = json.loads(output.with_name("rows.ndjson.manifest.json").read_text(encoding="utf-8")) + assert manifest["reward_config"]["w_coverage"] == 0.0 + assert manifest["reward_config"]["w_breadth"] == 0.48 diff --git a/tests/test_docs_cli_coverage.py b/tests/test_docs_cli_coverage.py index 758b57cf..556ccc3b 100644 --- a/tests/test_docs_cli_coverage.py +++ b/tests/test_docs_cli_coverage.py @@ -26,6 +26,7 @@ "build-fullmap": ("cli.md", "fullmap.md"), "build-kg": ("cli.md",), "distill-export": ("cli.md",), + "distill-weigh": ("cli.md",), "validate": ("cli.md",), "validate-kgx": ("cli.md",), } @@ -34,6 +35,7 @@ "build-fullmap": ("cli.md", "fullmap.md"), "build-kg": ("cli.md",), "distill-export": ("cli.md",), + "distill-weigh": ("cli.md",), "validate": ("cli.md",), "validate-kgx": ("cli.md",), }