[ASR] Wordboosting by phrase

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::

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(

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,19 +210,28 @@ 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
             used only when this graph applied for the keywords spotting system.
             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()
 

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