fix(speechlm2): keep EOS that falls in the delay_text_channel_by window by udsy19 · Pull Request #16165 · NVIDIA-NeMo/Speech

udsy19 · 2026-08-29T04:33:11Z

What does this PR do ?

Fixes prepare_text_and_asr_labels silently dropping the EOS token whenever delay_text_channel_by shifts it into the trailing delay window.

Collection: [speechlm2]

Changelog

nemo/collections/speechlm2/parts/label_prep.py:170: the EOS-rescue write targets
target_tokens[i, -(delay_by)], a position the very next line's
target_tokens[:, :-delay_by] slice drops. The relocated EOS was therefore
discarded along with the rest of the delay window every time. The slot that
actually survives the slice-then-pad is target_tokens[i, -(delay_by + 1)].
tests/collections/speechlm2/test_label_prep.py: new
test_prepare_text_and_asr_labels_delay_channel_keeps_eos_near_boundary,
covering an EOS at the last index still inside the delay window, an EOS at the
very last index, and a control EOS well outside the window.

Usage

import torch
from nemo.collections.speechlm2.parts.label_prep import prepare_text_and_asr_labels

PAD, BOS, EOS = 0, 1, 3
delay_by = 2
tokens = torch.tensor([[BOS, 10, 11, 12, 13, 14, EOS, PAD]])

out = prepare_text_and_asr_labels(
    batch={"target_token_lens": torch.tensor([7])},
    target_tokens=tokens,
    source_encoded=torch.zeros(1, tokens.shape[1], 4),
    cfg={"delay_text_channel_by": delay_by},
    text_pad_id=PAD, text_bos_id=BOS, text_eos_id=EOS,
)
print(out["text_labels"])

On main (6281c939aa): tensor([[ 0, 1, 10, 11, 12, 13, 14]]) — EOS gone. With this PR: EOS is present, shifted like every other token.

Negative control: with label_prep.py reverted to main and only the new test added, the file reports 1 failed, 2 passed — the new boundary test fails, the two pre-existing tests are unaffected. With the fix, all three tests in the file pass (3 passed).

nemo/collections/speechlm2/parts/label_prep.py has had exactly one commit since it was introduced (#15092); no test exercised this branch before this PR.

GitHub Actions CI

Requires a maintainer /ok to test <head-sha>.

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)

PR Type:

New Feature
Bugfix
Documentation

Additional Information

Fixes #16164

prepare_text_and_asr_labels's delay_text_channel_by path masks out any EOS in the last delay_by positions (they are about to be sliced off by the right-shift) and tries to relocate one of them into the slot that survives the shift by writing to target_tokens[i, -(delay_by)]. That index is itself one of the positions the following `target_tokens[:, :-delay_by]` slice drops, so the relocated EOS is discarded along with the positions it was rescued from. Every EOS that lands in the delay window is silently erased from the training labels, regardless of the rescue attempt. The slot that actually survives the slice-then-pad is target_tokens[i, -(delay_by + 1)]. label_prep.py has had exactly one commit since it was introduced (NVIDIA-NeMo#15092); the rescue branch has never been exercised by a test. Signed-off-by: Udaya Tejas <udayatejas2004@gmail.com>

copy-pr-bot · 2026-08-29T04:33:15Z

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.

pzelasko · 2026-08-31T16:04:34Z

@codex review

chatgpt-codex-connector · 2026-08-31T16:04:50Z

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-08-31T16:13:41.374445Z	`29c9bb6`	Manual request

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

pzelasko · 2026-08-31T16:05:01Z

/ok to test 29c9bb6

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 29c9bb60a4

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-08-31T16:13:38Z

        for i in range(target_tokens.size(0)):
            if eos_mask[i].any():
-                target_tokens[i, -(delay_by)] = text_eos_id
+                target_tokens[i, -(delay_by + 1)] = text_eos_id


Guard the full-width delay before indexing

When delay_text_channel_by equals target_tokens.size(1) and a row contains EOS, this index becomes -T-1, so PyTorch raises IndexError before performing the shift. The previous -delay_by index remained valid at this boundary, making this a new crash for short batches whose encoded width matches the configured delay; validate that the delay is smaller than the width or handle the full-width shift explicitly, with a boundary regression test.

