From ee514d0f1fa1544034e6c1bdc717c7a9ed0538a6 Mon Sep 17 00:00:00 2001 From: Artur Wierzbicki Date: Sun, 19 Jul 2026 16:36:04 +0400 Subject: [PATCH] Parakeet word timestamp verification Verify Parakeet TDT word timestamps against the NeMo reference: - dump reference word and token timings in the parakeet NeMo dumper (transcript.json gains word rows; token rows gain timing) - emit words/tokens in the CLI batch JSON, gated on the returned timestamp kind - compare word start/end against the reference in validate.py, within the new "timestamps" budget in tests/tolerances/.json (80 ms = one encoder frame for parakeet) - mirror NeMo's punctuation convention in the decoder: punctuation-only tokens are zero-length at the previous token's end, guarded for non-UTF-8 vocab pieces, with the punctuation set cached at model load. The first validation run caught a 560 ms word-end divergence this fixes; the re-run passes with max 80 ms and mean 3.6 ms on the validation clip, and parakeet-tdt-0.6b-v3 matches the reference exactly - update validate, model-family-testing, and model-card docs --- docs/model-family-testing.md | 11 +- docs/models/parakeet-tdt-0.6b-v2.md | 5 + docs/models/parakeet-tdt-0.6b-v3.md | 6 + docs/tools/validate.md | 10 +- examples/cli/main.cpp | 124 ++++++++- scripts/dump_reference_parakeet_nemo.py | 137 +++++++++- scripts/lib/ref_dump.py | 11 +- scripts/lib/test_validate.py | 182 +++++++++++++ scripts/validate.py | 340 ++++++++++++++++++++++-- src/arch/parakeet/model.cpp | 163 +++++++++++- src/arch/parakeet/parakeet.h | 10 +- tests/decoder_smoke.cpp | 19 +- tests/tolerances/parakeet.json | 4 + 13 files changed, 950 insertions(+), 72 deletions(-) create mode 100644 scripts/lib/test_validate.py diff --git a/docs/model-family-testing.md b/docs/model-family-testing.md index 4c5ebf03..27031b6b 100644 --- a/docs/model-family-testing.md +++ b/docs/model-family-testing.md @@ -127,7 +127,9 @@ Inputs: - Reference dumps come from the per-family `scripts/dump_reference_*.py`. - Tolerances come from `tests/tolerances/.json`. - Transcript artifacts are `transcript.json`; when the reference writes - one, validation requires exact C++ vs reference text equality. + one, validation requires exact C++ vs reference text equality. When the + reference artifact also carries word rows, validation additionally compares + each word's start and end against the family's `timestamps` tolerance. Do not grow bespoke C++ tensor-comparison smoke tests per family. C++ smokes should cover load/run/API behavior; the dump/compare workflow @@ -155,15 +157,20 @@ Tolerances are data, not C++ literals. The source of truth is: tests/tolerances/.json ``` -Use the current flat tensor-name mapping across families: +Use the current flat mapping across families, with tensor names and the optional +`timestamps` entry as top-level keys: ```json { "_comment": ["why the tolerances are what they are"], + "timestamps": {"max_abs_ms": 80.0}, "enc.final": {"max_abs": 1e-4, "mean_abs": 1e-6} } ``` +`timestamps.max_abs_ms` bounds each word's start and end deviation from the +reference in milliseconds; `validate.py` uses it for the word-timestamp compare. + Generic fallback tolerances belong in the tool, for example `compare_tensors.py --max-abs` and `--mean-abs`. Family-specific tensor tolerances belong in `tests/tolerances/.json`. diff --git a/docs/models/parakeet-tdt-0.6b-v2.md b/docs/models/parakeet-tdt-0.6b-v2.md index 87ada6d3..9576b391 100644 --- a/docs/models/parakeet-tdt-0.6b-v2.md +++ b/docs/models/parakeet-tdt-0.6b-v2.md @@ -11,6 +11,11 @@ a transcript with optional token-level timestamps. It is not a streaming model, does not translate, and v2 has no multilingual capability. For multilingual see v3. +Word timings are numerically validated against the NeMo reference within one +encoder frame (80 ms). Token tables are dumped on both sides for divergence +localization but are not compared; punctuation-only tokens follow NeMo's +convention and are zero-length at the previous token's end. + See NVIDIA's [model card](https://huggingface.co/nvidia/parakeet-tdt-0.6b-v2) for training data, intended use, and upstream evaluation methodology. diff --git a/docs/models/parakeet-tdt-0.6b-v3.md b/docs/models/parakeet-tdt-0.6b-v3.md index 567fd42a..9e1c5dd3 100644 --- a/docs/models/parakeet-tdt-0.6b-v3.md +++ b/docs/models/parakeet-tdt-0.6b-v3.md @@ -14,6 +14,12 @@ Estonian, Finnish, French, German, Greek, Hungarian, Italian, Latvian, Lithuanian, Maltese, Polish, Portuguese, Romanian, Russian, Slovak, Slovenian, Spanish, Swedish, Ukrainian. +Word timings are numerically validated against the NeMo reference within one +encoder frame (80 ms); on the validation clip v3 matches the reference +exactly. Token tables are dumped on both sides for divergence localization +but are not compared; punctuation-only tokens follow NeMo's convention and +are zero-length at the previous token's end. + See NVIDIA's [model card](https://huggingface.co/nvidia/parakeet-tdt-0.6b-v3) for training data, intended use, and upstream evaluation methodology. diff --git a/docs/tools/validate.md b/docs/tools/validate.md index d0fa9b4c..9cbab862 100644 --- a/docs/tools/validate.md +++ b/docs/tools/validate.md @@ -10,7 +10,7 @@ convention. ```text validate.py ref → runs dump_reference__*.py → writes build/validate////ref/ validate.py cpp → runs transcribe-cli with TRANSCRIBE_DUMP_DIR → writes build/validate/.../cpp/ - validate.py compare → runs compare_tensors.py with tests/tolerances/.json + validate.py compare → runs compare_tensors.py with tests/tolerances/.json; compares transcript text and, when the reference dump carries word rows, word timestamps validate.py all → ref + cpp + compare in sequence ``` @@ -42,9 +42,11 @@ uv run scripts/validate.py cpp --family cohere --gguf models/cohere-transcribe-0 ## Exit codes -- `0` — all contract tensors within tolerance. -- `1` — one or more contract tensors out of tolerance, missing, or - shape-mismatched. +- `0` — all contract tensors are within tolerance and transcript text matches; + when the reference dump carries word rows, all word timestamps are within the + family's `timestamps` budget. +- `1` — any tensor is out of tolerance, missing, or shape-mismatched; any + required transcript text or word timestamp is missing or mismatched. Exit-code driven, no interactive prompts (per project policy). diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 93c1c889..c57f62a5 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -12,6 +12,7 @@ #include "wav.h" #include +#include #include #include #include @@ -110,6 +111,113 @@ std::string batch_segments_json(const transcribe_session * ctx, int i) { return out; } +void append_word_json(std::string & out, const struct transcribe_word & wrd) { + char head[96]; + std::snprintf(head, sizeof(head), "{\"t0_ms\":%lld,\"t1_ms\":%lld,\"text\":\"", static_cast(wrd.t0_ms), + static_cast(wrd.t1_ms)); + out += head; + out += json_escape(wrd.text != nullptr ? wrd.text : ""); + out += "\"}"; +} + +void append_token_json(std::string & out, const struct transcribe_token & tok) { + char head[96]; + std::snprintf(head, sizeof(head), "{\"t0_ms\":%lld,\"t1_ms\":%lld,\"text\":\"", static_cast(tok.t0_ms), + static_cast(tok.t1_ms)); + out += head; + out += json_escape(tok.text != nullptr ? tok.text : ""); + char tail[96]; + if (std::isfinite(tok.p)) { + std::snprintf(tail, sizeof(tail), "\",\"id\":%d,\"p\":%.9g}", tok.id, static_cast(tok.p)); + } else { + std::snprintf(tail, sizeof(tail), "\",\"id\":%d,\"p\":null}", tok.id); + } + out += tail; +} + +std::string words_json(const transcribe_session * ctx) { + const transcribe_timestamp_kind kind = transcribe_returned_timestamp_kind(ctx); + if (kind != TRANSCRIBE_TIMESTAMPS_WORD && kind != TRANSCRIBE_TIMESTAMPS_TOKEN) { + return {}; + } + const int n_words = transcribe_n_words(ctx); + if (n_words <= 0) { + return {}; + } + std::string out = ",\"words\":["; + for (int w = 0; w < n_words; ++w) { + if (w > 0) { + out += ","; + } + struct transcribe_word wrd; + transcribe_word_init(&wrd); + (void) transcribe_get_word(ctx, w, &wrd); + append_word_json(out, wrd); + } + out += "]"; + return out; +} + +std::string tokens_json(const transcribe_session * ctx) { + if (transcribe_returned_timestamp_kind(ctx) != TRANSCRIBE_TIMESTAMPS_TOKEN) { + return {}; + } + const int n_tokens = transcribe_n_tokens(ctx); + if (n_tokens <= 0) { + return {}; + } + std::string out = ",\"tokens\":["; + for (int t = 0; t < n_tokens; ++t) { + if (t > 0) { + out += ","; + } + struct transcribe_token tok; + transcribe_token_init(&tok); + (void) transcribe_get_token(ctx, t, &tok); + append_token_json(out, tok); + } + out += "]"; + return out; +} + +std::string batch_words_json(const transcribe_session * ctx, int i) { + if (transcribe_batch_returned_timestamp_kind(ctx, i) < TRANSCRIBE_TIMESTAMPS_WORD) { + return {}; + } + const int n_words = transcribe_batch_n_words(ctx, i); + std::string out = ",\"words\":["; + for (int w = 0; w < n_words; ++w) { + if (w > 0) { + out += ","; + } + struct transcribe_word wrd; + transcribe_word_init(&wrd); + (void) transcribe_batch_get_word(ctx, i, w, &wrd); + append_word_json(out, wrd); + } + out += "]"; + return out; +} + +std::string batch_tokens_json(const transcribe_session * ctx, int i) { + if (transcribe_batch_returned_timestamp_kind(ctx, i) != TRANSCRIBE_TIMESTAMPS_TOKEN) { + return {}; + } + const int n_tokens = transcribe_batch_n_tokens(ctx, i); + std::string out = ",\"tokens\":["; + for (int t = 0; t < n_tokens; ++t) { + if (t > 0) { + out += ","; + } + struct transcribe_token tok; + transcribe_token_init(&tok); + (void) transcribe_batch_get_token(ctx, i, t, &tok); + append_token_json(out, tok); + } + out += "]"; + return out; +} + struct cli_args { std::string wav_path; std::string model_path; @@ -828,6 +936,8 @@ int main(int argc, char ** argv) { if (args.batch_jsonl) { const std::string escaped = json_escape(text); const std::string segments = batch_segments_json(ctx, static_cast(k)); + const std::string words = batch_words_json(ctx, static_cast(k)); + const std::string tokens = batch_tokens_json(ctx, static_cast(k)); std::string err_field; if (ust != TRANSCRIBE_OK) { err_field = ",\"error\":\""; @@ -842,11 +952,11 @@ int main(int argc, char ** argv) { transcribe_timings_init(&tm); (void) transcribe_batch_get_timings(ctx, static_cast(k), &tm); std::printf( - "{\"file\":\"%s\",\"text\":\"%s\"%s," + "{\"file\":\"%s\",\"text\":\"%s\"%s%s%s," "\"mel_ms\":%.1f,\"encode_ms\":%.1f," "\"decode_ms\":%.1f%s}\n", - wav.c_str(), escaped.c_str(), segments.c_str(), (double) tm.mel_ms, (double) tm.encode_ms, - (double) tm.decode_ms, err_field.c_str()); + wav.c_str(), escaped.c_str(), segments.c_str(), words.c_str(), tokens.c_str(), + (double) tm.mel_ms, (double) tm.encode_ms, (double) tm.decode_ms, err_field.c_str()); } else { std::printf("[%zu/%zu] %s", src_index[k] + 1, total, wav.c_str()); if (ust == TRANSCRIBE_OK) { @@ -971,6 +1081,8 @@ int main(int argc, char ** argv) { (void) transcribe_get_timings(ctx, &tm); const std::string escaped = json_escape(text); const std::string segments = segments_json(ctx); + const std::string words = words_json(ctx); + const std::string tokens = tokens_json(ctx); std::string err_field; if (run_st != TRANSCRIBE_OK) { err_field = ",\"error\":\""; @@ -978,11 +1090,11 @@ int main(int argc, char ** argv) { err_field += "\""; } std::printf( - "{\"file\":\"%s\",\"text\":\"%s\"%s," + "{\"file\":\"%s\",\"text\":\"%s\"%s%s%s," "\"mel_ms\":%.1f,\"encode_ms\":%.1f," "\"decode_ms\":%.1f%s}\n", - wav.c_str(), escaped.c_str(), segments.c_str(), (double) tm.mel_ms, (double) tm.encode_ms, - (double) tm.decode_ms, err_field.c_str()); + wav.c_str(), escaped.c_str(), segments.c_str(), words.c_str(), tokens.c_str(), + (double) tm.mel_ms, (double) tm.encode_ms, (double) tm.decode_ms, err_field.c_str()); } else { std::printf("[%zu/%zu] %s", i + 1, wav_paths.size(), wav.c_str()); if (run_st == TRANSCRIBE_OK) { diff --git a/scripts/dump_reference_parakeet_nemo.py b/scripts/dump_reference_parakeet_nemo.py index 4f77b0c2..23b655c7 100644 --- a/scripts/dump_reference_parakeet_nemo.py +++ b/scripts/dump_reference_parakeet_nemo.py @@ -108,6 +108,98 @@ def normalize_text(text: str) -> str: return " ".join(text.strip().lower().split()) +def enable_native_timestamps(model, arch: str) -> None: + """Enable NeMo's native alignment and timestamp post-processing.""" + if arch not in {"tdt", "hybrid"} or not hasattr(model.cfg, "decoding"): + return + + import copy + from omegaconf import OmegaConf, open_dict + + decoding_cfg = copy.deepcopy(model.cfg.decoding) + with open_dict(decoding_cfg): + decoding_cfg.compute_timestamps = True + if "greedy" not in decoding_cfg: + decoding_cfg.greedy = OmegaConf.create({}) + decoding_cfg.greedy.preserve_alignments = True + decoding_cfg.tdt_include_token_duration = True + model.change_decoding_strategy(decoding_cfg) + + +def nemo_timestamp_rows(hypothesis) -> tuple[list[dict[str, Any]], list[dict[str, Any]] | None]: + """Convert NeMo Hypothesis.timestamp rows to the reference artifact shape.""" + timestamps = getattr(hypothesis, "timestamp", None) + if not isinstance(timestamps, dict): + raise RuntimeError( + "NeMo returned no Hypothesis.timestamp mapping with timestamps=True" + ) + + def convert_rows( + values, + *, + level: str, + text_keys: tuple[str, ...], + ) -> list[dict[str, Any]]: + converted: list[dict[str, Any]] = [] + for index, value in enumerate(values or []): + if not isinstance(value, dict): + raise RuntimeError( + f"NeMo {level} timestamp row {index} is not an object: {value!r}" + ) + text_key = next((key for key in text_keys if key in value), None) + if text_key is None or "start" not in value or "end" not in value: + raise RuntimeError( + f"NeMo {level} timestamp row {index} lacks text/start/end: {value!r}" + ) + row_text = value[text_key] + if isinstance(row_text, (list, tuple)): + if len(row_text) != 1: + raise RuntimeError( + f"NeMo {level} timestamp row {index} has " + f"{len(row_text)} text values: {value!r}" + ) + row_text = row_text[0] + if not isinstance(row_text, str): + raise RuntimeError( + f"NeMo {level} timestamp row {index} text is not a string: " + f"{value!r}" + ) + converted.append({ + "start_s": float(value["start"]), + "end_s": float(value["end"]), + "text": row_text, + }) + return converted + + if "word" not in timestamps: + raise RuntimeError("NeMo returned no word timestamp table with timestamps=True") + words = convert_rows( + timestamps["word"], level="word", text_keys=("word", "text") + ) + + # NeMo 2.8.0rc0 at 6967f48fda2a exposes the finest decoder-token table only + # as `timestamp["char"]`, including for Parakeet's SentencePiece tokens. + # AbstractRNNTDecoding.compute_rnnt_timestamps forms a raw TDT token end + # from its emission offset plus Hypothesis.token_duration. Its TDT + # refinement then makes supported punctuation zero-width at the previous + # token's end. Word aggregation inherits that final, possibly clamped end. + token_values = timestamps.get("char") + tokens = None + if token_values is not None: + tokens = convert_rows( + token_values, level="token", text_keys=("char", "text") + ) + return words, tokens + + +def first_hypothesis(value): + while isinstance(value, (list, tuple)): + if not value: + raise RuntimeError("NeMo returned an empty transcription result") + value = value[0] + return value + + # --------------------------------------------------------------------------- # NeMo model loading + arch detection # --------------------------------------------------------------------------- @@ -1429,6 +1521,8 @@ def dump(name: str, t, stage: str) -> None: if not args.skip_transcript: print("\n Running full greedy transcription...") + tdt_timestamps = arch in {"tdt", "hybrid"} + enable_native_timestamps(model, arch) # Some .nemo archives (notably parakeet-unified-en-0.6b) ship without # a validation_ds config, which NeMo's transcribe() pipeline tries to # read for use_start_end_token. Stub it so transcribe() works. @@ -1442,28 +1536,49 @@ def dump(name: str, t, stage: str) -> None: # dataset throws on prompt_key='None'. Bypass it: run greedy # decoding directly on the prompt-conditioned encoder output we # already computed. + # rnnt_decoder_predictions_tensor returns frame offsets only; + # transcribe() adds seconds during post-processing. A future + # prompt-aware TDT/hybrid variant will fail loudly here until this + # branch converts those offsets. _, encoded_len = run_encoder_forward(model, audio_tensor, length_tensor) hyps = model.decoding.rnnt_decoder_predictions_tensor( encoder_output=encoded_for_joint, encoded_lengths=encoded_len, - return_hypotheses=False, + return_hypotheses=tdt_timestamps, ) - first = hyps[0] if isinstance(hyps, (list, tuple)) and len(hyps) > 0 else hyps + first = first_hypothesis(hyps) text = first if isinstance(first, str) else getattr(first, "text", str(first)) print(f" Transcription: {text}") - write_transcript(out_dir, text, source=source) + if tdt_timestamps: + words, tokens = nemo_timestamp_rows(first) + write_transcript( + out_dir, text, source=source, words=words, tokens=tokens + ) + else: + write_transcript(out_dir, text, source=source) return 0 - transcriptions = model.transcribe(audio=[str(audio_path)], batch_size=1) - if isinstance(transcriptions, (list, tuple)) and len(transcriptions) > 0: - first = transcriptions[0] - if isinstance(first, (list, tuple)): - first = first[0] - text = first if isinstance(first, str) else first.text + if tdt_timestamps: + transcriptions = model.transcribe( + audio=[str(audio_path)], + batch_size=1, + return_hypotheses=True, + timestamps=True, + ) else: - text = str(transcriptions) + transcriptions = model.transcribe( + audio=[str(audio_path)], batch_size=1 + ) + first = first_hypothesis(transcriptions) + text = first if isinstance(first, str) else getattr(first, "text", str(first)) print(f" Transcription: {text}") - write_transcript(out_dir, text, source=source) + if tdt_timestamps: + words, tokens = nemo_timestamp_rows(first) + write_transcript( + out_dir, text, source=source, words=words, tokens=tokens + ) + else: + write_transcript(out_dir, text, source=source) return 0 diff --git a/scripts/lib/ref_dump.py b/scripts/lib/ref_dump.py index c484ea76..08282631 100644 --- a/scripts/lib/ref_dump.py +++ b/scripts/lib/ref_dump.py @@ -101,18 +101,21 @@ def write_transcript( text: str, *, source: dict[str, Any], - tokens: list[int] | None = None, + tokens: list[int] | list[dict[str, Any]] | None = None, + words: list[dict[str, Any]] | None = None, ) -> None: """Write transcript.json alongside a decode dump. transcript.json is the behavioral artifact: the text the reference - framework produced for this audio, plus optional token IDs. Used by - validate.py to compare against the C++ transcript at the public - boundary (not just per-tensor). + framework produced for this audio, plus optional token IDs or timestamp + rows. Used by validate.py to compare against the C++ transcript at the + public boundary (not just per-tensor). """ out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) payload: dict[str, Any] = {"text": text, "source": source} + if words is not None: + payload["words"] = list(words) if tokens is not None: payload["tokens"] = list(tokens) (out_dir / "transcript.json").write_text(json.dumps(payload, indent=2) + "\n") diff --git a/scripts/lib/test_validate.py b/scripts/lib/test_validate.py new file mode 100644 index 00000000..edd5e34f --- /dev/null +++ b/scripts/lib/test_validate.py @@ -0,0 +1,182 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy"] +# /// +"""Focused unit tests for validate.py transcript and timestamp plumbing. + +Run standalone: uv run scripts/lib/test_validate.py +Run with pytest: uv run --with pytest --with numpy python -m pytest scripts/lib/test_validate.py +""" + +from __future__ import annotations + +import importlib.util +import json +import tempfile +from pathlib import Path +from types import SimpleNamespace + + +_VALIDATE_PATH = Path(__file__).resolve().parents[1] / "validate.py" +_SPEC = importlib.util.spec_from_file_location("transcribe_validate", _VALIDATE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +validate = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(validate) + +_PARAKEET_DUMPER_PATH = ( + Path(__file__).resolve().parents[1] / "dump_reference_parakeet_nemo.py" +) +_PARAKEET_DUMPER_SPEC = importlib.util.spec_from_file_location( + "transcribe_parakeet_dumper", _PARAKEET_DUMPER_PATH +) +assert _PARAKEET_DUMPER_SPEC is not None and _PARAKEET_DUMPER_SPEC.loader is not None +parakeet_dumper = importlib.util.module_from_spec(_PARAKEET_DUMPER_SPEC) +_PARAKEET_DUMPER_SPEC.loader.exec_module(parakeet_dumper) + + +def test_parse_cli_jsonl_result() -> None: + output = "\n".join([ + "backend log", + '{"type":"batch_header","load_ms":12.5}', + ( + '{"file":"jfk.wav","text":"hello","words":' + '[{"t0_ms":80,"t1_ms":240,"text":"hello"}],"tokens":' + '[{"t0_ms":80,"t1_ms":240,"text":" hello","id":7,"p":0.9}]}' + ), + ]) + result = validate.parse_cli_result(output) + assert result is not None + assert result["text"] == "hello" + assert result["words"][0]["t0_ms"] == 80 + assert result["tokens"][0]["id"] == 7 + + +def test_cpp_transcript_v2_carries_timestamp_rows() -> None: + words = [{"t0_ms": 80, "t1_ms": 240, "text": "hello"}] + tokens = [ + { + "t0_ms": 80, + "t1_ms": 240, + "text": " hello", + "id": 7, + "p": 0.9, + } + ] + with tempfile.TemporaryDirectory() as raw_dir: + out_dir = Path(raw_dir) + validate.write_cpp_transcript( + out_dir, + family="parakeet", + variant="unit", + case="jfk", + gguf=Path("unit.gguf"), + backend="cpu", + text="hello", + words=words, + tokens=tokens, + ) + payload = json.loads((out_dir / "transcript.json").read_text()) + assert payload["schema"] == "transcribe-cpp-transcript-v2" + assert payload["words"] == words + assert payload["tokens"] == tokens + + +def test_nemo_timestamp_rows_extracts_char_token_text() -> None: + hypothesis = SimpleNamespace( + timestamp={ + "word": [ + {"start": 0.0, "end": 0.8, "word": "And so,"}, + ], + "char": [ + {"start": 0.0, "end": 0.32, "char": ["And"]}, + {"start": 0.72, "end": 0.8, "char": ["so"]}, + {"start": 0.8, "end": 0.8, "char": [","]}, + ], + } + ) + + words, tokens = parakeet_dumper.nemo_timestamp_rows(hypothesis) + + assert words == [{"start_s": 0.0, "end_s": 0.8, "text": "And so,"}] + assert tokens == [ + {"start_s": 0.0, "end_s": 0.32, "text": "And"}, + {"start_s": 0.72, "end_s": 0.8, "text": "so"}, + {"start_s": 0.8, "end_s": 0.8, "text": ","}, + ] + + +def test_timestamp_compare_normalizes_and_reports_stats() -> None: + reference = { + "words": [ + {"start_s": 0.08, "end_s": 0.48, "text": "Hello,"}, + {"start_s": 0.56, "end_s": 1.0, "text": "WORLD"}, + ] + } + cpp = { + "words": [ + {"t0_ms": 120, "t1_ms": 520, "text": "hello"}, + {"t0_ms": 600, "t1_ms": 1040, "text": "world!"}, + ] + } + result = validate.case_timestamp_compare( + reference, cpp, text_mode="normalized", tolerance_ms=80.0 + ) + assert result["match"] is True + assert result["max_deviation_ms"] == 40.0 + assert result["mean_deviation_ms"] == 40.0 + + +def test_timestamp_compare_reports_count_mismatch() -> None: + reference = { + "words": [ + {"start_s": 0.0, "end_s": 0.08, "text": "one"}, + {"start_s": 0.08, "end_s": 0.16, "text": "two"}, + ] + } + cpp = {"words": [{"t0_ms": 0, "t1_ms": 80, "text": "one"}]} + result = validate.case_timestamp_compare( + reference, cpp, text_mode="exact", tolerance_ms=80.0 + ) + assert result["match"] is False + assert "word count mismatch: reference=2, C++=1" in result["reason"] + + +def test_timestamp_compare_enforces_endpoint_tolerance() -> None: + reference = { + "words": [{"start_s": 0.0, "end_s": 0.08, "text": "one"}] + } + cpp = {"words": [{"t0_ms": 0, "t1_ms": 161, "text": "one"}]} + result = validate.case_timestamp_compare( + reference, cpp, text_mode="exact", tolerance_ms=80.0 + ) + assert result["match"] is False + assert result["max_deviation_ms"] == 81.0 + assert "exceeds 80.000 ms" in result["reason"] + + +_TESTS = [ + test_parse_cli_jsonl_result, + test_cpp_transcript_v2_carries_timestamp_rows, + test_nemo_timestamp_rows_extracts_char_token_text, + test_timestamp_compare_normalizes_and_reports_stats, + test_timestamp_compare_reports_count_mismatch, + test_timestamp_compare_enforces_endpoint_tolerance, +] + + +def main() -> int: + failures = 0 + for test in _TESTS: + try: + test() + except AssertionError as exc: + failures += 1 + print(f"FAIL {test.__name__}: {exc}") + else: + print(f"ok {test.__name__}") + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/validate.py b/scripts/validate.py index a5987261..49f0ebc2 100644 --- a/scripts/validate.py +++ b/scripts/validate.py @@ -40,6 +40,7 @@ import datetime as dt import glob import json +import math import os import re import shutil @@ -178,6 +179,14 @@ def dediarize_text(text: str) -> str: return normalize_text_for_compare(stripped) +def text_for_compare(text: str, mode: str) -> str: + if mode == "exact": + return text + if mode == "dediarized": + return dediarize_text(text) + return normalize_text_for_compare(text) + + def find_gguf(repo: Path, family: str, slug: str | None = None, variant: str | None = None) -> Path: """Find a GGUF under models/. @@ -303,13 +312,30 @@ def run_cmd(cmd: list[str], repo: Path, label: str) -> None: ) -def parse_cli_transcript(output: str) -> str | None: +def parse_cli_result(output: str) -> dict[str, Any] | None: + for line in reversed(output.splitlines()): + try: + payload = json.loads(line) + except json.JSONDecodeError: + continue + if ( + isinstance(payload, dict) + and payload.get("type") != "batch_header" + and isinstance(payload.get("text"), str) + ): + return payload + for line in reversed(output.splitlines()): if line.startswith("text: "): - return line[len("text: "):] + return {"text": line[len("text: "):], "words": []} return None +def parse_cli_transcript(output: str) -> str | None: + payload = parse_cli_result(output) + return str(payload["text"]) if payload is not None else None + + def write_cpp_transcript( out_dir: Path, *, @@ -319,25 +345,169 @@ def write_cpp_transcript( gguf: Path, backend: str, text: str, + words: list[dict[str, Any]], + tokens: list[dict[str, Any]] | None, ) -> None: path = out_dir / "transcript.json" payload = { - "schema": "transcribe-cpp-transcript-v1", + "schema": "transcribe-cpp-transcript-v2", "family": family, "variant": variant, "case": case, "text": text, "normalized_text": normalize_text(text), + "words": words, "source": { "kind": "transcribe.cpp", "gguf": str(gguf), "backend": backend, }, } + if tokens is not None: + payload["tokens"] = tokens path.write_text(json.dumps(payload, indent=2) + "\n") print(f" wrote {path}", file=sys.stderr) +def timestamp_tolerance_ms(path: Path | None) -> float | None: + if path is None: + return None + payload = json.loads(path.read_text()) + entry = payload.get("timestamps") + if entry is None: + return None + if not isinstance(entry, dict) or "max_abs_ms" not in entry: + raise SystemExit( + f"error: {path}: timestamps must contain max_abs_ms" + ) + value = float(entry["max_abs_ms"]) + if not math.isfinite(value) or value < 0.0: + raise SystemExit( + f"error: {path}: timestamps.max_abs_ms must be finite and non-negative" + ) + return value + + +def case_timestamp_compare( + reference: dict[str, Any], + cpp: dict[str, Any], + *, + text_mode: str, + tolerance_ms: float | None, +) -> dict[str, Any]: + ref_words = reference.get("words") + cpp_words = cpp.get("words") + if not isinstance(ref_words, list): + return {"match": False, "reason": "reference words field is not a list"} + if not isinstance(cpp_words, list): + return {"match": False, "reason": "C++ words field is not a list"} + if tolerance_ms is None: + return { + "match": False, + "reason": "timestamp tolerance is missing from the tolerance file", + } + + def word_texts(rows: list[Any], side: str) -> tuple[list[str] | None, str | None]: + texts: list[str] = [] + for index, row in enumerate(rows): + if not isinstance(row, dict) or "text" not in row: + return None, f"{side} word {index} lacks a text field" + texts.append(str(row["text"])) + return texts, None + + ref_texts, error = word_texts(ref_words, "reference") + if error is not None: + return {"match": False, "reason": error} + cpp_texts, error = word_texts(cpp_words, "C++") + if error is not None: + return {"match": False, "reason": error} + assert ref_texts is not None and cpp_texts is not None + + ref_aligned = [text_for_compare(text, text_mode) for text in ref_texts] + cpp_aligned = [text_for_compare(text, text_mode) for text in cpp_texts] + if len(ref_words) != len(cpp_words): + divergent_index = next( + ( + index + for index, (ref_text, cpp_text) in enumerate( + zip(ref_aligned, cpp_aligned) + ) + if ref_text != cpp_text + ), + min(len(ref_aligned), len(cpp_aligned)), + ) + window_start = max(0, divergent_index - 2) + window_end = divergent_index + 3 + return { + "match": False, + "reason": ( + f"word count mismatch: reference={len(ref_words)}, " + f"C++={len(cpp_words)}; first divergent index={divergent_index}; " + f"reference[{window_start}:{min(window_end, len(ref_aligned))}]=" + f"{ref_aligned[window_start:window_end]!r}; " + f"C++[{window_start}:{min(window_end, len(cpp_aligned))}]=" + f"{cpp_aligned[window_start:window_end]!r}" + ), + } + + for index, (ref_text, cpp_text) in enumerate(zip(ref_aligned, cpp_aligned)): + if ref_text != cpp_text: + return { + "match": False, + "reason": ( + f"word text mismatch at index {index}: " + f"reference={ref_texts[index]!r} ({ref_text!r}), " + f"C++={cpp_texts[index]!r} ({cpp_text!r})" + ), + } + + starts: list[float] = [] + ends: list[float] = [] + for index, (ref_word, cpp_word) in enumerate(zip(ref_words, cpp_words)): + try: + ref_t0_ms = float(ref_word["start_s"]) * 1000.0 + ref_t1_ms = float(ref_word["end_s"]) * 1000.0 + cpp_t0_ms = float(cpp_word["t0_ms"]) + cpp_t1_ms = float(cpp_word["t1_ms"]) + except (KeyError, TypeError, ValueError) as exc: + return { + "match": False, + "reason": f"invalid timestamp fields at word {index}: {exc}", + } + values = (ref_t0_ms, ref_t1_ms, cpp_t0_ms, cpp_t1_ms) + if not all(math.isfinite(value) for value in values): + return { + "match": False, + "reason": f"non-finite timestamp at word {index}: {values!r}", + } + starts.append(abs(cpp_t0_ms - ref_t0_ms)) + ends.append(abs(cpp_t1_ms - ref_t1_ms)) + + deviations = starts + ends + max_deviation_ms = max(deviations, default=0.0) + mean_deviation_ms = ( + sum(deviations) / len(deviations) if deviations else 0.0 + ) + within_tolerance = max_deviation_ms <= tolerance_ms + result = { + "match": within_tolerance, + "n_words": len(ref_words), + "tolerance_ms": tolerance_ms, + "max_deviation_ms": max_deviation_ms, + "mean_deviation_ms": mean_deviation_ms, + "max_start_deviation_ms": max(starts, default=0.0), + "mean_start_deviation_ms": sum(starts) / len(starts) if starts else 0.0, + "max_end_deviation_ms": max(ends, default=0.0), + "mean_end_deviation_ms": sum(ends) / len(ends) if ends else 0.0, + } + if not result["match"]: + result["reason"] = ( + f"max endpoint deviation {max_deviation_ms:.3f} ms exceeds " + f"{tolerance_ms:.3f} ms" + ) + return result + + def cmd_ref(args: argparse.Namespace) -> int: repo = find_repo_root(Path(__file__).parent) manifest = load_manifest(repo, args.family, getattr(args, "variant", None)) @@ -440,6 +610,35 @@ def cmd_cpp(args: argparse.Namespace) -> int: env = os.environ.copy() env["TRANSCRIBE_DUMP_DIR"] = str(out_dir) + ref_transcript = ( + repo / "build" / "validate" / args.family / variant + / case_name / "ref" / "transcript.json" + ) + reference_data: dict[str, Any] = {} + if ref_transcript.exists(): + reference_data = json.loads(ref_transcript.read_text()) + else: + print( + " note: reference transcript artifact is absent; the CLI dump " + "is running without reference-informed timestamp mode, and the " + "later compare will report a word-count mismatch.", + file=sys.stderr, + ) + reference_words = reference_data.get("words") + reference_tokens = reference_data.get("tokens") + has_reference_words = isinstance(reference_words, list) + has_reference_tokens = ( + isinstance(reference_tokens, list) + and all( + isinstance(row, dict) + and {"start_s", "end_s", "text"}.issubset(row) + for row in reference_tokens + ) + ) + timestamp_request = None + if has_reference_words: + timestamp_request = "token" if has_reference_tokens else "word" + # Whisper: by default, exercise the production C++ MelFrontend so # the per-tensor compare covers the full mel→encoder→decoder # pipeline. The env-var ref-mel injection is preserved as an @@ -474,6 +673,9 @@ def cmd_cpp(args: argparse.Namespace) -> int: cmd += ["--language", language] if args.family == "whisper": cmd += ["--timestamps", "none"] + timestamp_request = None + elif timestamp_request is not None: + cmd += ["--timestamps", timestamp_request] if args.family in ("sensevoice", "parakeet"): # The reference dumper emits the raw token stream including # control / language tags (sensevoice: language / event / @@ -482,7 +684,15 @@ def cmd_cpp(args: argparse.Namespace) -> int: # strips these by default, so the validate dump must pass # --raw-tokens to keep them and match the reference exactly. cmd += ["--raw-tokens"] - cmd.append(str(audio)) + + batch_tmp = None + if timestamp_request is not None: + batch_tmp = tempfile.TemporaryDirectory(prefix="transcribe-validate-") + batch_list = Path(batch_tmp.name) / "audio.txt" + batch_list.write_text(f"{audio}\n") + cmd += ["--batch", str(batch_list), "--batch-jsonl"] + else: + cmd.append(str(audio)) print(f"\n{'=' * 60}", file=sys.stderr) print(f" cpp dump [{args.family}/{case_name}]", file=sys.stderr) @@ -490,15 +700,19 @@ def cmd_cpp(args: argparse.Namespace) -> int: print(f" {' '.join(cmd)}", file=sys.stderr) print(f"{'=' * 60}", file=sys.stderr) - result = subprocess.run( - cmd, - cwd=repo, - env=env, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - errors="replace", - ) + try: + result = subprocess.run( + cmd, + cwd=repo, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + errors="replace", + ) + finally: + if batch_tmp is not None: + batch_tmp.cleanup() if result.stdout: print(result.stdout, end="") if result.returncode != 0: @@ -506,11 +720,21 @@ def cmd_cpp(args: argparse.Namespace) -> int: f"error: cpp dump [{args.family}/{case_name}] failed " f"with exit code {result.returncode}" ) - transcript = parse_cli_transcript(result.stdout or "") - if transcript is None: + cli_result = parse_cli_result(result.stdout or "") + if cli_result is None: raise SystemExit( f"error: cpp dump [{args.family}/{case_name}] did not emit a transcript line" ) + words = cli_result.get("words", []) + tokens = cli_result.get("tokens") + if not isinstance(words, list): + raise SystemExit( + f"error: cpp dump [{args.family}/{case_name}] emitted a non-list words field" + ) + if tokens is not None and not isinstance(tokens, list): + raise SystemExit( + f"error: cpp dump [{args.family}/{case_name}] emitted a non-list tokens field" + ) write_cpp_transcript( out_dir, family=args.family, @@ -518,7 +742,9 @@ def cmd_cpp(args: argparse.Namespace) -> int: case=case_name, gguf=gguf, backend=args.backend, - text=transcript, + text=str(cli_result["text"]), + words=words, + tokens=tokens, ) return 0 @@ -547,11 +773,13 @@ def cmd_compare(args: argparse.Namespace) -> int: compare_script = repo / "scripts" / "compare_tensors.py" report_mode = getattr(args, "report", False) + timestamp_tolerance = timestamp_tolerance_ms(tolerances) cases = manifest.get("cases", ["jfk"]) all_passed = True compare_outputs: list[dict[str, Any]] = [] transcript_results: list[dict[str, Any]] = [] + timestamp_results: list[dict[str, Any]] = [] cmd_log: list[dict[str, Any]] = [] for case in cases: @@ -616,19 +844,18 @@ def cmd_compare(args: argparse.Namespace) -> int: "case": case_name, "match": False, "reason": "missing C++ transcript artifact", }) + if "words" in ref_data: + timestamp_results.append({ + "case": case_name, + "match": False, + "reason": "missing C++ transcript artifact", + }) continue cpp_data = json.loads(cpp_transcript.read_text()) cpp_text = str(cpp_data.get("text", "")) - if transcript_compare == "exact": - ref_compare = ref_text - cpp_compare = cpp_text - elif transcript_compare == "dediarized": - ref_compare = dediarize_text(ref_text) - cpp_compare = dediarize_text(cpp_text) - else: - ref_compare = normalize_text_for_compare(ref_text) - cpp_compare = normalize_text_for_compare(cpp_text) + ref_compare = text_for_compare(ref_text, transcript_compare) + cpp_compare = text_for_compare(cpp_text, transcript_compare) match = cpp_compare == ref_compare transcript_results.append({ "case": case_name, "match": match, @@ -646,6 +873,37 @@ def cmd_compare(args: argparse.Namespace) -> int: else: print(f"\n Transcript: ok ({transcript_compare}) {cpp_text!r}") + if "words" in ref_data: + timestamp_result = case_timestamp_compare( + ref_data, + cpp_data, + text_mode=transcript_compare, + tolerance_ms=timestamp_tolerance, + ) + timestamp_result = { + "case": case_name, + "mode": transcript_compare, + **timestamp_result, + } + timestamp_results.append(timestamp_result) + if timestamp_result["match"]: + print( + " Word timestamps: ok " + f"({timestamp_result['n_words']} words, " + f"max={timestamp_result['max_deviation_ms']:.3f} ms, " + f"mean={timestamp_result['mean_deviation_ms']:.3f} ms, " + f"tolerance={timestamp_result['tolerance_ms']:.3f} ms)" + ) + else: + print("\nFAIL word timestamp mismatch") + print(f" {timestamp_result['reason']}") + if "max_deviation_ms" in timestamp_result: + print( + f" max={timestamp_result['max_deviation_ms']:.3f} ms, " + f"mean={timestamp_result['mean_deviation_ms']:.3f} ms" + ) + all_passed = False + if report_mode: write_report_bundle( repo=repo, @@ -653,6 +911,7 @@ def cmd_compare(args: argparse.Namespace) -> int: variant=variant, compare_outputs=compare_outputs, transcript_results=transcript_results, + timestamp_results=timestamp_results, cmd_log=cmd_log, overall_passed=all_passed, ) @@ -765,16 +1024,17 @@ def write_report_bundle( variant: str, compare_outputs: list[dict[str, Any]], transcript_results: list[dict[str, Any]], + timestamp_results: list[dict[str, Any]], cmd_log: list[dict[str, Any]], overall_passed: bool, ) -> Path: """Write an ephemeral validation report bundle. The bundle captures what's not reproducible from git alone: the compare - stdout, exact command invocations, and transcript comparison at this - moment. The manifest, tolerance file, and compare_tensors.py live in the - repo — pin them via `validated_at_sha` (the git HEAD at bundle time) - rather than duplicating them here. + stdout, exact command invocations, transcript comparison, and timestamp + comparison at this moment. The manifest, tolerance file, and + compare_tensors.py live in the repo — pin them via `validated_at_sha` (the + git HEAD at bundle time) rather than duplicating them here. """ now = dt.datetime.now(dt.UTC) ts = now.strftime("%Y%m%dT%H%M%SZ") @@ -825,6 +1085,28 @@ def write_report_bundle( summary.append(f"- reference: `{tr.get('reference', '')!r}`") summary.append(f"- c++: `{tr.get('cpp', '')!r}`") summary.append("") + if timestamp_results: + summary += ["## Word timestamp comparison", ""] + for result in timestamp_results: + summary.append(f"### {result['case']}") + summary.append("") + summary.append( + f"- Match: **{'yes' if result['match'] else 'no'}**" + ) + if result["match"]: + summary.append(f"- Words: `{result['n_words']}`") + summary.append( + f"- Max endpoint deviation: `{result['max_deviation_ms']:.3f} ms`" + ) + summary.append( + f"- Mean endpoint deviation: `{result['mean_deviation_ms']:.3f} ms`" + ) + summary.append( + f"- Tolerance: `{result['tolerance_ms']:.3f} ms`" + ) + else: + summary.append(f"- Reason: {result['reason']}") + summary.append("") (bundle_dir / "summary.md").write_text("\n".join(summary)) repro = f"""# Reproducing this validation run diff --git a/src/arch/parakeet/model.cpp b/src/arch/parakeet/model.cpp index 735f84ae..3c258df3 100644 --- a/src/arch/parakeet/model.cpp +++ b/src/arch/parakeet/model.cpp @@ -24,6 +24,7 @@ #include "transcribe-log.h" #include "transcribe-mel.h" #include "transcribe-meta.h" +#include "transcribe-unicode.h" #include "transcribe/parakeet.h" #include "weights.h" @@ -382,6 +383,8 @@ void reset_streaming_decoder_state(ParakeetSession * pc, const ParakeetModel * p pc->stream_dec_state.initialized = true; } +static std::vector nemo_supported_punctuation(const transcribe::Tokenizer & tok); + // 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 **); @@ -444,6 +447,9 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par if (const transcribe_status st = read_parakeet_hparams(loader.gguf(), m->hparams); st != TRANSCRIBE_OK) { return st; } + if (m->hparams.head_kind == HeadKind::TDT) { + m->nemo_supported_punctuation = nemo_supported_punctuation(m->tok); + } // Derive supports_streaming from hparams: // ChunkedLimited + (L, R) >= 0 — cache-aware streaming @@ -633,6 +639,135 @@ static bool is_strippable_special(const transcribe::Tokenizer & tok, int id) { return tok.is_control(id) || is_lang_tag_piece(tok.token(id)); } +constexpr const char k_sentencepiece_marker[] = "\xE2\x96\x81"; +constexpr int k_sentencepiece_marker_len = 3; + +static bool starts_with_sentencepiece_marker(const std::string & piece) { + return piece.size() >= static_cast(k_sentencepiece_marker_len) && + std::memcmp(piece.data(), k_sentencepiece_marker, k_sentencepiece_marker_len) == 0; +} + +static bool is_valid_utf8(const std::string & text) { + size_t offset = 0; + while (offset < text.size()) { + const uint8_t lead = static_cast(text[offset]); + if (lead <= 0x7Fu) { + ++offset; + continue; + } + + size_t width = 0; + uint32_t cpt = 0; + uint32_t min = 0; + if ((lead & 0xE0u) == 0xC0u) { + width = 2; + cpt = lead & 0x1Fu; + min = 0x80u; + } else if ((lead & 0xF0u) == 0xE0u) { + width = 3; + cpt = lead & 0x0Fu; + min = 0x800u; + } else if ((lead & 0xF8u) == 0xF0u) { + width = 4; + cpt = lead & 0x07u; + min = 0x10000u; + } else { + return false; + } + if (width > text.size() - offset) { + return false; + } + for (size_t i = 1; i < width; ++i) { + const uint8_t next = static_cast(text[offset + i]); + if ((next & 0xC0u) != 0x80u) { + return false; + } + cpt = (cpt << 6) | (next & 0x3Fu); + } + if (cpt < min || cpt > 0x10FFFFu || (cpt >= 0xD800u && cpt <= 0xDFFFu)) { + return false; + } + offset += width; + } + return true; +} + +static bool is_nemo_punctuation_special_piece(const std::string & piece) { + const size_t n = piece.size(); + if ((n >= 2 && piece.front() == '[' && piece.back() == ']') || + (n >= 2 && piece.front() == '<' && piece.back() == '>') || (n >= 2 && piece[0] == '#' && piece[1] == '#') || + starts_with_sentencepiece_marker(piece)) { + return true; + } + + size_t offset = 0; + while (offset < n) { + const uint32_t cpt = unicode::cpt_from_utf8(piece, offset); + if (!unicode::flags_from_cpt(cpt).is_whitespace()) { + return false; + } + } + return true; +} + +static std::vector nemo_supported_punctuation(const transcribe::Tokenizer & tok) { + std::vector punctuation; + for (int id = 0; id < tok.n_tokens(); ++id) { + const std::string & piece = tok.token(id); + if (!is_valid_utf8(piece) || is_nemo_punctuation_special_piece(piece)) { + continue; + } + + size_t offset = 0; + while (offset < piece.size()) { + const uint32_t cpt = unicode::cpt_from_utf8(piece, offset); + if ((unicode::flags_from_cpt(cpt).bits & unicode::CptFlags::PUNCTUATION) == 0) { + continue; + } + bool seen = false; + for (uint32_t existing : punctuation) { + if (existing == cpt) { + seen = true; + break; + } + } + if (!seen) { + punctuation.push_back(cpt); + } + } + } + return punctuation; +} + +static bool is_nemo_supported_punctuation_token(const std::string & piece, + const std::string & text, + const std::vector & supported_punctuation) { + if (!is_valid_utf8(piece) || !is_valid_utf8(text)) { + return false; + } + + size_t offset = 0; + // SentencePiece removes a leading word-boundary marker when NeMo + // decodes one token in isolation; our decoder exposes it as a space. + if (starts_with_sentencepiece_marker(piece) && !text.empty() && text.front() == ' ') { + offset = 1; + } + if (offset >= text.size()) { + return false; + } + + const uint32_t cpt = unicode::cpt_from_utf8(text, offset); + if (offset != text.size()) { + return false; + } + for (uint32_t supported : supported_punctuation) { + if (supported == cpt) { + return true; + } + } + return false; +} + // Collapse runs of ASCII spaces to one and trim both ends — cleans up the // double/edge spaces a stripped tag leaves behind. Shared by both builders. static void normalize_transcript_whitespace(std::string & s) { @@ -725,21 +860,23 @@ static transcribe_status decode_and_populate(ParakeetSession * pc, // v2/v3). Shared with the streaming builder. const double frame_to_ms = parakeet_ms_per_enc_frame(pm->hparams); - // SentencePiece word-boundary marker U+2581 ("▁", UTF-8 E2 96 81): a - // raw piece starting with it opens a new word. We use the raw NeMo - // piece (its leading 3 bytes are unambiguous), not the post-decode - // "▁ → space" form. - constexpr const char k_sp_marker[] = "\xE2\x96\x81"; - constexpr int k_sp_marker_len = 3; - const transcribe::Tokenizer & tok = pm->tok; + // NeMo 2.8.0rc0 (6967f48fda2a) parts/utils/tokenizer_utils.py builds + // this set from Unicode P* code points. parts/submodules/rnnt_decoding.py + // makes a supported punctuation token zero-width at the previous end. + const std::vector & supported_punctuation = pm->nemo_supported_punctuation; + pc->tokens.reserve(pc->raw_tokens.size()); // Strip multilingual language tags by default (gated on // keep_special_tags / CLI --raw-tokens). Detection is // Tokenizer::is_control (CONTROL token_type); is_lang_tag_piece is a // transitional fallback for GGUFs predating that converter change. const bool strip_tags = (params == nullptr) ? true : !params->keep_special_tags; + // NeMo indexes the full non-blank sequence; this loop indexes tokens kept + // after control/language-tag stripping. They match for v2 (which has no + // strippable tokens) and validation's --raw-tokens path, but a future + // tag-emitting TDT model would diverge. for (const TdtToken & rt : pc->raw_tokens) { if (strip_tags && is_strippable_special(tok, rt.id)) { continue; @@ -752,6 +889,11 @@ static transcribe_status decode_and_populate(ParakeetSession * pc, const double end_frame = static_cast(rt.step_at_emit) + static_cast(rt.duration_frames); te.t1_ms = static_cast(std::llround(frame_to_ms * end_frame)); te.text = tok.decode(&rt.id, 1); + if (pm->host_decoder.head_kind == HostHeadKind::TDT && !pc->tokens.empty() && + is_nemo_supported_punctuation_token(tok.token(rt.id), te.text, supported_punctuation)) { + te.t0_ms = pc->tokens.back().t1_ms; + te.t1_ms = te.t0_ms; + } pc->tokens.push_back(std::move(te)); } @@ -787,8 +929,7 @@ static transcribe_status decode_and_populate(ParakeetSession * pc, for (size_t i = 0; i < pc->tokens.size(); ++i) { const auto & tk = pc->tokens[i]; const auto & raw_piece = tok.token(tk.id); // empty if id OOR - const bool starts_word = (i == 0) || (raw_piece.size() >= static_cast(k_sp_marker_len) && - std::memcmp(raw_piece.data(), k_sp_marker, k_sp_marker_len) == 0); + const bool starts_word = (i == 0) || starts_with_sentencepiece_marker(raw_piece); if (starts_word) { open_new_word(static_cast(i), tk); } @@ -2040,6 +2181,10 @@ void rebuild_streaming_result_text(ParakeetSession * pc, const ParakeetModel * p const double frame_to_ms = parakeet_ms_per_enc_frame(pm->hparams); + // The punctuation clamp intentionally stays in the offline decode path. + // Buffered-streaming timings keep raw TDT durations because NeMo's + // buffered-streaming path does not run TDT timestamp refinement either. + // Strip CONTROL / tag pieces from the public result, mirroring // decode_and_populate (gated on keep_special_tags). Only the public // projection is filtered; pc->raw_tokens stays whole. diff --git a/src/arch/parakeet/parakeet.h b/src/arch/parakeet/parakeet.h index 0574e544..4cb2c42c 100644 --- a/src/arch/parakeet/parakeet.h +++ b/src/arch/parakeet/parakeet.h @@ -73,10 +73,12 @@ void apply_family_invariants(transcribe_model & model); // data buffer; the destructor frees it, invalidating every borrowed // ggml_tensor* in `weights`. struct ParakeetModel final : public transcribe_model { - Tokenizer tok; - ParakeetHParams hparams; - ParakeetWeights weights; - ggml_context * ctx_meta = nullptr; + Tokenizer tok; + // Derived once from the tokenizer vocabulary during load. + std::vector nemo_supported_punctuation; + ParakeetHParams hparams; + ParakeetWeights weights; + ggml_context * ctx_meta = nullptr; // Runtime backend plan, resolved at load() via // load_common::init_backends (see transcribe-backend.h). plan's diff --git a/tests/decoder_smoke.cpp b/tests/decoder_smoke.cpp index 5b038c98..cb1b7677 100644 --- a/tests/decoder_smoke.cpp +++ b/tests/decoder_smoke.cpp @@ -233,8 +233,10 @@ int main() { const int n_tokens = transcribe_n_tokens(ctx); CHECK(n_tokens > 0); - int prev_t0 = -1; - int n_with_p = 0; + int prev_t0 = -1; + int prev_t1 = -1; + int n_with_p = 0; + int n_zero_width_punctuation = 0; for (int i = 0; i < n_tokens; ++i) { transcribe_token tok; transcribe_token_init(&tok); @@ -244,20 +246,31 @@ int main() { CHECK(tok.text != nullptr); CHECK(tok.t0_ms >= 0); CHECK(tok.t1_ms >= tok.t0_ms); - CHECK(tok.t0_ms >= prev_t0); // monotone + CHECK(tok.t0_ms >= prev_t0); // equality is valid at a zero-width punctuation boundary + CHECK(tok.t1_ms >= prev_t1); CHECK_EQ_INT(tok.seg_index, 0); CHECK(tok.word_index >= 0); CHECK(tok.p >= 0.0f && tok.p <= 1.0001f && !std::isnan(tok.p)); if (tok.p > 0.0f) { ++n_with_p; } + const bool is_jfk_punctuation = tok.text != nullptr && tok.text[0] != '\0' && tok.text[1] == '\0' && + std::strchr(",.", tok.text[0]) != nullptr; + if (is_jfk_punctuation) { + CHECK(i > 0); + CHECK_EQ_INT(tok.t0_ms, prev_t1); + CHECK_EQ_INT(tok.t1_ms, tok.t0_ms); + ++n_zero_width_punctuation; + } prev_t0 = static_cast(tok.t0_ms); + prev_t1 = static_cast(tok.t1_ms); } // At least most tokens should carry a positive confidence (the // entropy-based formula gives ~0 for a uniform distribution and // ~1 for a confident decode; on jfk.wav nearly every token is // emitted with very high confidence). CHECK(n_with_p > n_tokens / 2); + CHECK(n_zero_width_punctuation > 0); // ---- Per-word sanity ------------------------------------------- const int n_words = transcribe_n_words(ctx); diff --git a/tests/tolerances/parakeet.json b/tests/tolerances/parakeet.json index dea4f214..187190cb 100644 --- a/tests/tolerances/parakeet.json +++ b/tests/tolerances/parakeet.json @@ -120,6 +120,10 @@ "matches NeMo's [257,1024] shape, drift is well inside tolerance,", "and the JFK transcript matches exactly." ], + "timestamps": { + "_comment": "One encoder frame: 80 ms from 8x subsampling and a 0.01 s frontend hop. Pinned NeMo 2.8.0rc0 (6967f48fda2a) AbstractRNNTDecoding.compute_rnnt_timestamps computes a raw TDT token end as emission offset + token_duration. Its post-processing makes punctuation-only tokens in NeMo's supported set zero-length at the previous token's end, except for the first token; words inherit the clamped end.", + "max_abs_ms": 80.0 + }, "enc.mel.in": { "max_abs": 10.0, "mean_abs": 0.01