Source code for mbrs.decoders.rerank

 1from __future__ import annotations
 2
 3from mbrs import timer
 4
 5from . import DecoderReferenceless, register
 6
 7

[docs]
 8@register("rerank")
 9class DecoderRerank(DecoderReferenceless):
10    """Reranking decoder class.
11
12    - Time complexity: O(N)
13    - Space complexity: O(N)
14    """
15

[docs]
16    def decode(
17        self, hypotheses: list[str], source: str, nbest: int = 1
18    ) -> DecoderRerank.Output:
19        """Select the n-best hypotheses based on the strategy.
20
21        Args:
22            hypotheses (list[str]): Hypotheses.
23            source (str): A source.
24            nbest (int): Return the n-best hypotheses.
25
26        Returns:
27            DecoderRerank.Output: The n-best hypotheses.
28        """
29        with timer.measure("rerank"):
30            scores = self.metric.scores(hypotheses, sources=[source] * len(hypotheses))
31
32        selector_outputs = self.select(hypotheses, scores, nbest=nbest, source=source)
33        return (
34            self.Output(
35                idx=selector_outputs.idx,
36                sentence=selector_outputs.sentence,
37                score=selector_outputs.score,
38            )
39            | selector_outputs
40        )