class stoppingcriteriaabc addstartdocstringsstoppingcriteriainputsdocstring def callself inputids torch longtensor scores torch floattensor kwargs bool raise notimplementederrorstoppingcriteria needs to be subclassed class maxlengthcriteriastoppingcriteria def initself maxlength int maxpositionembeddings optionalint none self maxlength maxlength self maxpositionembeddings maxpositionembeddings addstartdocstringsstoppingcriteriainputsdocstring def callself inputids torch longtensor scores torch floattensor kwargs bool curlen inputids shape1 isdone curlen self maxlength if self maxpositionembeddings is not none and not isdone and curlen self maxpositionembeddings logger warningonce this is a friendly reminder the current text generation call will exceed the model s predefined fmaximum length self maxpositionembeddings depending on the model you may observe exceptions performance degradation or nothing at all return isdone class maxnewtokenscriteriastoppingcriteria def initself startlength int maxnewtokens int warnings warn the class maxnewtokenscriteria is deprecated fplease use maxlengthcriteriamaxlengthstartlength maxnewtokens with maxlength startlength maxnewtokens instead futurewarning self startlength startlength self maxnewtokens maxnewtokens self maxlength startlength maxnewtokens addstartdocstringsstoppingcriteriainputsdocstring def callself inputids torch longtensor scores torch floattensor kwargs bool return inputids shape1 self maxlength class maxtimecriteriastoppingcriteria def initself maxtime float initialtimestamp optionalfloat none self maxtime maxtime self initialtimestamp time time if initialtimestamp is none else initialtimestamp addstartdocstringsstoppingcriteriainputsdocstring def callself inputids torch longtensor scores torch floattensor kwargs bool return time time self initialtimestamp self maxtime class stoppingcriterialistlist addstartdocstringsstoppingcriteriainputsdocstring def callself inputids torch longtensor scores torch floattensor kwargs bool return anycriteriainputids scores for criteria in self property def maxlengthself optionalint for stoppingcriterium in self if isinstancestoppingcriterium maxlengthcriteria return stoppingcriterium maxlength elif isinstancestoppingcriterium maxnewtokenscriteria return stoppingcriterium maxlength return none def validatestoppingcriteriastoppingcriteria stoppingcriterialist maxlength int stoppingcriterialist stoppingmaxlength stoppingcriteria maxlength newstoppingcriteria deepcopystoppingcriteria if stoppingmaxlength is not none and stoppingmaxlength maxlength warnings warnyou set different maxlength for stopping criteria and maxlength parameter userwarning elif stoppingmaxlength is none newstoppingcriteria appendmaxlengthcriteriamaxlengthmaxlength return newstoppingcriteria args input_ids torch longtensor of shape batch_size sequence_length indices of input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids scores torch floattensor of shape batch_size config vocab_size prediction scores of a language modeling head these can be scores for each vocabulary token before softmax or scores for each vocabulary token after softmax if this stopping criteria depends on the scores input make sure you pass return_dict_in_generate true output_scores true to generate kwargs dict str any optional additional stopping criteria specific kwargs return bool false indicates we should continue true indicates we should stop abstract base class for all stopping criteria that can be applied during generation if your stopping criteria depends on the scores input make sure you pass return_dict_in_generate true output_scores true to generate this class can be used to stop generation whenever the full generated number of tokens exceeds max_length keep in mind for decoder only type of transformers this will include the initial prompted tokens args max_length int the maximum length that the output sequence can have in number of tokens max_position_embeddings int optional the maximum model length as defined by the model s config max_position_embeddings attribute this class can be used to stop generation whenever the generated number of tokens exceeds max_new_tokens keep in mind for decoder only type of transformers this will not include the initial prompted tokens this is very close to maxlengthcriteria but ignores the number of initial tokens args start_length int the number of initial tokens max_new_tokens int the maximum number of tokens to generate this class can be used to stop generation whenever the full generation exceeds some amount of time by default the time will start being counted when you initialize this function you can override this by passing an initial_time args max_time float the maximum allowed time in seconds for the generation initial_time float optional defaults to time time the start of the generation allowed time
import time
import warnings
from abc import ABC
from copy import deepcopy
from typing import Optional

import torch

from ..utils import add_start_docstrings, logging


logger = logging.get_logger(__name__)


STOPPING_CRITERIA_INPUTS_DOCSTRING = r


class StoppingCriteria(ABC):
    

    @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        raise NotImplementedError("StoppingCriteria needs to be subclassed")


class MaxLengthCriteria(StoppingCriteria):
    

    def __init__(self, max_length: int, max_position_embeddings: Optional[int] = None):
        self.max_length = max_length
        self.max_position_embeddings = max_position_embeddings

    @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        cur_len = input_ids.shape[-1]
        is_done = cur_len >= self.max_length
        if self.max_position_embeddings is not None and not is_done and cur_len >= self.max_position_embeddings:
            logger.warning_once(
                "This is a friendly reminder - the current text generation call will exceed the model's predefined "
                f"maximum length ({self.max_position_embeddings}). Depending on the model, you may observe "
                "exceptions, performance degradation, or nothing at all."
            )
        return is_done


class MaxNewTokensCriteria(StoppingCriteria):
    

    def __init__(self, start_length: int, max_new_tokens: int):
        warnings.warn(
            "The class `MaxNewTokensCriteria` is deprecated. "
            f"Please use `MaxLengthCriteria(max_length={start_length + max_new_tokens})` "
            "with `max_length = start_length + max_new_tokens` instead.",
            FutureWarning,
        )
        self.start_length = start_length
        self.max_new_tokens = max_new_tokens
        self.max_length = start_length + max_new_tokens

    @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        return input_ids.shape[-1] >= self.max_length


class MaxTimeCriteria(StoppingCriteria):
    

    def __init__(self, max_time: float, initial_timestamp: Optional[float] = None):
        self.max_time = max_time
        self.initial_timestamp = time.time() if initial_timestamp is None else initial_timestamp

    @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        return time.time() - self.initial_timestamp > self.max_time


class StoppingCriteriaList(list):
    @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        return any(criteria(input_ids, scores) for criteria in self)

    @property
    def max_length(self) -> Optional[int]:
        for stopping_criterium in self:
            if isinstance(stopping_criterium, MaxLengthCriteria):
                return stopping_criterium.max_length
            elif isinstance(stopping_criterium, MaxNewTokensCriteria):
                return stopping_criterium.max_length
        return None


def validate_stopping_criteria(stopping_criteria: StoppingCriteriaList, max_length: int) -> StoppingCriteriaList:
    stopping_max_length = stopping_criteria.max_length
    new_stopping_criteria = deepcopy(stopping_criteria)
    if stopping_max_length is not None and stopping_max_length != max_length:
        warnings.warn("You set different `max_length` for stopping criteria and `max_length` parameter", UserWarning)
    elif stopping_max_length is None:
        new_stopping_criteria.append(MaxLengthCriteria(max_length=max_length))
    return new_stopping_criteria