parakeet multitalker with diarization

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.**
 

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         = <diarizer stt.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())

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-<config>-<split>.manifest.jsonl   (ingest_ami.py)
+Writes samples/diar/ami-<config>-<split>-seglst/<meeting>.seglst.json
+       samples/diar/ami-<config>-<split>-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())