MCPcopy Create free account
hub / github.com/THUDM/GLM / MinLengthLogitsProcessor

Class MinLengthLogitsProcessor

generation_utils.py:416–441  ·  view source on GitHub ↗

r""" :class:`transformers.LogitsProcessor` enforcing a min-length by setting EOS probability to 0. Args: min_length (:obj:`int`): The minimum length below which the score of :obj:`eos_token_id` is set to :obj:`-float("Inf")`. eos_token_id (:obj:`int`):

Source from the content-addressed store, hash-verified

414
415
416class MinLengthLogitsProcessor(LogitsProcessor):
417 r"""
418 :class:`transformers.LogitsProcessor` enforcing a min-length by setting EOS probability to 0.
419
420 Args:
421 min_length (:obj:`int`):
422 The minimum length below which the score of :obj:`eos_token_id` is set to :obj:`-float("Inf")`.
423 eos_token_id (:obj:`int`):
424 The id of the `end-of-sequence` token.
425 """
426
427 def __init__(self, min_length: int, eos_token_id: int):
428 if not isinstance(min_length, int) or min_length < 0:
429 raise ValueError(f"`min_length` has to be a positive integer, but is {min_length}")
430
431 if not isinstance(eos_token_id, int) or eos_token_id < 0:
432 raise ValueError(f"`eos_token_id` has to be a positive integer, but is {eos_token_id}")
433
434 self.min_length = min_length
435 self.eos_token_id = eos_token_id
436
437 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
438 cur_len = input_ids.shape[-1]
439 if cur_len < self.min_length:
440 scores[:, self.eos_token_id] = -float("inf")
441 return scores
442
443
444class NoRepeatNGramLogitsProcessor(LogitsProcessor):

Callers 1

__init__Method · 0.90

Calls

no outgoing calls

Tested by

no test coverage detected