feat: nemotron-mlx ASR transducer backend (depends on #423) - #426
Open
clkao wants to merge 34 commits into
Open
Conversation
The model occasionally emits <｜hy_place▁holder▁no▁2｜> (U+FF5C pipes) and hallucinates free text after it. Display layers (tui, overlay) stripped the token, but the terminal transcript and file writer still showed the hallucinated tail. Port livecaption's source-level fix: _strip_hy_placeholder truncates at the FIRST placeholder occurrence (cut everything from it onward), applied to _translate_text output.
The model occasionally emits <｜hy_place▁holder▁no▁2｜> and hallucinates free text after it. The previous fix truncates the string after decode — correct output, but the hallucinated tail is still fully decoded. Now the stream_generate loop stops in-loop the moment the placeholder is emitted. The stop condition is resolved from the tokenizer at call time (_placeholder_stop_check): a single-id placeholder (Hy-MT2-1.8B: id 120020) stops exactly on that id; a fragmented placeholder (Hunyuan-MT-7B: ~13 byte-level BPE ids) uses a rolling id window. Tokenizers that fragment it beyond a 16-id window get no in-loop stop and keep the post-hoc string strip as the only mitigation. _strip_hy_placeholder stays as the defensive fallback for all consumers. The helper is module-level so the simul-MT subclass (PR QuentinFuxa#423 branch) can reuse the same mechanism at its commit/release sites. Tests: 4 new stub-stream tests (single-id stop, fragmented stop, no-stop-id fallback to string strip, clean output unaffected). 16/16 in tests/test_mlx_llm_mt.py.
- MlxLlmTranslation.new_session(target_language) creates a per-session client that shares the model cache but has fresh buffer/pending/metrics state; online_translation_factory and session_translation_factory both route through it so concurrent sessions don't cross-contaminate. - Define [tool.uv].conflicts for mlx-llm-mt vs qwen3-streaming/vllm (transformers>=5 vs ==4.57.6 pin); refresh uv.lock. - BenchmarkReport.has_wer is a @Property (was a method — 'if report.has_wer' was always truthy because the method object is truthy). - Remove mlx-qwen3-asr auto-detection from benchmark/compat.py (no such backend in this repo). - Ruff clean.
Add MlxLlmTranslationSimul subclass that drafts translation over the unstable ASR tail and commits target tokens via attention alignment (calibrated zh→en Hunyuan heads, top head L9/H5). Held tokens release from cached attention without a new MT call when ASR commits the tail. - simul_mt_capture.py: MLX Q/K capture (CapturedAttention) + commit policy - translation_mlx_llm_mt_simul.py: subclass with provisional draft + release - audio_processor.py: forward provisional buffer when no final translation - config.py: mlx_llm_mt_simultaneous flag - core.py: factory routes to Simul when flag set - parse_args.py: --simultaneous flag - cli.py: improved --simultaneous help text for wlk bench - test_mlx_llm_mt_simul.py: 21 tests (subclass, tail, commit, release, wiring) Benchmark (faster-whisper, zh_long, real-time): Base first_final=13.77s Simul first_provisional=10.48s Simul provisional EN arrives ~3.3s before base's first final. 32 tests pass (11 existing + 21 new).
…ivation Add a CALIBRATION_REGISTRY keyed by (model_repo, source_lang, target_lang) in simul_mt_capture.py. MlxLlmTranslationSimul looks up its tuple at init: found → install capture with calibrated heads; not found → silently deactivate (wants_hypothesis_tail=False, delegate to base class, log warning naming the missing tuple). 4bit zh→en is NOT seeded (calibration probe showed 48.9% argmax match vs 8bit; formal promotion gate could not run on MLX-format repo), so 4bit deactivates (translation still works via base). 10 new tests cover the 3-tuple matrix: calibrated (activates), uncalibrated (deactivates), 4bit (deactivates).
…lptext scope leak - CALIBRATION_REGISTRY key changed from fully-qualified repo (mlx-community/Hy-MT2-1.8B-8bit) to normalized model id (hy-mt2-1.8b) by stripping org prefix and quant suffix, so calibration entries are shareable across implementations (MLX, vLLM) and quantizations. - 4bit deactivation moved from missing-key to disabled_quants field on CalibrationEntry (model id matches, but quant is disabled). - Revert --target-language and --reference-translation helptext to PR1 wording; PR2 cli.py diff is now strictly the --simultaneous flag. - Add TODO comment for external-heads-loading refactor (AlignAtt4LLM translation_heads_<model>_<direction>.json pattern). - Add test for _normalize_model_id; update registry tests for new key shape.
The model occasionally emits <｜hy_place▁holder▁no▁2｜> (U+FF5C pipes) and hallucinates free text after it. Display layers (tui, overlay) stripped the token, but the terminal transcript and file writer still showed the hallucinated tail — and the simul commit policy committed draft tokens that included it. Port livecaption's source-level fix: _strip_hy_placeholder truncates at the FIRST placeholder occurrence (cut everything from it onward), applied to _translate_text, _translate_simul, and _release_held outputs.
Rebase result: the simul branch now sits on PR QuentinFuxa#422 head (f923d78) and inherits the module-level _placeholder_stop_check + the base engine's in-loop stop; the simul-side calibration/registry commits are intact. Wire the same stop predicate into the simul decode paths: - _translate_simul: the commit stream breaks the moment the placeholder is emitted (single-id exact stop / rolling id window, tokenizer-resolved at call time), and the token stream itself is truncated at the first placeholder id sequence so the commit policy and the stashed draft the release path reads never contain placeholder tokens; committed_len is clamped to the truncated stream. - _release_held: reads the clean stash (no new code needed); the post-hoc _strip_hy_placeholder calls at both sites stay as the fallback for tokenizers that fragment beyond the window cap. Tests: 3 new stub-stream tests (single-id stop, fragmented stop, clean stash release). Focused: 35 simul + 16 base = 51 passed. Full suite: only pre-existing qwen3-backend-shim failures (verified identical on the unmodified branch).
- simul_mt_capture.py: move mlx.core/mlx.nn imports from module level to inside CapturedAttention.__call__ and install_capture. CapturedAttention no longer inherits from nn.Module at class definition time — install_capture dynamically creates an nn.Module-backed subclass when MLX is available. Module is now collectable without MLX (Linux CI). - tests/test_simul_mt_capture_no_mlx.py: regression test importing simul_mt_capture with mlx modules removed from sys.modules. - Remove unused 'released' assignment (F841) in test_mlx_llm_mt_simul.py. - Fix import sort (I001) in translation_mlx_llm_mt_simul.py. ruff clean, 37/37 tests pass.
…-session The rebase onto PR QuentinFuxa#422 introduced new_session() per-session isolation in online_translation_factory. MlxLlmTranslationSimul inherited the base new_session which returned MlxLlmTranslation, losing the simul type and state. Override new_session to return MlxLlmTranslationSimul so each session gets fresh simul state (tail, committed tokens, draft) sharing the cached model. Update the factory test to check isinstance instead of identity (per-session creates a new instance, not the same object).
clkao
marked this pull request as draft
August 30, 2026 17:55
clkao
force-pushed
the
spacedock-ensign/nemotron-mlx-asr-backend-rebased
branch
from
August 30, 2026 23:34
2671610 to
04cbe9c
Compare
Mass commit mode: commit target tokens whose accessible attention mass on committed source tokens >= threshold (default 0.5), replacing the brittle argmax-only check. Measured best in livecaption A/B (more provisional content + less final lag). Token hysteresis: the MT-call hysteresis threshold is now in source BPE tokens (the MT's own unit) instead of chars, with a rolling chars-per-token ratio to estimate token growth without re-tokenizing. CJK and Latin converge to the same token budget (15 tokens ≈ one short sentence in both). For Latin source (en→zh), this dramatically reduces MT calls (13 vs 30 on a 30s demo) since the old 15-char threshold fired too often for English's ~5 chars/token ratio. Wired commit_mode + mass_threshold through config and core.
…/H5 top) en→zh: calibrated tencent/Hy-MT2-1.8B on 1138 Mxode en-zh pairs. Top head L9/H5 TS=0.86, 3/3 stability splits stable (max delta 0.0086). ja→zh: calibrated on 219 WikiMatrix ja-zh pairs. Top head L9/H5 TS=0.89, 3/3 stability splits stable (max delta 0.0244). All three directions (zh→en, en→zh, ja→zh) share L9/H5 as top head — strong evidence these are general alignment heads for hunyuan_v1_dense. 5/8 heads shared across all three directions. 4bit disabled for both (same attention-divergence reasoning as zh→en). With all three directions seeded, the simul-MT variant now covers the directions the eval harness needs (IWSLT en→zh anchor + zh→en/ja→zh flagship).
new_session dropped commit_mode/mass_threshold when constructing the per-session client, so server-context sessions defaulted to argmax even when mass was configured. Thread them so the configured policy reaches per-session clients.
clkao
marked this pull request as ready for review
August 30, 2026 23:40
…aft starvation fixed Two sources produced fragment finals and starved MT drafts (measured on zh_long: 12 fragment finals vs 6 sentence finals; 42/42 MT drafts released with an empty committed source): 1. Every short clause pause flushed the open utterance as a final — the translation loop flushed on Silence.is_starting regardless of pause length. The flush now fires only when a pause crosses pause_segmentation_seconds, checked at the pause's END, when the duration is known. 2. Every punctuation token closed a simul segment (a final per clause). Endpointing owns closure now (rule2-softmax / rule3): punctuation closes a segment only once it has run simul_soft_max_s (default 4.0); simul_hard_max_s (default 20.0) force-cuts a run-on utterance regardless of punctuation. New config knobs mlx_llm_mt_simul_soft_max_s / mlx_llm_mt_simul_hard_max_s. Measured: 12 fragment finals -> 4 coherent multi-sentence finals, each carrying complete sentences; starved MT drafts 42/42 -> 4/20 (early drafts only). The simul tests that encoded the old contract (a final at any punctuation) are updated to the new one.
clkao
force-pushed
the
spacedock-ensign/nemotron-mlx-asr-backend-rebased
branch
from
September 1, 2026 16:35
04cbe9c to
0126dc9
Compare
…from the MT source CL: 'hyperopia flashed into provisional' — after the soft_max close, the closed segment's draft cache survived; the next process() ran the release path against the STALE cached draft with a SHORTER new source and re-emitted the pre-final translation, so the display regressed: the final showed the complete sentence (with 'hyperopia'), the next draft reverted to it minus the last clause. The close now resets the draft cache like the silence-boundary flush does. Verified on the zh-en capture: the post-final drafts no longer revert (the pre-fix capture's drafts at +16.1s reverted; now they hold the full text). Also: a stale hypothesis tail (text already inside the committed prefix, or its audio range predating the commit boundary — variant spellings defeat exact containment) is dropped from the MT source; feeding it doubled the source (both script variants of the same sentence) and the draft translated the doubled text into a one-word fragment. Translation provisional events are deduped (identical consecutive drafts were re-emitted 5x by the release path).
…when the committed outgrows the cached span
CL: 'the first sentence has no translation draft' — two starvation
mechanisms, both verified with instrumented production runs:
1. The base warmup only covers _translate_text; the first simul draft
paid the Metal kernel compile + capture install inside the translation
loop, blocking every draft for the first sentence (~9s). The simul
warmup now runs one draft at init and resets the state.
2. The release path maps the committed text against the CACHED source
span — only valid while the committed text is a prefix of it. When the
ASR committed words beyond the cached draft's span (the tail grew with
new source words), the boundary mapping fails and the release emits
nothing: the display starved until the segment's final. A fresh draft
is now required whenever the committed text outgrows the cached span.
Measured on zh_long (canonical capture): the first draft lands at 2.27s
(was: nothing until the first final at 9.1s), drafts grow with the
speech ('In eye surgery, lasers...' -> the full sentence; the last
sentence drafts 'Future applications will be more widespread.' instead
of the one-word 'Future' fragment).
…HEAD head override
…-frontier tracking
…Torch head ranking
…n subsystem The live pipeline is not reproducible (ASR/VAD timing jitter), so fixes to the commit policy and the accessible frontier had no deterministic gate. scripts/simul_fixture.py records the translation engine's full input interface from one live run (insert_tokens / insert_silence / validate_buffer_and_reset / process, with exact tokens and timestamps) into tests/golden/simul_zh_long_calls.jsonl, then replays it through a fresh engine with greedy sampling. Verified: zero output diffs across two in-process replays AND across separate processes; the 4 replayed finals match the live captures. Baseline on the shipped policy: draft coverage 0.27 (pass >= 0.6) — the number the frontier + policy fixes must move.
…argmax commit policy Two fixes for simul-MT draft starvation, measured on the deterministic replay fixture (tests/golden/simul_zh_long_calls.jsonl, baseline 0.27): 1. committed_src_end_from_text froze the commit frontier at the first byte-split BPE token: decoding a token carrying half a UTF-8 character yields U+FFFD, which failed the startswith test and stopped the scan. cend stuck at 3 of 21 tokens while the committed text grew 17 -> 37 chars. U+FFFD-tolerant prefix matching unfreezes it; still rounds down to the last complete token. Coverage 0.27 -> 0.54. 2. mode="paper" implements the paper's decision rule (arxiv 2606.03967 §4.4): head-averaged rows over the calibrated head set, per-head prefix-online Welford z-normalization, width-7 median filter along the source axis, stabilized argmax vs the accessible frontier (argmax < cend + b, b=1), mass gates off (tau=0) as in the official operating points. First-failure scan emits the longest accepting prefix. Coverage 0.59 vs 0.54 for the mass gate; final 3 improved 0.10 -> 0.30. scripts/simul_fixture.py replay gains --commit-mode / --min-source-tokens for A/Bs. Lowering the re-draft hysteresis does not help (0.57): the residual gap is the last segment's tail sentence, which the ASR has not committed when the segment closes — the final's quality pass is by construction its first release. Tests: 58 passing (simul + caption events + overlay replay).
config, core wiring, and the simul engine default now use the paper's stabilized-argmax commit policy; argmax and mass remain selectable for comparison. The lc_terminal CLI knob lands with the PR that introduces that script.
The mlx-qwen3-asr backend (pure-MLX Qwen3-ASR, no torch/transformers/WebSocket sidecar) built through a new generalized wrapper layer. The wrapper factors the two jobs every non-transducer ASR duplicates: Job 1 (stable/unstable split for models that revise) in asr_commit.py, Job 2 (timestamp manufacture for forward-emit models) in asr_timestamps.py, composed via asr_wrapper.py. A second provider (Voxtral-MLX) shares the timestamp module.
…guage Finding 1 (AlignAtt seam): get_buffer returned the full rolling text\n(including the committed stable prefix), double-counting the prefix for\nconsumers that read get_buffer (display, and the AlignAtt translator\nwhich drafts over the tail via HypothesisTail). Return text[len(stable):]\n(the unstable tail only), matching WLK's contract (process_iter returns\nthe committed ASRToken list; get_buffer returns the unstable tail). Finding 2 (per-session language): SessionASRProxy delegates __getattr__\nto the shared ASR, so getattr(asr, 'language') returned the server-wide\ndefault, ignoring the per-session ?language= override. Prefer the proxy's\n_session_language when present so a multi-language WLK server transcribes\neach session in its own language.
…n wiring Blocker 1 — get_buffer contract conflict: StableCommitTransform read inner.get_buffer() as the full rolling hypothesis, but get_buffer returns only the unstable tail (WLK contract). The transform could not compute a stable prefix and emitted garbage deltas omitting the committed prefix. Add get_hypothesis() to MlxQwen3AsrOnlineProcessor returning the full rolling text (Transcript with text=self._text); update StableCommitTransform to call inner.get_hypothesis() for the full hypothesis. get_buffer stays the tail (WLK contract, for AlignAtt/HypothesisTail). Add a falsifiable test exercising the transform + get_hypothesis seam (fails when the transform reads get_buffer). Blocker 2 — out-of-scope translation wiring: Re-carve config.py/core.py/parse_args.py to drop the translation-backend wiring that belonged to a different PR (model field, import, translation-backend choice, flag, qwen3+NLLB guard removal). Restore the qwen3+NLLB guard. The PR3 diff is now ASR-only.
_finalize_utterance emitted the full re-decoded text at finalization, but the stable prefix was already emitted during streaming by StableCommitTransform → text appeared twice. Fix: StableCommitTransform now tracks the cumulative committed text on inner._emitted_stable; _finalize_utterance uses _compute_finalize_delta to emit only the uncommitted suffix (or full text when no streaming commit happened, or full corrected text when the re-decode corrected the prefix). 7 new tests including an integration test that exercises the transform + dedup path end-to-end and would fail on pre-fix code (verified: pre-fix produces 'alpha beta alpha beta gamma delta').
… override tests - Add [tool.uv].conflicts for mlx-qwen3-asr vs qwen3-streaming/qwen3-vllm/ qwen3-vllm-metal (transformers>=5 vs ==4.57.6), matching the mlx-llm-mt pattern. Update uv.lock conflicts section to match. - Add two-session ASR isolation regression test: two MlxQwen3AsrOnlineProcessor instances from shared model cache have independent per-session state (_state, _text, _stable_text, _emitted_stable, _utt_audio); model object is shared (loaded once). - Add per-session language override regression test: SessionASRProxy with language='ja' causes the processor to resolve to Japanese, not the server-wide default.
…P fallback The StableCommitTransform re-derived stability from the rolling hypothesis via LCP on split_text_units, which splits on whitespace only. CJK has no whitespace, so a 30s zh transcript collapsed into ONE unit, LCP was always 0, and the transform never committed during streaming (zh->en: 0 provisionals, 1 final at finish). Three changes: 1. Prefer the backend's NATIVE stable prefix (prefer_native_stable=True, default). Backends that expose a monotonic stable field (mlx-qwen3-asr's _stable_text) commit it directly — no LCP re-derivation. This is language-agnostic (the model's own field works on CJK and Latin alike). zh->en: 0 -> 11 streaming commits, first at 6s. en->zh: 7 streaming, first at 6s. 2. split_text_units: CJK char-split fallback. When native stable_text is unavailable, the LCP path now splits CJK runs into character units (one char = one unit ≈ one BPE token) instead of collapsing them. Detection is by Unicode range (not a lang switch), applying to zh/ja/ko equally. 3. Token-LCP variant (tokenize_fn): commit on the model's own token strings instead of text units. Normalizes CJK (1 char ≈ 1 token) and Latin (subword) to the same token budget. Faster first commit (4s) but fragments mid-word; kept as an option, not the default. No language-specific switches. The only lang-aware code is CJK Unicode-range detection in the fallback char-split.
- Pure MLX streaming via mlx-audio (Nemotron transducer, native per-token timestamps) - Adapts the transducer's streaming API to WLK's insert_audio_chunk/process_iter contract - Native audio-time boundary for AlignAtt time-based simul-MT frontier - Normalizes language tag before prompt_dictionary validation - uv conflicts for nemotron-mlx-asr vs qwen3-streaming/qwen3-vllm/qwen3-vllm-metal + diarization-diart vs diarization-sortformer (numpy/onnx win32 split)
clkao
force-pushed
the
spacedock-ensign/nemotron-mlx-asr-backend-rebased
branch
from
September 2, 2026 05:18
0126dc9 to
d7c1e6d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #423 (which is stacked on #422). Review the incremental diff against #423's head.
Adds a pure-MLX Nemotron transducer ASR backend (
mlx-audiopackage) with native per-token timestamps for the AlignAtt time-based simul-MT frontier.Backend (
asr_nemotron_mlx.py)mlx-audio(NemotronMLXASR+NemotronMLXOnlineProcessor)insert_audio_chunk/process_iter/get_buffer/finishcontractAlignedTokenswith realstart/endaudio time (append-only, no revision). This is the time-based source frontier AlignAtt commits against — qwen3-asr'sstable_textis a text-position proxy with no timestamps.Why native timestamps matter (measured)
E2E A/B (ASR + MT, thermal-safe interleaved, 3 measured trials each):
nemotron's native timestamps give 2-3x earlier simul-MT provisionals on both directions. qwen3's
stable_textproxy commits coarser and later. Full baseline:_work/en_zh_e2e_baseline.md.Config + args
--backend nemotron-mlx-asrchoice--nemotron-mlx-asr-model(defaultmlx-community/nemotron-3.5-asr-streaming-0.6b)[tool.uv].conflictsfornemotron-mlx-asrvsqwen3-streaming/qwen3-vllm/qwen3-vllm-metal+diarization-diartvsdiarization-sortformer(numpy/onnx win32 split)Tests
12 tests (
test_asr_nemotron_mlx.py): streaming, finalization, timestamp monotonicity, state reset. All pass with and without MLX installed (pytest --collect-onlysafe).Dependencies
mlx-audio>=0.4.4,<0.5(macOS arm64 only)qwen3-streaming/qwen3-vllm/qwen3-vllm-metal(transformers version pin)