ttucs 0.2.0__py3-none-any.whl

ttucs/__init__.py

@@ -0,0 +1,20 @@
+"""ttucs - Text to UCS.
+
+Classify spoken field-recording slates (speech-to-text output) into Universal Category
+System (UCS) CatIDs::
+
+    from ttucs import TTUCS
+
+    model = TTUCS.from_pretrained("models/ttucs-bow")
+    pred = model.predict("project nightfall, location tofino beach, waves crashing on rocks, take 3")
+    print(pred.cat_id, pred.confidence, pred.filename)
+"""
+
+from ttucs._version import __version__
+from ttucs.pipeline import TTUCS, Candidate, Prediction
+from ttucs.slate import SlateFields, SlateParser, parse_slate
+from ttucs.taxonomy import UCS_VERSION, UCSEntry, UCSFilename, UCSTaxonomy, parse_ucs_filename
+
+__all__ = ["TTUCS", "UCS_VERSION", "Candidate", "Prediction", "SlateFields", "SlateParser",
+           "UCSEntry", "UCSFilename", "UCSTaxonomy", "__version__", "parse_slate",
+           "parse_ucs_filename"]

ttucs/_extras.py

@@ -0,0 +1,33 @@
+"""Helpful errors for optional dependency groups (``pip install "ttucs[train]"``)."""
+
+from __future__ import annotations
+
+import importlib.util
+
+TRAIN_MODULES = {"sklearn": "scikit-learn", "matplotlib": "matplotlib", "pandas": "pandas"}
+
+
+class MissingExtraError(ImportError):
+    """Raised when a feature needs an optional dependency group that is not installed."""
+
+
+def missing(modules: dict[str, str]) -> list[str]:
+    return [dist for mod, dist in modules.items() if importlib.util.find_spec(mod) is None]
+
+
+def require_train_extras(feature: str = "training") -> None:
+    """Fail early (and clearly) when ``ttucs[train]`` is not installed."""
+    absent = missing(TRAIN_MODULES)
+    if absent:
+        raise MissingExtraError(
+            f"{feature} needs the optional training dependencies ({', '.join(absent)} not installed).\n"
+            '  Install them with:  pip install "ttucs[train]"\n'
+            "  Inference (ttucs predict / parse / info / download) works without them.")
+
+
+def require_mlflow(feature: str = "MLflow tracking") -> None:
+    if importlib.util.find_spec("mlflow") is None:
+        raise MissingExtraError(
+            f"{feature} needs MLflow, which is not installed.\n"
+            '  Install it with:  pip install "ttucs[train]"   (or "ttucs[mlflow]")\n'
+            "  or disable tracking for this run:  -o mlflow.enabled=false")

ttucs/_version.py

@@ -0,0 +1 @@
+__version__ = "0.2.0"

ttucs/cli.py

