Spaces:
Runtime error
Runtime error
| import torch | |
| from transformers import LogitsProcessor | |
| # HuggingFace's generate function does not yet support a `min_new_tokens`, so we need to add the functionality | |
| # ourselves by adding a custom logits processor. Adapted from: | |
| # https://huggingface.co/transformers/v4.1.1/_modules/transformers/generation_logits_process.html#MinLengthLogitsProcessor | |
| class MinNewTokensLogitsProcessor(LogitsProcessor): | |
| r""" | |
| A [`LogitsProcessor`] enforcing a minimum response length by setting the `EOS` probability to 0 until | |
| `min_new_tokens` new tokens have been generated since `input_length`. | |
| """ | |
| def __init__(self, min_new_tokens: int, eos_token_id: int, input_length: int): | |
| if not isinstance(min_new_tokens, int) or min_new_tokens < 0: | |
| raise ValueError(f"`min_new_tokens` has to be a positive integer, but is {min_new_tokens}") | |
| if not isinstance(eos_token_id, int) or eos_token_id < 0: | |
| raise ValueError(f"`eos_token_id` has to be a positive integer, but is {eos_token_id}") | |
| if not isinstance(input_length, int) or input_length < 0: | |
| raise ValueError(f"`input_length` has to be a positive integer, but is {input_length}") | |
| self.min_new_tokens = min_new_tokens | |
| self.eos_token_id = eos_token_id | |
| self.input_length = input_length | |
| def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: | |
| if not hasattr(self, "input_length"): | |
| raise ValueError("`save_input_length` has to be called before `__call__`") | |
| total_length = input_ids.shape[-1] | |
| response_len = total_length - self.input_length | |
| if response_len < self.min_new_tokens: | |
| scores[:, self.eos_token_id] = -float("inf") | |
| return scores | |