| 22 | """ |
| 23 | |
| 24 | class ConstrainedLogitsProcessor(LogitsProcessor): |
| 25 | |
| 26 | def __init__( |
| 27 | self, |
| 28 | prefix_allowed_tokens_fn: Callable[[int, torch.Tensor], List[int]], |
| 29 | num_beams: int, |
| 30 | base_model: str = None, |
| 31 | eos_token_id: int = None |
| 32 | ): |
| 33 | self._prefix_allowed_tokens_fn = prefix_allowed_tokens_fn |
| 34 | self._num_beams = num_beams |
| 35 | self.count=0 |
| 36 | self.base_model = base_model |
| 37 | self.eos_token_id = eos_token_id |
| 38 | if self.base_model.lower().find("gpt2") > -1: |
| 39 | self.prefix_index = 4 |
| 40 | else: |
| 41 | self.prefix_index = 3 |
| 42 | |
| 43 | |
| 44 | @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) |
| 45 | def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: |
| 46 | scores = torch.nn.functional.log_softmax(scores, dim=-1) |
| 47 | mask = torch.full_like(scores, float('-inf')) |
| 48 | |
| 49 | for batch_id, beam_sent in enumerate(input_ids.view(-1, self._num_beams, input_ids.shape[-1])): |
| 50 | for beam_id, sent in enumerate(beam_sent): |
| 51 | if self.count == 0: |
| 52 | hash_key = sent[-self.prefix_index:] |
| 53 | else: |
| 54 | hash_key=sent[-self.count:] |
| 55 | hash_key = hash_key.tolist() |
| 56 | prefix_allowed_tokens = self._prefix_allowed_tokens_fn(batch_id, hash_key) |
| 57 | |
| 58 | if len(prefix_allowed_tokens) == 0: |
| 59 | warnings.warn( |
| 60 | f"No valid tokens found for hash_key {hash_key} at step {self.count}. " |
| 61 | f"This indicates the model generated an unexpected token. " |
| 62 | ) |
| 63 | # Force EOS token to end invalid sequence |
| 64 | if self.eos_token_id is not None: |
| 65 | mask[batch_id * self._num_beams + beam_id, self.eos_token_id] = 0 |
| 66 | continue |
| 67 | |
| 68 | mask[batch_id * self._num_beams + beam_id, prefix_allowed_tokens] = 0 |
| 69 | |
| 70 | self.count += 1 |
| 71 | |
| 72 | scores = scores + mask |
| 73 | return scores |
no outgoing calls
no test coverage detected