diff --git a/docs/source/asr/asr_customization/word_boosting.rst b/docs/source/asr/asr_customization/word_boosting.rst index 62f3d2b45a47..d3b2a768c0f5 100644 --- a/docs/source/asr/asr_customization/word_boosting.rst +++ b/docs/source/asr/asr_customization/word_boosting.rst @@ -36,11 +36,30 @@ We recommend to provide the list of context phrases directly into ``speech_to_te List of the most important parameters: * ``strategy`` - The strategy to use for decoding depending on the model type (CTC - greedy_batch or beam_batch; RNN-T/TDT - greedy_batch or malsd_batch; AED - beam). -* ``model_path``, ``key_phrases_file``, ``key_phrases_list`` - The way to pass the context phrases into the decoding script. +* ``model_path``, ``key_phrases_file``, ``key_phrases_list``, ``key_phrase_items_list`` - The way to pass the context phrases into the decoding script. * ``context_score`` - The score for each arc transition in the context graph (1.0 is recommended). * ``depth_scaling`` - The scaling factor for the depth of the context graph (2.0 is recommended for CTC, RNN-T and TDT, 1.0 for Canary). * ``boosting_tree_alpha`` - Weight of the GPU-PB boosting tree during shallow fusion decoding (tune it according to your data). +Per-phrase boosting parameters +------------------------------ + +By default, all phrases share the global ``context_score`` and ``boosting_tree_alpha``. To boost individual phrases by different amounts, +use ``boosting_tree.key_phrase_items_list``, where each item can carry its own parameters: + +.. code-block:: + + boosting_tree.key_phrase_items_list='[{phrase:"nvlink",boosting_tree_alpha:2.0},{phrase:"omniverse cloud",context_score:1.5},{phrase:"gtc"}]' + +* ``context_score`` (per phrase) - Overrides the global ``context_score`` (and ``score_per_phrase``) for this phrase's arc transitions. +* ``boosting_tree_alpha`` (per phrase) - A relative boosting weight multiplier baked into the tree weights at build time. + The effective boost for a phrase is ``boosting_tree_alpha (decoding config) * boosting_tree_alpha (per phrase)``. + Note that the decode-time ``boosting_tree_alpha=0.0`` disables all boosting regardless of per-phrase values. +* ``lang`` (per phrase) - Language for the aggregate tokenizer; falls back to ``source_lang`` if omitted. + +Omitted fields fall back to the global values. Per-phrase parameters also work with ``build_gpu_boosting_tree.py`` +(the parameters are baked into the saved boosting tree). + **0.0. [Optional] Build the boosting tree for a specific ASR model:** .. code-block:: diff --git a/nemo/collections/asr/parts/context_biasing/boosting_graph_batched.py b/nemo/collections/asr/parts/context_biasing/boosting_graph_batched.py index a7cd54714bd0..b2349c86ddfb 100644 --- a/nemo/collections/asr/parts/context_biasing/boosting_graph_batched.py +++ b/nemo/collections/asr/parts/context_biasing/boosting_graph_batched.py @@ -35,8 +35,11 @@ @dataclass class PhraseItem: phrase: str # phrase itself - lang: str # per-phrase language (for aggregate tokenizer) - # custom weight can be further added + lang: Optional[str] = None # per-phrase language (for aggregate tokenizer); None -> cfg.source_lang + context_score: Optional[float] = None # per-phrase score for each arc transition; None -> global behavior + # per-phrase boosting weight multiplier applied on top of the decode-time boosting_tree_alpha, + # baked into the graph weights at build time; None -> 1.0 (global behavior) + boosting_tree_alpha: Optional[float] = None @dataclass @@ -51,8 +54,9 @@ class BoostingTreeModelConfig: None # The list of context-biasing phrases ['word1', 'word2', 'word3', ...] ) # The list of context-biasing phrases with custom options: - # [PhraseItem("word1", lang="en"), PhraseItem("frase dos", lang="es"), ...] - # in CLI: key_phrase_items_list='[{phrase:"word1",lang:en},{phrase:"frase dos",lang:es}]' + # [PhraseItem("word1", lang="en"), PhraseItem("word2", context_score=2.0, boosting_tree_alpha=4.0), ...] + # in CLI: key_phrase_items_list='[{phrase:"word1",lang:en},{phrase:"word2",context_score:2.0,boosting_tree_alpha:4.0}]' + # omitted per-phrase fields fall back to the global values key_phrase_items_list: list[PhraseItem] | None = None context_score: float = 1.0 # The score for each arc transition in the context graph depth_scaling: float = ( @@ -174,7 +178,7 @@ def _add_tbranches_first_order(self, tbranches: list): self.states[self.start_state]["order"] + 1, self.start_state, backoff_weight, - self.final_eos_score if tbranch.next_node.is_end else 0.0, + self.final_eos_score * tbranch.next_node.phrase_alpha if tbranch.next_node.is_end else 0.0, ) num_vocab_labels += 1 self._node_cache[tbranch.next_node.id] = next_state @@ -216,7 +220,7 @@ def _add_tbranches_next_order(self, tbranches: list): self.states[from_state]["order"] + 1, backoff_state, backoff_weight, - self.final_eos_score if tbranch.next_node.is_end else 0.0, + self.final_eos_score * tbranch.next_node.phrase_alpha if tbranch.next_node.is_end else 0.0, ) self._node_cache[tbranch.next_node.id] = next_state @@ -487,7 +491,7 @@ def dummy_boosting_tree( """ context_graph_trivial = ContextGraph(context_score=0.0, depth_scaling=0.0) - context_graph_trivial.build(token_ids=[[1]], phrases=["c"], scores=[0.0], uniform_weights=False) + context_graph_trivial.build(token_ids=[[1]], phrases=["c"], scores=None, uniform_weights=False) boosting_tree_trivial = GPUBoostingTreeModel.from_context_graph( context_graph=context_graph_trivial, @@ -534,6 +538,25 @@ def from_arpa( ) -> "NGramGPULanguageModel": raise NeMoBaseException("Boosting tree cannot be loaded from ARPA file") + @classmethod + def _validate_phrase_items(cls, phrase_items_list: list[PhraseItem]) -> None: + """ + Validate per-phrase boosting parameters (context_score and boosting_tree_alpha) + """ + for item in phrase_items_list: + if item.boosting_tree_alpha is not None: + if item.boosting_tree_alpha < 0: + raise ValueError( + f"Negative boosting_tree_alpha {item.boosting_tree_alpha} for phrase '{item.phrase}' " + "is not allowed (it would turn boosting into penalization)" + ) + if item.boosting_tree_alpha == 0.0: + logging.warning( + f"boosting_tree_alpha is 0.0 for phrase '{item.phrase}': the phrase is effectively disabled" + ) + if item.context_score is not None and item.context_score < 0: + logging.warning(f"Negative context_score {item.context_score} for phrase '{item.phrase}'") + @classmethod def from_config(cls, cfg: BoostingTreeModelConfig, tokenizer: TokenizerSpec) -> "GPUBoostingTreeModel": """ @@ -559,7 +582,10 @@ def from_config(cls, cfg: BoostingTreeModelConfig, tokenizer: TokenizerSpec) -> else: raise ValueError("No key phrases file or list specified") - # 2. tokenize key phrases + cls._validate_phrase_items(phrase_items_list) + + # 2. tokenize key phrases (deduplicated by phrase text, last occurrence wins) + items_by_phrase: dict[str, PhraseItem] = {item.phrase: item for item in phrase_items_list} phrases_dict = {} is_aggregate_tokenizer = isinstance(tokenizer, AggregateTokenizer) @@ -572,31 +598,37 @@ def from_config(cls, cfg: BoostingTreeModelConfig, tokenizer: TokenizerSpec) -> use_bpe_dropout = False spm.set_random_generator_seed(1234) # fix random seed for reproducibility of BPE dropout - for phrase_item in phrase_items_list: - phrase = phrase_item.phrase + for phrase, phrase_item in items_by_phrase.items(): if use_bpe_dropout: phrases_dict[phrase] = cls.get_alternative_transcripts(cfg, tokenizer, phrase) else: if is_aggregate_tokenizer: - phrases_dict[phrase] = tokenizer.text_to_ids(phrase, phrase_item.lang) + phrases_dict[phrase] = [tokenizer.text_to_ids(phrase, phrase_item.lang or cfg.source_lang)] else: - phrases_dict[phrase] = tokenizer.text_to_ids(phrase) - - # 3. build pythoncontext graph - contexts, scores, phrases = [], [], [] - for phrase in phrases_dict: - if use_bpe_dropout: - for transcript in phrases_dict[phrase]: - contexts.append(transcript) - scores.append(round(cfg.score_per_phrase / len(phrase), 2)) - phrases.append(phrase) + phrases_dict[phrase] = [tokenizer.text_to_ids(phrase)] + + # 3. build python context graph + contexts, scores, phrases, alphas = [], [], [], [] + for phrase, transcripts in phrases_dict.items(): + phrase_item = items_by_phrase[phrase] + # per-phrase context_score takes precedence over score_per_phrase; + # None means using the global cfg.context_score inside ContextGraph.build + if phrase_item.context_score is not None: + phrase_score = phrase_item.context_score + elif cfg.score_per_phrase != 0.0: + phrase_score = round(cfg.score_per_phrase / len(phrase), 2) else: - contexts.append(phrases_dict[phrase]) - scores.append(round(cfg.score_per_phrase / len(phrase), 2)) + phrase_score = None + for transcript in transcripts: + contexts.append(transcript) + scores.append(phrase_score) + alphas.append(phrase_item.boosting_tree_alpha) phrases.append(phrase) context_graph = ContextGraph(context_score=cfg.context_score, depth_scaling=cfg.depth_scaling) - context_graph.build(token_ids=contexts, scores=scores, phrases=phrases, uniform_weights=cfg.uniform_weights) + context_graph.build( + token_ids=contexts, scores=scores, phrases=phrases, alphas=alphas, uniform_weights=cfg.uniform_weights + ) # 4. build GPU boosting tree model from python context graph boosting_tree_model = GPUBoostingTreeModel.from_context_graph( diff --git a/nemo/collections/asr/parts/context_biasing/context_graph_universal.py b/nemo/collections/asr/parts/context_biasing/context_graph_universal.py index 5671f5bdd90a..4c700abb6c11 100644 --- a/nemo/collections/asr/parts/context_biasing/context_graph_universal.py +++ b/nemo/collections/asr/parts/context_biasing/context_graph_universal.py @@ -53,6 +53,7 @@ def __init__( level: int, phrase: str = "", ac_threshold: float = 1.0, + phrase_alpha: float = 1.0, ): """Create a ContextState. @@ -82,6 +83,9 @@ def __init__( The acoustic threshold (probability) of current context phrase, the value is valid only when current state is end state (is_end == True). Note: ac_threshold only used in keywords spotting. + phrase_alpha: + The per-phrase boosting weight multiplier of current context phrase, the + value is valid only when current state is end state (is_end == True). """ self.id = id self.token = token @@ -93,6 +97,7 @@ def __init__( self.next = {} self.phrase = phrase self.ac_threshold = ac_threshold + self.phrase_alpha = phrase_alpha self.fail = None self.output = None @@ -182,9 +187,10 @@ def build( self, token_ids: List[List[int]], phrases: Optional[List[str]] = None, - scores: Optional[List[float]] = None, + scores: Optional[List[Optional[float]]] = None, ac_thresholds: Optional[List[float]] = None, uniform_weights: Optional[bool] = False, + alphas: Optional[List[Optional[float]]] = None, ): """Build the ContextGraph from a list of token list. It first build a trie from the given token lists, then fill the fail arc @@ -204,9 +210,10 @@ def build( length of `phrases` MUST be equal to the length of `token_ids`. scores: The customize boosting score(token level) for each word/phrase, - 0 means using the default value (i.e. self.context_score). - It is a list of floats, and the length of `scores` MUST be equal to - the length of `token_ids`. + None means using the default value (i.e. self.context_score); + an explicit 0.0 is honored (zero-weight arcs for this phrase). + It is a list of optional floats, and the length of `scores` MUST be + equal to the length of `token_ids`. ac_thresholds: The customize trigger acoustic threshold (probability) for each phrase, 0 means using the default value (i.e. self.ac_threshold). It is @@ -214,9 +221,17 @@ def build( The length of `ac_threshold` MUST be equal to the length of `token_ids`. uniform_weights: If True, the weights will be distributed uniformly for all tokens as in Icefall. + alphas: + The customize boosting weight multiplier for each word/phrase, None means + the default value 1.0. The multiplier scales the whole token_score of the + phrase (including depth scaling terms), acting as a per-phrase counterpart + of the decode-time boosting_tree_alpha (the effective boost is + decode-time alpha * per-phrase alpha * base score). The length of `alphas` + MUST be equal to the length of `token_ids`. Note: The phrases would have shared states, the score of the shared states is - the MAXIMUM value among all the tokens sharing this state. + the MAXIMUM value among all the tokens sharing this state (effective, + i.e. alpha-scaled, scores are compared). """ num_phrases = len(token_ids) if phrases is not None: @@ -225,15 +240,18 @@ def build( assert len(scores) == num_phrases, (len(scores), num_phrases) if ac_thresholds is not None: assert len(ac_thresholds) == num_phrases, (len(ac_thresholds), num_phrases) + if alphas is not None: + assert len(alphas) == num_phrases, (len(alphas), num_phrases) for index, tokens in enumerate(token_ids): phrase = "" if phrases is None else phrases[index] - score = 0.0 if scores is None else scores[index] + score = None if scores is None else scores[index] ac_threshold = 0.0 if ac_thresholds is None else ac_thresholds[index] + alpha = 1.0 if alphas is None or alphas[index] is None else alphas[index] node = self.root # If has customized score using the customized token score, otherwise # using the default score - context_score = self.context_score if score == 0.0 else score + context_score = self.context_score if score is None else score threshold = self.ac_threshold if ac_threshold == 0.0 else ac_threshold for i, token in enumerate(tokens): if token not in node.next: @@ -243,6 +261,7 @@ def build( ) # depth scaling is used to give a larger score for all tokens after the first one else: token_score = context_score + token_score *= alpha self.num_nodes += 1 is_end = i == len(tokens) - 1 node_score = node.node_score + token_score @@ -256,10 +275,11 @@ def build( level=i + 1, phrase=phrase if is_end else "", ac_threshold=threshold if is_end else 0.0, + phrase_alpha=alpha if is_end else 1.0, ) else: # node exists, get the score of shared state. - token_score = max(context_score, node.next[token].token_score) + token_score = max(alpha * context_score, node.next[token].token_score) node.next[token].token_score = token_score node_score = node.node_score + token_score node.next[token].node_score = node_score @@ -269,6 +289,7 @@ def build( if i == len(tokens) - 1: node.next[token].phrase = phrase node.next[token].ac_threshold = threshold + node.next[token].phrase_alpha = max(alpha, node.next[token].phrase_alpha) node = node.next[token] self._fill_fail_output() diff --git a/scripts/asr_context_biasing/build_gpu_boosting_tree.py b/scripts/asr_context_biasing/build_gpu_boosting_tree.py index 852b61cee65b..3bde9392ae6d 100644 --- a/scripts/asr_context_biasing/build_gpu_boosting_tree.py +++ b/scripts/asr_context_biasing/build_gpu_boosting_tree.py @@ -33,6 +33,10 @@ class BuildWordBoostingTreeConfig(BoostingTreeModelConfig): """ Build GPU-accelerated phrase boosting tree (btree) to be used with greedy and beam search decoders of ASR models. + + Per-phrase boosting parameters (context_score and boosting_tree_alpha) can be passed via key_phrase_items_list, + e.g. key_phrase_items_list='[{phrase:"word1",context_score:2.0,boosting_tree_alpha:4.0},{phrase:"word2"}]'; + omitted fields fall back to the global values, and the parameters are baked into the saved boosting tree. """ asr_pretrained_name: Optional[str] = None # Name of a pretrained model diff --git a/tests/collections/asr/test_boosting_tree.py b/tests/collections/asr/test_boosting_tree.py index dff2124ff3c3..d97cc7f2513a 100644 --- a/tests/collections/asr/test_boosting_tree.py +++ b/tests/collections/asr/test_boosting_tree.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import numpy as np import pytest import torch from lightning.pytorch import Trainer @@ -21,6 +22,7 @@ from nemo.collections.asr.parts.context_biasing.boosting_graph_batched import ( BoostingTreeModelConfig, GPUBoostingTreeModel, + PhraseItem, ) from nemo.collections.asr.parts.context_biasing.context_graph_universal import ContextGraph @@ -34,9 +36,8 @@ def test_context_graph(): phrases = ["abc", "abd", "c"] phrases_ids = [[1, 2, 3], [1, 2, 4], [3]] - scores = [0.0, 0.0, 0.0] context_graph = ContextGraph(context_score=1.0, depth_scaling=1.0) - context_graph.build(token_ids=phrases_ids, phrases=phrases, scores=scores, uniform_weights=False) + context_graph.build(token_ids=phrases_ids, phrases=phrases, scores=None, uniform_weights=False) return context_graph @@ -185,12 +186,11 @@ def test_boosting_tree_model_from_config(self, conformer_ctc_bpe_model, tmp_path boosting_tree_cfg = BoostingTreeModelConfig() phrases = ["abc", "abd", "c"] phrases_ids = [conformer_ctc_bpe_model.tokenizer.text_to_ids(phrase) for phrase in phrases] - scores = [0.0, 0.0, 0.0] context_graph = ContextGraph( context_score=boosting_tree_cfg.context_score, depth_scaling=boosting_tree_cfg.depth_scaling ) context_graph.build( - token_ids=phrases_ids, phrases=phrases, scores=scores, uniform_weights=boosting_tree_cfg.uniform_weights + token_ids=phrases_ids, phrases=phrases, scores=None, uniform_weights=boosting_tree_cfg.uniform_weights ) test_boosting_tree = GPUBoostingTreeModel.from_context_graph( context_graph=context_graph, @@ -228,3 +228,185 @@ def test_boosting_tree_model_from_config(self, conformer_ctc_bpe_model, tmp_path assert torch.allclose( boosting_tree_from_model_path.arcs_weights, boosting_tree_from_key_phrases_list.arcs_weights ) + + +class TestPerPhraseBoostingParams: + @pytest.mark.unit + def test_per_phrase_context_score_affects_only_that_phrase(self): + """Per-phrase context_score changes arc scores of that phrase only (shared prefixes take the max)""" + phrases = ["abc", "abd", "c"] + phrases_ids = [[1, 2, 3], [1, 2, 4], [3]] + context_graph = ContextGraph(context_score=1.0, depth_scaling=1.0) + context_graph.build(token_ids=phrases_ids, phrases=phrases, scores=[None, 2.0, None], uniform_weights=False) + + root = context_graph.root + # shared prefix "ab" is raised to the max of the two phrase scores + assert root.next[1].token_score == pytest.approx(2.0) + assert root.next[1].next[2].token_score == pytest.approx(2.0) + # "abd" leaf uses its custom score, "abc" leaf keeps the global one + assert root.next[1].next[2].next[4].token_score == pytest.approx(2.0 + np.log(3)) + assert root.next[1].next[2].next[3].token_score == pytest.approx(1.0 + np.log(3)) + # "c" is unaffected + assert root.next[3].token_score == pytest.approx(1.0) + + @pytest.mark.unit + @pytest.mark.parametrize("device", DEVICES) + def test_per_phrase_alpha_scales_whole_phrase(self, device): + """Uniform per-phrase alpha scales all graph scores (arcs and backoffs) by the same factor""" + phrases = ["abc", "abd", "c"] + phrases_ids = [[1, 2, 3], [1, 2, 4], [3]] + alpha = 2.0 + + trees = [] + for alphas in (None, [alpha] * len(phrases)): + context_graph = ContextGraph(context_score=1.0, depth_scaling=1.0) + context_graph.build( + token_ids=phrases_ids, phrases=phrases, scores=None, alphas=alphas, uniform_weights=False + ) + trees.append( + GPUBoostingTreeModel.from_context_graph( + context_graph=context_graph, + vocab_size=5, + unk_score=0.0, + final_eos_score=0.0, + use_triton=True, + uniform_weights=False, + ).to(device) + ) + baseline_tree, scaled_tree = trees + + # node-level check: every token/node score is scaled by alpha + assert scaled_tree.num_states == baseline_tree.num_states + assert torch.allclose(scaled_tree.arcs_weights, alpha * baseline_tree.arcs_weights) + assert torch.allclose(scaled_tree.backoff_weights, alpha * baseline_tree.backoff_weights) + + # end-to-end check (includes backoff transitions): scores are exactly alpha * baseline + sentences_ids = [[1, 2, 3, 2, 1], [2, 2, 1, 2, 4], [3, 1, 2, 1]] + labels = pad_sequence([torch.LongTensor(s) for s in sentences_ids], batch_first=True).to(device) + lengths = torch.LongTensor([len(s) for s in sentences_ids]).to(device) + baseline_scores = baseline_tree(labels=labels, labels_lengths=lengths, bos=False, eos=False) + scaled_scores = scaled_tree(labels=labels, labels_lengths=lengths, bos=False, eos=False) + assert torch.allclose(scaled_scores, alpha * baseline_scores, atol=1e-4) + + @pytest.mark.unit + def test_per_phrase_alpha_shared_suffix_backoff(self): + """Backoff transitions between phrases with different alphas keep alpha-consistent cumulative scores""" + phrases = ["ab", "c"] + phrases_ids = [[1, 2], [3]] + context_graph = ContextGraph(context_score=1.0, depth_scaling=1.0) + context_graph.build(token_ids=phrases_ids, phrases=phrases, scores=None, alphas=[2.0, 3.0]) + boosting_tree = GPUBoostingTreeModel.from_context_graph( + context_graph=context_graph, vocab_size=5, unk_score=0.0, final_eos_score=0.0, use_triton=True + ) + + # [1, 3]: "a" arc (2.0), then backoff from non-final "a" state (-2.0) + "c" arc (3.0) = 1.0 + # [1, 2, 3]: "a" (2.0), "b" (2 * (1 + log(2))), then "c" from final state without penalty (3.0) + sentences_ids = [[1, 3, 0], [1, 2, 3]] + boosting_scores = boosting_tree( + labels=torch.LongTensor(sentences_ids), + labels_lengths=torch.LongTensor([2, 3]), + bos=False, + eos=False, + ) + expected = torch.tensor( + [ + [2.0, 1.0, 0.0], + [2.0, 2.0 * (1.0 + np.log(2)), 3.0], + ], + dtype=boosting_scores.dtype, + ) + assert torch.allclose(boosting_scores, expected, atol=1e-4) + + @pytest.mark.unit + def test_per_phrase_alpha_scales_final_eos_score(self): + """final_eos_score of an end state is scaled by the alpha of the phrase ending there""" + phrases = ["abc", "abd", "c"] + phrases_ids = [[1, 2, 3], [1, 2, 4], [3]] + context_graph = ContextGraph(context_score=1.0, depth_scaling=1.0) + context_graph.build(token_ids=phrases_ids, phrases=phrases, scores=None, alphas=[2.0, None, 3.0]) + boosting_tree = GPUBoostingTreeModel.from_context_graph( + context_graph=context_graph, vocab_size=5, unk_score=0.0, final_eos_score=1.0, use_triton=True + ) + + # three end states scaled by their phrase alphas (2.0, 1.0, 3.0), all other states are not final + final_weights = boosting_tree.final_weights[: boosting_tree.num_states] + assert sorted(w for w in final_weights.tolist() if w != 0.0) == pytest.approx([1.0, 2.0, 3.0]) + + @pytest.mark.unit + def test_explicit_zero_context_score_is_honored(self): + """An explicit per-phrase context_score of 0.0 is not treated as 'use the default' anymore""" + context_graph = ContextGraph(context_score=1.0, depth_scaling=1.0) + context_graph.build(token_ids=[[1, 2]], phrases=["ab"], scores=[0.0]) + + assert context_graph.root.next[1].token_score == pytest.approx(0.0) + # only the depth term remains for tokens after the first one + assert context_graph.root.next[1].next[2].token_score == pytest.approx(np.log(2)) + + @pytest.mark.unit + def test_negative_alpha_raises(self): + """Negative per-phrase boosting_tree_alpha is rejected""" + with pytest.raises(ValueError, match="Negative boosting_tree_alpha"): + GPUBoostingTreeModel._validate_phrase_items([PhraseItem("abc", boosting_tree_alpha=-1.0)]) + + @pytest.mark.unit + def test_per_phrase_params_none_matches_baseline(self, conformer_ctc_bpe_model): + """key_phrase_items_list with all-default params builds the same graph as key_phrases_list""" + phrases = ["abc", "abd", "c"] + baseline_tree = GPUBoostingTreeModel.from_config( + BoostingTreeModelConfig(key_phrases_list=phrases), tokenizer=conformer_ctc_bpe_model.tokenizer + ) + items_tree = GPUBoostingTreeModel.from_config( + BoostingTreeModelConfig(key_phrase_items_list=[PhraseItem(phrase) for phrase in phrases]), + tokenizer=conformer_ctc_bpe_model.tokenizer, + ) + + assert torch.allclose(baseline_tree.arcs_weights, items_tree.arcs_weights) + assert torch.allclose(baseline_tree.backoff_weights, items_tree.backoff_weights) + assert torch.allclose(baseline_tree.final_weights, items_tree.final_weights) + + @pytest.mark.unit + def test_per_phrase_context_score_precedence_over_score_per_phrase(self, conformer_ctc_bpe_model): + """Per-phrase context_score wins over score_per_phrase for that phrase only""" + tokenizer = conformer_ctc_bpe_model.tokenizer + cfg = BoostingTreeModelConfig( + key_phrase_items_list=[PhraseItem("abc", context_score=2.0), PhraseItem("abd")], + score_per_phrase=4.0, + ) + items_tree = GPUBoostingTreeModel.from_config(cfg, tokenizer=tokenizer) + + # reference graph: "abc" uses the per-phrase score, "abd" uses score_per_phrase / len(phrase) + reference_graph = ContextGraph(context_score=cfg.context_score, depth_scaling=cfg.depth_scaling) + reference_graph.build( + token_ids=[tokenizer.text_to_ids(phrase) for phrase in ("abc", "abd")], + phrases=["abc", "abd"], + scores=[2.0, round(4.0 / len("abd"), 2)], + ) + reference_tree = GPUBoostingTreeModel.from_context_graph( + context_graph=reference_graph, + vocab_size=tokenizer.vocab_size, + unk_score=cfg.unk_score, + final_eos_score=cfg.final_eos_score, + ) + + assert torch.allclose(items_tree.arcs_weights, reference_tree.arcs_weights) + assert torch.allclose(items_tree.backoff_weights, reference_tree.backoff_weights) + + @pytest.mark.unit + def test_per_phrase_params_with_bpe_dropout(self, conformer_ctc_bpe_model): + """All alternative BPE tokenizations of a phrase get that phrase's alpha. + + Note: two consecutive builds with BPE dropout sample different alternative tokenizations + (sentencepiece seeds its random generator only on first use), so the check is done + within a single build via the final weights of the end states. + """ + cfg = BoostingTreeModelConfig( + key_phrase_items_list=[PhraseItem("nvidia", boosting_tree_alpha=2.0), PhraseItem("omniverse")], + use_bpe_dropout=True, + final_eos_score=1.0, + ) + boosting_tree = GPUBoostingTreeModel.from_config(cfg, tokenizer=conformer_ctc_bpe_model.tokenizer) + + # every end state of a "nvidia" tokenization carries final_eos_score * 2.0, + # every end state of an "omniverse" tokenization keeps final_eos_score * 1.0 + final_weights = boosting_tree.final_weights[: boosting_tree.num_states] + assert set(w for w in final_weights.tolist() if w != 0.0) == {1.0, 2.0}