(
self,
batch_size: int,
max_length: int,
num_beams: int,
device: torch.device,
length_penalty: Optional[float] = 1.0,
do_early_stopping: Optional[bool] = False,
num_beam_hyps_to_keep: Optional[int] = 1,
)
| 167 | """ |
| 168 | |
| 169 | def __init__( |
| 170 | self, |
| 171 | batch_size: int, |
| 172 | max_length: int, |
| 173 | num_beams: int, |
| 174 | device: torch.device, |
| 175 | length_penalty: Optional[float] = 1.0, |
| 176 | do_early_stopping: Optional[bool] = False, |
| 177 | num_beam_hyps_to_keep: Optional[int] = 1, |
| 178 | ): |
| 179 | self.max_length = max_length |
| 180 | self.num_beams = num_beams |
| 181 | self.device = device |
| 182 | self.length_penalty = length_penalty |
| 183 | self.do_early_stopping = do_early_stopping |
| 184 | self.num_beam_hyps_to_keep = num_beam_hyps_to_keep |
| 185 | |
| 186 | self._is_init = False |
| 187 | self._beam_hyps = [ |
| 188 | BeamHypotheses( |
| 189 | num_beams=self.num_beams, |
| 190 | max_length=self.max_length, |
| 191 | length_penalty=self.length_penalty, |
| 192 | early_stopping=self.do_early_stopping, |
| 193 | ) |
| 194 | for _ in range(batch_size) |
| 195 | ] |
| 196 | self._done = torch.tensor([False for _ in range(batch_size)], dtype=torch.bool, device=self.device) |
| 197 | |
| 198 | # if not isinstance(num_beams, int) or num_beams <= 1: |
| 199 | # raise ValueError( |
| 200 | # f"`num_beams` has to be an integer strictly greater than 1, but is {num_beams}. For `num_beams` == 1, one should make use of `greedy_search` instead." |
| 201 | # ) |
| 202 | |
| 203 | @property |
| 204 | def is_done(self) -> bool: |
nothing calls this directly
no test coverage detected