fix(asr): use each session's own timestamps in SpeakerTaggedASR._add_speaker_transcriptions

[kzos]

What does this PR do ?

Makes SpeakerTaggedASR._add_speaker_transcriptions read each session's own word and segment timestamps. At the moment it reads session 0's for every session.

Collection: ASR

Changelog

• nemo/collections/asr/parts/utils/multispk_transcribe_utils.py, _add_speaker_transcriptions: trans_hyp[0] → trans_hyp[sess_idx] in the word loop (L998) and the segment loop (L1004).
• tests/collections/speaker_tasks/utils/test_spk_tagged_asr_utils.py: a new TestAddSpeakerTranscriptions class with one test, parametrized over two two-session cases and one three-session case.

The problem

Inside for sess_idx, (uniq_id, _) in enumerate(test_manifest_dict.items()), lines 996, 1002 and 1018 use trans_hyp[sess_idx]. The two loops that read the timestamps use index 0:

for word_idx, trans_wdict in enumerate(trans_hyp[0].timestamp['word']):        # L998
    ...
for word_idx, trans_segdict in enumerate(trans_hyp[0].timestamp['segment']):   # L1004

trans_hyp is a flat list with one hypothesis per session. merge_transcript_and_speakers builds word_and_ts_seq[uniq_id]['words'] for session idx from asr_hypotheses[idx] (L846), and _add_speaker_transcriptions then indexes it by word_idx and w_count. Index 0 is therefore session 0, not a nested per-session container, and session i ends up with session 0's words carrying session i's speaker labels.

I checked this with a CPU script that passes hand-built hypotheses and diarization tensors through the real merge_transcript_and_speakers → _add_speaker_transcriptions path. No model is involved. Session 0 is sess_a and session 1 is sess_b:

session 1 compared with session 0	current main
different words	ValueError: Word mismatch: 'hello' != 'good' at session 1, word count 0.
fewer words (1 vs 3)	IndexError: list index out of range
same words, different times	no error, but session 1 gets session 0's timestamps and session 0's segment speaker is overwritten

The script's output for the third case, exactly as printed apart from the # lines I added:

# input: sess_a has one segment 'yes okay'; sess_b has two segments, 'yes' and 'okay'
  input  sess_a: words=[('yes', 0.08, 0.4), ('okay', 0.48, 0.8)] diar speakers per word=['speaker_0', 'speaker_0']
  input  sess_b: words=[('yes', 1.6, 1.92), ('okay', 2.56, 3.04)] diar speakers per word=['speaker_1', 'speaker_2']
# current main
  output sess_a: word   =[('yes', 0.08, 0.4, 'speaker_0'), ('okay', 0.48, 0.8, 'speaker_0')]
  output sess_a: segment=[('yes okay', 0.08, 0.8, 'speaker_1')]
  output sess_b: word   =[('yes', 0.08, 0.4, 'speaker_1'), ('okay', 0.48, 0.8, 'speaker_2')]
  output sess_b: segment=[('yes okay', 0.08, 0.8, 'speaker_1')]
# with this PR
  output sess_a: word   =[('yes', 0.08, 0.4, 'speaker_0'), ('okay', 0.48, 0.8, 'speaker_0')]
  output sess_a: segment=[('yes okay', 0.08, 0.8, 'speaker_0')]
  output sess_b: word   =[('yes', 1.6, 1.92, 'speaker_1'), ('okay', 2.56, 3.04, 'speaker_2')]
  output sess_b: segment=[('yes', 1.6, 1.92, 'speaker_1'), ('okay', 2.56, 3.04, 'speaker_2')]

On main, session 0's segment speaker changes from speaker_0 to speaker_1 because the segment loop writes trans_segdict['speaker'] into the dicts it iterates over, and on main those dicts belong to session 0.

With this change all three cases give each session its own timestamps and speakers. A single session is not affected, because sess_idx is always 0. The script's single-session output is identical before and after (compared as JSON).

The script is not part of this PR. The new test covers the same three situations by calling _add_speaker_transcriptions directly, so it serves as the runnable reproduction.

