Beam search implementation.
| 32 | |
| 33 | |
| 34 | class BeamSearch(torch.nn.Module): |
| 35 | """Beam search implementation.""" |
| 36 | |
| 37 | def __init__( |
| 38 | self, |
| 39 | scorers: Dict[str, ScorerInterface], |
| 40 | weights: Dict[str, float], |
| 41 | beam_size: int, |
| 42 | vocab_size: int, |
| 43 | sos: int, |
| 44 | eos: int, |
| 45 | token_list: List[str] = None, |
| 46 | pre_beam_ratio: float = 1.5, |
| 47 | pre_beam_score_key: str = None, |
| 48 | ): |
| 49 | """Initialize beam search. |
| 50 | |
| 51 | Args: |
| 52 | scorers (dict[str, ScorerInterface]): Dict of decoder modules |
| 53 | e.g., Decoder, CTCPrefixScorer, LM |
| 54 | The scorer will be ignored if it is `None` |
| 55 | weights (dict[str, float]): Dict of weights for each scorers |
| 56 | The scorer will be ignored if its weight is 0 |
| 57 | beam_size (int): The number of hypotheses kept during search |
| 58 | vocab_size (int): The number of vocabulary |
| 59 | sos (int): Start of sequence id |
| 60 | eos (int): End of sequence id |
| 61 | token_list (list[str]): List of tokens for debug log |
| 62 | pre_beam_score_key (str): key of scores to perform pre-beam search |
| 63 | pre_beam_ratio (float): beam size in the pre-beam search |
| 64 | will be `int(pre_beam_ratio * beam_size)` |
| 65 | |
| 66 | """ |
| 67 | super().__init__() |
| 68 | # set scorers |
| 69 | self.weights = weights |
| 70 | self.scorers = dict() |
| 71 | self.full_scorers = dict() |
| 72 | self.part_scorers = dict() |
| 73 | # this module dict is required for recursive cast |
| 74 | # `self.to(device, dtype)` in `recog.py` |
| 75 | self.nn_dict = torch.nn.ModuleDict() |
| 76 | for k, v in scorers.items(): |
| 77 | w = weights.get(k, 0) |
| 78 | if w == 0 or v is None: |
| 79 | continue |
| 80 | assert isinstance( |
| 81 | v, ScorerInterface |
| 82 | ), f"{k} ({type(v)}) does not implement ScorerInterface" |
| 83 | self.scorers[k] = v |
| 84 | if isinstance(v, PartialScorerInterface): |
| 85 | self.part_scorers[k] = v |
| 86 | else: |
| 87 | self.full_scorers[k] = v |
| 88 | if isinstance(v, torch.nn.Module): |
| 89 | self.nn_dict[k] = v |
| 90 | |
| 91 | # set configurations |
no outgoing calls
no test coverage detected
searching dependent graphs…