142 lines
6.5 KiB
Python
142 lines
6.5 KiB
Python
"""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])))
|