icefall/context_graph.py

# Copyright    2023  Xiaomi Corp.        (authors: Wei Kang)
#
# See ../LICENSE for clarification regarding multiple authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import shutil
from collections import deque
from typing import Dict, List, Optional, Tuple, Union


class ContextState:
    """The state in ContextGraph"""

    def __init__(
        self,
        id: int,
        token: int,
        token_score: float,
        node_score: float,
        output_score: float,
        is_end: bool,
        level: int,
        phrase: str = "",
        ac_threshold: float = 1.0,
    ):
        """Create a ContextState.

        Args:
          id:
            The node id, only for visualization now. A node is in [0, graph.num_nodes).
            The id of the root node is always 0.
          token:
            The token id.
          token_score:
            The bonus for each token during decoding, which will hopefully
            boost the token up to survive beam search.
          node_score:
            The accumulated bonus from root of graph to current node, it will be
            used to calculate the score for fail arc.
          output_score:
            The total scores of matched phrases, sum of the node_score of all
            the output node for current node.
          is_end:
            True if current token is the end of a context.
          level:
            The distance from current node to root.
          phrase:
            The context phrase of current state, the value is valid only when
            current state is end state (is_end == True).
          ac_threshold:
            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.
        """
        self.id = id
        self.token = token
        self.token_score = token_score
        self.node_score = node_score
        self.output_score = output_score
        self.is_end = is_end
        self.level = level
        self.next = {}
        self.phrase = phrase
        self.ac_threshold = ac_threshold
        self.fail = None
        self.output = None


