






'''simple docstring'''



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


SCREAMING_SNAKE_CASE						=	logging.get_logger(__name__)


SCREAMING_SNAKE_CASE						=	r'\n    Args:\n        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary.\n\n            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n            [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`):\n            Prediction scores of a language modeling head. These can be scores for each vocabulary token before SoftMax\n            or scores for each vocabulary token after SoftMax.\n        kwargs (`Dict[str, Any]`, *optional*):\n            Additional stopping criteria specific kwargs.\n\n    Return:\n        `bool`. `False` indicates we should continue, `True` indicates we should stop.\n\n'
class UpperCAmelCase_					(    __A				):




     """simple docstring"""




     @add_start_docstrings(UpperCAmelCase		)
     def __call__( self			:		Any     ,      UpperCAmelCase			:		Optional[Any]     ,      UpperCAmelCase			:		Tuple     ,      **UpperCAmelCase			:		Tuple		)  ->  bool:


      '''simple docstring'''

      raise NotImplementedError('''StoppingCriteria needs to be subclassed'''		)





class UpperCAmelCase_					(    __A				):




     """simple docstring"""




     def __init__( self			:		Any     ,      UpperCAmelCase			:		Tuple     ,      UpperCAmelCase			:		str = None		)  ->  Optional[int]:


      '''simple docstring'''

      lowercase					:   Any				  =max_length
      lowercase					:   int				  =max_position_embeddings






     @add_start_docstrings(UpperCAmelCase		)
     def __call__( self			:		int     ,      UpperCAmelCase			:		List[str]     ,      UpperCAmelCase			:		List[str]     ,      **UpperCAmelCase			:		Optional[int]		)  ->  bool:


      '''simple docstring'''

      lowercase					:   Tuple				  =input_ids.shape[-1]
      lowercase					:   Any				  =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 UpperCAmelCase_					(    __A				):




     """simple docstring"""




     def __init__( self			:		List[str]     ,      UpperCAmelCase			:		Optional[int]     ,      UpperCAmelCase			:		List[str]		)  ->  Tuple:


      '''simple docstring'''

      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.'''     ,      UpperCAmelCase     ,      )
      lowercase					:   Optional[int]				  =start_length
      lowercase					:   List[Any]				  =max_new_tokens
      lowercase					:   Optional[Any]				  =start_length + max_new_tokens






     @add_start_docstrings(UpperCAmelCase		)
     def __call__( self			:		str     ,      UpperCAmelCase			:		int     ,      UpperCAmelCase			:		Tuple     ,      **UpperCAmelCase			:		Any		)  ->  bool:


      '''simple docstring'''

      return input_ids.shape[-1] >= self.max_length





class UpperCAmelCase_					(    __A				):




     """simple docstring"""




     def __init__( self			:		str     ,      UpperCAmelCase			:		Union[str, Any]     ,      UpperCAmelCase			:		Dict = None		)  ->  int:


      '''simple docstring'''

      lowercase					:   Union[str, Any]				  =max_time
      lowercase					:   str				  =time.time() if initial_timestamp is None else initial_timestamp






     @add_start_docstrings(UpperCAmelCase		)
     def __call__( self			:		Tuple     ,      UpperCAmelCase			:		int     ,      UpperCAmelCase			:		Tuple     ,      **UpperCAmelCase			:		Optional[Any]		)  ->  bool:


      '''simple docstring'''

      return time.time() - self.initial_timestamp > self.max_time





class UpperCAmelCase_					(    __A				):




     """simple docstring"""




     @add_start_docstrings(UpperCAmelCase		)
     def __call__( self			:		str     ,      UpperCAmelCase			:		List[str]     ,      UpperCAmelCase			:		List[str]     ,      **UpperCAmelCase			:		Tuple		)  ->  bool:


      '''simple docstring'''

      return any(criteria(UpperCAmelCase     ,      UpperCAmelCase		) for criteria in self		)






     @property
     def 				A__      ( self			:		List[str]		)  ->  Optional[int]:


      '''simple docstring'''

      for stopping_criterium in self:
       if isinstance(UpperCAmelCase     ,      UpperCAmelCase		):
        return stopping_criterium.max_length
       elif isinstance(UpperCAmelCase     ,      UpperCAmelCase		):
        return stopping_criterium.max_length
      return None





def 							lowercase_				(    __A :       StoppingCriteriaList					,      __A :       int			)       ->      Optional[Any]:



 """simple docstring"""




 lowercase					:   Tuple				  =stopping_criteria.max_length
 lowercase					:   List[str]				  =deepcopy(__A			)
 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'''					,      __A			)
 elif stopping_max_length is None:
  new_stopping_criteria.append(MaxLengthCriteria(max_length=__A			)			)
 return new_stopping_criteria











'''simple docstring'''



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


SCREAMING_SNAKE_CASE						=	{
    'configuration_mega': ['MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MegaConfig', 'MegaOnnxConfig'],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE						=	[
          'MEGA_PRETRAINED_MODEL_ARCHIVE_LIST',
          'MegaForCausalLM',
          'MegaForMaskedLM',
          'MegaForMultipleChoice',
          'MegaForQuestionAnswering',
          'MegaForSequenceClassification',
          'MegaForTokenClassification',
          'MegaModel',
          'MegaPreTrainedModel',
      ]

if TYPE_CHECKING:
      from .configuration_mega import MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP, MegaConfig, MegaOnnxConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_mega import (
                MEGA_PRETRAINED_MODEL_ARCHIVE_LIST,
                MegaForCausalLM,
                MegaForMaskedLM,
                MegaForMultipleChoice,
                MegaForQuestionAnswering,
                MegaForSequenceClassification,
                MegaForTokenClassification,
                MegaModel,
                MegaPreTrainedModel,
            )

else:
      import sys

      SCREAMING_SNAKE_CASE						=	_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