Impact, stated honestly

This is a latent defect in a helper, and nothing in the repository reaches it today. The only caller is SpeakerTaggedASR.perform_offline_stt_spk (L1043). Neither speech_to_text_multitalker_streaming_infer.py nor the multitalker tutorial calls that method; both use the streaming paths.

On current main, perform_offline_stt_spk failed before reaching this function in every case I tried:

• It unpacks best_hyp, _ = transcriptions (L1036). That matches the (hypotheses, all_hypotheses) tuple that EncDecRNNTModel._transcribe_output_processing returned before changed asr models outputs to be consistent #11818; it now returns a list. I built a small, randomly initialised EncDecRNNTBPEModel locally, with a SentencePiece tokenizer trained on a few lines of text and noise audio. Its transcribe() returns a flat list of Hypothesis, and with a stub diarizer:
  • a 1-file manifest raises ValueError: not enough values to unpack (expected 2, got 1);
  • a 3-file manifest raises ValueError: too many values to unpack (expected 2);
  • a 2-file manifest gets as far as merge_transcript_and_speakers, where asr_hypotheses[idx] raises TypeError: 'Hypothesis' object is not subscriptable (L841).
• It passes self.cfg.dataset_manifest to transcribe() (L1033). The example script's MultitalkerTranscriptionConfig has manifest_file but no dataset_manifest, so accessing it raises ConfigAttributeError.

I did not run the released multitalker checkpoints. From reading the code, EncDecMultiTalkerRNNTBPEModel overrides _transcribe_forward and _setup_transcribe_dataloader but not transcribe or _transcribe_output_processing, so I expect it to return the same list.

I kept this PR to the indexing fix. Whether the offline entry point should be updated to the current transcribe() output or removed is for you to decide, and I'm happy to follow up either way. From reading the code, if it is updated, this fix is needed for any manifest with more than one file.

Tests

The new test is CPU-only and needs no model, tokenizer or download. It calls _add_speaker_transcriptions directly with a SimpleNamespace as self, the way the existing tests call perform_parallel_streaming_stt_spk, and asserts that each session's word and segment timestamps and speakers come back as its own. The cases:

• different_words: on main this hits the ValueError above.
• same_words_different_times: on main this runs without error but returns the wrong values. The first mismatch is session 0's segment speaker, speaker_1 instead of speaker_0.
• three_sessions_fewer_words: session 1 has one word and session 0 has three, so on main this hits the IndexError above. The third session also makes the test reject an index that agrees with sess_idx only for the first two sessions: I checked trans_hyp[-sess_idx], which passes the two two-session cases and fails this one.

TestAddSpeakerTranscriptions   without this change : 3 failed
                               with this change    : 3 passed

test_spk_tagged_asr_utils.py   without this change : 3 failed, 108 passed, 12 errors
                               with this change    : 111 passed, 12 errors

Fixing only one of the two lines still fails all three cases.

The 12 errors are the same tests in both runs, and each one is fixture 'test_data_dir' not found. tests/conftest.py downloads the test-data archive, and I ran with --noconftest to avoid that download (registering conftest's three markers so --strict-markers still applied). These tests use the RNNT asr_model fixture from test_asr_rnnt_encoder_model_bpe.py, whose tokenizer comes from that archive, so they could not be set up here. None of them reaches the changed code: under tests/, _add_speaker_transcriptions and perform_offline_stt_spk are referenced only by the new class.

black 24.10.0 and isort 5.13.2 clean at line length 119.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc) — no

PR Type:

• New Feature
• Bugfix
• Documentation

Additional Information

• Both trans_hyp[0] lines date from the file's first version in Merge updates of Multi-Talker Parakeet Model, Modules, Dataloader and Utils PR 01 #14905.
• No open PR changes _add_speaker_transcriptions. Draft [Streaming SpeechLM] Add audomodel support #16045, which targets a feature branch rather than main, modifies this file and its test file in other places; both its head and its base branch still have the same two trans_hyp[0] lines.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.