diff --git a/docs/models/multitalker-parakeet-streaming-0.6b-v1.md b/docs/models/multitalker-parakeet-streaming-0.6b-v1.md index 44d1fc0c..25c8c67e 100644 --- a/docs/models/multitalker-parakeet-streaming-0.6b-v1.md +++ b/docs/models/multitalker-parakeet-streaming-0.6b-v1.md @@ -12,13 +12,15 @@ RNN-T decoding. Outputs cased, punctuated transcripts. Token- and word-level timestamps are available. Upstream this is a **multitalker (speaker-attributed)** checkpoint: it can -transcribe several overlapping speakers into per-speaker channels. That -path depends on machinery this port does **not** ship (see -[Known limitations](#known-limitations)). transcribe.cpp exposes only the -model's `single_speaker_mode` ASR path, which disables the multitalker -machinery and runs the checkpoint as a cache-aware streaming RNN-T with -the base model's frontend, encoder backbone, decoder, and tokenizer plus -the checkpoint's always-on layer-0 speaker-kernel injection. +transcribe several overlapping speakers into per-speaker channels. This +port ships that path too, via **bundle GGUFs** that embed the +[`nvidia/diar_streaming_sortformer_4spk-v2.1`](https://huggingface.co/nvidia/diar_streaming_sortformer_4spk-v2.1) +streaming diarizer alongside the ASR model. A plain (non-bundle) GGUF runs +the model's `single_speaker_mode` ASR path — a cache-aware streaming RNN-T +with the checkpoint's always-on layer-0 speaker-kernel injection. A bundle +GGUF with `--diarize` runs the full multitalker pipeline and emits a +speaker-tagged transcript (see +[Multitalker](#multitalker-speaker-attributed-asr)). This port runs the model in both **offline** and **cache-aware streaming** modes. @@ -165,20 +167,64 @@ CPU same-length tensor parity is bit-exact (max_abs 0.0) at batch=4 on `jfk.wav`, diverse-length flash tensor parity is bit-exact across arbitrary length mixes, and full test-clean batch-8 WER equals batch-1 (2.19%). +## Multitalker (speaker-attributed ASR) + +A **bundle GGUF** (composed by `scripts/compose-multitalker-bundle.py`) +embeds the streaming Sortformer diarizer. Running it with `--diarize` +produces a speaker-tagged transcript (speaker ids on segments/words; up to +4 speakers): the embedded diarizer streams over the clip at the reference +operating point (14-frame chunks, spkcache/FIFO 188), per-chunk cache +gating selects which speakers are active, and each active speaker gets a +private cache-aware streaming encoder+decoder instance fed only its active +chunks — memory is O(1) in clip length (~5 MB per speaker), so hour-scale +meetings run in bounded memory. + +Two supervision modes, selected via `TRANSCRIBE_MULTITALKER_MODE`: + +- **kernel** (default): trained speaker-kernel injection at conformer + layer 0 (`masked_asr=False` upstream). Audio is unmodified; the model + suppresses non-target speech. +- **masked**: mel feature masking (`masked_asr=True`, the NeMo example + default). Non-target frames are zero-masked before pre-encode. + +cpWER on ami-ihm-test (16 meetings, meeteval, refs from `edinburghcstr/ami` +segment transcripts; NeMo baseline = the pinned reference pipeline on +L40S, batch 1): + +| system | kernel | masked | +| --- | ---: | ---: | +| transcribe.cpp (CUDA F32) | **19.35%** | 23.73% | +| NeMo reference | 21.39% | 24.00% | +| NVIDIA self-reported | 21.26% | — | + +Kernel mode is the default because it wins on both implementations. The +C++ kernel number is *better* than the reference for a verified reason: +the reference harness has a background-slot indexing bug +(`get_active_speakers_info` builds `inactive_speaker_ids` from +`range(len(active))` minus a slot id), confirmed at tensor level from the +reference's own supervision dumps. With +`TRANSCRIBE_MULTITALKER_REF_BG_COMPAT=1` (parity tooling only) the port +reproduces that bug and scores 21.23% — matching the reference (and +NVIDIA's published 21.26%). Under matched configuration both modes are +transcript-exact vs the reference on CPU F32 (0 differing words on a +3-minute 4-speaker AMI slice); on long meetings a measured ~3e-5 F32 pred +drift eventually flips near-tie top-k picks inside the diarizer's +spkcache compression (a discontinuous selection both sides implement +bit-identically vs a neutral arbiter), after which trajectories diverge +benignly (~0.4 cpWER on the affected meeting). + +`TRANSCRIBE_MULTITALKER_OFFLINE=1` selects the older offline whole-clip +replay instead of the streaming pass (parity/dump tooling; attention +buffers grow quadratically with clip length — do not use on meetings +longer than ~15 minutes). + ## Known limitations -- **Single-speaker only. The multitalker / speaker-attributed ASR path is - not ported.** Upstream, this checkpoint transcribes several overlapping - speakers into per-speaker channels. That path requires (1) an external - streaming speaker-diarization model - ([`nvidia/diar_streaming_sortformer_4spk-v2.1`](https://huggingface.co/nvidia/diar_streaming_sortformer_4spk-v2.1), - a separate Sortformer checkpoint that this repo does not ship), (2) - per-frame speaker-kernel injection driven by that diarizer's supervision, - (3) one encoder+decoder instance per speaker (up to 4), and (4) a - speaker-tagged SegLST output contract. None of these exist in - transcribe.cpp today, so the port runs `single_speaker_mode` and produces - a single flat transcript. It does **not** separate speakers, diarize, or - emit speaker turns. +- **Multitalker is offline-API only.** `--diarize` on a bundle runs the + streaming pipeline internally, but the public streaming API + (`stream_feed`) still exposes single-speaker transcription only; a + per-speaker streaming text ABI is future work. `run_batch` on a bundle + warns and runs single-speaker. - **English only.** The model is English-only by training. - **No translation and no language detection.** diff --git a/samples/multitalker-2spk-mix.wav b/samples/multitalker-2spk-mix.wav new file mode 100644 index 00000000..35ea11c4 Binary files /dev/null and b/samples/multitalker-2spk-mix.wav differ diff --git a/scripts/compose-multitalker-bundle.py b/scripts/compose-multitalker-bundle.py new file mode 100644 index 00000000..7f7f8398 --- /dev/null +++ b/scripts/compose-multitalker-bundle.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +""" +compose-multitalker-bundle.py - merge a multitalker Parakeet ASR GGUF and a +Sortformer diarizer GGUF into one self-contained bundle GGUF. + +The multitalker checkpoint CONSUMES external speaker-activity supervision at +inference time (layer-0 speaker-kernel injection); NVIDIA validated it paired +with streaming Sortformer. Shipping the pair as one file removes the +mismatched-pair support surface: one artifact, one tested behavior, and the +model-level DIARIZATION capability stays a property of the file. + +Layout of the bundle: + + KV every ASR KV verbatim (arch stays 'parakeet'), plus: + stt.parakeet.diarizer.embedded = true + stt.parakeet.diarizer.variant = + stt.parakeet.diarizer.tensor_prefix = 'sortformer.' + stt.sortformer.* copied verbatim + The diarizer's frontend KVs are NOT copied: this pairing shares + one mel frontend (both are 128-mel, 16 kHz, hann/512/400/160, + normalize=none, same dither/pre-emphasis) and the compose step + VERIFIES that agreement instead of duplicating the keys. + tensors every ASR tensor verbatim, then every diarizer tensor renamed + with the 'sortformer.' prefix (the diarizer reuses Parakeet's + NEST encoder, so its enc.* names collide with the ASR's own + enc.* catalog without the prefix). + +The C++ parakeet loader materializes and uploads every tensor in the file +and resolves its catalog by name, so a bundle loads and runs the shipped +single-speaker path unchanged even before the multitalker composition code +exists; the sortformer.* tensors ride along for it to claim. + +Dtype policy mirrors the shipped matrices: the ASR half may be any shipped +quant; the diarizer half should be F32/F16/Q8_0 only (sortformer k-quants +are withdrawn). Pick the diarizer file explicitly per bundle tier. + +Usage (any env with gguf, e.g. the moonshine or sortformer env): + uv run --project scripts/envs/moonshine scripts/compose-multitalker-bundle.py \ + models/multitalker-parakeet-streaming-0.6b-v1/multitalker-parakeet-streaming-0.6b-v1-F32.gguf \ + models/diar_streaming_sortformer_4spk-v2.1/diar_streaming_sortformer_4spk-v2.1-F32.gguf \ + -o models/multitalker-parakeet-streaming-0.6b-v1/bundle/multitalker-parakeet-streaming-0.6b-v1-F32.gguf +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +from gguf import GGUFReader, GGUFValueType + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from lib.gguf_common import gguf_writer # noqa: E402 + +DIAR_TENSOR_PREFIX = "sortformer." + +# KVs the two halves must agree on for the bundle to share one mel frontend. +FRONTEND_AGREEMENT_KEYS = [ + "stt.frontend.sample_rate", + "stt.frontend.num_mels", + "stt.frontend.n_fft", + "stt.frontend.hop_length", + "stt.frontend.win_length", + "stt.frontend.window", + "stt.frontend.normalize", + "stt.frontend.pre_emphasis", + "stt.frontend.dither", +] + +# Keys the GGUFReader synthesizes (virtual) or the writer emits itself. +SKIP_COPY_KEYS = {"general.architecture"} + + +def copy_kv(writer, field) -> None: + vtype = field.types[0] + sub_type = field.types[-1] if vtype == GGUFValueType.ARRAY else None + writer.add_key_value(field.name, field.contents(), vtype, sub_type=sub_type) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("asr_gguf", help="multitalker Parakeet GGUF (any shipped quant)") + ap.add_argument("diar_gguf", help="Sortformer diarizer GGUF (F32/F16/Q8_0)") + ap.add_argument("-o", "--out", required=True, help="output bundle path") + args = ap.parse_args() + + asr = GGUFReader(args.asr_gguf) + diar = GGUFReader(args.diar_gguf) + + asr_arch = asr.fields["general.architecture"].contents() + diar_arch = diar.fields["general.architecture"].contents() + if asr_arch != "parakeet": + print(f"error: ASR input arch is '{asr_arch}', expected 'parakeet'", file=sys.stderr) + return 1 + if diar_arch != "sortformer": + print(f"error: diarizer input arch is '{diar_arch}', expected 'sortformer'", file=sys.stderr) + return 1 + if "stt.parakeet.encoder.spk_kernel_layers" not in asr.fields: + print("error: ASR input has no speaker-kernel tensors; not a multitalker checkpoint", file=sys.stderr) + return 1 + + # The bundle carries a single stt.frontend.* block (the ASR's). Refuse to + # compose a pairing whose diarizer disagrees with it rather than silently + # feeding the diarizer the wrong mel geometry. + for key in FRONTEND_AGREEMENT_KEYS: + a = asr.fields[key].contents() if key in asr.fields else None + d = diar.fields[key].contents() if key in diar.fields else None + if a != d: + print(f"error: frontend mismatch on {key}: asr={a!r} diar={d!r}", file=sys.stderr) + return 1 + + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + writer = gguf_writer(str(out_path), "parakeet") + + for name, field in asr.fields.items(): + if name.startswith("GGUF.") or name in SKIP_COPY_KEYS: + continue + copy_kv(writer, field) + + writer.add_bool("stt.parakeet.diarizer.embedded", True) + writer.add_string("stt.parakeet.diarizer.variant", diar.fields["stt.variant"].contents()) + writer.add_string("stt.parakeet.diarizer.tensor_prefix", DIAR_TENSOR_PREFIX) + for name, field in diar.fields.items(): + if name.startswith("stt.sortformer."): + copy_kv(writer, field) + + for t in asr.tensors: + writer.add_tensor(t.name, t.data, raw_dtype=t.tensor_type) + for t in diar.tensors: + writer.add_tensor(DIAR_TENSOR_PREFIX + t.name, t.data, raw_dtype=t.tensor_type) + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + n_asr, n_diar = len(asr.tensors), len(diar.tensors) + print(f"wrote {out_path}: {n_asr} ASR + {n_diar} diarizer tensors ({n_asr + n_diar} total)") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/diar/ingest_ami_seglst.py b/scripts/diar/ingest_ami_seglst.py new file mode 100644 index 00000000..92b562d1 --- /dev/null +++ b/scripts/diar/ingest_ami_seglst.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +""" +ingest_ami_seglst.py - build SegLST reference transcripts for the AMI +acceptance manifest. + +cpWER needs speaker-attributed reference TRANSCRIPTS; the +diarizers-community/ami dataset behind ingest_ami.py ships only RTTM +(speaker turns, no words). This pulls the per-segment manual transcripts +of edinburghcstr/ami via the HF datasets-server rows API (JSON, text +columns only — the parquet shards embed segmented audio and would cost +tens of GB of cache), groups them by meeting, and writes one SegLST JSON +per meeting, plus a manifest joining audio + rttm + seglst. + +Reads samples/diar/ami--.manifest.jsonl (ingest_ami.py) +Writes samples/diar/ami---seglst/.seglst.json + samples/diar/ami---mt.manifest.jsonl + +Usage: + uv run scripts/diar/ingest_ami_seglst.py --config ihm --split test +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +import urllib.parse +import urllib.request +from collections import defaultdict +from pathlib import Path + +REPO = Path(__file__).resolve().parent.parent.parent +API = "https://datasets-server.huggingface.co/rows" +PAGE = 100 + + +def hf_token() -> str | None: + import os + + tok = os.environ.get("HF_TOKEN") + if tok: + return tok + p = Path.home() / ".cache" / "huggingface" / "token" + return p.read_text().strip() if p.exists() else None + + +def fetch_page(config: str, split: str, offset: int, cache_dir: Path, token: str | None) -> dict: + cache = cache_dir / f"{config}-{split}-{offset}.json" + if cache.exists(): + return json.loads(cache.read_text()) + q = urllib.parse.urlencode( + {"dataset": "edinburghcstr/ami", "config": config, "split": split, + "offset": offset, "length": PAGE} + ) + req = urllib.request.Request(f"{API}?{q}") + if token: + req.add_header("Authorization", f"Bearer {token}") + for attempt in range(10): + try: + with urllib.request.urlopen(req, timeout=60) as r: + payload = json.load(r) + cache.write_text(json.dumps(payload)) + return payload + except urllib.error.HTTPError as exc: + if attempt == 9: + raise + retry_after = exc.headers.get("Retry-After") if exc.headers else None + wait = float(retry_after) if retry_after else min(60.0, 2.0 ** attempt) + print(f"HTTP {exc.code} offset={offset}, waiting {wait:.0f}s", file=sys.stderr) + time.sleep(wait) + except Exception as exc: + if attempt == 9: + raise + time.sleep(min(60.0, 2.0 ** attempt)) + print(f"retry offset={offset}: {exc}", file=sys.stderr) + raise RuntimeError("unreachable") + + +def fetch_rows(config: str, split: str, cache_dir: Path): + token = hf_token() + offset = 0 + while True: + payload = fetch_page(config, split, offset, cache_dir, token) + rows = payload.get("rows", []) + if not rows: + return + for r in rows: + yield r["row"] + offset += len(rows) + if offset >= payload.get("num_rows_total", 0): + return + time.sleep(0.75) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--config", default="ihm", choices=["ihm", "sdm"]) + ap.add_argument("--split", default="test") + args = ap.parse_args() + + base = REPO / "samples" / "diar" + src_manifest = base / f"ami-{args.config}-{args.split}.manifest.jsonl" + if not src_manifest.exists(): + print(f"error: {src_manifest} missing (run ingest_ami.py first)", file=sys.stderr) + return 1 + entries = [json.loads(l) for l in src_manifest.read_text().splitlines() if l.strip()] + meetings = {e["id"].split(".")[0]: e for e in entries} + + cache_dir = base / ".rows-cache" + cache_dir.mkdir(parents=True, exist_ok=True) + per_meeting: dict[str, list[dict]] = defaultdict(list) + n_rows = 0 + for row in fetch_rows(args.config, args.split, cache_dir): + mid = row.get("meeting_id") + text = (row.get("text") or "").strip() + if mid not in meetings or not text: + continue + per_meeting[mid].append( + { + "session_id": mid, + "speaker": row["speaker_id"], + "start_time": float(row["begin_time"]), + "end_time": float(row["end_time"]), + "words": text, + } + ) + n_rows += 1 + if n_rows % 2000 == 0: + print(f"...{n_rows} segments", file=sys.stderr) + + missing = sorted(set(meetings) - set(per_meeting)) + if missing: + print(f"warning: no transcript rows for {missing}", file=sys.stderr) + + out_dir = base / f"ami-{args.config}-{args.split}-seglst" + out_dir.mkdir(parents=True, exist_ok=True) + out_manifest = base / f"ami-{args.config}-{args.split}-mt.manifest.jsonl" + with out_manifest.open("w") as mf: + for mid, entry in sorted(meetings.items()): + segs = sorted(per_meeting.get(mid, []), key=lambda s: s["start_time"]) + seglst_path = out_dir / f"{mid}.seglst.json" + seglst_path.write_text(json.dumps(segs, indent=1) + "\n") + row = dict(entry) + row["seglst"] = str(seglst_path.relative_to(REPO)) + mf.write(json.dumps(row) + "\n") + print(f"{mid}: {len(segs)} reference segments") + + print(f"wrote {out_manifest} ({len(meetings)} meetings, {n_rows} segments total)") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/diar/modal_multitalker_cpp.py b/scripts/diar/modal_multitalker_cpp.py new file mode 100644 index 00000000..09ce3de6 --- /dev/null +++ b/scripts/diar/modal_multitalker_cpp.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python3 +""" +modal_multitalker_cpp.py - run the C++ multitalker bundle over the AMI +acceptance manifest on Modal CUDA and emit SegLST hypotheses. + +The long AMI meetings need more attention-buffer memory than the local +16 GB M4 can give (a ~50-min meeting wants a ~5.5 GB F32 mask; a local +attempt hard-crashed the machine), so the C++ side runs on an L40S with +a CUDA build. The image bakes a TRANSCRIBE_CUDA=ON build of the repo +source; bundle GGUF + audio come from the `transcribe-ami-cpp` volume: + + modal volume create transcribe-ami-cpp + modal volume put transcribe-ami-cpp \ + models/multitalker-parakeet-streaming-0.6b-v1/bundle/multitalker-parakeet-streaming-0.6b-v1-F32.gguf \ + /model/bundle-F32.gguf + modal volume put transcribe-ami-cpp samples/diar/ami-ihm-test /audio + +One container per (meeting, mode) cell, batch size 1 (multitalker is a +transcribe_run path). Mode is selected via TRANSCRIBE_MULTITALKER_MODE, +same contract as run_cpp_multitalker.py. Hypotheses land in +reports/cpwer/cpp-/.seglst.json for score_cpwer.py. + +Usage: + modal run scripts/diar/modal_multitalker_cpp.py --modes both + modal run scripts/diar/modal_multitalker_cpp.py --modes kernel --meetings IS1009a +""" + +from __future__ import annotations + +import json +import os +import subprocess +from pathlib import Path + +import modal + +REPO = Path(__file__).resolve().parent.parent.parent + +app = modal.App("transcribe-multitalker-cpp") +vol = modal.Volume.from_name("transcribe-ami-cpp") + +image = ( + modal.Image.from_registry("nvidia/cuda:12.4.1-devel-ubuntu22.04", add_python="3.11") + .apt_install("build-essential", "git", "ninja-build") + .pip_install("cmake") + .add_local_file(REPO / "CMakeLists.txt", "/repo/CMakeLists.txt", copy=True) + .add_local_dir(REPO / "include", "/repo/include", copy=True) + .add_local_dir(REPO / "cmake", "/repo/cmake", copy=True) + .add_local_dir(REPO / "ggml", "/repo/ggml", copy=True) + .add_local_dir(REPO / "src", "/repo/src", copy=True) + .add_local_dir(REPO / "examples", "/repo/examples", copy=True) + .run_commands( + "cmake -S /repo -B /repo/build -G Ninja -DCMAKE_BUILD_TYPE=Release " + "-DTRANSCRIBE_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=89 " + "-DTRANSCRIBE_BUILD_TESTS=OFF -DTRANSCRIBE_BUILD_EXAMPLES=ON", + "cmake --build /repo/build -j 16 --target transcribe-cli", + ) +) + + +@app.function(image=image, gpu="L40S", volumes={"/data": vol}, timeout=3600) +def run_meeting(meeting_wav: str, mode: str, ref_bg_compat: bool = False) -> dict: + mid = Path(meeting_wav).name.split(".")[0] + batch_list = "/tmp/batch.list" + Path(batch_list).write_text(f"/data/audio/{meeting_wav}\n") + + env = dict(os.environ) + if ref_bg_compat: + env["TRANSCRIBE_MULTITALKER_REF_BG_COMPAT"] = "1" + env["TRANSCRIBE_MULTITALKER_MODE"] = mode + + cmd = [ + "/repo/build/bin/transcribe-cli", "-q", "--diarize", "--backend", "cuda", + "--batch", batch_list, "--batch-jsonl", + "-m", "/data/model/bundle-F32.gguf", + ] + proc = subprocess.run(cmd, capture_output=True, text=True, env=env) + if proc.returncode != 0: + return {"id": mid, "mode": mode, "error": f"exit {proc.returncode}: {proc.stderr[-2000:]}"} + + row = None + for line in proc.stdout.splitlines(): + line = line.strip() + if line.startswith("{"): + try: + cand = json.loads(line) + except json.JSONDecodeError: + continue + if cand.get("segments"): + row = cand + if row is None: + return {"id": mid, "mode": mode, "error": "no JSONL row with segments"} + + seglst = [] + for seg in row["segments"]: + text = (seg.get("text") or "").strip() + if not text or seg.get("speaker_id", 0) <= 0: + continue + seglst.append( + { + "session_id": mid, + "speaker": f"S{seg['speaker_id']}", + "start_time": seg["t0_ms"] / 1000.0, + "end_time": seg["t1_ms"] / 1000.0, + "words": text, + } + ) + return {"id": mid, "mode": mode, "seglst": seglst} + + +@app.local_entrypoint() +def main(modes: str = "both", meetings: str = "", ref_bg_compat: bool = False): + manifest = REPO / "samples" / "diar" / "ami-ihm-test.manifest.jsonl" + entries = [json.loads(l) for l in manifest.read_text().splitlines() if l.strip()] + want = {m.strip() for m in meetings.split(",") if m.strip()} + mode_list = ["masked", "kernel"] if modes == "both" else [modes] + + cells = [] + for e in entries: + mid = e["id"].split(".")[0] + if want and mid not in want: + continue + for mode in mode_list: + cells.append((Path(e["audio"]).name, mode, ref_bg_compat)) + + n_ok = 0 + for res in run_meeting.starmap(cells, order_outputs=False): + if res.get("error"): + print(f"FAIL {res['id']} [{res['mode']}]: {res['error']}") + continue + suffix = "-refbug" if ref_bg_compat else "" + out_dir = REPO / "reports" / "cpwer" / f"cpp-{res['mode']}{suffix}" + out_dir.mkdir(parents=True, exist_ok=True) + out = out_dir / f"{res['id']}.seglst.json" + out.write_text(json.dumps(res["seglst"], indent=1) + "\n") + n_ok += 1 + print(f"ok {res['id']} [{res['mode']}]: {len(res['seglst'])} segments -> {out.relative_to(REPO)}") + + print(f"{n_ok}/{len(cells)} cells succeeded") diff --git a/scripts/diar/modal_multitalker_ref.py b/scripts/diar/modal_multitalker_ref.py new file mode 100644 index 00000000..b375d6ce --- /dev/null +++ b/scripts/diar/modal_multitalker_ref.py @@ -0,0 +1,127 @@ +""" +modal_multitalker_ref.py - NeMo multitalker baseline (ASR + streaming +Sortformer) on Modal L40S for the AMI-IHM cpWER acceptance run. + +Fans out one container per (meeting, supervision-mode) cell, each running +scripts/dump_reference_multitalker_nemo.py (--device cuda --light) over +one staged AMI meeting, and collects the SegLST hypotheses locally under +reports/cpwer/nemo-/.seglst.json for scoring with +scripts/diar/score_cpwer.py. + +Batch size is 1 by design: one meeting per container, the pipeline's +audio_file mode. + +Usage (from the repo root): + # smoke: one meeting, both modes + modal run scripts/diar/modal_multitalker_ref.py --meetings ES2004a --modes both + + # full acceptance sweep (meetings from the staged manifest) + modal run scripts/diar/modal_multitalker_ref.py --modes both +""" + +from __future__ import annotations + +import json +import pathlib + +import modal + +REPO = pathlib.Path(__file__).resolve().parent.parent.parent +NEMO_REV = "6967f48fda2a776a68bbaed6a71d7fea78ccc3f6" + +app = modal.App("transcribe-multitalker-cpwer") + +image = ( + modal.Image.from_registry( + "nvidia/cuda:12.4.1-runtime-ubuntu22.04", + add_python="3.11", + ) + .apt_install("git", "ca-certificates", "ffmpeg", "libsndfile1", "build-essential") + .pip_install("Cython", "packaging", "huggingface_hub") + .pip_install(f"nemo_toolkit[asr] @ git+https://github.com/NVIDIA-NeMo/NeMo.git@{NEMO_REV}") + .add_local_file( + REPO / "scripts" / "dump_reference_multitalker_nemo.py", + "/src/dump_reference_multitalker_nemo.py", + ) + .add_local_dir(REPO / "samples" / "diar" / "ami-ihm-test", "/ami") +) + +hf_vol = modal.Volume.from_name("hf-cache", create_if_missing=True) + + +@app.function( + image=image, + gpu="L40S", + timeout=4 * 3600, + volumes={"/root/.cache/huggingface": hf_vol}, + secrets=[modal.Secret.from_name("huggingface-secret")], +) +def run_meeting(meeting: str, masked: bool, strict_fp32: bool = False) -> dict: + import subprocess + + from huggingface_hub import hf_hub_download + + asr = hf_hub_download( + "nvidia/multitalker-parakeet-streaming-0.6b-v1", + "multitalker-parakeet-streaming-0.6b-v1.nemo", + ) + diar = hf_hub_download( + "nvidia/diar_streaming_sortformer_4spk-v2.1", + "diar_streaming_sortformer_4spk-v2.1.nemo", + ) + hf_vol.commit() + + wav = f"/ami/{meeting}.Mix-Headset.wav" + out = f"/tmp/out/{meeting}" + mode = "true" if masked else "false" + cmd = [ + "python", "/src/dump_reference_multitalker_nemo.py", + "--audio", wav, + "--asr-model", asr, + "--diar-model", diar, + "--masked-asr", mode, + "--device", "cuda", + "--light", + "--out", out, + ] + if strict_fp32: + cmd.append("--strict-fp32") + proc = subprocess.run(cmd, capture_output=True, text=True) + tail = "\n".join((proc.stdout + "\n" + proc.stderr).splitlines()[-30:]) + if proc.returncode != 0: + return {"meeting": meeting, "masked": masked, "ok": False, "log": tail} + seglst = json.loads(pathlib.Path(out, "seglst.json").read_text()) + return {"meeting": meeting, "masked": masked, "ok": True, "seglst": seglst, "log": tail} + + +@app.local_entrypoint() +def main(meetings: str = "", modes: str = "both", strict_fp32: bool = False): + if meetings: + ids = [m.strip() for m in meetings.split(",") if m.strip()] + else: + manifest = REPO / "samples" / "diar" / "ami-ihm-test.manifest.jsonl" + ids = sorted( + {json.loads(l)["id"].split(".")[0] for l in manifest.read_text().splitlines() if l.strip()} + ) + mode_list = [True, False] if modes == "both" else [modes == "masked"] + + cells = [(m, masked, strict_fp32) for m in ids for masked in mode_list] + print(f"dispatching {len(cells)} cells ({len(ids)} meetings x {len(mode_list)} modes) on L40S") + + n_ok = 0 + for res in run_meeting.starmap(cells, return_exceptions=True): + if isinstance(res, Exception): + print(f"FAIL (exception): {res}") + continue + mode_name = "masked" if res["masked"] else "kernel" + if not res["ok"]: + print(f"FAIL {res['meeting']} [{mode_name}]:\n{res['log']}") + continue + suffix = "-fp32" if strict_fp32 else "" + out_dir = REPO / "reports" / "cpwer" / f"nemo-{mode_name}{suffix}" + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / f"{res['meeting']}.seglst.json" + out_path.write_text(json.dumps(res["seglst"], indent=1) + "\n") + n_ok += 1 + print(f"ok {res['meeting']} [{mode_name}]: {len(res['seglst'])} segments -> {out_path.relative_to(REPO)}") + print(f"{n_ok}/{len(cells)} cells succeeded") diff --git a/scripts/diar/run_cpp_multitalker.py b/scripts/diar/run_cpp_multitalker.py new file mode 100644 index 00000000..485361ae --- /dev/null +++ b/scripts/diar/run_cpp_multitalker.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +""" +run_cpp_multitalker.py - drive the C++ multitalker bundle over an AMI +manifest and emit SegLST hypotheses for cpWER scoring. + +One transcribe-cli invocation per meeting (batch size 1: multitalker is a +transcribe_run path), diarize=ON, --batch-jsonl output parsed into +/.seglst.json. Supervision mode is selected via +TRANSCRIBE_MULTITALKER_MODE (masked default | kernel), forwarded from +--mode. + +Usage: + uv run scripts/diar/run_cpp_multitalker.py \ + --model models/multitalker-parakeet-streaming-0.6b-v1/bundle/multitalker-parakeet-streaming-0.6b-v1-F32.gguf \ + --manifest samples/diar/ami-ihm-test.manifest.jsonl \ + --mode kernel \ + --out-dir reports/cpwer/cpp-kernel \ + [--meetings IS1009a,ES2004a] [--backend cpu|metal|cuda] [--cli PATH] +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import tempfile +import time +from pathlib import Path + +REPO = Path(__file__).resolve().parent.parent.parent + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", required=True) + ap.add_argument("--manifest", required=True) + ap.add_argument("--out-dir", required=True) + ap.add_argument("--mode", choices=["masked", "kernel"], default="masked") + ap.add_argument("--backend", default="auto") + ap.add_argument("--cli", default=str(REPO / "build" / "bin" / "transcribe-cli")) + ap.add_argument("--meetings", default="", help="comma-separated meeting-id filter") + args = ap.parse_args() + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + want = {m.strip() for m in args.meetings.split(",") if m.strip()} + + env = dict(os.environ) + if args.mode == "kernel": + env["TRANSCRIBE_MULTITALKER_MODE"] = "kernel" + else: + env.pop("TRANSCRIBE_MULTITALKER_MODE", None) + + entries = [json.loads(l) for l in Path(args.manifest).read_text().splitlines() if l.strip()] + n_ok = 0 + for e in entries: + mid = e["id"].split(".")[0] + if want and mid not in want: + continue + audio = e["audio"] + if not Path(audio).is_absolute(): + audio = str(REPO / audio) + + with tempfile.NamedTemporaryFile("w", suffix=".list", delete=False) as bf: + bf.write(audio + "\n") + batch_list = bf.name + cmd = [ + args.cli, "-q", "--diarize", "--backend", args.backend, + "--batch", batch_list, "--batch-jsonl", + "-m", args.model, + ] + t0 = time.time() + proc = subprocess.run(cmd, capture_output=True, text=True, env=env) + os.unlink(batch_list) + if proc.returncode != 0: + print(f"FAIL {mid}: exit {proc.returncode}\n{proc.stderr[-2000:]}", file=sys.stderr) + continue + + row = None + for line in proc.stdout.splitlines(): + line = line.strip() + if line.startswith("{"): + try: + cand = json.loads(line) + except json.JSONDecodeError: + continue + if cand.get("segments"): + row = cand + if row is None: + print(f"FAIL {mid}: no JSONL row with segments in CLI output", file=sys.stderr) + continue + + seglst = [] + for seg in row["segments"]: + text = (seg.get("text") or "").strip() + if not text or seg.get("speaker_id", 0) <= 0: + continue + seglst.append( + { + "session_id": mid, + "speaker": f"S{seg['speaker_id']}", + "start_time": seg["t0_ms"] / 1000.0, + "end_time": seg["t1_ms"] / 1000.0, + "words": text, + } + ) + (out_dir / f"{mid}.seglst.json").write_text(json.dumps(seglst, indent=1) + "\n") + n_ok += 1 + print(f"ok {mid}: {len(seglst)} segments in {time.time() - t0:.0f}s") + + print(f"{n_ok} meetings -> {out_dir}") + return 0 if n_ok > 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/diar/score_cpwer.py b/scripts/diar/score_cpwer.py new file mode 100644 index 00000000..bf852902 --- /dev/null +++ b/scripts/diar/score_cpwer.py @@ -0,0 +1,126 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = ["meeteval>=0.4"] +# /// +""" +score_cpwer.py - cpWER for speaker-attributed multitalker hypotheses. + +Scores hypothesis SegLST files against the reference SegLSTs from +ingest_ami_seglst.py using meeteval's cpwer (concatenated +minimum-permutation WER — the metric NVIDIA reports for the multitalker +parakeet models). Text is normalized with a whisper-style English +normalizer to match the repo's WER convention. + +Inputs: a -mt manifest (audio + rttm + seglst columns) and a hypothesis +directory containing .seglst.json files (from the C++ or NeMo +multitalker runners). + +Usage: + uv run scripts/diar/score_cpwer.py \ + --manifest samples/diar/ami-ihm-test-mt.manifest.jsonl \ + --hyp-dir reports/cpwer// \ + --out reports/cpwer/.score.json +""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from pathlib import Path + +import meeteval.io +from meeteval.wer import cpwer + +REPO = Path(__file__).resolve().parent.parent.parent + +_norm_re = re.compile(r"[^a-z0-9' ]+") + + +def normalize(text: str) -> str: + # Lowercase, strip punctuation, collapse whitespace. Matches the + # de-PnC convention score.py applies before WER. + text = text.lower().replace("-", " ") + text = _norm_re.sub(" ", text) + return " ".join(text.split()) + + +def load_seglst(path: Path, session_id: str) -> meeteval.io.SegLST: + rows = json.loads(path.read_text()) + cleaned = [] + for r in rows: + words = normalize(r.get("words", "")) + if not words: + continue + cleaned.append( + { + "session_id": session_id, + "speaker": str(r["speaker"]), + "start_time": float(r.get("start_time", 0.0)), + "end_time": float(r.get("end_time", 0.0)), + "words": words, + } + ) + return meeteval.io.SegLST.new(cleaned) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--manifest", required=True) + ap.add_argument("--hyp-dir", required=True) + ap.add_argument("--out", required=True) + args = ap.parse_args() + + hyp_dir = Path(args.hyp_dir) + per_meeting = {} + tot_err = tot_len = 0 + for line in Path(args.manifest).read_text().splitlines(): + if not line.strip(): + continue + e = json.loads(line) + mid = e["id"].split(".")[0] + ref_path = REPO / e["seglst"] + hyp_path = hyp_dir / f"{mid}.seglst.json" + if not hyp_path.exists(): + print(f"warning: missing hyp for {mid}, skipping", file=sys.stderr) + continue + ref = load_seglst(ref_path, mid) + hyp = load_seglst(hyp_path, mid) + r = cpwer(ref, hyp)[mid] + per_meeting[mid] = { + "errors": r.errors, + "length": r.length, + "cpwer_pct": round(100.0 * r.error_rate, 3) if r.error_rate is not None else None, + "insertions": r.insertions, + "deletions": r.deletions, + "substitutions": r.substitutions, + "missed_speaker": r.missed_speaker, + "falarm_speaker": r.falarm_speaker, + "scored_speaker": r.scored_speaker, + } + tot_err += r.errors + tot_len += r.length + + if tot_len == 0: + print("error: nothing scored", file=sys.stderr) + return 1 + summary = { + "metric": "cpWER", + "cpwer_pct": round(100.0 * tot_err / tot_len, 3), + "errors": tot_err, + "length": tot_len, + "n_meetings": len(per_meeting), + "per_meeting": per_meeting, + } + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(summary, indent=2) + "\n") + print(f"cpWER: {summary['cpwer_pct']}% over {len(per_meeting)} meetings ({tot_err}/{tot_len})") + print(f"wrote {out}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/dump_reference_multitalker_nemo.py b/scripts/dump_reference_multitalker_nemo.py new file mode 100644 index 00000000..0a138a7e --- /dev/null +++ b/scripts/dump_reference_multitalker_nemo.py @@ -0,0 +1,320 @@ +#!/usr/bin/env python3 +""" +dump_reference_multitalker_nemo.py - NeMo reference dump for the multitalker +Parakeet + streaming Sortformer pipeline. + +Replicates the pinned NeMo example +(examples/asr/asr_cache_aware_streaming/speech_to_text_multitalker_streaming_infer.py +at rev 6967f48fda2a, single-audio-file mode, parallel speaker strategy) and +instruments it so the C++ port has an oracle for every hand-off: + + - per-step diar predictions (cumulative stream and per-chunk slice) + - per-step active-speaker selection (cache gating) + - per-step ASR supervision: feature masks (masked_asr=true) or + spk/bg speaker-kernel targets (masked_asr=false) + - final SegLST and per-speaker word streams + - the effective streaming configs actually in force (sortformer module + attrs, ASR streaming_cfg) so the C++ side mirrors reality, not defaults + +All defaults follow MultitalkerTranscriptionConfig at the pin; the only +deliberate overrides are CPU placement and use_amp=False (parity runs are +fp32 CPU, per the porting convention). + +Usage: + uv run --project scripts/envs/parakeet scripts/dump_reference_multitalker_nemo.py \ + --audio samples/multitalker-2spk-mix.wav \ + --asr-model ~/.cache/huggingface/.../multitalker-parakeet-streaming-0.6b-v1.nemo \ + --diar-model ~/.cache/huggingface/.../diar_streaming_sortformer_4spk-v2.1.nemo \ + --masked-asr true \ + --out build/validate/parakeet/multitalker-parakeet-streaming-0.6b-v1/multitalker/masked +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch +from omegaconf import OmegaConf + + +def build_cfg(args) -> OmegaConf: + # Field-for-field mirror of MultitalkerTranscriptionConfig defaults at + # the pinned rev; overrides limited to input/output paths, CPU, no AMP, + # and the masked_asr toggle under test. + cfg = { + "diar_model": str(args.diar_model), + "diar_pretrained_name": None, + "max_num_of_spks": 4, + "parallel_speaker_strategy": True, + "masked_asr": args.masked_asr, + "mask_preencode": False, + "cache_gating": True, + "cache_gating_buffer_size": 2, + "single_speaker_mode": False, + "feat_len_sec": 0.01, + "session_len_sec": -1, + "num_workers": 0, + "random_seed": None, + "log": False, + "streaming_mode": True, + "spkcache_len": 188, + "spkcache_refresh_rate": 0, + "fifo_len": 188, + "chunk_len": 0, + "chunk_left_context": 0, + "chunk_right_context": 0, + "cuda": None, + "allow_mps": False, + "matmul_precision": "highest", + "asr_model": str(args.asr_model), + "device": "cpu", + "audio_file": str(args.audio), + "manifest_file": None, + "att_context_size": [70, 13], + "use_amp": False, + "debug_mode": False, + "deploy_mode": False, + "batch_size": 1, + "chunk_size": -1, + "shift_size": -1, + "left_chunks": 2, + "online_normalization": False, + "output_path": None, + "pad_and_drop_preencoded": False, + "generate_realtime_scripts": False, + "spk_supervision": "diar", + "binary_diar_preds": False, + "verbose": False, + "word_window": 50, + "sent_break_sec": 30.0, + "fix_prev_words_count": 5, + "update_prev_words_sentence": 5, + "left_frame_shift": -1, + "right_frame_shift": 0, + "min_sigmoid_val": 1e-2, + "discarded_frames": 8, + "print_time": False, + "print_sample_indices": [0], + "colored_text": False, + "real_time_mode": False, + "print_path": None, + "ignored_initial_frame_steps": 5, + "finetune_realtime_ratio": 0.01, + } + return OmegaConf.create(cfg) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--audio", required=True) + ap.add_argument("--asr-model", required=True, help="path to the multitalker .nemo") + ap.add_argument("--diar-model", required=True, help="path to the sortformer .nemo") + ap.add_argument("--masked-asr", choices=["true", "false"], required=True) + ap.add_argument("--out", required=True, help="output directory") + ap.add_argument("--device", choices=["cpu", "cuda"], default="cpu", + help="cpu (parity dumps) or cuda (long-form baseline runs)") + ap.add_argument("--dump-compress", action="store_true", + help="record per-call _compress_spkcache selection (compress-ref/ dumps)") + ap.add_argument("--strict-fp32", action="store_true", + help="disable TF32 everywhere (cudnn convs + cuda matmul); " + "isolates precision when comparing against true-F32 backends") + ap.add_argument("--light", action="store_true", + help="skip per-step instrumentation dumps (seglst + config echo only); " + "required for long meetings where per-chunk npz would be huge") + args = ap.parse_args() + args.masked_asr = args.masked_asr == "true" + + out_dir = Path(args.out) + out_dir.mkdir(parents=True, exist_ok=True) + + import nemo.collections.asr as nemo_asr + from nemo.collections.asr.models.sortformer_diar_models import SortformerEncLabelModel + from nemo.collections.asr.parts.utils.multispk_transcribe_utils import SpeakerTaggedASR + from nemo.collections.asr.parts.utils.streaming_utils import CacheAwareStreamingAudioBuffer + + cfg = build_cfg(args) + torch.set_float32_matmul_precision(cfg.matmul_precision) + if args.strict_fp32: + torch.backends.cudnn.allow_tf32 = False + torch.backends.cuda.matmul.allow_tf32 = False + device = torch.device(args.device) + + diar_model = SortformerEncLabelModel.restore_from(restore_path=cfg.diar_model, map_location=device) + diar_model = diar_model.eval().to(device) + compress_calls = None + if args.dump_compress: + import sys as _sys + from pathlib import Path as _P + _sys.path.insert(0, str(_P(__file__).resolve().parent)) + from dump_reference_sortformer_nemo import _install_compress_hooks, _write_compress_dump + compress_calls = _install_compress_hooks(diar_model) + diar_model.streaming_mode = cfg.streaming_mode + diar_model.sortformer_modules.chunk_len = cfg.chunk_len + diar_model.sortformer_modules.spkcache_len = cfg.spkcache_len + diar_model.sortformer_modules.chunk_left_context = cfg.chunk_left_context + diar_model.sortformer_modules.chunk_right_context = cfg.chunk_right_context + diar_model.sortformer_modules.fifo_len = cfg.fifo_len + diar_model.sortformer_modules.log = cfg.log + diar_model.sortformer_modules.spkcache_refresh_rate = cfg.spkcache_refresh_rate + + asr_model = nemo_asr.models.ASRModel.restore_from(restore_path=cfg.asr_model, map_location=device) + asr_model.encoder.set_default_att_context_size(att_context_size=list(cfg.att_context_size)) + asr_model = asr_model.to(device).eval() + + # ------------------------------------------------------------------ # + # Instrumentation: record every hand-off the C++ port must reproduce # + # ------------------------------------------------------------------ # + steps: list[dict] = [] + instrument = not args.light + + orig_fss = diar_model.forward_streaming_step + + def rec_fss(*a, **kw): + state, preds = orig_fss(*a, **kw) + steps.append({"kind": "diar_preds", "step": len([s for s in steps if s["kind"] == "diar_preds"]), + "total_preds": preds.detach().cpu().numpy().copy()}) + return state, preds + + if instrument: + diar_model.forward_streaming_step = rec_fss + + orig_sst = asr_model.set_speaker_targets + + def rec_sst(spk_targets=None, bg_spk_targets=None): + steps.append({"kind": "spk_targets", + "spk": None if spk_targets is None else spk_targets.detach().cpu().numpy().copy(), + "bg": None if bg_spk_targets is None else bg_spk_targets.detach().cpu().numpy().copy()}) + return orig_sst(spk_targets, bg_spk_targets) + + if instrument: + asr_model.set_speaker_targets = rec_sst + + streamer = SpeakerTaggedASR(cfg, asr_model, diar_model) + + orig_find = streamer._find_active_speakers + + def rec_find(diar_preds, n_active_speakers_per_stream): + res = orig_find(diar_preds, n_active_speakers_per_stream) + steps.append({"kind": "active_speakers", "active": [list(map(int, r)) for r in res]}) + return res + + if instrument: + streamer._find_active_speakers = rec_find + + orig_mask_feat = streamer.mask_features + + def rec_mask_feat(chunk_audio, mask, threshold=0.5, mask_value=-16.6355): + steps.append({"kind": "feature_mask", "mask": mask.detach().cpu().numpy().copy()}) + return orig_mask_feat(chunk_audio, mask, threshold, mask_value) + + if instrument: + streamer.mask_features = rec_mask_feat + + # ------------------------------------------------------------------ # + # Streaming loop (mirrors launch_parallel_streaming at the pin) # + # ------------------------------------------------------------------ # + streaming_buffer = CacheAwareStreamingAudioBuffer( + model=asr_model, + online_normalization=cfg.online_normalization, + pad_and_drop_preencoded=cfg.pad_and_drop_preencoded, + ) + streaming_buffer.append_audio_file(audio_filepath=cfg.audio_file, stream_id=-1) + + autocast = torch.amp.autocast("cpu", enabled=cfg.use_amp) + for step_num, (chunk_audio, chunk_lengths) in enumerate(iter(streaming_buffer)): + drop_extra_pre_encoded = ( + 0 if step_num == 0 and not cfg.pad_and_drop_preencoded + else asr_model.encoder.streaming_cfg.drop_extra_pre_encoded + ) + with torch.inference_mode(), autocast, torch.no_grad(): + streamer.perform_parallel_streaming_stt_spk( + step_num=step_num, + chunk_audio=chunk_audio, + chunk_lengths=chunk_lengths, + is_buffer_empty=streaming_buffer.is_buffer_empty(), + drop_extra_pre_encoded=drop_extra_pre_encoded, + ) + + samples = [{"audio_filepath": cfg.audio_file}] + seglst = streamer.generate_seglst_dicts_from_parallel_streaming(samples=samples) + + # ------------------------------------------------------------------ # + # Dump # + # ------------------------------------------------------------------ # + # Accumulated diar preds are tiny and always worth keeping (light + # mode included): they are the supervision ground truth for parity. + diar_states = streamer.instance_manager.diar_states + final_preds = diar_states.diar_pred_out_stream + np.save(out_dir / "diar_total_preds.npy", final_preds.detach().cpu().numpy()) + if compress_calls is not None: + _write_compress_dump(compress_calls, out_dir / "compress-ref") + + npz: dict[str, np.ndarray] = {} + counters: dict[str, int] = {} + step_index = [] + for s in steps: + i = counters.get(s["kind"], 0) + counters[s["kind"]] = i + 1 + if s["kind"] == "diar_preds": + npz[f"diar_preds_{i:03d}"] = s["total_preds"] + elif s["kind"] == "spk_targets": + if s["spk"] is not None: + npz[f"spk_targets_{i:03d}"] = s["spk"] + if s["bg"] is not None: + npz[f"bg_targets_{i:03d}"] = s["bg"] + elif s["kind"] == "feature_mask": + npz[f"feature_mask_{i:03d}"] = s["mask"] + elif s["kind"] == "active_speakers": + step_index.append(s["active"]) + np.savez(out_dir / "per_step.npz", **npz) + + scfg = asr_model.encoder.streaming_cfg + echo = { + "masked_asr": cfg.masked_asr, + "binary_diar_preds": cfg.binary_diar_preds, + "cache_gating": cfg.cache_gating, + "cache_gating_buffer_size": cfg.cache_gating_buffer_size, + "single_speaker_mode": cfg.single_speaker_mode, + "att_context_size": list(cfg.att_context_size), + "nframes_per_chunk": streamer._nframes_per_chunk, + "sent_break_sec": streamer._sent_break_sec, + "active_speakers_per_step": step_index, + "sortformer_modules": { + k: getattr(diar_model.sortformer_modules, k) + for k in ("chunk_len", "spkcache_len", "fifo_len", "chunk_left_context", + "chunk_right_context", "spkcache_refresh_rate") + }, + "asr_streaming_cfg": { + "chunk_size": list(scfg.chunk_size) if hasattr(scfg.chunk_size, "__len__") else scfg.chunk_size, + "shift_size": list(scfg.shift_size) if hasattr(scfg.shift_size, "__len__") else scfg.shift_size, + "pre_encode_cache_size": list(scfg.pre_encode_cache_size) + if hasattr(scfg.pre_encode_cache_size, "__len__") else scfg.pre_encode_cache_size, + "drop_extra_pre_encoded": scfg.drop_extra_pre_encoded, + "valid_out_len": scfg.valid_out_len, + }, + } + (out_dir / "config_echo.json").write_text(json.dumps(echo, indent=2, default=str) + "\n") + (out_dir / "seglst.json").write_text(json.dumps(seglst, indent=2) + "\n") + + per_speaker: dict[str, str] = {} + for seg in seglst: + per_speaker.setdefault(seg["speaker"], []) + for seg in sorted(seglst, key=lambda s: s["start_time"]): + per_speaker[seg["speaker"]] = per_speaker.get(seg["speaker"], []) + per_speaker[seg["speaker"]].append(seg["words"]) + per_speaker_text = {k: " ".join(v) for k, v in per_speaker.items()} + (out_dir / "per_speaker_text.json").write_text(json.dumps(per_speaker_text, indent=2) + "\n") + + print(f"steps recorded: {counters}") + print(f"seglst segments: {len(seglst)}") + for spk, text in per_speaker_text.items(): + print(f" {spk}: {text[:120]}") + print(f"wrote {out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/envs/sortformer/pyproject.toml b/scripts/envs/sortformer/pyproject.toml index 0e402747..7d16c7aa 100644 --- a/scripts/envs/sortformer/pyproject.toml +++ b/scripts/envs/sortformer/pyproject.toml @@ -14,6 +14,7 @@ dependencies = [ "numpy>=1.26", "sentencepiece>=0.2", "gguf>=0.10.0", + "meeteval>=0.4.3", ] # NeMo main maps torch by platform: darwin->pypi (CPU), linux->cu129, diff --git a/scripts/envs/sortformer/uv.lock b/scripts/envs/sortformer/uv.lock index 80608bd5..b41cb9ce 100644 --- a/scripts/envs/sortformer/uv.lock +++ b/scripts/envs/sortformer/uv.lock @@ -406,6 +406,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, ] +[[package]] +name = "cython" +version = "3.2.9" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/de/db48b8870e766cfea809986cc50c1e986c663a9ab7bafd0ac1a2512c4a26/cython-3.2.9.tar.gz", hash = "sha256:d249c9022ab13286b17bd66f30609e800c5f95efeecb06168990c7a66cecde6c", size = 3293493, upload-time = "2026-07-24T06:21:21.867Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/97/2f477d30fc6ca0ff333552233aa6dee0e57323a55913f84f24da9cac26fc/cython-3.2.9-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e8ec2d5a7b798c84d6779e7ca5318fd928b8fe9dd021106d714dc2e77cdc7555", size = 2992696, upload-time = "2026-07-24T06:21:44.931Z" }, + { url = "https://files.pythonhosted.org/packages/0a/bf/0f280f7960c129957d2ce650e7fd7dec8e706c26d61a428e93927f4c9904/cython-3.2.9-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cac792b0bde1c86d8f832e513cd061b7e682181cc5a6d843d487e6c8ae9d5fcb", size = 3307656, upload-time = "2026-07-24T06:21:46.649Z" }, + { url = "https://files.pythonhosted.org/packages/ac/cd/6d7c35a1065f1a07c9fa5ec784ee562f32f36a8fd05748e5f10694887eb7/cython-3.2.9-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8c843dd85857cd0d0da055489f448d032866120cfb094ce16205a1ff54d06353", size = 3459259, upload-time = "2026-07-24T06:21:48.881Z" }, + { url = "https://files.pythonhosted.org/packages/fd/37/c74d842306c8fe381c415b37460d5e3086a820fac72b8ff5cb48513ccfcd/cython-3.2.9-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:114b2dee0fa1daa48a59574d848da0ff1b6bdb725a755e9b92fad14962e1ff8d", size = 3009571, upload-time = "2026-07-24T06:21:52.534Z" }, + { url = "https://files.pythonhosted.org/packages/15/c0/f7c42b161edd585e3ae556fd62a2c72cd80a6ed527a9907f0c5c6fb060de/cython-3.2.9-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a5cd9c5f138cb052130b40ad3b6976d2180c35348410995812678f4636bd8f94", size = 3183562, upload-time = "2026-07-24T06:21:54.62Z" }, + { url = "https://files.pythonhosted.org/packages/56/1b/c04520ac7f3157aa12a69b632c16261170dac9fab6c48608cc004b8f1b17/cython-3.2.9-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:23e80bc885c599e72072e18d0746df82d394b73100c1e153cda7359e6e59fe09", size = 3354811, upload-time = "2026-07-24T06:21:56.72Z" }, + { url = "https://files.pythonhosted.org/packages/0a/4a/342312c5fe021c8e0c386e1915d138e0902c48ae179b0374ab04773a8831/cython-3.2.9-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:944dc8747f640b3527649c566a5fc75ee0c15e80642ea2fdae4fe6378e1a9d4a", size = 2899729, upload-time = "2026-07-24T06:22:24.877Z" }, + { url = "https://files.pythonhosted.org/packages/f0/62/ea919ee426cb4d435ec8155e1ee6bcbb46b20d8f070527191b59769d4e7f/cython-3.2.9-cp39-abi3-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:4b871ad97dd7fb1cbf56f6238c54423febd310afc1d9d9bc70c69c89b7ce57fc", size = 3226650, upload-time = "2026-07-24T06:22:26.947Z" }, + { url = "https://files.pythonhosted.org/packages/18/02/057b4f63e2ced8c3cf217c4e9fb544bfe48145f493347c7ca3f51607526c/cython-3.2.9-cp39-abi3-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e9b6ebc6c74b4318eaa4e51e520dc8b95ebc7b262953c3ecb24131104681f14e", size = 2881919, upload-time = "2026-07-24T06:22:29.318Z" }, + { url = "https://files.pythonhosted.org/packages/64/e4/e158793ee3de7e4417ba17e7ff1015d6e2cf557cb485ad270b2446c9d1c7/cython-3.2.9-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:92989da161a7d18a7ad4baebc49289b2b77556d5a94916f90140ba26aecf6892", size = 3004702, upload-time = "2026-07-24T06:22:31.535Z" }, + { url = "https://files.pythonhosted.org/packages/f9/c0/d5bbbd743ab4feddb24a7e823b34c3ec4ebab91ff503d16743a4e7ce106b/cython-3.2.9-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:e75ec625d8f8781ced690b7a2f5c2d138067711cf24bb8fb68c872c30c2fefe5", size = 2902695, upload-time = "2026-07-24T06:22:33.525Z" }, + { url = "https://files.pythonhosted.org/packages/d6/2f/80850817395985259f135baa510d9186d3a325df81cd1862060bba977029/cython-3.2.9-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:2b1756ddc3bc0cd4341a515fc420c3e25e13c249f5537159b3fb0bff8d19e55c", size = 3241554, upload-time = "2026-07-24T06:22:35.667Z" }, + { url = "https://files.pythonhosted.org/packages/4c/ce/6be776814f6cb81751f3da737ed537385738148e1ea99f89fb4637799198/cython-3.2.9-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7d41baea51ea00f9237f75af498577827493bca5e9b45bbd4e351543727e589a", size = 3124337, upload-time = "2026-07-24T06:22:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/00/ec/e61deec9bcfbb0e1b36f8b5ba75cb44644419b4bfd0fdd666bffd21d9579/cython-3.2.9-py3-none-any.whl", hash = "sha256:a2b0e87f6b80790c929308ca0831d686f7a180feab684fe8cd4a4380bd96aaca", size = 1259272, upload-time = "2026-07-24T06:21:18.95Z" }, +] + [[package]] name = "cytoolz" version = "1.1.0" @@ -1286,6 +1308,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e3/a0/0d55c59ea8a5f1b13b3eb931e1c0eab9c2a07322cad79cb51a596d2d2a5c/mediapy-1.1.6-py3-none-any.whl", hash = "sha256:c74370808b445666f95272bfdf0eb5707a43b7e05e5527f2dd0830e6892f976f", size = 24955, upload-time = "2023-02-24T13:08:40.53Z" }, ] +[[package]] +name = "meeteval" +version = "0.4.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cython", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "kaldialign", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "packaging", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and sys_platform == 'darwin') or (python_full_version < '3.12' and sys_platform == 'linux')" }, + { name = "scipy", version = "1.18.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and sys_platform == 'darwin') or (python_full_version >= '3.12' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d4/13/10edb36a8881b24690e27213bfb2395e8857a70eeef61d241a463cb5522c/meeteval-0.4.3.tar.gz", hash = "sha256:02d3a359f375d39c67dfb8fe1c061e7dac19d6fc1fb89ee72d793a5813dafeb2", size = 843558, upload-time = "2025-09-08T11:57:30.158Z" } + [[package]] name = "ml-dtypes" version = "0.5.4" @@ -2886,6 +2922,7 @@ source = { virtual = "." } dependencies = [ { name = "gguf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "meeteval", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "nemo-toolkit", extra = ["asr"], marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "sentencepiece", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -2898,6 +2935,7 @@ dependencies = [ requires-dist = [ { name = "gguf", specifier = ">=0.10.0" }, { name = "huggingface-hub", specifier = ">=0.20" }, + { name = "meeteval", specifier = ">=0.4.3" }, { name = "nemo-toolkit", extras = ["asr"], git = "https://github.com/NVIDIA-NeMo/NeMo.git?rev=6967f48fda2a" }, { name = "numpy", specifier = ">=1.26" }, { name = "sentencepiece", specifier = ">=0.2" }, diff --git a/scripts/gen_multitalker_oracle_audio.py b/scripts/gen_multitalker_oracle_audio.py new file mode 100644 index 00000000..563a032f --- /dev/null +++ b/scripts/gen_multitalker_oracle_audio.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Deterministically synthesize a 2-speaker ENGLISH multitalker oracle clip. + +The multitalker Parakeet + Sortformer pipeline needs an English +multi-speaker fixture (the sortformer 2spk mix is cross-language, which is +right for a diarizer oracle but wrong for scoring English ASR hyps). This +builds a reproducible 2-speaker English mixture from two committed clips +with distinct voices, plus a ground-truth RTTM for the authored timeline. + +Speaker A: samples/jfk.wav (JFK, 1961) +Speaker B: samples/whole-earth.wav (different male voice) + +Timeline: mostly turn-taking with one authored 0.5 s overlap at +[5.0, 5.5] s so the overlap path is exercised without dominating the clip. + +Outputs (16 kHz mono, deterministic): + samples/multitalker-2spk-mix.wav + tests/golden/parakeet/multitalker-2spk-mix.rttm + +Run: + uv run scripts/gen_multitalker_oracle_audio.py +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +import numpy as np +import soundfile as sf + +SR = 16000 +REPO = Path(__file__).resolve().parent.parent +SPK_A_SRC = REPO / "samples" / "jfk.wav" +SPK_B_SRC = REPO / "samples" / "whole-earth.wav" +OUT_WAV = REPO / "samples" / "multitalker-2spk-mix.wav" +OUT_RTTM = REPO / "tests" / "golden" / "parakeet" / "multitalker-2spk-mix.rttm" +CLIP_ID = "multitalker-2spk-mix" + +# Authored timeline in seconds. Each speaker's source audio is consumed +# sequentially into that speaker's windows, so window boundaries cut the +# source mid-sentence — fine for parity fixtures (we compare hyps against +# the NeMo reference, not against a ground-truth transcript). +SPK_A_WINDOWS = [(0.0, 5.5), (9.5, 14.0)] +SPK_B_WINDOWS = [(5.0, 9.5)] +TOTAL_SEC = 14.0 + + +def _load_mono_16k(path: Path) -> np.ndarray: + audio, sr = sf.read(str(path), dtype="float32", always_2d=False) + if audio.ndim > 1: + audio = audio.mean(axis=1) + if sr != SR: + n_out = int(round(len(audio) * SR / sr)) + x_old = np.linspace(0.0, 1.0, num=len(audio), endpoint=False) + x_new = np.linspace(0.0, 1.0, num=n_out, endpoint=False) + audio = np.interp(x_new, x_old, audio).astype(np.float32) + peak = float(np.max(np.abs(audio))) or 1.0 + return (audio / peak * 0.9).astype(np.float32) + + +def _lay_windows_sequential(src: np.ndarray, windows: list[tuple[float, float]], total_n: int) -> np.ndarray: + """Consume src sequentially into each window (no tiling restart per + window), so speech flows naturally across a speaker's turns.""" + track = np.zeros(total_n, dtype=np.float32) + cursor = 0 + for start, end in windows: + i0 = int(round(start * SR)) + i1 = int(round(end * SR)) + n = i1 - i0 + seg = src[cursor : cursor + n] + if len(seg) < n: # source exhausted: wrap deterministically + reps = int(np.ceil(n / max(1, len(src)))) + seg = np.concatenate([seg, np.tile(src, reps)])[:n] + track[i0:i1] = seg + cursor += n + return track + + +def main() -> int: + for p in (SPK_A_SRC, SPK_B_SRC): + if not p.exists(): + print(f"error: missing source clip {p}", file=sys.stderr) + return 1 + total_n = int(round(TOTAL_SEC * SR)) + a = _lay_windows_sequential(_load_mono_16k(SPK_A_SRC), SPK_A_WINDOWS, total_n) + b = _lay_windows_sequential(_load_mono_16k(SPK_B_SRC), SPK_B_WINDOWS, total_n) + mix = a + b + peak = float(np.max(np.abs(mix))) or 1.0 + mix = (mix / peak * 0.9).astype(np.float32) + + OUT_WAV.parent.mkdir(parents=True, exist_ok=True) + OUT_RTTM.parent.mkdir(parents=True, exist_ok=True) + sf.write(str(OUT_WAV), mix, SR, subtype="PCM_16") + + lines = [] + for spk, windows in (("spk_A", SPK_A_WINDOWS), ("spk_B", SPK_B_WINDOWS)): + for start, end in windows: + lines.append(f"SPEAKER {CLIP_ID} 1 {start:.3f} {end - start:.3f} {spk} ") + OUT_RTTM.write_text("\n".join(lines) + "\n") + + dur = len(mix) / SR + print(f"wrote {OUT_WAV.relative_to(REPO)} ({dur:.2f}s, 16kHz mono)") + print(f"wrote {OUT_RTTM.relative_to(REPO)} ({len(lines)} turns, 2 speakers, overlap [5.0,5.5]s)") + sha = subprocess.run(["shasum", "-a", "256", str(OUT_WAV)], capture_output=True, text=True) + print("sha256:", sha.stdout.split()[0] if sha.returncode == 0 else "n/a") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 87fcbfc1..5e4a7e92 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -34,6 +34,7 @@ add_library(transcribe causal_lm/causal_lm.cpp granite_conformer/shaw_attn.cpp arch/parakeet/model.cpp + arch/parakeet/multitalker.cpp arch/parakeet/capabilities.cpp arch/parakeet/weights.cpp arch/parakeet/encoder.cpp diff --git a/src/arch/parakeet/encoder.cpp b/src/arch/parakeet/encoder.cpp index 16db9d62..a26af118 100644 --- a/src/arch/parakeet/encoder.cpp +++ b/src/arch/parakeet/encoder.cpp @@ -258,13 +258,26 @@ ggml_tensor * build_spk_kernel_ff(ggml_context * ctx, ggml_tensor * x, const Par return y; } -// Single-speaker mode applies the speaker kernel to x and the background -// kernel to zeros. The latter reduces to a per-channel constant broadcast -// over time. +// Speaker-kernel layer-0 injection (NeMo SpeakerKernelMixin pre-layer hook). +// +// Single-speaker mode (spk_mask == nullptr) mirrors spk_targets=None: +// the speaker kernel sees x unmasked (mask defaults to all-ones) and the +// background kernel sees zeros (bg default_value=0.0), which reduces to a +// per-channel constant broadcast over time. This path is byte-identical +// to the shipped single-speaker graph. +// +// Multitalker mode (spk_mask/bg_mask ne=[1, T] f32) mirrors +// mask_with_speaker_targets: x_spk = FF_spk(x * m_spk) added to x, then +// x_bg = FF_bg(x' * m_bg) added on the UPDATED x' (NeMo applies the bg +// kernel after the spk residual), with m_spk the target speaker's raw +// per-frame sigmoids and m_bg the binarized union of the other active +// speakers. ggml_tensor * apply_spk_kernel_injection(ggml_context * ctx, ggml_tensor * x, const ParakeetWeights & w, - const ParakeetHParams & hp) { + const ParakeetHParams & hp, + ggml_tensor * spk_mask = nullptr, + ggml_tensor * bg_mask = nullptr) { if (!hp.has_spk_kernel) { return x; } @@ -272,14 +285,23 @@ ggml_tensor * apply_spk_kernel_injection(ggml_context * ctx, if (kernel.layer != 0) { continue; // loader rejects non-zero layers; defensive } - ggml_tensor * x_spk = build_spk_kernel_ff(ctx, x, kernel.spk); - x = ggml_add(ctx, x, x_spk); - if (kernel.has_bg) { - // W3_bg * relu(b0_bg) + b3_bg is broadcast over time. - ggml_tensor * hb = ggml_relu(ctx, kernel.bg.lin0_b); - ggml_tensor * c_bg = ggml_mul_mat(ctx, kernel.bg.lin3_w, hb); - c_bg = ggml_add(ctx, c_bg, kernel.bg.lin3_b); - x = ggml_add(ctx, x, c_bg); + if (spk_mask == nullptr) { + ggml_tensor * x_spk = build_spk_kernel_ff(ctx, x, kernel.spk); + x = ggml_add(ctx, x, x_spk); + if (kernel.has_bg) { + // W3_bg * relu(b0_bg) + b3_bg is broadcast over time. + ggml_tensor * hb = ggml_relu(ctx, kernel.bg.lin0_b); + ggml_tensor * c_bg = ggml_mul_mat(ctx, kernel.bg.lin3_w, hb); + c_bg = ggml_add(ctx, c_bg, kernel.bg.lin3_b); + x = ggml_add(ctx, x, c_bg); + } + } else { + ggml_tensor * x_spk = build_spk_kernel_ff(ctx, ggml_mul(ctx, x, spk_mask), kernel.spk); + x = ggml_add(ctx, x, x_spk); + if (kernel.has_bg && bg_mask != nullptr) { + ggml_tensor * x_bg = build_spk_kernel_ff(ctx, ggml_mul(ctx, x, bg_mask), kernel.bg); + x = ggml_add(ctx, x, x_bg); + } } } return x; @@ -295,7 +317,9 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const char * backend_name, const BufferedStreamMaskOverride * buf_mask, int n_batch, - bool batch_var_len) { + bool batch_var_len, + bool spk_supervision, + int mt_keep_frames) { if (n_batch < 1) { n_batch = 1; } @@ -386,9 +410,40 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, x = conf::named(x, "enc.pre_encode.xscaled"); } - // NeMo applies the optional layer-0 injection after xscaling. + // NeMo applies the optional layer-0 injection after xscaling. With + // spk_supervision (multitalker bundle, n_batch == 1), the per-frame + // speaker/background masks become graph inputs the driver fills from + // the embedded diarizer's predictions. + // Kernel-mode chunk gating: keep only the target speaker's active + // chunks (reference cache_gating — inactive chunks never reach that + // speaker's conformer/decoder, and its attention context spans its + // own kept history, matching the per-speaker streaming caches). + // Applied after the continuous pre_encode + xscaling, before the + // layer-0 injection, exactly where the reference gathers instances. + if (spk_supervision && n_batch == 1 && mt_keep_frames > 0) { + ggml_tensor * keep = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, mt_keep_frames); + ggml_set_name(keep, "mt_keep.in"); + ggml_set_input(keep); + eb.mt_keep_in = keep; + x = ggml_get_rows(ctx, x, keep); + x = conf::named(x, "enc.pre_encode.mt_gathered"); + } + if (hp.has_spk_kernel) { - x = apply_spk_kernel_injection(ctx, x, w, hp); + ggml_tensor * spk_mask = nullptr; + ggml_tensor * bg_mask = nullptr; + if (spk_supervision && n_batch == 1) { + const int64_t T_inj = x->ne[1]; + spk_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, T_inj); + ggml_set_name(spk_mask, "spk_mask.in"); + ggml_set_input(spk_mask); + bg_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, T_inj); + ggml_set_name(bg_mask, "bg_mask.in"); + ggml_set_input(bg_mask); + eb.spk_mask_in = spk_mask; + eb.bg_mask_in = bg_mask; + } + x = apply_spk_kernel_injection(ctx, x, w, hp, spk_mask, bg_mask); x = conf::named(x, "enc.pre_encode.spk_injected"); } @@ -644,7 +699,8 @@ EncoderBuild build_encoder_graph_streaming(ggml_context * ctx, int drop_extra_pre_encoded, StreamingEncoderCacheIO & cache_io, ggml_type kv_type, - const char * backend_name) { + const char * backend_name, + bool spk_supervision) { conf::ConvPolicy policy{}; policy.direct_pw = conf::detect_direct_pw(backend_name); policy.direct_dw_in_block = detect_direct_dw_in_block(backend_name); @@ -722,8 +778,24 @@ EncoderBuild build_encoder_graph_streaming(ggml_context * ctx, } // Apply the same optional layer-0 injection to each streaming chunk. + // With spk_supervision (multitalker streaming pass), the per-chunk + // speaker/background masks become graph inputs the driver fills from + // the diarizer predictions for this chunk's frames. if (hp.has_spk_kernel) { - x = apply_spk_kernel_injection(ctx, x, w, hp); + ggml_tensor * spk_mask = nullptr; + ggml_tensor * bg_mask = nullptr; + if (spk_supervision) { + const int64_t T_inj = x->ne[1]; + spk_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, T_inj); + ggml_set_name(spk_mask, "stream.spk_mask.in"); + ggml_set_input(spk_mask); + bg_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, T_inj); + ggml_set_name(bg_mask, "stream.bg_mask.in"); + ggml_set_input(bg_mask); + eb.spk_mask_in = spk_mask; + eb.bg_mask_in = bg_mask; + } + x = apply_spk_kernel_injection(ctx, x, w, hp, spk_mask, bg_mask); } const int64_t T_q_new = x->ne[1]; diff --git a/src/arch/parakeet/encoder.h b/src/arch/parakeet/encoder.h index 23f94270..7105e9b1 100644 --- a/src/arch/parakeet/encoder.h +++ b/src/arch/parakeet/encoder.h @@ -134,6 +134,22 @@ struct EncoderBuild { ggml_tensor * pre_encode_mask_s2_in = nullptr; // after relu3 ggml_tensor * pre_encode_mask_s3_in = nullptr; // after relu6 + // Multitalker speaker-kernel supervision inputs, ne=[1, T_enc] f32 + // each. Null unless build_encoder_graph was called with + // spk_supervision on a speaker-kernel variant (n_batch == 1). The + // driver fills spk_mask_in with the target speaker's per-frame + // activity (raw diarizer sigmoids) and bg_mask_in with the binarized + // union of the other active speakers, both aligned to encoder frames. + ggml_tensor * spk_mask_in = nullptr; + ggml_tensor * bg_mask_in = nullptr; + + // Kernel-mode chunk gating (multitalker): I32 ne=[T_keep] gather + // indices into the pre_encode output, filled by the driver with the + // target speaker's active-chunk encoder frames. Null unless + // build_encoder_graph was called with mt_keep_frames > 0. When + // present, spk_mask_in/bg_mask_in/out are sized T_keep, not T_enc. + ggml_tensor * mt_keep_in = nullptr; + // Encoder forward output, ne=[d_model, T_enc, 1, 1] f32. Equal // to dumps.final_out; provided as a separate field so callers // that don't care about intermediates can ignore the dumps @@ -189,16 +205,27 @@ EncoderBuild build_encoder_graph(ggml_context * compute_ctx, const ParakeetWeights & weights, const ParakeetHParams & hp, int n_mel_frames, - ggml_type kv_type = GGML_TYPE_COUNT, - const char * backend_name = "", - const BufferedStreamMaskOverride * buf_mask = nullptr, - int n_batch = 1, + ggml_type kv_type = GGML_TYPE_COUNT, + const char * backend_name = "", + const BufferedStreamMaskOverride * buf_mask = nullptr, + int n_batch = 1, // When true and n_batch > 1, allocate the // variable-length batch masks (attn_pad_mask_in // + conv_pad_mask_in sized for the batch) and // wire them into every conformer block. The // driver fills them from per-utterance lengths. - bool batch_var_len = false); + bool batch_var_len = false, + // Multitalker bundle (n_batch == 1 only): expose + // per-frame speaker/background supervision masks + // as graph inputs (EncoderBuild::spk_mask_in / + // bg_mask_in) consumed by the layer-0 + // speaker-kernel injection. + bool spk_supervision = false, + // Kernel-mode chunk gating: > 0 gathers that many + // pre_encode frames (driver-filled indices in + // EncoderBuild::mt_keep_in) before injection and + // the conformer blocks. 0 = no gather. + int mt_keep_frames = 0); // Per-layer streaming cache I/O for the streaming encoder graph. // The inputs are persistent backend tensors (allocated outside the @@ -269,7 +296,12 @@ EncoderBuild build_encoder_graph_streaming(ggml_context * compute_ctx int n_mel_chunk_frames, int drop_extra_pre_encoded, StreamingEncoderCacheIO & cache_io, - ggml_type kv_type = GGML_TYPE_COUNT, - const char * backend_name = ""); + ggml_type kv_type = GGML_TYPE_COUNT, + const char * backend_name = "", + // Multitalker streaming pass: expose the + // per-chunk speaker/background supervision + // masks (ne=[1, T_q_new]) as graph inputs + // for the layer-0 kernel injection. + bool spk_supervision = false); } // namespace transcribe::parakeet diff --git a/src/arch/parakeet/model.cpp b/src/arch/parakeet/model.cpp index 793efd7c..242d9931 100644 --- a/src/arch/parakeet/model.cpp +++ b/src/arch/parakeet/model.cpp @@ -118,6 +118,15 @@ ParakeetModel::~ParakeetModel() { safe_buffer_free(conv_pw_f32_buffer); conv_pw_f32_buffer = nullptr; } + // Embedded-diarizer fused BN (multitalker bundle only). + if (diar_bn_ctx != nullptr) { + ggml_free(diar_bn_ctx); + diar_bn_ctx = nullptr; + } + if (diar_bn_buffer != nullptr) { + safe_buffer_free(diar_bn_buffer); + diar_bn_buffer = nullptr; + } if (ctx_meta != nullptr) { ggml_free(ctx_meta); ctx_meta = nullptr; @@ -276,6 +285,8 @@ transcribe_status promote_conv_pw_to_f32_on_cpu(ParakeetModel & m) { // Default variant string when the GGUF did not carry stt.variant. constexpr const char k_default_variant[] = "tdt-0.6b-v2"; +} // namespace + // Allocate the streaming encoder caches (cache_last_channel, // cache_last_time), lazily on the first stream_begin for ChunkedLimited // variants. Layout matches NeMo's get_initial_cache_state (see @@ -382,6 +393,8 @@ void reset_streaming_decoder_state(ParakeetSession * pc, const ParakeetModel * p pc->stream_dec_state.initialized = true; } +namespace { + // Forward declarations for the Arch trait below. extern transcribe_status load(Loader &, const transcribe_model_load_params *, transcribe_model **); extern transcribe_status init_context(transcribe_model *, const transcribe_session_params *, transcribe_session **); @@ -542,6 +555,27 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par return st; } + // Multitalker bundle: claim the embedded Sortformer diarizer. The + // bundle GGUF (scripts/compose-multitalker-bundle.py) carries the + // diarizer's tensors under a name prefix next to the ASR's own, plus + // the stt.sortformer.* hparams. Weight resolution needs gguf_data (KVs) + // + ctx_meta (tensors), so it runs before gguf_free; the BN fusion + // below needs uploaded tensor DATA, so it runs after this point too. + { + const int64_t k_embedded = gguf_find_key(gguf_data, "stt.parakeet.diarizer.embedded"); + if (k_embedded >= 0 && gguf_get_val_bool(gguf_data, k_embedded)) { + const int64_t k_prefix = gguf_find_key(gguf_data, "stt.parakeet.diarizer.tensor_prefix"); + const char * prefix = (k_prefix >= 0) ? gguf_get_val_str(gguf_data, k_prefix) : "sortformer."; + m->diar = std::make_unique(); + if (const transcribe_status st = + transcribe::sortformer::init_embedded_diarizer(gguf_data, m->ctx_meta, prefix, *m->diar); + st != TRANSCRIBE_OK) { + gguf_free(gguf_data); + return st; + } + } + } + gguf_free(gguf_data); // Fuse BatchNorm into scale + bias (replaces 4 elementwise ops per @@ -565,6 +599,31 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par return st; } + // Multitalker bundle post-load: fuse the diarizer conformer's BN (needs + // uploaded data), build its ceil-framed mel frontend, and publish the + // DIARIZATION capability — for a bundle it is a property of the file. + if (m->diar != nullptr) { + if (const transcribe_status st = transcribe::sortformer::fuse_embedded_diar_bn( + *m->diar, m->plan.scheduler_list.back(), &m->diar_bn_ctx, &m->diar_bn_buffer); + st != TRANSCRIBE_OK) { + return st; + } + { + transcribe::MelConfig cfg{}; + cfg.sample_rate = m->diar->hp.fe_sample_rate; + cfg.num_mels = m->diar->hp.fe_num_mels; + cfg.n_fft = m->diar->hp.fe_n_fft; + cfg.win_length = m->diar->hp.fe_win_length; + cfg.hop_length = m->diar->hp.fe_hop_length; + cfg.pre_emphasis = m->diar->hp.fe_pre_emphasis; + cfg.normalize = m->diar->hp.fe_normalize; + cfg.pad_mode = "constant"; + cfg.nemo_seq_len_ceil = true; // sortformer port framing + m->diar_mel.emplace(cfg); + } + transcribe::set_feature(m.get(), TRANSCRIBE_FEATURE_DIARIZATION, true); + } + m->t_load_us = ggml_time_us() - t_load_start; // The caller now owns the model and must call transcribe_model_free. @@ -717,9 +776,19 @@ static transcribe_status decode_and_populate(ParakeetSession * pc, } } pc->t_decode_us = ggml_time_us() - t_dec_start; + return build_result_from_raw_tokens(pc, pm, params); +} - // ----- Build the public result hierarchy ----- - // +} // namespace + +// ----- Build the public result hierarchy from pc->raw_tokens ----- +// +// External linkage: the multitalker streaming pass accumulates raw +// tokens per speaker (with absolute-frame step_at_emit) and reuses this +// to materialize tokens/words/segments/text. +transcribe_status build_result_from_raw_tokens(ParakeetSession * pc, + ParakeetModel * pm, + const transcribe_run_params * params) { // step_at_emit is an encoder frame index; one frame is // subsampling_factor * hop_length / sample_rate seconds (80 ms on // v2/v3). Shared with the streaming builder. @@ -896,11 +965,14 @@ static transcribe_status decode_and_populate(ParakeetSession * pc, return TRANSCRIBE_OK; } -static transcribe_status run_one_shot_inner(ParakeetSession * pc, - ParakeetModel * pm, - const float * pcm, - int n_samples, - const transcribe_run_params * params) { +// External linkage: also driven per speaker by run_multitalker +// (multitalker.cpp) with a supervision pass. +transcribe_status run_one_shot_inner(ParakeetSession * pc, + ParakeetModel * pm, + const float * pcm, + int n_samples, + const transcribe_run_params * params, + const MultitalkerPass * mt) { // Pre-run abort check (the single observation point on this path). if (pc->poll_abort()) { return TRANSCRIBE_ERR_ABORTED; @@ -925,6 +997,29 @@ static transcribe_status run_one_shot_inner(ParakeetSession * pc, } pc->t_mel_us = ggml_time_us() - t_mel_start; + // Multitalker masked mode (NeMo mask_features): binarize the target + // speaker's per-diar-frame activity, expand 8x to feature frames from + // t=0, multiply the mel; feature frames past the supervision (or with + // an inactive target) zero out, and original exact-zero features take + // the log-floor mask_value. Mirrors the reference quirk exactly. + if (mt != nullptr && !mt->use_kernel) { + const int sub = pm->hparams.enc_subsampling_factor > 0 ? pm->hparams.enc_subsampling_factor : 8; + const size_t T_diar = mt->spk.size(); + for (int m = 0; m < mel_n_mels; ++m) { + float * row = pc->mel_buf.data() + static_cast(m) * mel_n_frames; + for (int f = 0; f < mel_n_frames; ++f) { + const size_t d = static_cast(f / sub); + const bool active = d < T_diar && mt->spk[d] > mt->threshold; + const float v = row[f]; + if (v == 0.0f) { + row[f] = mt->mask_value; + } else if (!active) { + row[f] = 0.0f; + } + } + } + } + // ----- Reset per-call compute state ----- // Free the previous run's compute_ctx; encoder_out is invalidated. if (pc->compute_ctx != nullptr) { @@ -959,11 +1054,34 @@ static transcribe_status run_one_shot_inner(ParakeetSession * pc, resolved_kv = GGML_TYPE_F16; } - EncoderBuild eb = - build_encoder_graph(pc->compute_ctx, pm->weights, pm->hparams, mel_n_frames, resolved_kv, pm->backend.c_str()); + const bool spk_supervision = (mt != nullptr && mt->use_kernel); + // Kernel-mode chunk gating: sanitize the orchestrator's keep list + // against the actual pre_encode output length (the orchestrator + // builds its chunk grid from the diar frame count, which can trail + // the encoder length by a frame at the clip tail). + std::vector mt_keep; + if (spk_supervision && !mt->keep_enc_frames.empty()) { + const bool causal = (pm->hparams.enc_att_context_style == ParakeetHParams::AttContextStyle::ChunkedLimited); + const int T_pred = + pre_encode_time_out(pre_encode_time_out(pre_encode_time_out(mel_n_frames, causal), causal), causal); + mt_keep.reserve(mt->keep_enc_frames.size()); + for (const int32_t i : mt->keep_enc_frames) { + if (i >= 0 && i < T_pred) { + mt_keep.push_back(i); + } + } + } + const int mt_keep_frames = static_cast(mt_keep.size()); + EncoderBuild eb = build_encoder_graph(pc->compute_ctx, pm->weights, pm->hparams, mel_n_frames, resolved_kv, + pm->backend.c_str(), /*buf_mask=*/nullptr, /*n_batch=*/1, + /*batch_var_len=*/false, spk_supervision, mt_keep_frames); if (eb.mel_in == nullptr || eb.out == nullptr || eb.graph == nullptr) { return TRANSCRIBE_ERR_GGUF; // build_encoder_graph logged } + if (spk_supervision && (eb.spk_mask_in == nullptr || eb.bg_mask_in == nullptr)) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet run: speaker supervision requested but model has no spk kernel"); + return TRANSCRIBE_ERR_INVALID_ARG; + } // ----- Allocate compute tensors via the multi-backend scheduler. // Created lazily, persists across calls. ----- @@ -988,6 +1106,47 @@ static transcribe_status run_one_shot_inner(ParakeetSession * pc, transcribe::debug::dump_tensor("enc.mel.in", eb.mel_in, "encoder.mel"); + // Multitalker kernel-mode supervision upload. Length alignment mirrors + // NeMo's solve_length_mismatch: pad LEFT with the default value (1.0 + // for the target-speaker mask, 0.0 for background) when the + // supervision is short, keep the LAST T frames when it is long. + if (eb.spk_mask_in != nullptr && mt != nullptr) { + // Ungated encoder length (pre-gather); the injection masks are + // sized to the gathered length when chunk gating is active. + const int T_full = static_cast(eb.dumps.pre_encode_out->ne[1]); + const int T_inj = static_cast(eb.spk_mask_in->ne[1]); + auto align_full = [&](const std::vector & src, float default_value) { + std::vector buf(static_cast(T_full), default_value); + const int n = static_cast(src.size()); + if (n >= T_full) { + std::copy(src.end() - T_full, src.end(), buf.begin()); + } else { + std::copy(src.begin(), src.end(), buf.begin() + (T_full - n)); + } + return buf; + }; + auto fill_mask = [&](ggml_tensor * dst, const std::vector & src, float default_value) { + std::vector full = align_full(src, default_value); + if (mt_keep_frames > 0) { + std::vector gathered(static_cast(T_inj), default_value); + for (int g = 0; g < T_inj && g < mt_keep_frames; ++g) { + const int32_t i = mt_keep[static_cast(g)]; + if (i >= 0 && i < T_full) { + gathered[static_cast(g)] = full[static_cast(i)]; + } + } + ggml_backend_tensor_set(dst, gathered.data(), 0, gathered.size() * sizeof(float)); + } else { + ggml_backend_tensor_set(dst, full.data(), 0, full.size() * sizeof(float)); + } + }; + fill_mask(eb.spk_mask_in, mt->spk, 1.0f); + fill_mask(eb.bg_mask_in, mt->bg, 0.0f); + if (eb.mt_keep_in != nullptr) { + ggml_backend_tensor_set(eb.mt_keep_in, mt_keep.data(), 0, mt_keep.size() * sizeof(int32_t)); + } + } + // Prompt one-hot input (multilingual variants only): fill // prompt_one_hot_in with the resolved language's one-hot replicated // across T_enc frames (see resolve_prompt_id). @@ -1185,6 +1344,8 @@ static transcribe_status run_one_shot_inner(ParakeetSession * pc, return decode_and_populate(pc, pm, params, pc->enc_host.data(), T_enc, d_enc, /*utt_index=*/-1, enc_dump_name); } +namespace { + // One-shot entry point. Validates session/pm, clears the previous result // snapshot, then forwards to the shared inference helper (also reused by // the streaming hooks). @@ -1206,6 +1367,15 @@ transcribe_status run(transcribe_session * session, } pc->clear_result(); + + // Multitalker bundle: diarize=ON routes through the embedded-sortformer + // orchestrator (per-speaker passes + merged speaker-tagged result). + // Without an embedded diarizer the dispatcher has already WARNed and we + // proceed single-speaker, preserving the shipped behavior. + if (params != nullptr && params->diarize == TRANSCRIBE_DIARIZE_MODE_ON && pm->diar != nullptr) { + return run_multitalker(pc, pm, pcm, n_samples, params); + } + return run_one_shot_inner(pc, pm, pcm, n_samples, params); } @@ -1486,6 +1656,15 @@ transcribe_status run_batch(transcribe_session * session, } transcribe::debug::init(); + // Multitalker is a transcribe_run concern for now: batch keeps the + // shipped single-speaker semantics even on a bundle model. Warn so a + // diarize=ON batch caller isn't silently downgraded. + if (params != nullptr && params->diarize == TRANSCRIBE_DIARIZE_MODE_ON && pm->diar != nullptr) { + log_msg(TRANSCRIBE_LOG_LEVEL_WARN, + "parakeet: diarize=ON is not yet supported in transcribe_run_batch; " + "utterances are decoded single-speaker (use transcribe_run for multitalker)"); + } + // Compute each utterance's mel in parallel (pure host, no // cross-utterance state). A malformed utterance falls the whole call // back to the per-utterance path (keeps the batch tensor rectangular). @@ -1763,8 +1942,6 @@ void compute_chunked_limited_with_rc_mask(float * out_buf, } } -namespace { - // Build, run, and post-process a single streaming encoder chunk. // // mel_chunk_data row-major [n_mels, n_mel_chunk_frames] f32. @@ -1774,14 +1951,23 @@ namespace { // mel_frames_advance mel frames this chunk consumes from the input // (NOT n_mel_chunk_frames when a cache prepend // is in play). +// mt_spk / mt_bg / mt_mask_len optional multitalker per-chunk +// supervision (kernel injection masks over the +// chunk's new encoder frames); both null for +// the single-speaker paths. // On success, appends TdtTokens to pc->raw_tokens and advances the cache -// + decoder state. +// + decoder state. External linkage: the multitalker streaming pass in +// multitalker.cpp drives this per active speaker with swapped-in +// per-speaker cache/decoder state. transcribe_status emit_streaming_chunk(ParakeetSession * pc, ParakeetModel * pm, const float * mel_chunk_data, int n_mel_chunk_frames, int drop_extra_pre_encoded, - int mel_frames_advance) { + int mel_frames_advance, + const float * mt_spk, + const float * mt_bg, + int mt_mask_len) { const auto & hp = pm->hparams; const int n_layers = static_cast(pm->weights.blocks.size()); @@ -1853,8 +2039,10 @@ transcribe_status emit_streaming_chunk(ParakeetSession * pc, cache_io.pos_proj = pc->stream_caches.pos_proj; cache_io.pos_proj_len = pc->stream_caches.pos_proj_len; + const bool mt_supervised = (mt_spk != nullptr && mt_bg != nullptr); EncoderBuild eb = build_encoder_graph_streaming(pc->compute_ctx, pm->weights, hp, n_mel_chunk_frames, - drop_extra_pre_encoded, cache_io, resolved_kv, pm->backend.c_str()); + drop_extra_pre_encoded, cache_io, resolved_kv, pm->backend.c_str(), + /*spk_supervision=*/mt_supervised); if (eb.out == nullptr || eb.graph == nullptr) { return TRANSCRIBE_ERR_GGUF; } @@ -1884,6 +2072,23 @@ transcribe_status emit_streaming_chunk(ParakeetSession * pc, const int T_cache = hp.enc_att_context_left; const int T_q_new = T_virtual - T_cache; + // Multitalker per-chunk supervision upload (kernel mode). Values + // cover this chunk's new encoder frames; pad with the mask defaults + // (spk 1.0 / bg 0.0) when the chunk's frame count trails the + // supplied length at the clip tail. + if (mt_supervised && eb.spk_mask_in != nullptr) { + const int T_inj = static_cast(eb.spk_mask_in->ne[1]); + auto fill = [&](ggml_tensor * dst, const float * src, float default_value) { + std::vector buf(static_cast(T_inj), default_value); + for (int t = 0; t < T_inj && t < mt_mask_len; ++t) { + buf[static_cast(t)] = src[t]; + } + ggml_backend_tensor_set(dst, buf.data(), 0, buf.size() * sizeof(float)); + }; + fill(eb.spk_mask_in, mt_spk, 1.0f); + fill(eb.bg_mask_in, mt_bg, 0.0f); + } + // Build & upload pos_emb. The zero-offset row sits at T_virtual - 1. // Null when every block consumes the memoized rel-pos projection. if (eb.pos_emb_in != nullptr) { @@ -2031,6 +2236,8 @@ transcribe_status emit_streaming_chunk(ParakeetSession * pc, return TRANSCRIBE_OK; } +namespace { + // Rebuild the public result vectors (tokens, full_text) from the // committed raw_tokens; idempotent. Words/segments are NOT populated // during streaming (deferred to finalize). result_kind is TOKEN whenever @@ -2685,6 +2892,15 @@ transcribe_status stream_begin(transcribe_session * session, // or a back-to-back stream_begin carries tokens. pc->stream_run_params = *run_params; + // Multitalker is a transcribe_run concern for now: the incremental + // streaming API keeps the shipped single-speaker semantics even on a + // bundle model. Warn so a diarize=ON stream isn't silently downgraded. + if (run_params->diarize == TRANSCRIBE_DIARIZE_MODE_ON && pm->diar != nullptr) { + log_msg(TRANSCRIBE_LOG_LEVEL_WARN, + "parakeet: diarize=ON is not yet supported on the streaming API; " + "the stream is decoded single-speaker (use transcribe_run for multitalker)"); + } + // Allocate streaming caches on first stream_begin (idempotent); zero // contents and reset cursors on every begin. if (const transcribe_status st = init_streaming_caches(pc, pm); st != TRANSCRIBE_OK) { diff --git a/src/arch/parakeet/multitalker.cpp b/src/arch/parakeet/multitalker.cpp new file mode 100644 index 00000000..09d7e3b8 --- /dev/null +++ b/src/arch/parakeet/multitalker.cpp @@ -0,0 +1,700 @@ +// arch/parakeet/multitalker.cpp - multitalker orchestration for bundle +// GGUFs (multitalker Parakeet + embedded streaming Sortformer). +// +// Mirrors NeMo's SpeakerTaggedASR parallel streaming pipeline +// (speech_to_text_multitalker_streaming_infer.py at the pinned rev), +// replayed offline inside one transcribe_run: +// +// 1. The embedded Sortformer runs its streaming AOSC/FIFO forward over +// the whole clip at the multitalker operating point (chunk = ASR +// chunk cadence att_context_right+1 = 14 diar frames, lc/rc = 0, +// spkcache 188, fifo 188, update period 188 — the example's config +// on the v2.1 checkpoint defaults). The accumulated per-chunk preds +// equal the causal per-chunk supervision the reference loop feeds +// each ASR instance. +// 2. Active speakers = columns whose activity exceeds 0.5 anywhere +// (the offline analog of the reference's per-chunk cache gating, +// which tops-k by max over the last 2 chunks). +// 3. One ASR decode pass per active speaker over the same audio, with +// that speaker's supervision: masked mode (NeMo masked_asr=True +// default) masks mel features; kernel mode (masked_asr=False, +// TRANSCRIBE_MULTITALKER_MODE=kernel) drives the layer-0 +// speaker-kernel injection. +// 4. Per-speaker results merge into one speaker-tagged result: segments +// ordered by start time with speaker_id set, words/tokens re-indexed +// under them, full_text joined in segment order (SegLST rendering), +// and the transcript-independent speaker_segment rows emitted from +// the diarizer preds. +// +// Deliberate deviation from the reference (documented in the family doc): +// the ASR pass is the validated OFFLINE chunked-attention graph over the +// whole clip, not a replay of the cache-aware streaming loop, so encoder +// state at chunk boundaries can differ at numerical (not transcript) +// scale. The streaming API multitalker path will be chunk-faithful by +// construction. + +#include "../sortformer/sortformer.h" +#include "ggml.h" +#include "parakeet.h" +#include "transcribe-debug.h" +#include "transcribe-log.h" + +#include +#include +#include +#include +#include +#include + +namespace transcribe::parakeet { + +namespace sf = transcribe::sortformer; + +namespace { + +struct SpeakerDecode { + int speaker = 0; // 0-based diar column + std::vector tokens; + std::vector words; + std::vector segments; + std::string full_text; + std::string raw_text; +}; + +} // namespace + +// ---- Streaming multitalker passes (bounded memory) ---- +// +// Chunk-faithful mirror of the reference loop (conformer_stream_step + +// instance manager): one global 14-enc-frame chunk grid; per active +// speaker a private streaming cache + decoder-state instance that only +// advances on chunks where cache-gating selects that speaker. Memory is +// O(1) in clip length (~5 MB per speaker + one chunk graph), unlike the +// offline replay whose attention buffers grow quadratically. The +// per-speaker instances are swapped into the session around each +// emit_streaming_chunk call so the single-speaker streaming driver runs +// unchanged. +static transcribe_status run_streaming_passes(ParakeetSession * pc, + ParakeetModel * pm, + const float * pcm, + int n_samples, + const transcribe_run_params * params, + const std::vector & probs, + int T_diar, + int n_spk, + const std::vector & active, + bool use_kernel, + std::vector & per_spk, + int64_t & mel_us, + int64_t & enc_us, + int64_t & dec_us) { + const auto & hp = pm->hparams; + + // ASR mel over the whole clip (row-major [n_mels, n_frames]; each + // mel row contiguous over frames). ~170 MB at 55 min — bounded. + const int64_t t_mel_start = ggml_time_us(); + std::vector mel; + int n_mels = 0, n_frames = 0; + if (const transcribe_status st = pm->mel->compute(pcm, static_cast(n_samples), mel, n_mels, n_frames); + st != TRANSCRIBE_OK) { + return st; + } + mel_us += ggml_time_us() - t_mel_start; + + const int chunk_enc = hp.enc_att_context_right + 1; // 14 + const int sub = hp.enc_subsampling_factor; // 8 + const int mel_first = hp.enc_stream_sampling_frames_first + sub * hp.enc_att_context_right; // 105 + const int mel_chunk = chunk_enc * sub; // 112 + const int mel_cache = hp.enc_stream_pre_encode_cache_size; // 9 + const int drop_sub = hp.enc_stream_drop_extra_pre_encoded; // 2 + + // Reference cache_gating: speaker s is active in chunk k iff the max + // of its accumulated preds over the last 2 chunks (incl. current) + // exceeds 0.5. Applies to BOTH supervision modes (the reference + // instance selection is mode-independent). + auto chunk_active = [&](int s, int k) { + const int begin = std::max(0, (k - 1) * chunk_enc); + const int end = std::min((k + 1) * chunk_enc, T_diar); + float mx = 0.0f; + for (int t = begin; t < end; ++t) { + mx = std::max(mx, probs[static_cast(t) * n_spk + s]); + } + return mx > 0.5f; + }; + + struct Inst { + ParakeetStreamingCaches caches; + ParakeetStreamingDecoderState dec; + std::vector raw; + std::vector fed; // absolute enc frame per gathered frame + std::vector spk; // [T_diar] supervision values + std::vector bg; + bool started = false; + }; + + std::vector inst(active.size()); + + auto swap_in = [&](Inst & I) { + std::swap(pc->stream_caches, I.caches); + std::swap(pc->stream_dec_state, I.dec); + std::swap(pc->raw_tokens, I.raw); + }; + auto free_inst = [&](Inst & I) { + if (I.caches.buffer != nullptr) { + transcribe::safe_buffer_free(I.caches.buffer); + I.caches.buffer = nullptr; + } + if (I.caches.ctx != nullptr) { + ggml_free(I.caches.ctx); + I.caches.ctx = nullptr; + } + if (I.caches.pos_proj_buf != nullptr) { + transcribe::safe_buffer_free(I.caches.pos_proj_buf); + I.caches.pos_proj_buf = nullptr; + } + if (I.caches.pos_proj_ctx != nullptr) { + ggml_free(I.caches.pos_proj_ctx); + I.caches.pos_proj_ctx = nullptr; + } + }; + + // Per-speaker supervision value arrays (same construction as the + // offline pass: raw sigmoids for the target, binarized union of the + // other active speakers for the background). + for (size_t i = 0; i < active.size(); ++i) { + const int s = active[i]; + inst[i].spk.resize(static_cast(T_diar)); + inst[i].bg.assign(static_cast(T_diar), 0.0f); + for (int t = 0; t < T_diar; ++t) { + inst[i].spk[static_cast(t)] = probs[static_cast(t) * n_spk + s]; + } + for (const int o : active) { + if (o == s) { + continue; + } + for (int t = 0; t < T_diar; ++t) { + if (probs[static_cast(t) * n_spk + o] > 0.5f) { + inst[i].bg[static_cast(t)] = 1.0f; + } + } + } + } + + const int64_t t_enc_start = ggml_time_us(); + std::vector chunk_buf; + transcribe_status st = TRANSCRIBE_OK; + + for (int k = 0;; ++k) { + const int c0_new = (k == 0) ? 0 : mel_first + (k - 1) * mel_chunk; // first NEW mel col + if (c0_new >= n_frames) { + break; + } + const int c0 = (k == 0) ? 0 : c0_new - mel_cache; // history prepend + const int c1 = std::min(c0_new + ((k == 0) ? mel_first : mel_chunk), n_frames); + const int n_feed = c1 - c0; + const int advance = c1 - c0_new; + if (n_feed <= 0 || advance <= 0) { + break; + } + + for (size_t i = 0; i < active.size() && st == TRANSCRIBE_OK; ++i) { + const int s = active[i]; + if (!chunk_active(s, k)) { + continue; + } + Inst & I = inst[i]; + + // Slice this chunk's mel columns (shared grid; the 9-column + // history prepend reproduces the continuous pre_encode). + chunk_buf.resize(static_cast(n_mels) * static_cast(n_feed)); + for (int m = 0; m < n_mels; ++m) { + std::memcpy(chunk_buf.data() + static_cast(m) * n_feed, + mel.data() + static_cast(m) * n_frames + c0, + static_cast(n_feed) * sizeof(float)); + } + + // Masked mode: reference mask_features semantics — + // masked = mel * mask (masked columns -> 0.0, NOT the + // log-silence value); cells whose ORIGINAL mel was exactly + // 0 -> mask_value; the upsampled mask is left-aligned to + // the chunk's NEW columns (col j -> chunk frame j/8) and + // the short side is left-PADDED WITH 0, so the history + // prepend columns are always masked. + if (!use_kernel) { + const int hist = c0_new - c0; // 9 (0 for the first chunk) + for (int c = 0; c < n_feed; ++c) { + bool keep = false; + if (c >= hist) { + const int f = k * chunk_enc + (c - hist) / sub; + keep = f < T_diar && I.spk[static_cast(f)] > 0.5f; + } + for (int m = 0; m < n_mels; ++m) { + float & v = chunk_buf[static_cast(m) * n_feed + c]; + if (v == 0.0f) { + v = -16.6355f; + } else if (!keep) { + v = 0.0f; + } + } + } + } + + // Kernel mode: per-chunk supervision over the NEW enc frames. + const int t_enc0 = k * chunk_enc; + const int mask_len = std::max(0, std::min(chunk_enc, T_diar - t_enc0)); + const float * spk_ptr = nullptr; + const float * bg_ptr = nullptr; + std::vector bg_compat; + if (use_kernel && mask_len > 0) { + spk_ptr = I.spk.data() + t_enc0; + bg_ptr = I.bg.data() + t_enc0; + // TRANSCRIBE_MULTITALKER_REF_BG_COMPAT=1 reproduces the + // reference harness's background-slot indexing bug + // (multispk_transcribe_utils.get_active_speakers_info + // takes range(len(active)) minus the SLOT id, so with + // active slots like [0,2] speaker 0's background tracks + // the silent slot 1 instead of slot 2). Verified against + // the reference's own bg_targets dumps. Parity tooling + // only — the default is the intended union of the other + // ACTIVE slots. + static const bool ref_bg_compat = (std::getenv("TRANSCRIBE_MULTITALKER_REF_BG_COMPAT") != nullptr); + if (ref_bg_compat) { + std::vector act_k; + for (const int o : active) { + if (chunk_active(o, k)) { + act_k.push_back(o); + } + } + bg_compat.assign(static_cast(mask_len), 0.0f); + for (int slot = 0; slot < static_cast(act_k.size()); ++slot) { + if (slot == s) { + continue; // reference compares POSITION to SLOT id + } + for (int t = 0; t < mask_len; ++t) { + if (probs[static_cast(t_enc0 + t) * n_spk + slot] > 0.5f) { + bg_compat[static_cast(t)] = 1.0f; + } + } + } + bg_ptr = bg_compat.data(); + } + } + + swap_in(I); + if (!I.started) { + // (fields now live on pc after the swap) + if ((st = init_streaming_caches(pc, pm)) == TRANSCRIBE_OK) { + reset_streaming_decoder_state(pc, pm); + pc->stream_caches.att_context_left = hp.enc_att_context_left; + pc->stream_caches.att_context_right = hp.enc_att_context_right; + pc->stream_caches.channel_len = 0; + pc->stream_caches.chunk_step = 0; + } + } + if (st == TRANSCRIBE_OK) { + const int64_t off_before = pc->stream_dec_state.frame_offset; + st = emit_streaming_chunk(pc, pm, chunk_buf.data(), n_feed, (k == 0) ? 0 : drop_sub, advance, spk_ptr, + bg_ptr, use_kernel ? mask_len : 0); + if (st == TRANSCRIBE_OK) { + const int n_new = static_cast(pc->stream_dec_state.frame_offset - off_before); + for (int f = 0; f < n_new; ++f) { + I.fed.push_back(t_enc0 + f); + } + } + } + swap_in(I); // swap back + I.started = true; + } + if (st != TRANSCRIBE_OK) { + break; + } + } + enc_us += ggml_time_us() - t_enc_start; + + // Build per-speaker results from the accumulated raw tokens, with + // step_at_emit remapped from the gathered timeline to absolute + // encoder frames through the fed-chunk list. + const int64_t t_dec_start = ggml_time_us(); + for (size_t i = 0; i < active.size(); ++i) { + Inst & I = inst[i]; + if (st == TRANSCRIBE_OK && !I.raw.empty()) { + for (auto & rt : I.raw) { + const size_t f = std::min(static_cast(std::max(rt.step_at_emit, 0)), I.fed.size() - 1); + rt.step_at_emit = I.fed[f]; + } + transcribe_run_params sp = *params; + sp.timestamps = TRANSCRIBE_TIMESTAMPS_AUTO; + pc->clear_result(); + pc->raw_tokens = std::move(I.raw); + const transcribe_status bst = build_result_from_raw_tokens(pc, pm, &sp); + if (bst != TRANSCRIBE_OK) { + st = bst; + } else { + SpeakerDecode r; + r.speaker = active[i]; + r.tokens = std::move(pc->tokens); + r.words = std::move(pc->words); + r.segments = std::move(pc->segments); + r.full_text = std::move(pc->full_text); + r.raw_text = std::move(pc->raw_text); + per_spk.push_back(std::move(r)); + } + } + free_inst(I); + } + dec_us += ggml_time_us() - t_dec_start; + return st; +} + +transcribe_status run_multitalker(ParakeetSession * pc, + ParakeetModel * pm, + const float * pcm, + int n_samples, + const transcribe_run_params * params) { + if (pm->diar == nullptr || !pm->diar_mel.has_value() || !pm->mel.has_value()) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + if (pc->poll_abort()) { + return TRANSCRIBE_ERR_ABORTED; + } + transcribe::debug::init(); + + // Supervision mode. Default is the checkpoint's trained speaker-kernel + // path: on ami-ihm-test it scores 19.35% cpWER vs masked's 24.61% + // (and vs the NeMo reference's own 21.39% kernel / 24.00% masked). + // TRANSCRIBE_MULTITALKER_MODE=masked selects mel feature masking + // (the NeMo example default) until the family run-ext grows a knob. + bool use_kernel = true; + if (const char * mode = std::getenv("TRANSCRIBE_MULTITALKER_MODE")) { + use_kernel = (std::strcmp(mode, "masked") != 0); + } + + const sf::SortformerEmbedded & diar = *pm->diar; + + // ---- 1. Embedded diarizer streaming forward ---- + std::vector diar_mel; + int dm_mels = 0, dm_frames = 0; + if (const transcribe_status st = + pm->diar_mel->compute(pcm, static_cast(n_samples), diar_mel, dm_mels, dm_frames); + st != TRANSCRIBE_OK) { + return st; + } + + if (pc->sched == nullptr) { + pc->sched = ggml_backend_sched_new(pm->plan.scheduler_list.data(), nullptr, + static_cast(pm->plan.scheduler_list.size()), + /*graph_size=*/8192, /*parallel=*/false, /*op_offload=*/true); + if (pc->sched == nullptr) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet multitalker: ggml_backend_sched_new failed"); + return TRANSCRIBE_ERR_GGUF; + } + } + + sf::SortformerStreamParams P = sf::resolve_stream_params(diar.hp, TRANSCRIBE_SORTFORMER_PRESET_DEFAULT); + P.chunk_len = pm->hparams.enc_att_context_right + 1; // ASR chunk cadence (14 at [70,13]) + P.chunk_left_context = 0; + P.chunk_right_context = 0; + P.spkcache_len = 188; + P.fifo_len = 188; + P.spkcache_update_period = 188; + // Reference-exact windowing: the NeMo loop feeds the diarizer the + // ASR streaming buffer's chunks (chunk_size_first = sampling_frames + + // sub * R mel frames first, then pre_encode_cache-frame overlapped + // windows with drop_extra_pre_encoded outputs discarded). + P.feat_first_chunk = pm->hparams.enc_stream_sampling_frames_first + + pm->hparams.enc_subsampling_factor * pm->hparams.enc_att_context_right; + P.feat_cache = pm->hparams.enc_stream_pre_encode_cache_size; + P.drop_pre_encode = pm->hparams.enc_stream_drop_extra_pre_encoded; + + const int64_t t_diar_start = ggml_time_us(); + transcribe::debug::push_name_prefix("mt.diar"); + const transcribe_status dst = sf::run_diar_streaming_core( + pc->diar_scratch, diar.hp, diar.conformer_hp, diar.conformer, diar.weights, pm->backend.c_str(), pc->sched, + pc->n_threads, P, diar_mel.data(), dm_mels, dm_frames, pc); + transcribe::debug::pop_name_prefix(); + if (dst != TRANSCRIBE_OK) { + return dst; + } + const int64_t t_diar_us = ggml_time_us() - t_diar_start; + + const int sub = diar.hp.enc_subsampling_factor > 0 ? diar.hp.enc_subsampling_factor : 8; + const int n_spk = diar.hp.max_speakers; + const int T_diar = std::min(pc->diar_scratch.stream.total_n, (dm_frames + sub - 1) / sub); + const std::vector & probs = pc->diar_scratch.stream.total_preds; // row-major [total_n, n_spk] + + if (transcribe::debug::enabled()) { + const long long shape[2] = { T_diar, n_spk }; + transcribe::debug::dump_host_f32("mt.diar.probs", probs.data(), static_cast(T_diar) * n_spk, shape, + 2, "diarize"); + } + + // ---- 2. Active speakers ---- + std::vector active; + for (int s = 0; s < n_spk; ++s) { + float mx = 0.0f; + for (int t = 0; t < T_diar; ++t) { + mx = std::max(mx, probs[static_cast(t) * n_spk + s]); + } + if (mx > 0.5f) { + active.push_back(s); + } + } + + const double diar_ms_per_frame = + 1000.0 * static_cast(diar.hp.frame_hop) / static_cast(diar.hp.fe_sample_rate); + + if (active.empty()) { + // No speech activity detected anywhere (reference behavior: every + // chunk skipped). Empty-but-valid result. + pc->t_mel_us = 0; + pc->t_encode_us = t_diar_us; + pc->t_decode_us = 0; + pc->result_kind = TRANSCRIBE_TIMESTAMPS_NONE; + pc->has_result = true; + return TRANSCRIBE_OK; + } + + // ---- 3. Per-speaker decode passes ---- + // + // Default is the bounded-memory streaming pass (chunk-faithful to the + // reference loop). TRANSCRIBE_MULTITALKER_OFFLINE=1 selects the + // offline whole-clip replay (O(T^2) attention buffers; parity/dump + // tooling only). + std::vector per_spk; + per_spk.reserve(active.size()); + int64_t mel_us = 0, enc_us = 0, dec_us = 0; + + const bool offline_replay = (std::getenv("TRANSCRIBE_MULTITALKER_OFFLINE") != nullptr); + if (!offline_replay) { + if (const transcribe_status st = run_streaming_passes(pc, pm, pcm, n_samples, params, probs, T_diar, n_spk, + active, use_kernel, per_spk, mel_us, enc_us, dec_us); + st != TRANSCRIBE_OK) { + return st; + } + } else { + for (const int s : active) { + MultitalkerPass pass; + pass.use_kernel = use_kernel; + pass.spk.resize(static_cast(T_diar)); + pass.bg.assign(static_cast(T_diar), 0.0f); + for (int t = 0; t < T_diar; ++t) { + pass.spk[static_cast(t)] = probs[static_cast(t) * n_spk + s]; + } + for (const int o : active) { + if (o == s) { + continue; + } + for (int t = 0; t < T_diar; ++t) { + if (probs[static_cast(t) * n_spk + o] > 0.5f) { + pass.bg[static_cast(t)] = 1.0f; + } + } + } + + // Kernel-mode chunk gating (reference cache_gating=true, + // cache_gating_buffer_size=2): a speaker's conformer/decoder only + // ever sees chunks where _find_active_speakers selected it — max + // of the accumulated preds over the last 2 chunks (incl. current) + // > 0.5. Without this, a mostly-inactive speaker's pass + // transcribes everyone else's speech (the trained kernels only + // ever saw gated chunks). Masked mode hard-silences non-target + // features instead, so it stays ungated here. + if (use_kernel) { + const int chunk = P.chunk_len; // 14 = nframes_per_chunk + const int n_chunks = (T_diar + chunk - 1) / chunk; + // +1 covers the encoder length trailing the diar frame count + // at the tail; run_one_shot_inner clamps to the true length. + pass.keep_enc_frames.reserve(static_cast(T_diar) + 1); + for (int k = 0; k < n_chunks; ++k) { + const int end = std::min((k + 1) * chunk, T_diar); + const int begin = std::max(0, (k - 1) * chunk); // 2-chunk window + float mx = 0.0f; + for (int t = begin; t < end; ++t) { + mx = std::max(mx, probs[static_cast(t) * n_spk + s]); + } + if (mx > 0.5f) { + const int enc_end = (k == n_chunks - 1) ? end + 1 : end; + for (int t = k * chunk; t < enc_end; ++t) { + pass.keep_enc_frames.push_back(t); + } + } + } + if (pass.keep_enc_frames.empty()) { + continue; // globally active but never chunk-selected + } + } + + // Decode at full timestamp resolution; the merged result is elided + // to the caller's request at the end. + transcribe_run_params sp = *params; + sp.timestamps = TRANSCRIBE_TIMESTAMPS_AUTO; + + pc->clear_result(); + char prefix[32]; + std::snprintf(prefix, sizeof(prefix), "mt.spk%d", s); + transcribe::debug::push_name_prefix(prefix); + const transcribe_status st = run_one_shot_inner(pc, pm, pcm, n_samples, &sp, &pass); + transcribe::debug::pop_name_prefix(); + if (st != TRANSCRIBE_OK) { + return st; + } + mel_us += pc->t_mel_us; + enc_us += pc->t_encode_us; + dec_us += pc->t_decode_us; + + SpeakerDecode r; + r.speaker = s; + r.tokens = std::move(pc->tokens); + r.words = std::move(pc->words); + r.segments = std::move(pc->segments); + r.full_text = std::move(pc->full_text); + r.raw_text = std::move(pc->raw_text); + + // Gated pass timestamps are in the gathered timeline; map each + // frame back through the keep list to absolute clip time. + if (!pass.keep_enc_frames.empty()) { + const auto & keep = pass.keep_enc_frames; + const double f_ms = diar_ms_per_frame; // 80 ms encoder cadence + auto remap_t0 = [&](int64_t t_ms) { + size_t f = static_cast(std::max(0, t_ms) / static_cast(f_ms)); + f = std::min(f, keep.size() - 1); + return static_cast(keep[f]) * static_cast(f_ms); + }; + auto remap_t1 = [&](int64_t t_ms) { + size_t f = static_cast(std::max(0, t_ms - 1) / static_cast(f_ms)); + f = std::min(f, keep.size() - 1); + return static_cast(keep[f] + 1) * static_cast(f_ms); + }; + for (auto & tk : r.tokens) { + tk.t0_ms = remap_t0(tk.t0_ms); + tk.t1_ms = remap_t1(tk.t1_ms); + } + for (auto & wd : r.words) { + wd.t0_ms = remap_t0(wd.t0_ms); + wd.t1_ms = remap_t1(wd.t1_ms); + } + for (auto & sg : r.segments) { + sg.t0_ms = remap_t0(sg.t0_ms); + sg.t1_ms = remap_t1(sg.t1_ms); + } + } + per_spk.push_back(std::move(r)); + } + } + + // ---- 4. Merge into one speaker-tagged result ---- + pc->clear_result(); + + struct SegRef { + const SpeakerDecode * r; + int seg_idx; + }; + + std::vector order; + for (const SpeakerDecode & r : per_spk) { + for (size_t i = 0; i < r.segments.size(); ++i) { + order.push_back({ &r, static_cast(i) }); + } + } + std::stable_sort(order.begin(), order.end(), [](const SegRef & a, const SegRef & b) { + return a.r->segments[static_cast(a.seg_idx)].t0_ms < + b.r->segments[static_cast(b.seg_idx)].t0_ms; + }); + + std::string merged_text; + std::string merged_raw; + for (const SegRef & ref : order) { + const SpeakerDecode & r = *ref.r; + const transcribe_session::SegmentEntry & src = r.segments[static_cast(ref.seg_idx)]; + + const int new_seg_idx = static_cast(pc->segments.size()); + const int word_base = static_cast(pc->words.size()); + const int token_base = static_cast(pc->tokens.size()); + + for (int j = 0; j < src.n_tokens; ++j) { + transcribe_session::TokenEntry te = r.tokens[static_cast(src.first_token + j)]; + te.seg_index = new_seg_idx; + te.word_index = te.word_index - src.first_word + word_base; + pc->tokens.push_back(std::move(te)); + } + for (int j = 0; j < src.n_words; ++j) { + transcribe_session::WordEntry we = r.words[static_cast(src.first_word + j)]; + we.seg_index = new_seg_idx; + we.first_token = we.first_token - src.first_token + token_base; + pc->words.push_back(std::move(we)); + } + + transcribe_session::SegmentEntry seg = src; + seg.first_word = word_base; + seg.first_token = token_base; + seg.speaker_id = r.speaker + 1; // 1-based + pc->segments.push_back(std::move(seg)); + + if (!src.text.empty()) { + if (!merged_text.empty()) { + merged_text += ' '; + } + merged_text += src.text; + } + // One raw-text contribution per speaker, in first-segment order. + if (ref.seg_idx == 0 && !r.raw_text.empty()) { + if (!merged_raw.empty()) { + merged_raw += ' '; + } + merged_raw += r.raw_text; + } + } + pc->full_text = std::move(merged_text); + pc->raw_text = std::move(merged_raw); + + // Transcript-independent "who spoke when" rows from the diarizer preds + // (same emission as the standalone sortformer family). + sf::probs_to_speaker_segments(pc, probs, T_diar, n_spk, diar_ms_per_frame, /*threshold=*/0.5f); + + // ---- Elide to the caller's requested granularity (same convention + // as decode_and_populate). ---- + transcribe_timestamp_kind eff = params->timestamps; + if (eff == TRANSCRIBE_TIMESTAMPS_AUTO) { + eff = pm->caps.max_timestamp_kind; // TOKEN for parakeet + } + if (eff == TRANSCRIBE_TIMESTAMPS_NONE) { + pc->tokens.clear(); + pc->words.clear(); + for (auto & s : pc->segments) { + s.t0_ms = 0; + s.t1_ms = 0; + s.first_word = 0; + s.n_words = 0; + s.first_token = 0; + s.n_tokens = 0; + } + } else if (eff == TRANSCRIBE_TIMESTAMPS_SEGMENT) { + pc->tokens.clear(); + pc->words.clear(); + for (auto & s : pc->segments) { + s.first_word = 0; + s.n_words = 0; + s.first_token = 0; + s.n_tokens = 0; + } + } else if (eff == TRANSCRIBE_TIMESTAMPS_WORD) { + pc->tokens.clear(); + for (auto & w : pc->words) { + w.first_token = 0; + w.n_tokens = 0; + } + for (auto & s : pc->segments) { + s.first_token = 0; + s.n_tokens = 0; + } + } + + pc->t_mel_us = mel_us; + pc->t_encode_us = enc_us + t_diar_us; + pc->t_decode_us = dec_us; + pc->result_kind = eff; + pc->has_result = true; + return TRANSCRIBE_OK; +} + +} // namespace transcribe::parakeet diff --git a/src/arch/parakeet/parakeet.h b/src/arch/parakeet/parakeet.h index 0574e544..131c7528 100644 --- a/src/arch/parakeet/parakeet.h +++ b/src/arch/parakeet/parakeet.h @@ -7,6 +7,7 @@ #pragma once +#include "../sortformer/sortformer.h" // embedded diarizer (multitalker bundle) #include "decoder.h" #include "transcribe-backend.h" #include "transcribe-mel.h" @@ -16,6 +17,7 @@ #include "weights.h" #include +#include #include #include #include @@ -112,6 +114,21 @@ struct ParakeetModel final : public transcribe_model { // without a per-run backend readback. HostDecoderWeights host_decoder; + // Embedded Sortformer diarizer (multitalker bundle GGUFs). Present iff + // stt.parakeet.diarizer.embedded = true in the GGUF. The weight slots + // are borrowed views into this model's ctx_meta (resolved under the + // bundle's tensor prefix); the diarizer conformer's fused-BN params + // live in their own ctx + buffer, freed in the dtor. + std::unique_ptr diar; + ggml_context * diar_bn_ctx = nullptr; + ggml_backend_buffer_t diar_bn_buffer = nullptr; + + // Diarizer mel front-end. Same acoustic settings as the ASR frontend + // (the compose step enforces agreement) but with the sortformer port's + // ceil framing, so the embedded forward sees exactly what the + // standalone sortformer port validated against. + std::optional diar_mel; + ParakeetModel() = default; ~ParakeetModel() override; @@ -239,6 +256,10 @@ struct ParakeetSession final : public transcribe_session { std::vector enc_host; std::vector raw_tokens; + // Multitalker bundle scratch: streaming-diarizer AOSC/FIFO state for + // the embedded sortformer forward (diarize=ON runs only). + transcribe::sortformer::DiarStreamScratch diar_scratch; + // Per-call timings and the TDT decode result (tokens / words / // segments / full_text / result_kind / has_result) live on the base // transcribe_session; no per-family copies here. @@ -298,4 +319,72 @@ struct ParakeetSession final : public transcribe_session { ~ParakeetSession() override; }; +// ---- Multitalker (bundle) internals ----------------------------------- // + +// Per-speaker supervision for one multitalker decode pass, at the diar +// frame rate (80 ms). Mirrors NeMo's SpeakerTaggedASR hand-off: +// masked mode (use_kernel=false, the NeMo example default): the mel is +// masked per feature frame with the 8x-expanded binarized target +// activity (mask_features: masked = mel * mask; original-zero +// features -> mask_value). +// kernel mode (use_kernel=true): spk (raw sigmoids) and bg (binarized +// union of other actives) drive the layer-0 speaker-kernel injection +// as graph inputs (set_speaker_targets path). +struct MultitalkerPass { + bool use_kernel = false; + float threshold = 0.5f; + float mask_value = -16.6355f; + std::vector spk; // [T_diar] target-speaker activity + std::vector bg; // [T_diar] binarized union of other actives + // Kernel-mode chunk gating (reference cache_gating): encoder-frame + // indices of this speaker's ACTIVE 14-frame chunks, in order. The + // conformer + decoder see only these frames (gathered after the + // continuous pre_encode), mirroring the reference per-speaker cache + // that only advances on active chunks. Empty = no gating (masked + // mode / single-speaker). + std::vector keep_enc_frames; +}; + +// Shared one-utterance inference helper (mel -> encoder -> host decode -> +// result build on pc). Defined in model.cpp; also driven per speaker by +// the multitalker orchestrator with a supervision pass. +transcribe_status run_one_shot_inner(ParakeetSession * pc, + ParakeetModel * pm, + const float * pcm, + int n_samples, + const transcribe_run_params * params, + const MultitalkerPass * mt = nullptr); + +// Build tokens/words/segments/full_text on pc from pc->raw_tokens +// (defined in model.cpp; shared by the offline decode and the +// multitalker streaming pass). +transcribe_status build_result_from_raw_tokens(ParakeetSession * pc, + ParakeetModel * pm, + const transcribe_run_params * params); + +// Streaming primitives (defined in model.cpp) shared with the +// multitalker streaming pass, which drives one cache/decoder-state +// instance per speaker by swapping it into pc around each call. +transcribe_status init_streaming_caches(ParakeetSession * pc, ParakeetModel * pm); +void reset_streaming_decoder_state(ParakeetSession * pc, const ParakeetModel * pm); +transcribe_status emit_streaming_chunk(ParakeetSession * pc, + ParakeetModel * pm, + const float * mel_chunk_data, + int n_mel_chunk_frames, + int drop_extra_pre_encoded, + int mel_frames_advance, + const float * mt_spk = nullptr, + const float * mt_bg = nullptr, + int mt_mask_len = 0); + +// Multitalker offline run (diarize=ON on a bundle model): embedded +// sortformer streaming forward -> active-speaker selection -> one decode +// pass per active speaker -> merged speaker-tagged result. Defined in +// multitalker.cpp. +transcribe_status run_multitalker(ParakeetSession * pc, + ParakeetModel * pm, + const float * pcm, + int n_samples, + const transcribe_run_params * params); + } // namespace transcribe::parakeet diff --git a/src/arch/sortformer/model.cpp b/src/arch/sortformer/model.cpp index e59c746b..ef989c44 100644 --- a/src/arch/sortformer/model.cpp +++ b/src/arch/sortformer/model.cpp @@ -64,6 +64,12 @@ SortformerModel::~SortformerModel() { plan.primary = nullptr; } +DiarStreamScratch::~DiarStreamScratch() { + if (compute_ctx != nullptr) { + ggml_free(compute_ctx); + } +} + SortformerSession::~SortformerSession() { if (sched != nullptr) { transcribe::safe_sched_free(sched); @@ -111,12 +117,20 @@ void fill_conformer_hp(const SortformerHParams & s, pk::ParakeetHParams & p) { // Populate ONLY the pre_encode + block slots of a ParakeetWeights from the // GGUF (the conformer). Predictor / joint / head stay empty. Mirrors the // encoder portion of build_parakeet_weights (identical tensor names). -transcribe_status load_conformer_weights(ggml_context * ctx, const pk::ParakeetHParams & hp, pk::ParakeetWeights & w) { - char name[128]; - auto get = [&](const char * nm) -> ggml_tensor * { - ggml_tensor * t = ggml_get_tensor(ctx, nm); +// `prefix` (nullable, standalone default "") namespaces every tensor name +// for the multitalker bundle case. +transcribe_status load_conformer_weights(ggml_context * ctx, + const pk::ParakeetHParams & hp, + pk::ParakeetWeights & w, + const char * prefix = nullptr) { + const char * pfx = (prefix != nullptr) ? prefix : ""; + char name[160]; + auto get = [&](const char * nm) -> ggml_tensor * { + char full[192]; + std::snprintf(full, sizeof(full), "%s%s", pfx, nm); + ggml_tensor * t = ggml_get_tensor(ctx, full); if (t == nullptr) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "sortformer: missing conformer tensor %s", nm); + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "sortformer: missing conformer tensor %s", full); } return t; }; @@ -197,32 +211,37 @@ transcribe_status load_conformer_weights(ggml_context * ctx, const pk::ParakeetH // Fuse the conformer BatchNorm into scale + bias tensors (parakeet's // build_encoder_graph consumes conv_bn_fused_scale/bias, not the raw BN). -transcribe_status fuse_conformer_bn(SortformerModel & m) { - const size_t n_blocks = m.conformer.blocks.size(); +// The fused tensors are allocated on `alloc_backend`; the caller owns and +// frees *out_ctx / *out_buffer. Must run AFTER tensor data is uploaded. +transcribe_status fuse_conformer_bn_core(std::vector & blocks, + int64_t d, + ggml_backend_t alloc_backend, + ggml_context ** out_ctx, + ggml_backend_buffer_t * out_buffer) { + const size_t n_blocks = blocks.size(); if (n_blocks == 0) { return TRANSCRIBE_OK; } - const int64_t d = m.conformer_hp.enc_d_model; - const size_t tensor_bytes = static_cast(d) * sizeof(float); + const size_t tensor_bytes = static_cast(d) * sizeof(float); const size_t ctx_size = n_blocks * 2 * ggml_tensor_overhead() + 256; ggml_init_params params = { ctx_size, nullptr, /*no_alloc=*/true }; - m.bn_fused_ctx = ggml_init(params); - if (m.bn_fused_ctx == nullptr) { + *out_ctx = ggml_init(params); + if (*out_ctx == nullptr) { return TRANSCRIBE_ERR_BACKEND; } for (size_t i = 0; i < n_blocks; ++i) { - auto & b = m.conformer.blocks[i]; - b.conv_bn_fused_scale = ggml_new_tensor_1d(m.bn_fused_ctx, GGML_TYPE_F32, d); - b.conv_bn_fused_bias = ggml_new_tensor_1d(m.bn_fused_ctx, GGML_TYPE_F32, d); + auto & b = blocks[i]; + b.conv_bn_fused_scale = ggml_new_tensor_1d(*out_ctx, GGML_TYPE_F32, d); + b.conv_bn_fused_bias = ggml_new_tensor_1d(*out_ctx, GGML_TYPE_F32, d); } - m.bn_fused_buffer = ggml_backend_alloc_ctx_tensors(m.bn_fused_ctx, m.plan.scheduler_list.back()); - if (m.bn_fused_buffer == nullptr) { + *out_buffer = ggml_backend_alloc_ctx_tensors(*out_ctx, alloc_backend); + if (*out_buffer == nullptr) { return TRANSCRIBE_ERR_BACKEND; } std::vector bn_w(d), bn_b(d), rm(d), rv(d), fused_s(d), fused_b(d); for (size_t i = 0; i < n_blocks; ++i) { - auto & b = m.conformer.blocks[i]; + auto & b = blocks[i]; ggml_backend_tensor_get(b.conv_bn_w, bn_w.data(), 0, tensor_bytes); ggml_backend_tensor_get(b.conv_bn_b, bn_b.data(), 0, tensor_bytes); ggml_backend_tensor_get(b.conv_bn_rm, rm.data(), 0, tensor_bytes); @@ -238,6 +257,11 @@ transcribe_status fuse_conformer_bn(SortformerModel & m) { return TRANSCRIBE_OK; } +transcribe_status fuse_conformer_bn(SortformerModel & m) { + return fuse_conformer_bn_core(m.conformer.blocks, m.conformer_hp.enc_d_model, m.plan.scheduler_list.back(), + &m.bn_fused_ctx, &m.bn_fused_buffer); +} + // ---- diar-head graph helpers ---- ggml_tensor * layer_norm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * w, ggml_tensor * b) { @@ -412,14 +436,18 @@ struct PreEncodeBuild { ggml_tensor * out = nullptr; }; -PreEncodeBuild build_pre_encode_graph(ggml_context * ctx, SortformerModel & m, int M) { +PreEncodeBuild build_pre_encode_graph(ggml_context * ctx, + const pk::ParakeetHParams & chp, + const pk::ParakeetPreEncode & pe, + const char * backend, + int M) { PreEncodeBuild b{}; - const int n_mels = m.conformer_hp.fe_num_mels; + const int n_mels = chp.fe_num_mels; ggml_tensor * mel_in = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, M, n_mels, 1, 1); ggml_set_name(mel_in, "chunk.mel.in"); ggml_set_input(mel_in); - const conf::ConvPolicy policy = sf_conv_policy(m.backend.c_str()); - ggml_tensor * out = conf::build_pre_encode(ctx, sf_pre_view(m.conformer.pre_encode), mel_in, policy, + const conf::ConvPolicy policy = sf_conv_policy(backend); + ggml_tensor * out = conf::build_pre_encode(ctx, sf_pre_view(pe), mel_in, policy, /*name_prefix=*/"chunk.pre_encode", /*error_tag=*/"sortformer", nullptr); if (out == nullptr) { return b; @@ -441,9 +469,15 @@ struct StreamInferBuild { ggml_tensor * preds = nullptr; }; -StreamInferBuild build_stream_infer_graph(ggml_context * ctx, SortformerModel & m, int T_concat) { +StreamInferBuild build_stream_infer_graph(ggml_context * ctx, + const SortformerHParams & hp, + const pk::ParakeetHParams & chp, + const pk::ParakeetWeights & conformer, + const SortformerWeights & w, + const char * backend, + int T_concat) { StreamInferBuild b{}; - const int ed = m.conformer_hp.enc_d_model; + const int ed = chp.enc_d_model; ggml_tensor * concat_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ed, T_concat); ggml_set_name(concat_in, "stream.concat.in"); ggml_set_input(concat_in); @@ -458,13 +492,13 @@ StreamInferBuild build_stream_infer_graph(ggml_context * ctx, SortformerModel & conf::BlockParams bp{}; bp.d_model = ed; - bp.n_head = m.conformer_hp.enc_n_heads; - bp.conv_kernel = m.conformer_hp.enc_conv_kernel; + bp.n_head = chp.enc_n_heads; + bp.conv_kernel = chp.enc_conv_kernel; bp.kv_type = GGML_TYPE_F32; bool enc_use_flash = true, dec_use_flash = false; transcribe::flash::apply_env_overrides(enc_use_flash, dec_use_flash); bp.use_flash = enc_use_flash; - bp.policy = sf_conv_policy(m.backend.c_str()); + bp.policy = sf_conv_policy(backend); bp.att_context_left = -1; // full (offline) attention over the concat bp.att_context_right = -1; bp.att_context_style = conf::BlockParams::AttContextStyle::Regular; @@ -472,24 +506,24 @@ StreamInferBuild build_stream_infer_graph(ggml_context * ctx, SortformerModel & bp.conv_context_right = -1; bp.conv_norm_type = conf::BlockParams::ConvNormType::BatchNorm; - for (int i = 0; i < m.conformer_hp.enc_n_layers; ++i) { - x = conf::build_conformer_block(ctx, x, pos_emb_in, sf_block_view(m.conformer.blocks[static_cast(i)]), + for (int i = 0; i < chp.enc_n_layers; ++i) { + x = conf::build_conformer_block(ctx, x, pos_emb_in, sf_block_view(conformer.blocks[static_cast(i)]), bp); } // encoder_proj (enc_d_model -> tf_d_model) then 18x post-LN transformer. - ggml_tensor * t = linear(ctx, m.weights.enc_proj_w, x, m.weights.enc_proj_b); - const int d = m.hparams.tf_d_model; - for (int i = 0; i < m.hparams.tf_n_layers; ++i) { - t = tf_block(ctx, m.weights.tf_blocks[static_cast(i)], t, d, m.hparams.tf_n_heads, T_concat); + ggml_tensor * t = linear(ctx, w.enc_proj_w, x, w.enc_proj_b); + const int d = hp.tf_d_model; + for (int i = 0; i < hp.tf_n_layers; ++i) { + t = tf_block(ctx, w.tf_blocks[static_cast(i)], t, d, hp.tf_n_heads, T_concat); } // Diar head (forward_speaker_sigmoids). All frames valid in sync mode, so // the encoder_mask multiply is a no-op. ggml_tensor * h = ggml_relu(ctx, t); - h = linear(ctx, m.weights.fc1_w, h, m.weights.fc1_b); + h = linear(ctx, w.fc1_w, h, w.fc1_b); h = ggml_relu(ctx, h); - ggml_tensor * s = linear(ctx, m.weights.single_spk_head_w, h, m.weights.single_spk_head_b); + ggml_tensor * s = linear(ctx, w.single_spk_head_w, h, w.single_spk_head_b); ggml_tensor * preds = ggml_sigmoid(ctx, s); // [n_spk, T_concat] ggml_cgraph * graph = ggml_new_graph_custom(ctx, 8192, false); @@ -602,6 +636,33 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par return TRANSCRIBE_OK; } +// ---- Embedded-diarizer surface (multitalker bundle; see sortformer.h) ---- + +transcribe_status init_embedded_diarizer(const gguf_context * gguf, + ggml_context * ctx_meta, + const char * tensor_prefix, + SortformerEmbedded & out) { + if (gguf == nullptr || ctx_meta == nullptr) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + if (const transcribe_status st = read_sortformer_hparams(gguf, out.hp); st != TRANSCRIBE_OK) { + return st; + } + fill_conformer_hp(out.hp, out.conformer_hp); + if (const transcribe_status st = load_conformer_weights(ctx_meta, out.conformer_hp, out.conformer, tensor_prefix); + st != TRANSCRIBE_OK) { + return st; + } + return build_sortformer_weights(ctx_meta, out.hp, out.weights, tensor_prefix); +} + +transcribe_status fuse_embedded_diar_bn(SortformerEmbedded & e, + ggml_backend_t backend, + ggml_context ** out_ctx, + ggml_backend_buffer_t * out_buffer) { + return fuse_conformer_bn_core(e.conformer.blocks, e.conformer_hp.enc_d_model, backend, out_ctx, out_buffer); +} + transcribe_status init_context(transcribe_model * model, const transcribe_session_params * params, transcribe_session ** out_ctx) { @@ -619,13 +680,14 @@ transcribe_status init_context(transcribe_model * model, // Threshold-based probs -> speaker segments (simple offline segmentation; // the DER-grade dihard3 post-processing is a later, streaming-checkpoint // concern). probs is row-major [T, n_spk]. Emits one contiguous run per -// speaker as a speaker_segment row. -static void probs_to_speaker_segments(transcribe_session * pc, - const std::vector & probs, - int T, - int n_spk, - double ms_per_frame, - float threshold) { +// speaker as a speaker_segment row. Non-static: the parakeet multitalker +// bundle path emits its transcript-independent view through this too. +void probs_to_speaker_segments(transcribe_session * pc, + const std::vector & probs, + int T, + int n_spk, + double ms_per_frame, + float threshold) { for (int s = 0; s < n_spk; ++s) { int run_start = -1; for (int t = 0; t <= T; ++t) { @@ -720,8 +782,9 @@ static transcribe_status run_offline_forward(SortformerSession * pc, SortformerM if (eb.pos_emb_in != nullptr) { const int d_model = pm->conformer_hp.enc_d_model; const int pos_len = static_cast(eb.pos_emb_in->ne[1]); - fill_rel_pos_emb(pc->pos_buf, pc->pos_div_term, pos_len, d_model); - ggml_backend_tensor_set(eb.pos_emb_in, pc->pos_buf.data(), 0, pc->pos_buf.size() * sizeof(float)); + fill_rel_pos_emb(pc->scratch.pos_buf, pc->scratch.pos_div_term, pos_len, d_model); + ggml_backend_tensor_set(eb.pos_emb_in, pc->scratch.pos_buf.data(), 0, + pc->scratch.pos_buf.size() * sizeof(float)); } transcribe::configure_sched_n_threads(pc->sched, pc->n_threads); @@ -737,155 +800,221 @@ static transcribe_status run_offline_forward(SortformerSession * pc, SortformerM return TRANSCRIBE_OK; } -// Streaming AOSC/FIFO forward (the product path). Chunks the mel per NeMo's -// streaming_feat_loader, runs Graph A (pre_encode) then Graph B (blocks + -// proj + transformer + head) over the [spkcache|fifo|chunk] concat, and drives -// the host streaming_update_sync / _compress_spkcache state machine. Emits the -// accumulated T x n_spk probs as diar.probs and the speaker segments. -static transcribe_status run_streaming(SortformerSession * pc, - SortformerModel * pm, - int mel_n_mels, - int mel_n_frames, - double ms_per_frame, - transcribe_sortformer_preset preset) { - const SortformerStreamParams P = resolve_stream_params(pm->hparams, preset); - const int ed = pm->conformer_hp.enc_d_model; - const int n_spk = pm->hparams.max_speakers; - const int sub = pm->hparams.enc_subsampling_factor; - const int feat_len = mel_n_frames; - - if (sub <= 0 || P.chunk_len <= 0) { +// Streaming AOSC/FIFO forward core (the product path). Chunks the mel per +// NeMo's streaming_feat_loader, runs Graph A (pre_encode) then Graph B +// (blocks + proj + transformer + head) over the [spkcache|fifo|chunk] +// concat, and drives the host streaming_update_sync / _compress_spkcache +// state machine. Accumulates the T x n_spk probs in sc.stream.total_preds. +// Family-agnostic: also driven by the parakeet multitalker bundle path. +transcribe_status run_diar_streaming_core(DiarStreamScratch & sc, + const SortformerHParams & hp, + const pk::ParakeetHParams & chp, + const pk::ParakeetWeights & conformer, + const SortformerWeights & w, + const char * backend, + ggml_backend_sched_t sched, + int n_threads, + const SortformerStreamParams & P, + const float * mel_buf, + int mel_n_mels, + int mel_n_frames, + transcribe_session * abort_session) { + const int ed = chp.enc_d_model; + const int n_spk = hp.max_speakers; + const int sub = hp.enc_subsampling_factor; + const int feat_len = mel_n_frames; + + if (sub <= 0 || P.chunk_len <= 0 || sched == nullptr || mel_buf == nullptr) { return TRANSCRIBE_ERR_INVALID_ARG; } - pc->stream.reset(ed); - if (const transcribe_status st = ensure_sched(pc, pm); st != TRANSCRIBE_OK) { - return st; + sc.stream.reset(ed); + + const bool external_cadence = P.feat_first_chunk > 0; + if (external_cadence && (P.chunk_left_context != 0 || P.chunk_right_context != 0)) { + return TRANSCRIBE_ERR_INVALID_ARG; } - const int64_t t_stream_start = ggml_time_us(); - int stt = 0, end = 0, chunk_idx = 0; + int stt = 0, end = 0, chunk_idx = 0; while (end < feat_len) { - if (pc->poll_abort()) { + if (abort_session != nullptr && abort_session->poll_abort()) { return TRANSCRIBE_ERR_ABORTED; } - const int left_offset = std::min(P.chunk_left_context * sub, stt); - end = std::min(stt + P.chunk_len * sub, feat_len); - const int right_offset = std::min(P.chunk_right_context * sub, feat_len - end); - const int win_lo = stt - left_offset; - const int win_hi = end + right_offset; - const int M = win_hi - win_lo; - // Diar-frame context: left_offset is a multiple of sub -> round is - // exact; right_offset may be short on the final chunk -> ceil. - const int lc = (left_offset + sub / 2) / sub; - const int rc = (right_offset + sub - 1) / sub; - - if (pc->compute_ctx != nullptr) { - ggml_free(pc->compute_ctx); - pc->compute_ctx = nullptr; + int win_lo = 0, win_hi = 0, lc = 0, rc = 0, drop = 0; + if (external_cadence) { + // NeMo CacheAwareStreamingAudioBuffer windows: first chunk is + // feat_first_chunk fresh frames; later chunks prepend + // feat_cache frames of real left context to chunk_len*sub new + // frames and drop the duplicated pre-encode outputs below. + if (chunk_idx == 0) { + end = std::min(P.feat_first_chunk, feat_len); + win_lo = 0; + } else { + end = std::min(stt + P.chunk_len * sub, feat_len); + win_lo = std::max(0, stt - P.feat_cache); + drop = P.drop_pre_encode; + } + win_hi = end; + } else { + const int left_offset = std::min(P.chunk_left_context * sub, stt); + end = std::min(stt + P.chunk_len * sub, feat_len); + const int right_offset = std::min(P.chunk_right_context * sub, feat_len - end); + win_lo = stt - left_offset; + win_hi = end + right_offset; + // Diar-frame context: left_offset is a multiple of sub -> round + // is exact; right_offset may be short on the final chunk -> ceil. + lc = (left_offset + sub / 2) / sub; + rc = (right_offset + sub - 1) / sub; + } + const int M = win_hi - win_lo; + + if (sc.compute_ctx != nullptr) { + ggml_free(sc.compute_ctx); + sc.compute_ctx = nullptr; } { ggml_init_params ip{}; - ip.mem_size = 32 * 1024 * 1024; - ip.mem_buffer = nullptr; - ip.no_alloc = true; - pc->compute_ctx = ggml_init(ip); - if (pc->compute_ctx == nullptr) { + ip.mem_size = 32 * 1024 * 1024; + ip.mem_buffer = nullptr; + ip.no_alloc = true; + sc.compute_ctx = ggml_init(ip); + if (sc.compute_ctx == nullptr) { return TRANSCRIBE_ERR_GGUF; } } - ggml_context * ctx = pc->compute_ctx; + ggml_context * ctx = sc.compute_ctx; transcribe::debug::push_name_prefix("stream.chunk"); // ---- Graph A: pre_encode over the mel window ---- - PreEncodeBuild A = build_pre_encode_graph(ctx, *pm, M); + PreEncodeBuild A = build_pre_encode_graph(ctx, chp, conformer.pre_encode, backend, M); if (A.graph == nullptr) { transcribe::debug::pop_name_prefix(); return TRANSCRIBE_ERR_GGUF; } - const int T_diar = static_cast(A.out->ne[1]); + int T_diar = static_cast(A.out->ne[1]); // Mel window [n_mels, M] from full mel [n_mels, feat_len], cols // [win_lo, win_hi). ggml ne=[M, n_mels] -> offset mel*M + col. - pc->chunk_mel_buf.resize(static_cast(mel_n_mels) * M); + sc.chunk_mel_buf.resize(static_cast(mel_n_mels) * M); for (int mel = 0; mel < mel_n_mels; ++mel) { - const float * src = pc->mel_buf.data() + static_cast(mel) * feat_len + win_lo; - std::copy(src, src + M, pc->chunk_mel_buf.data() + static_cast(mel) * M); + const float * src = mel_buf + static_cast(mel) * feat_len + win_lo; + std::copy(src, src + M, sc.chunk_mel_buf.data() + static_cast(mel) * M); } - ggml_backend_sched_reset(pc->sched); - if (!ggml_backend_sched_alloc_graph(pc->sched, A.graph)) { + ggml_backend_sched_reset(sched); + if (!ggml_backend_sched_alloc_graph(sched, A.graph)) { transcribe::debug::pop_name_prefix(); return TRANSCRIBE_ERR_GGUF; } - ggml_backend_tensor_set(A.mel_in, pc->chunk_mel_buf.data(), 0, pc->chunk_mel_buf.size() * sizeof(float)); - transcribe::configure_sched_n_threads(pc->sched, pc->n_threads); - if (ggml_backend_sched_graph_compute(pc->sched, A.graph) != GGML_STATUS_SUCCESS) { + ggml_backend_tensor_set(A.mel_in, sc.chunk_mel_buf.data(), 0, sc.chunk_mel_buf.size() * sizeof(float)); + transcribe::configure_sched_n_threads(sched, n_threads); + if (ggml_backend_sched_graph_compute(sched, A.graph) != GGML_STATUS_SUCCESS) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "sortformer streaming: pre_encode compute failed"); transcribe::debug::pop_name_prefix(); return TRANSCRIBE_ERR_GGUF; } - pc->chunk_embs_host.resize(static_cast(T_diar) * ed); - ggml_backend_tensor_get(A.out, pc->chunk_embs_host.data(), 0, pc->chunk_embs_host.size() * sizeof(float)); + sc.chunk_embs_host.resize(static_cast(T_diar) * ed); + ggml_backend_tensor_get(A.out, sc.chunk_embs_host.data(), 0, sc.chunk_embs_host.size() * sizeof(float)); + + // External cadence: discard the duplicated / partially-contexted + // pre-encode outputs (NeMo drop_extra_pre_encoded, applied right + // after pre_encode and before the encoder). + if (drop > 0) { + if (T_diar <= drop) { + // Degenerate tail chunk: nothing valid to append. + transcribe::debug::pop_name_prefix(); + stt = end; + ++chunk_idx; + continue; + } + sc.chunk_embs_host.erase(sc.chunk_embs_host.begin(), + sc.chunk_embs_host.begin() + static_cast(drop) * ed); + T_diar -= drop; + } // ---- host concat [spkcache | fifo | chunk_embs] ---- - const int S = pc->stream.spkcache_n; - const int F = pc->stream.fifo_n; + const int S = sc.stream.spkcache_n; + const int F = sc.stream.fifo_n; const int T_concat = S + F + T_diar; - pc->concat_host.resize(static_cast(T_concat) * ed); - std::copy(pc->stream.spkcache.begin(), pc->stream.spkcache.begin() + static_cast(S) * ed, - pc->concat_host.begin()); - std::copy(pc->stream.fifo.begin(), pc->stream.fifo.begin() + static_cast(F) * ed, - pc->concat_host.begin() + static_cast(S) * ed); - std::copy(pc->chunk_embs_host.begin(), pc->chunk_embs_host.end(), - pc->concat_host.begin() + static_cast(S + F) * ed); + sc.concat_host.resize(static_cast(T_concat) * ed); + std::copy(sc.stream.spkcache.begin(), sc.stream.spkcache.begin() + static_cast(S) * ed, + sc.concat_host.begin()); + std::copy(sc.stream.fifo.begin(), sc.stream.fifo.begin() + static_cast(F) * ed, + sc.concat_host.begin() + static_cast(S) * ed); + std::copy(sc.chunk_embs_host.begin(), sc.chunk_embs_host.end(), + sc.concat_host.begin() + static_cast(S + F) * ed); // ---- Graph B: xscale + 17 blocks + proj + 18 tf + head ---- - StreamInferBuild B = build_stream_infer_graph(ctx, *pm, T_concat); + StreamInferBuild B = build_stream_infer_graph(ctx, hp, chp, conformer, w, backend, T_concat); if (B.graph == nullptr) { transcribe::debug::pop_name_prefix(); return TRANSCRIBE_ERR_GGUF; } - fill_rel_pos_emb(pc->pos_buf, pc->pos_div_term, 2 * T_concat - 1, ed); + fill_rel_pos_emb(sc.pos_buf, sc.pos_div_term, 2 * T_concat - 1, ed); - ggml_backend_sched_reset(pc->sched); - if (!ggml_backend_sched_alloc_graph(pc->sched, B.graph)) { + ggml_backend_sched_reset(sched); + if (!ggml_backend_sched_alloc_graph(sched, B.graph)) { transcribe::debug::pop_name_prefix(); return TRANSCRIBE_ERR_GGUF; } - ggml_backend_tensor_set(B.concat_in, pc->concat_host.data(), 0, pc->concat_host.size() * sizeof(float)); - ggml_backend_tensor_set(B.pos_emb_in, pc->pos_buf.data(), 0, pc->pos_buf.size() * sizeof(float)); - transcribe::configure_sched_n_threads(pc->sched, pc->n_threads); - if (ggml_backend_sched_graph_compute(pc->sched, B.graph) != GGML_STATUS_SUCCESS) { + ggml_backend_tensor_set(B.concat_in, sc.concat_host.data(), 0, sc.concat_host.size() * sizeof(float)); + ggml_backend_tensor_set(B.pos_emb_in, sc.pos_buf.data(), 0, sc.pos_buf.size() * sizeof(float)); + transcribe::configure_sched_n_threads(sched, n_threads); + if (ggml_backend_sched_graph_compute(sched, B.graph) != GGML_STATUS_SUCCESS) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "sortformer streaming: infer compute failed"); transcribe::debug::pop_name_prefix(); return TRANSCRIBE_ERR_GGUF; } - pc->stream_preds_host.resize(static_cast(T_concat) * n_spk); - ggml_backend_tensor_get(B.preds, pc->stream_preds_host.data(), 0, pc->stream_preds_host.size() * sizeof(float)); + sc.stream_preds_host.resize(static_cast(T_concat) * n_spk); + ggml_backend_tensor_get(B.preds, sc.stream_preds_host.data(), 0, sc.stream_preds_host.size() * sizeof(float)); transcribe::debug::pop_name_prefix(); // ---- host streaming update (FIFO + AOSC compress) ---- - streaming_update_sync(pc->stream, P, n_spk, ed, pc->chunk_embs_host, T_diar, pc->stream_preds_host, T_concat, - lc, rc); + streaming_update_sync(sc.stream, P, n_spk, ed, sc.chunk_embs_host, T_diar, sc.stream_preds_host, T_concat, lc, + rc); stt = end; ++chunk_idx; } + (void) chunk_idx; + return TRANSCRIBE_OK; +} + +// Family wrapper: resolve the operating point, run the core, trim, dump, +// and emit speaker segments on the session. +static transcribe_status run_streaming(SortformerSession * pc, + SortformerModel * pm, + int mel_n_mels, + int mel_n_frames, + double ms_per_frame, + transcribe_sortformer_preset preset) { + const SortformerStreamParams P = resolve_stream_params(pm->hparams, preset); + if (const transcribe_status st = ensure_sched(pc, pm); st != TRANSCRIBE_OK) { + return st; + } + + const int64_t t_stream_start = ggml_time_us(); + const transcribe_status st = run_diar_streaming_core(pc->scratch, pm->hparams, pm->conformer_hp, pm->conformer, + pm->weights, pm->backend.c_str(), pc->sched, pc->n_threads, P, + pc->mel_buf.data(), mel_n_mels, mel_n_frames, pc); + if (st != TRANSCRIBE_OK) { + return st; + } pc->t_encode_us = ggml_time_us() - t_stream_start; // Trim padding tail: NeMo total_preds[:, :ceil(feat_len/sub)]. - const int n_frames = (feat_len + sub - 1) / sub; - const int used_n = std::min(pc->stream.total_n, n_frames); + const int sub = pm->hparams.enc_subsampling_factor; + const int n_spk = pm->hparams.max_speakers; + const int n_frames = (mel_n_frames + sub - 1) / sub; + const int used_n = std::min(pc->scratch.stream.total_n, n_frames); if (transcribe::debug::enabled()) { const long long shape[2] = { used_n, n_spk }; - transcribe::debug::dump_host_f32("diar.probs", pc->stream.total_preds.data(), + transcribe::debug::dump_host_f32("diar.probs", pc->scratch.stream.total_preds.data(), static_cast(used_n) * n_spk, shape, 2, "diarize"); } - probs_to_speaker_segments(pc, pc->stream.total_preds, used_n, n_spk, ms_per_frame, /*threshold=*/0.5f); - (void) chunk_idx; + probs_to_speaker_segments(pc, pc->scratch.stream.total_preds, used_n, n_spk, ms_per_frame, /*threshold=*/0.5f); return TRANSCRIBE_OK; } diff --git a/src/arch/sortformer/sortformer.h b/src/arch/sortformer/sortformer.h index 9d044c34..3edcd749 100644 --- a/src/arch/sortformer/sortformer.h +++ b/src/arch/sortformer/sortformer.h @@ -27,8 +27,11 @@ struct ggml_context; struct ggml_tensor; +struct ggml_backend; struct ggml_backend_buffer; struct ggml_backend_sched; +struct gguf_context; +typedef struct ggml_backend * ggml_backend_t; typedef struct ggml_backend_buffer * ggml_backend_buffer_t; typedef struct ggml_backend_sched * ggml_backend_sched_t; @@ -90,6 +93,19 @@ struct SortformerStreamParams { float weak_boost_rate = 1.5f; float min_pos_scores_rate = 0.5f; int max_index = 99999; + + // External-cadence chunking (multitalker bundle only; 0 = off, the + // standalone uniform-window path). When feat_first_chunk > 0 the core + // mirrors NeMo's CacheAwareStreamingAudioBuffer windows instead of the + // uniform streaming_feat_loader ones: the first chunk consumes + // feat_first_chunk mel frames, each later chunk consumes + // chunk_len*sub new frames with feat_cache frames of real left + // context prepended, and drop_pre_encode pre-encode output frames are + // discarded per later chunk (the duplicated / partially-contexted + // outputs). chunk_left/right_context must be 0 in this mode. + int feat_first_chunk = 0; + int feat_cache = 0; + int drop_pre_encode = 0; }; // Host-side streaming state (AOSC speaker cache + FIFO), mirroring NeMo's @@ -150,17 +166,14 @@ void streaming_update_sync(SortformerStreamState & st, // spkcache_len frames, matching sortformer_modules._compress_spkcache. void compress_spkcache(SortformerStreamState & st, const SortformerStreamParams & p, int n_spk, int emb_dim); -// Concrete context. Owns a per-call compute context and a persistent -// multi-backend scheduler. -struct SortformerSession final : public transcribe_session { - ggml_context * compute_ctx = nullptr; - ggml_backend_sched_t sched = nullptr; - - // Per-context scratch reused across runs. - std::vector mel_buf; +// Scratch for one streaming diarization forward, owned by whichever session +// drives it: the sortformer session (standalone family) or the parakeet +// session (multitalker bundle). compute_ctx is (re)created per chunk by +// run_diar_streaming_core and freed by the destructor. +struct DiarStreamScratch { + ggml_context * compute_ctx = nullptr; std::vector pos_buf; std::vector pos_div_term; - std::vector probs_host; // [n_spk * T], read back from diar.preds // Streaming scratch (AOSC/FIFO path). SortformerStreamState stream; @@ -169,8 +182,88 @@ struct SortformerSession final : public transcribe_session { std::vector concat_host; // [T_concat * enc_d_model] Graph B input std::vector stream_preds_host; // [T_concat * n_spk] Graph B readback + DiarStreamScratch() = default; + ~DiarStreamScratch(); +}; + +// Concrete context. Owns a per-call compute context and a persistent +// multi-backend scheduler. +struct SortformerSession final : public transcribe_session { + ggml_context * compute_ctx = nullptr; // offline-forward graph ctx + ggml_backend_sched_t sched = nullptr; + + // Per-context scratch reused across runs. + std::vector mel_buf; + std::vector probs_host; // [n_spk * T], read back from diar.preds + + // Streaming scratch (AOSC/FIFO path + rel-pos tables, shared with the + // offline forward's pos_emb fill). + DiarStreamScratch scratch; + SortformerSession() = default; ~SortformerSession() override; }; +// ---- Embedded-diarizer surface (multitalker bundle) -------------------- // +// +// A multitalker Parakeet bundle GGUF carries this diarizer's tensors under +// a name prefix (stt.parakeet.diarizer.tensor_prefix, "sortformer.") next +// to the ASR's own tensors, and its hparams under the usual +// stt.sortformer.* KVs. The host family owns the storage (its ctx_meta / +// backend buffers); this surface only resolves borrowed tensor pointers +// and runs the forward on the host's scheduler. + +// Weight + hparam core of a diarizer embedded in another family's GGUF. +struct SortformerEmbedded { + SortformerHParams hp; + transcribe::parakeet::ParakeetHParams conformer_hp; + transcribe::parakeet::ParakeetWeights conformer; // pre_encode + blocks only + SortformerWeights weights; +}; + +// Read stt.sortformer.* hparams from `gguf` and resolve every diarizer +// tensor (conformer + sortformer-specific) in `ctx_meta` under +// `tensor_prefix`. Defined in model.cpp. +transcribe_status init_embedded_diarizer(const gguf_context * gguf, + ggml_context * ctx_meta, + const char * tensor_prefix, + SortformerEmbedded & out); + +// Fuse the embedded conformer's BatchNorm into scale/bias tensors allocated +// on `backend` (must run AFTER tensor data is uploaded). The caller owns and +// frees *out_ctx / *out_buffer. Defined in model.cpp. +transcribe_status fuse_embedded_diar_bn(SortformerEmbedded & e, + ggml_backend_t backend, + ggml_context ** out_ctx, + ggml_backend_buffer_t * out_buffer); + +// Full streaming AOSC/FIFO diarization forward over a precomputed mel +// buffer (row-major [n_mels, n_frames]), on the caller's scheduler. On +// success sc.stream.total_preds holds the accumulated row-major +// [total_n, n_spk] probs (trim to ceil(n_frames / subsampling) frames; +// see run_streaming in model.cpp). abort_session is polled per chunk when +// non-null. Defined in model.cpp. +transcribe_status run_diar_streaming_core(DiarStreamScratch & sc, + const SortformerHParams & hp, + const transcribe::parakeet::ParakeetHParams & chp, + const transcribe::parakeet::ParakeetWeights & conformer, + const SortformerWeights & w, + const char * backend, + ggml_backend_sched_t sched, + int n_threads, + const SortformerStreamParams & P, + const float * mel_buf, + int mel_n_mels, + int mel_n_frames, + transcribe_session * abort_session); + +// Threshold-based probs -> speaker_segment rows (probs row-major [T, n_spk], +// speaker_id 1-based, p = NaN). Shared with the multitalker bundle path. +void probs_to_speaker_segments(transcribe_session * session, + const std::vector & probs, + int T, + int n_spk, + double ms_per_frame, + float threshold); + } // namespace transcribe::sortformer diff --git a/src/arch/sortformer/weights.cpp b/src/arch/sortformer/weights.cpp index bcd2de4e..708aaee6 100644 --- a/src/arch/sortformer/weights.cpp +++ b/src/arch/sortformer/weights.cpp @@ -138,18 +138,23 @@ ggml_tensor * get_checked(ggml_context * ctx, const char * name, int64_t ne0, in } // namespace -transcribe_status build_sortformer_weights(ggml_context * ctx, const SortformerHParams & hp, SortformerWeights & w) { +transcribe_status build_sortformer_weights(ggml_context * ctx, + const SortformerHParams & hp, + SortformerWeights & w, + const char * tensor_prefix) { const int64_t d = hp.tf_d_model; const int64_t dff = hp.tf_d_ff; const int64_t ed = hp.enc_d_model; const int64_t spk = hp.max_speakers; - char name[128]; - -#define GET(dst, nm, e0, e1) \ - do { \ - (dst) = get_checked(ctx, (nm), (e0), (e1)); \ - if ((dst) == nullptr) \ - return TRANSCRIBE_ERR_GGUF; \ + const char * pfx = (tensor_prefix != nullptr) ? tensor_prefix : ""; + char name[160]; + +#define GET(dst, nm, e0, e1) \ + do { \ + std::snprintf(name, sizeof(name), "%s%s", pfx, (nm)); \ + (dst) = get_checked(ctx, name, (e0), (e1)); \ + if ((dst) == nullptr) \ + return TRANSCRIBE_ERR_GGUF; \ } while (0) GET(w.enc_proj_w, "diar.encoder_proj.weight", ed, d); @@ -158,10 +163,11 @@ transcribe_status build_sortformer_weights(ggml_context * ctx, const SortformerH w.tf_blocks.resize(static_cast(hp.tf_n_layers)); for (int i = 0; i < hp.tf_n_layers; ++i) { SortformerTfBlock & b = w.tf_blocks[static_cast(i)]; -#define GETB(dst, suffix, e0, e1) \ - do { \ - std::snprintf(name, sizeof(name), "tf.blocks.%d.%s", i, suffix); \ - GET(dst, name, e0, e1); \ + char blk[96]; +#define GETB(dst, suffix, e0, e1) \ + do { \ + std::snprintf(blk, sizeof(blk), "tf.blocks.%d.%s", i, suffix); \ + GET(dst, blk, e0, e1); \ } while (0) GETB(b.norm1_w, "norm_1.weight", d, -1); GETB(b.norm1_b, "norm_1.bias", d, -1); diff --git a/src/arch/sortformer/weights.h b/src/arch/sortformer/weights.h index 9b655e96..c1b8f2be 100644 --- a/src/arch/sortformer/weights.h +++ b/src/arch/sortformer/weights.h @@ -112,8 +112,14 @@ transcribe_status read_sortformer_hparams(const gguf_context * gguf, SortformerH // head) by name in ctx_meta, validate shapes against hp, store borrowed // pointers. Returns TRANSCRIBE_ERR_GGUF (naming the tensor) on any // missing / mis-shaped tensor. +// +// tensor_prefix (nullable, standalone default "") is prepended to every +// tensor name: the multitalker bundle GGUF embeds these tensors under a +// "sortformer." namespace so they cannot collide with the host family's +// own enc.* catalog. transcribe_status build_sortformer_weights(ggml_context * ctx_meta, const SortformerHParams & hp, - SortformerWeights & weights); + SortformerWeights & weights, + const char * tensor_prefix = nullptr); } // namespace transcribe::sortformer diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e4b12330..42b9e303 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -669,6 +669,25 @@ if(TRANSCRIBE_BUILD_REAL_MODEL_TESTS) COMMAND transcribe_parakeet_real_smoke) set_tests_properties(transcribe_parakeet_real_smoke PROPERTIES SKIP_RETURN_CODE 77) + + # Multitalker bundle end-to-end smoke (diarize=ON on the composed + # parakeet + sortformer bundle; both supervision modes). Gated on + # TRANSCRIBE_MULTITALKER_BUNDLE_GGUF at runtime, exit 77 when unset. + add_executable(transcribe_parakeet_multitalker_e2e_smoke + parakeet_multitalker_e2e_smoke.cpp) + + target_link_libraries(transcribe_parakeet_multitalker_e2e_smoke + PRIVATE transcribe transcribe-common-example) + + target_compile_definitions(transcribe_parakeet_multitalker_e2e_smoke PRIVATE + "TRANSCRIBE_TEST_SAMPLES_DIR=\"${CMAKE_SOURCE_DIR}/samples\"") + + transcribe_apply_warnings(transcribe_parakeet_multitalker_e2e_smoke) + + add_test(NAME transcribe_parakeet_multitalker_e2e_smoke + COMMAND transcribe_parakeet_multitalker_e2e_smoke) + set_tests_properties(transcribe_parakeet_multitalker_e2e_smoke PROPERTIES + SKIP_RETURN_CODE 77) endif() # ----------------------------------------------------------------------------- diff --git a/tests/golden/parakeet/multitalker-2spk-mix.rttm b/tests/golden/parakeet/multitalker-2spk-mix.rttm new file mode 100644 index 00000000..b30d4e69 --- /dev/null +++ b/tests/golden/parakeet/multitalker-2spk-mix.rttm @@ -0,0 +1,3 @@ +SPEAKER multitalker-2spk-mix 1 0.000 5.500 spk_A +SPEAKER multitalker-2spk-mix 1 9.500 4.500 spk_A +SPEAKER multitalker-2spk-mix 1 5.000 4.500 spk_B diff --git a/tests/parakeet_multitalker_e2e_smoke.cpp b/tests/parakeet_multitalker_e2e_smoke.cpp new file mode 100644 index 00000000..dac2e83e --- /dev/null +++ b/tests/parakeet_multitalker_e2e_smoke.cpp @@ -0,0 +1,170 @@ +// parakeet_multitalker_e2e_smoke.cpp - real-model gated end-to-end test +// for the multitalker bundle (parakeet + embedded sortformer). +// +// Gating (same pattern as parakeet_real_smoke): +// - Built only under TRANSCRIBE_BUILD_REAL_MODEL_TESTS. +// - The bundle path comes from TRANSCRIBE_MULTITALKER_BUNDLE_GGUF; if +// unset, exits 77 (CTest "skipped"). +// +// What we assert, on samples/multitalker-2spk-mix.wav (the deterministic +// 2-speaker fixture whose NeMo reference lives under +// build/validate/parakeet/.../multitalker/): +// +// 1. The bundle loads and reports TRANSCRIBE_FEATURE_DIARIZATION. +// 2. diarize=OFF: single-speaker semantics — a result with every +// segment's speaker_id == 0 and no speaker_segment rows. +// 3. diarize=ON (masked mode, the default): >= 2 segments carrying at +// least two distinct 1-based speaker_ids, speaker_segment rows +// present, and each fixture voice's anchor phrase attributed to a +// single speaker ("fellow Americans" for the JFK track, +// "publication" for the Whole Earth track). +// 4. Same invariants under TRANSCRIBE_MULTITALKER_MODE=kernel. + +#include "transcribe.h" +#include "wav.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +int g_failures = 0; + +void check(bool ok, const char * what) { + if (!ok) { + std::fprintf(stderr, "FAIL: %s\n", what); + ++g_failures; + } +} + +// speaker_id owning the segment whose text contains `needle`; 0 if absent +// or ambiguous (present under more than one speaker). +int32_t phrase_speaker(struct transcribe_session * ctx, const char * needle) { + std::set owners; + const int n = transcribe_n_segments(ctx); + for (int i = 0; i < n; ++i) { + struct transcribe_segment seg; + transcribe_segment_init(&seg); + if (transcribe_get_segment(ctx, i, &seg) != TRANSCRIBE_OK || seg.text == nullptr) { + continue; + } + if (std::strstr(seg.text, needle) != nullptr) { + owners.insert(seg.speaker_id); + } + } + return owners.size() == 1 ? *owners.begin() : 0; +} + +void run_multitalker_checks(struct transcribe_session * ctx, const std::vector & pcm, const char * label) { + struct transcribe_run_params rp; + transcribe_run_params_init(&rp); + rp.diarize = TRANSCRIBE_DIARIZE_MODE_ON; + + const transcribe_status st = transcribe_run(ctx, pcm.data(), static_cast(pcm.size()), &rp); + std::fprintf(stderr, "[%s] run: %s\n", label, transcribe_status_string(st)); + check(st == TRANSCRIBE_OK, "diarize=ON run returns OK"); + if (st != TRANSCRIBE_OK) { + return; + } + + const int n_segments = transcribe_n_segments(ctx); + check(n_segments >= 2, "diarize=ON yields >= 2 segments"); + + std::set speakers; + for (int i = 0; i < n_segments; ++i) { + struct transcribe_segment seg; + transcribe_segment_init(&seg); + if (transcribe_get_segment(ctx, i, &seg) == TRANSCRIBE_OK) { + std::fprintf(stderr, "[%s] S%d [%lld, %lld] %s\n", label, seg.speaker_id, + static_cast(seg.t0_ms), static_cast(seg.t1_ms), + seg.text != nullptr ? seg.text : ""); + speakers.insert(seg.speaker_id); + } + } + check(speakers.size() >= 2, "segments carry >= 2 distinct speaker_ids"); + check(speakers.count(0) == 0, "every segment is speaker-attributed (no speaker_id 0)"); + + check(transcribe_n_speaker_segments(ctx) > 0, "speaker_segment rows populated"); + + const int32_t spk_a = phrase_speaker(ctx, "fellow Americans"); + const int32_t spk_b = phrase_speaker(ctx, "publication"); + check(spk_a > 0, "JFK anchor phrase attributed to exactly one speaker"); + check(spk_b > 0, "Whole Earth anchor phrase attributed to exactly one speaker"); + check(spk_a == 0 || spk_b == 0 || spk_a != spk_b, "the two voices land on different speakers"); +} + +} // namespace + +int main() { + const char * gguf = std::getenv("TRANSCRIBE_MULTITALKER_BUNDLE_GGUF"); + if (gguf == nullptr || gguf[0] == '\0') { + std::fprintf(stderr, + "SKIP: set TRANSCRIBE_MULTITALKER_BUNDLE_GGUF to a multitalker bundle GGUF " + "(scripts/compose-multitalker-bundle.py output)\n"); + return 77; + } + const std::string wav_path = std::string(TRANSCRIBE_TEST_SAMPLES_DIR) + "/multitalker-2spk-mix.wav"; + + struct transcribe_model * model = nullptr; + { + const transcribe_status st = transcribe_model_load_file(gguf, nullptr, &model); + if (st != TRANSCRIBE_OK || model == nullptr) { + std::fprintf(stderr, "FAIL: model load: %s\n", transcribe_status_string(st)); + return EXIT_FAILURE; + } + } + + check(transcribe_model_supports(model, TRANSCRIBE_FEATURE_DIARIZATION), + "bundle reports TRANSCRIBE_FEATURE_DIARIZATION"); + + std::vector pcm; + std::string load_err; + if (!transcribe_cli::load_wav_mono_16k(wav_path.c_str(), pcm, load_err)) { + std::fprintf(stderr, "FAIL: wav load: %s\n", load_err.c_str()); + transcribe_model_free(model); + return EXIT_FAILURE; + } + + struct transcribe_session * ctx = nullptr; + if (transcribe_session_init(model, nullptr, &ctx) != TRANSCRIBE_OK || ctx == nullptr) { + std::fprintf(stderr, "FAIL: session init\n"); + transcribe_model_free(model); + return EXIT_FAILURE; + } + + // diarize=OFF: shipped single-speaker semantics. + { + const transcribe_status st = transcribe_run(ctx, pcm.data(), static_cast(pcm.size()), nullptr); + check(st == TRANSCRIBE_OK, "diarize=OFF run returns OK"); + if (st == TRANSCRIBE_OK) { + check(transcribe_n_segments(ctx) >= 1, "diarize=OFF yields a result"); + check(transcribe_n_speaker_segments(ctx) == 0, "diarize=OFF has no speaker_segment rows"); + struct transcribe_segment seg; + transcribe_segment_init(&seg); + if (transcribe_get_segment(ctx, 0, &seg) == TRANSCRIBE_OK) { + check(seg.speaker_id == 0, "diarize=OFF segments are unattributed"); + } + } + } + + // diarize=ON in both supervision modes. + unsetenv("TRANSCRIBE_MULTITALKER_MODE"); + run_multitalker_checks(ctx, pcm, "masked"); + setenv("TRANSCRIBE_MULTITALKER_MODE", "kernel", /*overwrite=*/1); + run_multitalker_checks(ctx, pcm, "kernel"); + unsetenv("TRANSCRIBE_MULTITALKER_MODE"); + + transcribe_session_free(ctx); + transcribe_model_free(model); + + if (g_failures != 0) { + std::fprintf(stderr, "%d check(s) failed\n", g_failures); + return EXIT_FAILURE; + } + std::fprintf(stderr, "parakeet_multitalker_e2e_smoke: all checks passed\n"); + return EXIT_SUCCESS; +}