feat: publish downloadable wakeword trainer source

2026-08-14 05:03:00 -05:00
parent 000671eb93
commit aab9ce0fe7
30 changed files with 6757 additions and 2 deletions
@@ -0,0 +1,141 @@
+"""Materialize approved inputs then call only released LiveKit CLI stages."""
+
+from __future__ import annotations
+
+import json
+import shutil
+import sys
+from pathlib import Path
+
+from stage_evidence import run_stage
+from wakeword_training_image import InputMaterializer
+
+RUN_DIR = Path("/run")
+
+
+def _load_json(path: Path) -> dict[str, object]:
+    value = json.loads(path.read_text(encoding="utf-8"))
+    if not isinstance(value, dict):
+        raise ValueError(f"{path} must contain a JSON object")
+    return value
+
+
+def _copy_as_clips(source_paths: list[Path], target: Path) -> None:
+    target.mkdir(parents=True, exist_ok=True)
+    for index, source in enumerate(source_paths):
+        if source.suffix.lower() != ".wav":
+            raise ValueError(f"recorded source must be a WAV: {source.name}")
+        shutil.copyfile(source, target / f"clip_{index:06d}.wav")
+
+
+def _stage_inputs(recorded_manifest: Path, config: dict[str, object]) -> None:
+    manifest = _load_json(recorded_manifest)
+    if manifest.get("classification") != "commercial_review_required":
+        raise ValueError("commercial pipeline refuses a non-commercial source manifest")
+    model_name = config.get("model_name")
+    if not isinstance(model_name, str) or not model_name:
+        raise ValueError("resolved config must declare model_name")
+    model_dir = RUN_DIR / "output" / model_name
+    stage_data = RUN_DIR / "stage-data"
+    sources = manifest.get("sources")
+    if not isinstance(sources, list):
+        raise ValueError("recorded manifest sources must be a list")
+    general_negative_features: dict[str, Path] = {}
+    for source in sources:
+        if not isinstance(source, dict):
+            raise ValueError("recorded manifest source must be an object")
+        purpose = source.get("purpose")
+        partition = source.get("partition")
+        files = source.get("files")
+        if (
+            not isinstance(purpose, str)
+            or not isinstance(partition, str)
+            or not isinstance(files, list)
+        ):
+            raise ValueError("recorded manifest source is incomplete")
+        paths = []
+        for file in files:
+            if not isinstance(file, dict) or not isinstance(file.get("materialized_path"), str):
+                raise ValueError("recorded manifest file is incomplete")
+            paths.append(RUN_DIR / str(file["materialized_path"]))
+        if purpose == "positive":
+            split = "positive_train" if partition == "train" else "positive_test"
+            _copy_as_clips(paths, model_dir / split)
+        elif purpose in {"adversarial_negative", "near_miss"}:
+            split = "negative_train" if partition == "train" else "negative_test"
+            _copy_as_clips(paths, model_dir / split)
+        elif purpose == "background":
+            target = stage_data / "backgrounds"
+            target.mkdir(parents=True, exist_ok=True)
+            for source_path in paths:
+                shutil.copyfile(source_path, target / source_path.name)
+        elif purpose == "rir":
+            target = stage_data / "rirs"
+            target.mkdir(parents=True, exist_ok=True)
+            for source_path in paths:
+                shutil.copyfile(source_path, target / source_path.name)
+        elif purpose == "general_negative_speech":
+            if len(paths) != 1 or paths[0].suffix != ".npy":
+                raise ValueError(
+                    "general-negative source must provide exactly one .npy feature bundle"
+                )
+            if partition not in {"train", "validation"}:
+                raise ValueError("general-negative feature source must use train or validation")
+            if partition in general_negative_features:
+                raise ValueError(f"duplicate general-negative feature bundle for {partition}")
+            compatibility_name = {
+                "train": "openwakeword_features_ACAV100M_2000_hrs_16bit.npy",
+                "validation": "validation_set_features.npy",
+            }[partition]
+            target = stage_data / "features" / compatibility_name
+            target.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copyfile(paths[0], target)
+            general_negative_features[partition] = target
+
+    if set(general_negative_features) != {"train", "validation"}:
+        raise ValueError(
+            "commercial pipeline requires train and validation general-negative features"
+        )
+
+    config["data_dir"] = "/run/stage-data"
+    config["output_dir"] = "/run/output"
+    batch_sizes = config.setdefault("batch_n_per_class", {})
+    if not isinstance(batch_sizes, dict):
+        raise ValueError("batch_n_per_class must be an object")
+    # Preserve LiveKit's dataset key. The source lock/recorded manifest is the
+    # provenance boundary that says these bytes are approved general-negative
+    # speech, not ACAV100M data.
+    batch_sizes["ACAV100M_sample"] = batch_sizes.pop("general_negative_speech", 1024)
+    augmentation = config.setdefault("augmentation", {})
+    if not isinstance(augmentation, dict):
+        raise ValueError("augmentation must be an object")
+    augmentation["background_paths"] = ["/run/stage-data/backgrounds"]
+    augmentation["rir_paths"] = ["/run/stage-data/rirs"]
+    model = config.get("model")
+    if not isinstance(model, dict) or model.get("model_type", "conv_attention") != "conv_attention":
+        raise ValueError("commercial pipeline requires the FP32 conv_attention model")
+
+
+def main(config_path: Path) -> int:
+    request = _load_json(RUN_DIR / "request.json")
+    request_payload = request.get("request")
+    if not isinstance(request_payload, dict) or request_payload.get("run_kind") != "commercial":
+        raise ValueError("commercial pipeline refuses a quarantined or malformed request")
+    config = _load_json(config_path)
+    materialized = InputMaterializer(
+        RUN_DIR / "sources.lock.json", Path("/inputs"), RUN_DIR, run_kind="commercial"
+    ).materialize()
+    _stage_inputs(materialized.recorded_manifest_path, config)
+    effective_config_path = RUN_DIR / "effective-config.yaml"
+    effective_config_path.write_text(json.dumps(config, sort_keys=True) + "\n", encoding="utf-8")
+    run_stage("provenance", ["python", "/opt/jr-wakeword/record-image-provenance.py"])
+    for stage in ("augment", "train", "export", "eval"):
+        run_stage(stage, ["livekit-wakeword", stage, str(effective_config_path)])
+    run_stage(
+        "finalize", ["python", "/opt/jr-wakeword/finalize-run.py", str(effective_config_path)]
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main(Path(sys.argv[1])))