AGENTS.md reference: AGENTS.md:L58-L60

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-08-31T16:13:38Z

 import torch

-from nemo.collections.speechlm2.parts.label_prep import maybe_prepend_prompt_tokens
+from nemo.collections.speechlm2.parts.label_prep import maybe_prepend_prompt_tokens, prepare_text_and_asr_labels


Add the required DCO sign-off

The reviewed commit object 599996a2ccd1ea717ce60751fb6dab3e948bf807 contains no Signed-off-by trailer, so it does not satisfy the repository's mandatory DCO requirement and cannot pass contribution compliance as submitted. Recreate the commit with the appropriate sign-off trailer.

AGENTS.md reference: AGENTS.md:L52-L52

Useful? React with 👍 / 👎.

chtruong814 · 2026-08-31T23:10:37Z

/ok to test 7dbdf41

chtruong814 · 2026-09-01T06:37:21Z

/ok to test 2639b6a

nemo-automation-bot Bot added the community-request label Aug 29, 2026

svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Aug 31, 2026

pzelasko self-requested a review August 31, 2026 16:04

copy-pr-bot Bot deployed to public August 31, 2026 16:06 Active

copy-pr-bot Bot deployed to test August 31, 2026 16:06 Active

copy-pr-bot Bot deployed to public August 31, 2026 16:09 Active

copy-pr-bot Bot deployed to public August 31, 2026 16:10 Active

chatgpt-codex-connector Bot reviewed Aug 31, 2026

View reviewed changes

svcnvidia-nemo-ci removed the waiting-on-maintainers Waiting on maintainers to respond label Aug 31, 2026

Merge branch 'main' into fix/duplex-stt-delay-eos-clamp

7dbdf41

copy-pr-bot Bot deployed to public August 31, 2026 23:11 Active

copy-pr-bot Bot deployed to test August 31, 2026 23:12 Active

copy-pr-bot Bot deployed to public August 31, 2026 23:15 Active

Merge branch 'main' into fix/duplex-stt-delay-eos-clamp

2639b6a

copy-pr-bot Bot deployed to public September 1, 2026 06:38 Active

copy-pr-bot Bot deployed to test September 1, 2026 06:39 Active

copy-pr-bot Bot deployed to public September 1, 2026 06:41 Active

copy-pr-bot Bot deployed to public September 1, 2026 06:42 Active

Navigation Menu

Uh oh!

fix(speechlm2): keep EOS that falls in the delay_text_channel_by window - #16165

fix(speechlm2): keep EOS that falls in the delay_text_channel_by window#16165
udsy19 wants to merge 3 commits into
NVIDIA-NeMo:mainfrom
udsy19:fix/duplex-stt-delay-eos-clamp

udsy19 commented Aug 29, 2026

copy-pr-bot Bot commented Aug 29, 2026

pzelasko commented Aug 31, 2026

chatgpt-codex-connector Bot commented Aug 31, 2026 •

edited

Loading

pzelasko commented Aug 31, 2026

chatgpt-codex-connector Bot left a comment

chatgpt-codex-connector Bot Aug 31, 2026

chatgpt-codex-connector Bot Aug 31, 2026

chtruong814 commented Aug 31, 2026

chtruong814 commented Sep 1, 2026

Labels

4 participants

Uh oh!

Conversation

udsy19 commented Aug 29, 2026

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Additional Information

copy-pr-bot Bot commented Aug 29, 2026

pzelasko commented Aug 31, 2026

chatgpt-codex-connector Bot commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codex Review Summary

pzelasko commented Aug 31, 2026

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

chatgpt-codex-connector Bot Aug 31, 2026

Choose a reason for hiding this comment

chatgpt-codex-connector Bot Aug 31, 2026

Choose a reason for hiding this comment

chtruong814 commented Aug 31, 2026

chtruong814 commented Sep 1, 2026

Labels

4 participants

chatgpt-codex-connector Bot commented Aug 31, 2026 •

edited

Loading