class ContextGraph:
    """The ContextGraph is modified from Aho-Corasick which is mainly
    a Trie with a fail arc for each node.
    See https://en.wikipedia.org/wiki/Aho%E2%80%93Corasick_algorithm for more details
    of Aho-Corasick algorithm.

    A ContextGraph contains some words / phrases that we expect to boost their
    scores during decoding. If the substring of a decoded sequence matches the word / phrase
    in the ContextGraph, we will give the decoded sequence a bonus to make it survive
    beam search.
    """

    def __init__(self, context_score: float, ac_threshold: float = 1.0):
        """Initialize a ContextGraph with the given ``context_score``.

        A root node will be created (**NOTE:** the token of root is hardcoded to -1).

        Args:
          context_score:
            The bonus score for each token(note: NOT for each word/phrase, it means longer
            word/phrase will have larger bonus score, they have to be matched though).
            Note: This is just the default score for each token, the users can manually
            specify the context_score for each word/phrase (i.e. different phrase might
            have different token score).
          ac_threshold:
            The acoustic threshold (probability) to trigger the word/phrase, this argument
            is used only when applying the graph to keywords spotting system.
        """
        self.context_score = context_score
        self.ac_threshold = ac_threshold
        self.num_nodes = 0
        self.root = ContextState(
            id=self.num_nodes,
            token=-1,
            token_score=0,
            node_score=0,
            output_score=0,
            is_end=False,
            level=0,
        )
        self.root.fail = self.root

    def _fill_fail_output(self):
        """This function fills the fail arc for each trie node, it can be computed
        in linear time by performing a breadth-first search starting from the root.
        See https://en.wikipedia.org/wiki/Aho%E2%80%93Corasick_algorithm for the
        details of the algorithm.
        """
        queue = deque()
        for token, node in self.root.next.items():
            node.fail = self.root
            queue.append(node)
        while queue:
            current_node = queue.popleft()
            for token, node in current_node.next.items():
                fail = current_node.fail
                if token in fail.next:
                    fail = fail.next[token]
                else:
                    fail = fail.fail
                    while token not in fail.next:
                        fail = fail.fail
                        if fail.token == -1:  # root
                            break
                    if token in fail.next:
                        fail = fail.next[token]
                node.fail = fail
                # fill the output arc
                output = node.fail
                while not output.is_end:
                    output = output.fail
                    if output.token == -1:  # root
                        output = None
                        break
                node.output = output
                node.output_score += 0 if output is None else output.output_score
                queue.append(node)

    def build(
        self,
        token_ids: List[List[int]],
        phrases: Optional[List[str]] = None,
        scores: Optional[List[float]] = None,
        ac_thresholds: Optional[List[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
        for each trie node.

        See https://en.wikipedia.org/wiki/Trie for how to build a trie.

        Args:
          token_ids:
            The given token lists to build the ContextGraph, it is a list of
            token list, the token list contains the token ids
            for a word/phrase. The token id could be an id of a char
            (modeling with single Chinese char) or an id of a BPE
            (modeling with BPEs).
          phrases:
            The given phrases, they are the original text of the token_ids, the
            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`.
          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`.

        Note: The phrases would have shared states, the score of the shared states is
              the MAXIMUM value among all the tokens sharing this state.
        """
        num_phrases = len(token_ids)
        if phrases is not None:
            assert len(phrases) == num_phrases, (len(phrases), num_phrases)
        if scores is not None:
            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)

        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]
            ac_threshold = 0.0 if ac_thresholds is None else ac_thresholds[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
            threshold = self.ac_threshold if ac_threshold == 0.0 else ac_threshold
            for i, token in enumerate(tokens):
                node_next = {}
                if token not in node.next:
                    self.num_nodes += 1
                    is_end = i == len(tokens) - 1
                    node_score = node.node_score + context_score
                    node.next[token] = ContextState(
                        id=self.num_nodes,
                        token=token,
                        token_score=context_score,
                        node_score=node_score,
                        output_score=node_score if is_end else 0,
                        is_end=is_end,
                        level=i + 1,
                        phrase=phrase if is_end else "",
                        ac_threshold=threshold if is_end else 0.0,
                    )
                else:
                    # node exists, get the score of shared state.
                    token_score = max(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
                    is_end = i == len(tokens) - 1 or node.next[token].is_end
                    node.next[token].output_score = node_score if is_end else 0
                    node.next[token].is_end = is_end
                    if i == len(tokens) - 1:
                        node.next[token].phrase = phrase
                        node.next[token].ac_threshold = threshold
                node = node.next[token]
        self._fill_fail_output()

    def forward_one_step(
        self, state: ContextState, token: int, strict_mode: bool = True
    ) -> Tuple[float, ContextState, ContextState]:
        """Search the graph with given state and token.

        Args:
          state:
            The given token containing trie node to start.
          token:
            The given token.
          strict_mode:
            If the `strict_mode` is True, it can match multiple phrases simultaneously,
            and will continue to match longer phrase after matching a shorter one.
            If the `strict_mode` is False, it can only match one phrase at a time,
            when it matches a phrase, then the state will fall back to root state
            (i.e. forgetting all the history state and starting a new match). If
            the matched state have multiple outputs (node.output is not None), the
            longest phrase will be return.
            For example, if the phrases are `he`, `she` and `shell`, the query is
            `like shell`, when `strict_mode` is True, the query will match `he` and
            `she` at token `e` and `shell` at token `l`, while when `strict_mode`
            if False, the query can only match `she`(`she` is longer than `he`, so
            `she` not `he`) at token `e`.
            Caution: When applying this graph for keywords spotting system, the
            `strict_mode` MUST be True.

        Returns:
          Return a tuple of boosting score for current state, next state and matched
          state (if any). Note: Only returns the matched state with longest phrase of
          current state, even if there are multiple matches phrases. If no phrase
          matched, the matched state is None.
        """
        node = None
        score = 0
        # token matched
        if token in state.next:
            node = state.next[token]
            score = node.token_score
        else:
            # token not matched
            # We will trace along the fail arc until it matches the token or reaching
            # root of the graph.
            node = state.fail
            while token not in node.next:
                node = node.fail
                if node.token == -1:  # root
                    break

            if token in node.next:
                node = node.next[token]

            # The score of the fail path
            score = node.node_score - state.node_score
        assert node is not None

        # The matched node of current step, will only return the node with
        # longest phrase if there are multiple phrases matches this step.
        # None if no matched phrase.
        matched_node = (
            node if node.is_end else (node.output if node.output is not None else None)
        )
        if not strict_mode and node.output_score != 0:
            # output_score != 0 means at least on phrase matched
            assert matched_node is not None
            output_score = (
                node.node_score
                if node.is_end
                else (
                    node.node_score if node.output is None else node.output.node_score
                )
            )
            return (score + output_score - node.node_score, self.root, matched_node)
        assert (node.output_score != 0 and matched_node is not None) or (
            node.output_score == 0 and matched_node is None
        ), (
            node.output_score,
            matched_node,
        )
        return (score + node.output_score, node, matched_node)

    def is_matched(self, state: ContextState) -> Tuple[bool, ContextState]:
        """Whether current state matches any phrase (i.e. current state is the
        end state or the output of current state is not None.

        Args:
          state:
            The given state(trie node).

        Returns:
          Return a tuple of status and matched state.
        """
        if state.is_end:
            return True, state
        else:
            if state.output is not None:
                return True, state.output
            return False, None

    def finalize(self, state: ContextState) -> Tuple[float, ContextState]:
        """When reaching the end of the decoded sequence, we need to finalize
        the matching, the purpose is to subtract the added bonus score for the
        state that is not the end of a word/phrase.

        Args:
          state:
            The given state(trie node).

        Returns:
          Return a tuple of score and next state. If state is the end of a word/phrase
          the score is zero, otherwise the score is the score of a implicit fail arc
          to root. The next state is always root.
        """
        # The score of the fail arc
        score = -state.node_score
        return (score, self.root)

    def draw(
        self,
        title: Optional[str] = None,
        filename: Optional[str] = "",
        symbol_table: Optional[Dict[int, str]] = None,
    ) -> "Digraph":  # noqa
        """Visualize a ContextGraph via graphviz.

        Render ContextGraph as an image via graphviz, and return the Digraph object;
        and optionally save to file `filename`.
        `filename` must have a suffix that graphviz understands, such as
        `pdf`, `svg` or `png`.

        Note:
          You need to install graphviz to use this function::

            pip install graphviz

        Args:
           title:
              Title to be displayed in image, e.g. 'A simple FSA example'
           filename:
              Filename to (optionally) save to, e.g. 'foo.png', 'foo.svg',
              'foo.png'  (must have a suffix that graphviz understands).
           symbol_table:
              Map the token ids to symbols.
        Returns:
          A Diagraph from grahpviz.
        """

        try:
            import graphviz
        except Exception:
            print("You cannot use `to_dot` unless the graphviz package is installed.")
            raise

        graph_attr = {
            "rankdir": "LR",
            "size": "8.5,11",
            "center": "1",
            "orientation": "Portrait",
            "ranksep": "0.4",
            "nodesep": "0.25",
        }
        if title is not None:
            graph_attr["label"] = title

        default_node_attr = {
            "shape": "circle",
            "style": "bold",
            "fontsize": "14",
        }

        final_state_attr = {
            "shape": "doublecircle",
            "style": "bold",
            "fontsize": "14",
        }

        final_state = -1
        dot = graphviz.Digraph(name="Context Graph", graph_attr=graph_attr)

        seen = set()
        queue = deque()
        queue.append(self.root)
        # root id is always 0
        dot.node("0", label="0", **default_node_attr)
        dot.edge("0", "0", color="red")
        seen.add(0)

        while len(queue):
            current_node = queue.popleft()
            for token, node in current_node.next.items():
                if node.id not in seen:
                    node_score = f"{node.node_score:.2f}".rstrip("0").rstrip(".")
                    output_score = f"{node.output_score:.2f}".rstrip("0").rstrip(".")
                    label = f"{node.id}/({node_score}, {output_score})"
                    if node.is_end:
                        dot.node(str(node.id), label=label, **final_state_attr)
                    else:
                        dot.node(str(node.id), label=label, **default_node_attr)
                    seen.add(node.id)
                weight = f"{node.token_score:.2f}".rstrip("0").rstrip(".")
                label = str(token) if symbol_table is None else symbol_table[token]
                dot.edge(str(current_node.id), str(node.id), label=f"{label}/{weight}")
                dot.edge(
                    str(node.id),
                    str(node.fail.id),
                    color="red",
                )
                if node.output is not None:
                    dot.edge(
                        str(node.id),
                        str(node.output.id),
                        color="green",
                    )
                queue.append(node)

        if filename:
            _, extension = os.path.splitext(filename)
            if extension == "" or extension[0] != ".":
                raise ValueError(
                    "Filename needs to have a suffix like .png, .pdf, .svg: {}".format(
                        filename
                    )
                )

            import tempfile

            with tempfile.TemporaryDirectory() as tmp_dir:
                temp_fn = dot.render(
                    filename="temp",
                    directory=tmp_dir,
                    format=extension[1:],
                    cleanup=True,
                )

                shutil.move(temp_fn, filename)

        return dot


def _test(queries, score, strict_mode):
    contexts_str = [
        "S",
        "HE",
        "SHE",
        "SHELL",
        "HIS",
        "HERS",
        "HELLO",
        "THIS",
        "THEM",
    ]

    # test default score (1)
    contexts = []
    scores = []
    phrases = []
    for s in contexts_str:
        contexts.append([ord(x) for x in s])
        scores.append(round(score / len(s), 2))
        phrases.append(s)

    context_graph = ContextGraph(context_score=1)
    context_graph.build(token_ids=contexts, scores=scores, phrases=phrases)

    symbol_table = {}
    for contexts in contexts_str:
        for s in contexts:
            symbol_table[ord(s)] = s

    context_graph.draw(
        title="Graph for: " + " / ".join(contexts_str),
        filename=f"context_graph_{score}.pdf",
        symbol_table=symbol_table,
    )

    for query, expected_score in queries.items():
        total_scores = 0
        state = context_graph.root
        for q in query:
            score, state, phrase = context_graph.forward_one_step(
                state, ord(q), strict_mode
            )
            total_scores += score
        score, state = context_graph.finalize(state)
        assert state.token == -1, state.token
        total_scores += score
        assert round(total_scores, 2) == expected_score, (
            total_scores,
            expected_score,
            query,
        )


if __name__ == "__main__":
    # test default score
    queries = {
        "HEHERSHE": 14,  # "HE", "HE", "HERS", "S", "SHE", "HE"
        "HERSHE": 12,  # "HE", "HERS", "S", "SHE", "HE"
        "HISHE": 9,  # "HIS", "S", "SHE", "HE"
        "SHED": 6,  # "S", "SHE", "HE"
        "SHELF": 6,  # "S", "SHE", "HE"
        "HELL": 2,  # "HE"
        "HELLO": 7,  # "HE", "HELLO"
        "DHRHISQ": 4,  # "HIS", "S"
        "THEN": 2,  # "HE"
    }
    _test(queries, 0, True)

    queries = {
        "HEHERSHE": 7,  # "HE", "HE", "S", "HE"
        "HERSHE": 5,  # "HE", "S", "HE"
        "HISHE": 5,  # "HIS", "HE"
        "SHED": 3,  # "S", "HE"
        "SHELF": 3,  # "S", "HE"
        "HELL": 2,  # "HE"
        "HELLO": 2,  # "HE"
        "DHRHISQ": 3,  # "HIS"
        "THEN": 2,  # "HE"
    }
    _test(queries, 0, False)

    # test custom score
    # S : 5
    # HE : 5 (2.5 + 2.5)
    # SHE : 8.34 (5 + 1.67 + 1.67)
    # SHELL : 10.34 (5 + 1.67 + 1.67 + 1 + 1)
    # HIS : 5.84 (2.5 + 1.67 + 1.67)
    # HERS : 7.5 (2.5 + 2.5 + 1.25 + 1.25)
    # HELLO : 8 (2.5 + 2.5 + 1 + 1 + 1)
    # THIS : 5 (1.25 + 1.25 + 1.25 + 1.25)
    queries = {
        "HEHERSHE": 35.84,  # "HE", "HE", "HERS", "S", "SHE", "HE"
        "HERSHE": 30.84,  # "HE", "HERS", "S", "SHE", "HE"
        "HISHE": 24.18,  # "HIS", "S", "SHE", "HE"
        "SHED": 18.34,  # "S", "SHE", "HE"
        "SHELF": 18.34,  # "S", "SHE", "HE"
        "HELL": 5,  # "HE"
        "HELLO": 13,  # "HE", "HELLO"
        "DHRHISQ": 10.84,  # "HIS", "S"
        "THEN": 5,  # "HE"
    }

    _test(queries, 5, True)

    queries = {
        "HEHERSHE": 20,  # "HE", "HE", "S", "HE"
        "HERSHE": 15,  # "HE", "S", "HE"
        "HISHE": 10.84,  # "HIS", "HE"
        "SHED": 10,  # "S", "HE"
        "SHELF": 10,  # "S", "HE"
        "HELL": 5,  # "HE"
        "HELLO": 5,  # "HE"
        "DHRHISQ": 5.84,  # "HIS"
        "THEN": 5,  # "HE"
    }
    _test(queries, 5, False)