#16247·Speech

Early-interruption augmentation can drop the relocated EOS token entirely

Author: udsy19Created Sep 12, 2026Updated Sep 16, 2026
Labelscommunity-request

Describe the bug

DuplexSTTDataset._apply_early_interruption_augmentation (nemo/collections/speechlm2/data/duplex_stt_dataset.py) simulates an early user interruption by relocating a turn's EOS token earlier in the sequence, then padding the vacated tail. It writes the new EOS token before computing the tail-shift, then unconditionally overwrites the trailing frames_to_remove positions with the pad id:

python
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 the turn's original EOS sits close enough to the end of the sequence that no tail-shift runs (cont_start_pos >= seq_len - 1), the final pad-fill range [-frames_to_remove:] starts at or before new_eos_pos and immediately overwrites the EOS token that was just written — the augmented turn ends up with no EOS token at all.

Steps/Code to reproduce bug

python
import torch
target_tokens = torch.full((1, 25), 0, dtype=torch.long)
target_tokens[0, 0] = 1        # BOS
target_tokens[0, 1:13] = torch.arange(10, 22)  # content
target_tokens[0, 24] = 2       # original EOS at the last index

# with overlap_tokens=5, any cutoff in [1, 12] gives:
#   cont_start_pos = 24 + 5 = 29 >= seq_len (25) -> no tail-shift runs
#   frames_to_remove = 24 - cutoff -> pad-fill range covers new_eos_pos

Calling _apply_early_interruption_augmentation with this input leaves target_tokens[0] with zero occurrences of the EOS id — the relocated EOS is clobbered by the pad-fill.

Expected behavior

The relocated EOS token should survive the augmentation regardless of how close the original EOS sits to the end of the sequence.

Environment overview (please complete the following information)

  • Environment location: Bare-metal
  • Method of NeMo install: from source (editable install of a fresh clone of main)

Environment details

  • OS version: Ubuntu 24.04.4 LTS
  • PyTorch version: 2.14.0+cpu
  • Python version: 3.12.3

Additional context

early_interruption_prob is a documented, opt-in DuplexSTTDataset config knob (default 0.0) read directly from the model config at construction time; no shipped examples/speechlm2/conf/*.yaml currently enables it, but any training run that does silently loses the EOS token for the augmented turn whenever the original EOS is within overlap_tokens positions of the end of a short-enough sequence — the model would be trained on a turn with no end-of-turn signal. git log -S shows the function has been touched by exactly one commit since its introduction (#15092), so this is an oversight rather than a deliberate choice, and no test previously exercised this branch.