fix(speechlm2): preserve relocated EOS token during early interruption augmentation by ManoharPaturi · Pull Request #16268 · NVIDIA-NeMo/Speech · GitHub
Skip to content

fix(speechlm2): preserve relocated EOS token during early interruption augmentation - #16268

Open
ManoharPaturi wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
ManoharPaturi:fix/speechlm2-preserve-relocated-eos
Open

ManoharPaturi wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
ManoharPaturi:fix/speechlm2-preserve-relocated-eos

Conversation

@ManoharPaturi

Copy link
Copy Markdown
Contributor

Fixes #16247.

Problem
In DuplexSTTDataset._apply_early_interruption_augmentation (nemo/collections/speechlm2/data/duplex_stt_dataset.py), early interruption is simulated by placing eos_id at new_eos_pos, shifting the tail if applicable, and then overwriting trailing frames with pad_id:

target_tokens[batch_idx, new_eos_pos] = eos_id
seq_len = target_tokens.shape[1]
cont_start_pos = original_eos_pos + overlap_tokens
tail_length = seq_len - (cont_start_pos + 1)
if tail_length > 0:
    target_tokens[batch_idx, new_eos_pos + 1 : new_eos_pos + 1 + tail_length] = target_tokens[
        batch_idx, cont_start_pos + 1 : cont_start_pos + 1 + tail_length
    ].clone()
target_tokens[batch_idx, -frames_to_remove:] = pad_id

When original_eos_pos sits near the end of the sequence such that cont_start_pos >= seq_len - 1 (no tail-shift runs), frames_to_remove = original_eos_pos - cutoff. Because new_eos_pos = cutoff, the pad-fill range [-frames_to_remove:] starts at or before new_eos_pos, immediately overwriting the relocated eos_id with pad_id. The augmented turn ends up with no EOS token at all.

Fix
Calculate the pad start boundary as pad_start = new_eos_pos + 1 + max(0, tail_length) and set target_tokens[batch_idx, pad_start:] = pad_id, and ensure target_tokens[batch_idx, new_eos_pos] = eos_id is placed so that the relocated EOS token is strictly preserved.

Testing

  • Added unit test in tests/collections/speechlm2/test_duplex_early_interruption_eos.py exercising the boundary condition (seq_len=25, original_eos_pos=24, overlap_tokens=5) and asserting that the relocated EOS token is preserved.
  • Verified test passes on CPU.

…n augmentation

Fixes NVIDIA-NeMo#16247.

In DuplexSTTDataset._apply_early_interruption_augmentation, target_tokens
places eos_id at new_eos_pos before shifting the tail and padding the vacated end.
When original_eos_pos sits close enough to the end of the sequence that no tail-shift
runs (cont_start_pos >= seq_len - 1), the padding slice [-frames_to_remove:]
starts at or before new_eos_pos, immediately overwriting the newly placed EOS token.

This change calculates pad_start as new_eos_pos + 1 + max(0, tail_length) and places
eos_id at new_eos_pos, guaranteeing that the relocated EOS token is never clobbered
by pad_id. Added unit regression test.

Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com>
Copilot AI lite review requested due to automatic review settings September 15, 2026 17:25
@copy-pr-bot

copy-pr-bot Bot commented Sep 15, 2026

Copy link
Copy Markdown

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Sep 17, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-request waiting-on-maintainers Waiting on maintainers to respond

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Early-interruption augmentation can drop the relocated EOS token entirely

3 participants