@@ -0,0 +1,304 @@
+"""Command line interface: ``ttucs predict|parse|info|download|generate-data|train|eval|report``."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import sys
+from pathlib import Path
+
+from ttucs._version import __version__
+
+MODEL_HELP = ("model directory or published model name (default: 'bow', resolved offline-first: "
+              "installed ttucs-model-bow package -> cache -> download unless TTUCS_OFFLINE=1)")
+
+
+def _add_config(p: argparse.ArgumentParser) -> None:
+    p.add_argument("-c", "--config", default=None, help="YAML config (defaults are used if omitted)")
+    p.add_argument("-o", "--override", action="append", default=[], metavar="KEY=VALUE",
+                   help="dotted config override, e.g. -o train.epochs=3 (repeatable)")
+
+
+def _texts(args: argparse.Namespace) -> list[str]:
+    texts = list(args.text or [])
+    if args.file:
+        with open(args.file, encoding="utf-8") as fh:
+            texts += [line.strip() for line in fh if line.strip()]
+    if not texts and not sys.stdin.isatty():
+        texts = [line.strip() for line in sys.stdin if line.strip()]
+    if not texts:
+        raise SystemExit("no input text: pass TEXT arguments, --file or pipe lines on stdin")
+    return texts
+
+
+# ----------------------------------------------------------------------------- commands
+def cmd_generate_data(args: argparse.Namespace) -> int:
+    from ttucs.config import load_config
+    from ttucs.train import generate_data
+
+    cfg = load_config(args.config, args.override)
+    splits = generate_data(cfg, args.out)
+    for k, v in splits.items():
+        print(f"{k:<20} {len(v):>8} examples  {len({e.label for e in v}):>4} labels")
+    return 0
+
+
+def cmd_train(args: argparse.Namespace) -> int:
+    from ttucs._extras import require_mlflow, require_train_extras
+
+    require_train_extras("`ttucs train`")
+    from ttucs.config import load_config
+    from ttucs.train import train
+
+    cfg = load_config(args.config, args.override)
+    if cfg.get("mlflow", {}).get("enabled", True):
+        require_mlflow("`ttucs train` with mlflow.enabled=true")
+    summary = train(cfg, rebuild_data=args.rebuild_data)
+    print(json.dumps({k: summary.get(k) for k in ("output_dir", "temperature", "mlflow_run_id",
+                                                     "model_uri", "registered_model_version")},
+                     indent=2))
+    return 0
+
+
+def cmd_eval(args: argparse.Namespace) -> int:
+    from ttucs._extras import require_train_extras
+
+    require_train_extras("`ttucs eval`")
+    from ttucs.config import load_config
+    from ttucs.datasets.build import read_jsonl
+    from ttucs.evaluate import (
+        SLATE_SPLITS,
+        evaluate_split,
+        format_table,
+        parser_field_accuracy,
+        predictions_csv,
+    )
+    from ttucs.pipeline import TTUCS
+
+    cfg = load_config(args.config, args.override)
+    pipe = TTUCS.from_pretrained(args.model)
+    splits: dict = {}
+    if args.data:
+        for path in args.data:
+            splits[Path(path).stem] = read_jsonl(path)
+    else:
+        proc = Path(cfg["data"].get("processed_dir", "data/processed"))
+        for f in sorted(proc.glob("test*.jsonl")):
+            splits[f.stem] = read_jsonl(f)
+        rs = cfg["data"].get("realistic_slates")
+        if rs and Path(rs).is_file():
+            splits["test_realistic"] = read_jsonl(rs)
+    if not splits:
+        raise SystemExit("no evaluation data found: run `ttucs generate-data` or pass --data FILE")
+    metrics = {}
+    for name, exs in splits.items():
+        parse = args.parse or name in SLATE_SPLITS or any("slate" in e.source for e in exs[:5])
+        m, _ = evaluate_split(pipe, exs, parse=parse)
+        if m:
+            metrics[name] = m
+    print(format_table(metrics))
+    parser_m = parser_field_accuracy(splits.get("test_slates_synth", []), pipe.parser)
+    if parser_m:
+        print("\nslate parser (synthetic slates):")
+        for k, v in parser_m.items():
+            print(f"  {k:<34} {v:.4f}")
+    if args.predictions and "test_realistic" in splits:
+        predictions_csv(pipe, splits["test_realistic"], args.predictions)
+        print(f"\nwrote {args.predictions}")
+    if args.json:
+        Path(args.json).write_text(json.dumps({"metrics": metrics, "parser": parser_m}, indent=2))
+    if args.mlflow:
+        import mlflow
+
+        mlflow.set_tracking_uri(cfg["mlflow"]["tracking_uri"])
+        mlflow.set_experiment(cfg["mlflow"]["experiment"])
+        with mlflow.start_run(run_name=f"eval:{Path(str(pipe.model_path)).name}"):
+            mlflow.set_tag("stage", "eval")
+            mlflow.log_param("model_dir", str(pipe.model_path))
+            for split, m in metrics.items():
+                mlflow.log_metrics({f"{split}.{k}": float(v) for k, v in m.items()})
+            if parser_m:
+                mlflow.log_metrics(parser_m)
+    return 0
+
+
+def cmd_predict(args: argparse.Namespace) -> int:
+    from ttucs.pipeline import TTUCS
+
+    pipe = TTUCS.from_pretrained(args.model, **({"unknown_threshold": args.threshold}
+                                                if args.threshold is not None else {}))
+    preds = pipe.predict_batch(_texts(args), top_k=args.top_k, creator_id=args.creator_id,
+                               source_id=args.source_id)
+    if args.json:
+        out = [p.to_dict() for p in preds]
+        print(json.dumps(out if len(out) > 1 else out[0], indent=2, ensure_ascii=False))
+        return 0
+    for p in preds:
+        print(f"\n> {p.text}")
+        print(f"  description : {p.description}")
+        print(f"  prediction  : {p.cat_id}  ({p.category_full})  conf {p.confidence:.3f}  [{p.decision}]")
+        print("  top-k       : " + ", ".join(f"{a.cat_id} {a.confidence:.3f}" for a in p.alternatives))
+        print(f"  filename    : {p.filename}")
+        print(f"  keywords    : {', '.join(p.keywords)}")
+        meta = {k: v for k, v in p.slate.items()
+                if v and k not in {"raw_text", "normalized_text", "description", "unparsed"}}
+        if meta:
+            print(f"  slate       : {json.dumps(meta, ensure_ascii=False)}")
+    return 0
+
+
+def cmd_parse(args: argparse.Namespace) -> int:
+    from ttucs.slate import parse_slate
+
+    out = [parse_slate(t).to_dict() for t in _texts(args)]
+    if not args.verbose:
+        out = [{k: v for k, v in d.items() if v} for d in out]
+    print(json.dumps(out if len(out) > 1 else out[0], indent=2, ensure_ascii=False))
+    return 0
+
+
+def cmd_report(args: argparse.Namespace) -> int:
+    from ttucs._extras import require_train_extras
+
+    require_train_extras("`ttucs report`")
+    from ttucs.config import load_config
+    from ttucs.reporting import distribution_report
+    from ttucs.taxonomy import UCSTaxonomy
+    from ttucs.train import load_or_build
+
+    cfg = load_config(args.config, args.override)
+    splits = load_or_build(cfg)
+    tax = UCSTaxonomy.load(cfg["ucs_version"])
+    path = distribution_report(splits["train"], tax, args.out,
+                               set(cfg["data"].get("exclude_categories") or []))
+    print(f"wrote {path}")
+    return 0
+
+
+def cmd_info(args: argparse.Namespace) -> int:
+    from ttucs import hub
+    from ttucs.taxonomy import UCS_VERSION
+
+    info: dict = {"ttucs_version": __version__, "ucs_version": UCS_VERSION,
+                  "offline_mode": hub.is_offline(), "cache_dir": str(hub.cache_root())}
+    try:
+        resolved = hub.resolve_model(args.model, allow_download=False)
+        info["model"] = hub.describe(resolved)
+        if args.verify:
+            hub.verify_manifest(resolved.path)
+            info["model"]["manifest_verified"] = True
+    except hub.ModelNotFoundError as exc:
+        info["model"] = None
+        info["error"] = str(exc)
+    if args.json:
+        print(json.dumps(info, indent=2))
+    else:
+        for k, v in info.items():
+            if isinstance(v, dict):
+                print(f"{k}:")
+                for kk, vv in v.items():
+                    print(f"  {kk:<20} {vv}")
+            else:
+                print(f"{k:<22} {v}")
+    return 0 if info["model"] else 1
+
+
+def cmd_download(args: argparse.Namespace) -> int:
+    from ttucs import hub
+
+    path = hub.download_model(args.model, dest=args.dest, url=args.url, sha256=args.sha256, force=args.force)
+    print(f"model {args.model!r} ready at {path}")
+    return 0
+
+
+# ------------------------------------------------------------------------------ parser
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(prog="ttucs", description="Text to UCS: classify spoken slates "
+                                "into Universal Category System CatIDs.")
+    p.add_argument("--version", action="version", version=f"ttucs {__version__}")
+    p.add_argument("-v", "--verbose-log", action="store_true", help="debug logging")
+    sub = p.add_subparsers(dest="command", required=True)
+
+    s = sub.add_parser("generate-data", help="build synthetic + real train/val/test splits")
+    _add_config(s)
+    s.add_argument("--out", default=None, help="output dir (default: data.processed_dir)")
+    s.set_defaults(func=cmd_generate_data)
+
+    s = sub.add_parser("train", help="train, calibrate, evaluate and log to MLflow")
+    _add_config(s)
+    s.add_argument("--rebuild-data", action="store_true", help="ignore cached generated data")
+    s.set_defaults(func=cmd_train)
+
+    s = sub.add_parser("eval", help="evaluate a trained model")
+    _add_config(s)
+    s.add_argument("-m", "--model", default=None, help=MODEL_HELP)
+    s.add_argument("--data", nargs="*", help="JSONL files with text/label (default: test splits)")
+    s.add_argument("--parse", action="store_true", help="run the slate parser on --data texts")
+    s.add_argument("--predictions", default=None, help="write realistic-slate predictions CSV")
+    s.add_argument("--json", default=None, help="write metrics JSON")
+    s.add_argument("--mlflow", action="store_true", help="log an evaluation run to MLflow")
+    s.set_defaults(func=cmd_eval)
+
+    for name, func, hlp in (("predict", cmd_predict, "classify slate transcripts"),
+                            ("parse", cmd_parse, "parse slate transcripts (no model needed)")):
+        s = sub.add_parser(name, help=hlp)
+        s.add_argument("text", nargs="*")
+        s.add_argument("-f", "--file", help="read one transcript per line")
+        if name == "predict":
+            s.add_argument("-m", "--model", default=None, help=MODEL_HELP)
+            s.add_argument("-k", "--top-k", type=int, default=5)
+            s.add_argument("--threshold", type=float, default=None, help="unknown threshold")
+            s.add_argument("--creator-id", default=None)
+            s.add_argument("--source-id", default=None)
+            s.add_argument("--json", action="store_true")
+        else:
+            s.add_argument("--verbose", action="store_true", help="include empty fields")
+        s.set_defaults(func=func)
+
+    s = sub.add_parser("info", help="show the resolved model, versions and offline mode (no network)")
+    s.add_argument("-m", "--model", default=None, help=MODEL_HELP)
+    s.add_argument("--verify", action="store_true", help="check weights against manifest.json sha256")
+    s.add_argument("--json", action="store_true")
+    s.set_defaults(func=cmd_info)
+
+    s = sub.add_parser("download", help="download + verify published weights into the cache")
+    s.add_argument("--model", default="bow", help="published model name (default: bow)")
+    s.add_argument("--dest", default=None, help="target directory (default: $TTUCS_HOME or user cache)")
+    s.add_argument("--url", default=None, help="override the release URL (mirror / local file)")
+    s.add_argument("--sha256", default=None, help="expected sha256 of the archive (with --url)")
+    s.add_argument("--force", action="store_true", help="re-download even if present")
+    s.set_defaults(func=cmd_download)
+
+    s = sub.add_parser("report", help="class distribution report + plots")
+    _add_config(s)
+    s.add_argument("--out", default="docs/assets")
+    s.set_defaults(func=cmd_report)
+    return p
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    logging.basicConfig(level=logging.DEBUG if args.verbose_log else logging.INFO,
+                        format="%(asctime)s %(levelname)s %(name)s: %(message)s")
+    for noisy in ("alembic", "urllib3", "matplotlib", "PIL", "sentence_transformers", "httpx"):
+        logging.getLogger(noisy).setLevel(logging.WARNING)
+    from ttucs._extras import MissingExtraError
+    from ttucs.hub import (
+        ChecksumError,
+        DownloadError,
+        IncompatibleModelError,
+        ModelNotFoundError,
+        OfflineError,
+    )
+
+    try:
+        return int(args.func(args) or 0)
+    except (MissingExtraError, ModelNotFoundError, OfflineError, ChecksumError,
+            IncompatibleModelError, DownloadError) as exc:
+        print(f"ttucs {args.command}: error: {exc}", file=sys.stderr)
+        return 2
+
+
+if __name__ == "__main__":  # pragma: no cover
+    raise SystemExit(main())

ttucs/config.py

@@ -0,0 +1,113 @@
+"""YAML experiment configuration with defaults and dotted command-line overrides.
+
+A tiny, dependency-free take on the Hydra/Lightning-CLI pattern::
+
+    cfg = load_config("configs/bow.yaml", ["train.epochs=3", "data.label_documents=false"])
+"""
+
+from __future__ import annotations
+
+import copy
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+DEFAULTS: dict[str, Any] = {
+    "seed": 42,
+    "ucs_version": "8.2.1",
+    "output_dir": "models/ttucs-bow",
+    "data": {
+        "exclude_categories": ["ARCHIVED"],
+        "synthetic": {
+            "enabled": True,
+            "target_per_class": 150,   # synthetic examples per CatID
+            "min_per_class": 60,       # lower bound if target_per_class is set lower
+            "max_synonym_df": 25,
+            "cache": "data/processed/synthetic.jsonl",
+        },
+        "label_documents": True,       # add UCS explanation/synonym docs as training text
+        "slate_augment_fraction": 0.25,
+        "split": {"val": 0.1, "test": 0.1},
+        "realistic_slates": "examples/realistic_slates.jsonl",
+        "processed_dir": "data/processed",
+    },
+    "model": {
+        "name": "bow",
+        "embedding_dim": 128,
+        "num_buckets": 262144,
+        "word_ngrams": 2,
+        "char_ngrams": [3, 5],
+        "dropout": 0.2,
+        "category_loss_weight": 0.3,
+        "encoder": "sentence-transformers/all-MiniLM-L6-v2",
+        "hidden_dim": 512,
+    },
+    "train": {
+        "epochs": 12,
+        "batch_size": 256,
+        "lr": 1.0e-2,
+        "weight_decay": 1.0e-5,
+        "label_smoothing": 0.05,
+        "patience": 3,
+        "num_threads": 8,
+        "calibrate": True,
+    },
+    "inference": {"unknown_threshold": 0.12, "category_threshold": 0.5, "top_k": 5},
+    "mlflow": {
+        "enabled": True,
+        "tracking_uri": "sqlite:///mlflow.db",
+        "experiment": "ttucs",
+        "register_model": True,
+        "registered_model_name": "ttucs",
+        "log_model": True,
+    },
+}
+
+
+def deep_merge(base: dict, override: dict) -> dict:
+    out = copy.deepcopy(base)
+    for k, v in (override or {}).items():
+        if isinstance(v, dict) and isinstance(out.get(k), dict):
+            out[k] = deep_merge(out[k], v)
+        else:
+            out[k] = copy.deepcopy(v)
+    return out
+
+
+def _parse_value(raw: str) -> Any:
+    return yaml.safe_load(raw)
+
+
+def apply_overrides(cfg: dict, overrides: list[str] | None) -> dict:
+    cfg = copy.deepcopy(cfg)
+    for ov in overrides or []:
+        if "=" not in ov:
+            raise ValueError(f"override must look like key.sub=value, got {ov!r}")
+        key, raw = ov.split("=", 1)
+        node = cfg
+        parts = key.strip().split(".")
+        for p in parts[:-1]:
+            node = node.setdefault(p, {})
+        node[parts[-1]] = _parse_value(raw)
+    return cfg
+
+
+def load_config(path: str | Path | None = None, overrides: list[str] | None = None) -> dict:
+    cfg = DEFAULTS
+    if path:
+        with open(path, encoding="utf-8") as fh:
+            cfg = deep_merge(cfg, yaml.safe_load(fh) or {})
+    return apply_overrides(cfg, overrides)
+
+
+def flatten(d: dict, prefix: str = "") -> dict[str, Any]:
+    """Flatten nested config for MLflow params: {'a': {'b': 1}} -> {'a.b': 1}."""
+    out: dict[str, Any] = {}
+    for k, v in d.items():
+        key = f"{prefix}{k}"
+        if isinstance(v, dict):
+            out.update(flatten(v, f"{key}."))
+        else:
+            out[key] = v
+    return out

ttucs/datasets/__init__.py

@@ -0,0 +1,5 @@
+"""Dataset generation and split assembly."""
+
+from ttucs.datasets.build import build_datasets, group_split, read_jsonl, write_jsonl
+
+__all__ = ["build_datasets", "group_split", "read_jsonl", "write_jsonl"]

ttucs/datasets/build.py

@@ -0,0 +1,146 @@
+"""Assemble balanced train/val/test splits from synthetic data and UCS label documents."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import random
+from collections import defaultdict
+from pathlib import Path
+
+from ttucs.slate.parser import parse_slate
+from ttucs.synth.generator import DescriptionGenerator, Example, SlateGenerator
+from ttucs.taxonomy import UCSTaxonomy
+
+log = logging.getLogger(__name__)
+
+
+# ------------------------------------------------------------------------------ io
+def write_jsonl(examples: list[Example], path: str | Path) -> Path:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as fh:
+        for ex in examples:
+            fh.write(json.dumps(ex.to_dict(), ensure_ascii=False) + "\n")
+    return path
+
+
+def read_jsonl(path: str | Path) -> list[Example]:
+    out = []
+    with open(path, encoding="utf-8") as fh:
+        for line in fh:
+            if line.strip():
+                d = json.loads(line)
+                out.append(Example(text=d["text"], label=d["label"], source=d.get("source", ""),
+                                   group=d.get("group", ""), meta=d.get("meta", {})))
+    return out
+
+
+# --------------------------------------------------------------------------- splits
+def split_bucket(group: str, seed: int) -> float:
+    """Deterministic uniform number in [0, 1) for a group key (stable across runs/OSes)."""
+    h = hashlib.sha1(f"{seed}:{group}".encode(), usedforsecurity=False).hexdigest()
+    return int(h[:12], 16) / float(16**12)
+
+
+def group_split(examples: list[Example], val: float, test: float, seed: int
+                ) -> tuple[list[Example], list[Example], list[Example]]:
+    """Split by ``group`` so near-duplicates (take 01/02/03) never straddle splits."""
+    tr, va, te = [], [], []
+    for ex in examples:
+        b = split_bucket(ex.group or ex.text, seed)
+        (te if b < test else va if b < test + val else tr).append(ex)
+    return tr, va, te
+
+
+# ------------------------------------------------------------------------ synthetic
+def generate_synthetic(taxonomy: UCSTaxonomy, counts: dict[str, int], seed: int,
+                       max_synonym_df: int = 25) -> list[Example]:
+    gen = DescriptionGenerator(taxonomy, seed=seed, max_synonym_df=max_synonym_df)
+    out: list[Example] = []
+    for cat_id, n in counts.items():
+        entry = taxonomy[cat_id]
+        for i, text in enumerate(gen.generate(entry, n)):
+            # groups of 4 consecutive samples share a split bucket (limits template leakage)
+            out.append(Example(text=text, label=cat_id, source="synthetic",
+                               group=f"synth:{cat_id}:{i // 4}"))
+    return out
+
+
+def label_documents(taxonomy: UCSTaxonomy, cat_ids: list[str]) -> list[Example]:
+    """UCS explanation / synonym documents as extra training text (cold start for rare classes)."""
+    out = []
+    for cid in cat_ids:
+        e = taxonomy[cid]
+        docs = [e.explanation, f"{e.category.lower()} {e.subcategory.lower()}"]
+        syn = list(e.synonyms)
+        for i in range(0, min(len(syn), 24), 8):
+            docs.append(" ".join(s.lower() for s in syn[i:i + 8]))
+        out += [Example(text=d, label=cid, source="ucs_doc", group=f"doc:{cid}") for d in docs if d]
+    return out
+
+
+def slate_augment(examples: list[Example], fraction: float, seed: int) -> list[Example]:
+    """Replace a fraction of texts by parser output on a synthetic slate around them."""
+    rng = random.Random(seed)
+    sg = SlateGenerator(seed=seed)
+    out = []
+    for ex in examples:
+        if ex.source != "ucs_doc" and rng.random() < fraction:
+            slate, _ = sg.wrap(ex.text)
+            parsed = parse_slate(slate).description
+            out.append(Example(text=parsed, label=ex.label, source=ex.source, group=ex.group,
+                               meta={**ex.meta, "slate": slate, "augmented": True}))
+        else:
+            out.append(ex)
+    return out
+
+
+def make_slate_testset(examples: list[Example], seed: int, n: int = 2000) -> list[Example]:
+    rng = random.Random(seed)
+    sg = SlateGenerator(seed=seed + 1)
+    picked = rng.sample(examples, min(n, len(examples)))
+    out = []
+    for ex in picked:
+        slate, truth = sg.wrap(ex.text)
+        out.append(Example(text=slate, label=ex.label, source=f"slate:{ex.source}", group=ex.group,
+                           meta={"truth": truth}))
+    return out
+
+
+# ---------------------------------------------------------------------------- build
+def build_datasets(cfg: dict, taxonomy: UCSTaxonomy) -> dict[str, list[Example]]:
+    """Build all splits described by ``cfg['data']``. Returns a dict of split -> examples."""
+    seed = int(cfg["seed"])
+    dcfg = cfg["data"]
+    exclude = set(dcfg.get("exclude_categories") or [])
+    labels = [e.cat_id for e in taxonomy.trainable(exclude)]
+    vf, tf = dcfg["split"]["val"], dcfg["split"]["test"]
+    splits: dict[str, list[Example]] = defaultdict(list)
+
+    # ---- synthetic descriptions (balanced: the same budget for every CatID)
+    scfg = dcfg["synthetic"]
+    if scfg.get("enabled", True):
+        counts = {cid: max(scfg["min_per_class"], scfg["target_per_class"]) for cid in labels}
+        synth = generate_synthetic(taxonomy, counts, seed, scfg.get("max_synonym_df", 25))
+        if scfg.get("cache"):
+            write_jsonl(synth, scfg["cache"])
+        str_, sva, ste = group_split(synth, vf, tf, seed)
+        splits["train"] += str_
+        splits["val"] += sva
+        splits["test_synth"] = ste
+        splits["test_slates_synth"] = make_slate_testset(ste, seed)
+
+    if dcfg.get("label_documents", True):
+        splits["train"] += label_documents(taxonomy, labels)
+    if dcfg.get("slate_augment_fraction", 0) > 0:
+        splits["train"] = slate_augment(splits["train"], dcfg["slate_augment_fraction"], seed)
+
+    rs = dcfg.get("realistic_slates")
+    if rs and Path(rs).is_file():
+        splits["test_realistic"] = read_jsonl(rs)
+    random.Random(seed).shuffle(splits["train"])
+    for k, v in splits.items():
+        log.info("split %-18s %7d examples, %4d labels", k, len(v), len({e.label for e in v}))
+    return dict(splits)