(self, input_ids: torch.LongTensor, scores: torch.FloatTensor)
| 457 | self.ngram_size = ngram_size |
| 458 | |
| 459 | def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: |
| 460 | num_batch_hypotheses = scores.shape[0] |
| 461 | cur_len = input_ids.shape[-1] |
| 462 | banned_batch_tokens = self._calc_banned_ngram_tokens(input_ids, num_batch_hypotheses, cur_len) |
| 463 | |
| 464 | for i, banned_tokens in enumerate(banned_batch_tokens): |
| 465 | scores[i, banned_tokens] = -float("inf") |
| 466 | |
| 467 | return scores |
| 468 | |
| 469 | def _calc_banned_ngram_tokens( |
| 470 | self, prev_input_ids: torch.Tensor, num_hypos: int, cur_len: int |
nothing calls this directly
no test coverage detected