diff --git a/nemo/collections/asr/parts/context_biasing/biasing_multi_model.py b/nemo/collections/asr/parts/context_biasing/biasing_multi_model.py index 694f6f190a24..d5af66da2df1 100644 --- a/nemo/collections/asr/parts/context_biasing/biasing_multi_model.py +++ b/nemo/collections/asr/parts/context_biasing/biasing_multi_model.py @@ -562,22 +562,19 @@ def remove_model(self, model_id: int): self.model2num_states[model_id] = 0 self.model2num_arcs[model_id] = 0 self.model2num_arcs_extended[model_id] = 0 - # shift model offsets + # Shift the offsets of the models that lived after the removed one. + # Only ACTIVE models take part in the shift. An unused slot holds offset 0, and 0 is also a + # legitimate arena address, so a shift applied to every slot underflows every unused slot to a + # negative offset whenever the removed model started at 0. `model2active[model_id]` is already + # False here, so the removed model is excluded from the mask and its own offsets are cleared last. + states_shift_mask = self.model2active & (self.model2states_offset >= start_state) + arcs_shift_mask = self.model2active & (self.model2arcs_offset >= start_arc) + self.model2states_offset[states_shift_mask] -= num_states + self.model2arcs_offset[arcs_shift_mask] -= num_arcs + + # clear the removed model's own offsets last self.model2states_offset[model_id] = 0 self.model2arcs_offset[model_id] = 0 - # shift states and arcs offsets - torch.where( - self.model2states_offset < start_state, - self.model2states_offset, - self.model2states_offset - num_states, - out=self.model2states_offset, - ) - torch.where( - self.model2arcs_offset < start_arc, - self.model2arcs_offset, - self.model2arcs_offset - num_arcs, - out=self.model2arcs_offset, - ) def get_init_states(self, batch_size: int, bos=True) -> torch.Tensor: """ diff --git a/tests/collections/asr/decoding/test_biasing_multi_model.py b/tests/collections/asr/decoding/test_biasing_multi_model.py index c4b0dd1b5cd8..6af9a41e3b1e 100644 --- a/tests/collections/asr/decoding/test_biasing_multi_model.py +++ b/tests/collections/asr/decoding/test_biasing_multi_model.py @@ -143,6 +143,67 @@ def test_add_then_remove_model(self, stt_en_conformer_transducer_small, device: assert multi_model.model2states_offset[model_id2].item() == 0 assert multi_model.model2arcs_offset[model_id2].item() == 0 + # Every offset must stay a valid, non-negative arena index, including the reserved slots that + # have never held a model. Before the ordering fix in remove_model, removing the model that + # started at offset 0 underflowed every zero-valued slot to -num_states / -num_arcs. + assert torch.all(multi_model.model2states_offset >= 0) + assert torch.all(multi_model.model2arcs_offset >= 0) + + # Inactive slots carry a canonical, empty descriptor. + inactive = ~multi_model.model2active + assert torch.all(multi_model.model2states_offset[inactive] == 0) + assert torch.all(multi_model.model2arcs_offset[inactive] == 0) + assert torch.all(multi_model.model2num_states[inactive] == 0) + assert torch.all(multi_model.model2num_arcs_extended[inactive] == 0) + + @pytest.mark.unit + @pytest.mark.with_downloads + @pytest.mark.parametrize("device", DEVICES) + @pytest.mark.parametrize("remove_index", [0, 1, 2]) + def test_offsets_stay_non_negative_for_any_removal_position( + self, stt_en_conformer_transducer_small, device: torch.device, remove_index: int + ): + """Removing a model from any position leaves every offset a valid arena index. + + Removing the model at the front of the arena is the interesting case: its start offset is 0, + which every unused reserved slot also holds, so a shift that does not exclude inactive slots + drives all of them negative. + """ + tokenizer = stt_en_conformer_transducer_small.tokenizer + vocab_size = tokenizer.vocab_size + + multi_model = GPUBiasingMultiModel(vocab_size=vocab_size).to(device) + phrase_lists = [["alpha", "beta"], ["gamma"], ["delta", "epsilon", "zeta"]] + model_ids = [ + multi_model.add_model(create_boosting_model(phrases, tokenizer, device), alpha=1.0) + for phrases in phrase_lists + ] + + # A surviving model's scores must be unchanged by an unrelated removal. + survivor = model_ids[(remove_index + 1) % len(model_ids)] + states = multi_model.get_init_states(batch_size=2, bos=True) + survivor_ids = torch.full((2,), survivor, dtype=torch.long, device=device) + scores_before, _ = multi_model.advance(states=states, model_ids=survivor_ids) + scores_before = scores_before.clone() + + multi_model.remove_model(model_ids[remove_index]) + + assert torch.all(multi_model.model2states_offset >= 0) + assert torch.all(multi_model.model2arcs_offset >= 0) + + inactive = ~multi_model.model2active + assert torch.all(multi_model.model2states_offset[inactive] == 0) + assert torch.all(multi_model.model2arcs_offset[inactive] == 0) + + active = multi_model.model2active + assert torch.all( + multi_model.model2states_offset[active] + multi_model.model2num_states[active] + <= multi_model.num_states_total + ) + + scores_after, _ = multi_model.advance(states=states, model_ids=survivor_ids) + assert torch.equal(scores_before, scores_after) + @pytest.mark.unit @pytest.mark.with_downloads @pytest.mark.parametrize("device", DEVICES)