

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


UpperCAmelCase_							:					List[str]		 =   logging.get_logger(__name__)


UpperCAmelCase_							:					Union[str, Any]		 =   r'''
    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.
        kwargs (`Dict[str, Any]`, *optional*):
            Additional stopping criteria specific kwargs.

    Return:
        `bool`. `False` indicates we should continue, `True` indicates we should stop.

'''






class      lowerCamelCase_      (	_lowercase	):




	@add_start_docstrings(__A							)
	def __call__(     self						:     str					,      __A						:     torch.LongTensor					,      __A						:     torch.FloatTensor					,      **__A						:     Optional[Any]							):
		raise NotImplementedError("""StoppingCriteria needs to be subclassed"""							)






class      lowerCamelCase_      (	_lowercase	):




	def __init__(     self						:     Union[str, Any]					,      __A						:     int					,      __A						:     Optional[int] = None							):
		__A :			Optional[int]		    =			max_length
		__A :			Optional[int]		    =			max_position_embeddings




	@add_start_docstrings(__A							)
	def __call__(     self						:     Union[str, Any]					,      __A						:     torch.LongTensor					,      __A						:     torch.FloatTensor					,      **__A						:     Optional[int]							):
		__A :			Optional[Any]		    =			input_ids.shape[-1]
		__A :			Union[str, 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      lowerCamelCase_      (	_lowercase	):




	def __init__(     self						:     List[str]					,      __A						:     int					,      __A						:     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."""					,      __A					,      )
		__A :			Dict		    =			start_length
		__A :			Optional[int]		    =			max_new_tokens
		__A :			Tuple		    =			start_length + max_new_tokens




	@add_start_docstrings(__A							)
	def __call__(     self						:     Tuple					,      __A						:     torch.LongTensor					,      __A						:     torch.FloatTensor					,      **__A						:     str							):
		return input_ids.shape[-1] >= self.max_length






class      lowerCamelCase_      (	_lowercase	):




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




	@add_start_docstrings(__A							)
	def __call__(     self						:     int					,      __A						:     torch.LongTensor					,      __A						:     torch.FloatTensor					,      **__A						:     Optional[int]							):
		return time.time() - self.initial_timestamp > self.max_time






class      lowerCamelCase_      (	_lowercase	):




	@add_start_docstrings(__A							)
	def __call__(     self						:     List[str]					,      __A						:     torch.LongTensor					,      __A						:     torch.FloatTensor					,      **__A						:     List[str]							):
		return any(criteria(__A					,      __A							) for criteria in self							)




	@property
	def        lowerCAmelCase_       (     self						:     int							):
		for stopping_criterium in self:
			if isinstance(__A					,      __A							):
				return stopping_criterium.max_length
			elif isinstance(__A					,      __A							):
				return stopping_criterium.max_length
		return None






def 							__SCREAMING_SNAKE_CASE     (   a__    :   StoppingCriteriaList			,a__    :   int       )					->				StoppingCriteriaList:
	__A :			int		    =			stopping_criteria.max_length
	__A :			Optional[int]		    =			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













# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
     import torch




class       snake_case__ (				TensorFormatter[Mapping, "torch.Tensor", Mapping]):

       '''simple docstring'''



       def __init__(				self      ,    a__=None      ,    **a__     )			->   Tuple:






            '''simple docstring'''





            super().__init__(features=a__     )
            __snake_case      :int 							=    torch_tensor_kwargs
            import torch  # noqa import torch at initialization



       def        __lowercase    (				self      ,    a__     )			->   List[str]:






            '''simple docstring'''





            import torch

            if isinstance(a__      ,    a__     ) and column:
                 if all(
                     isinstance(a__      ,    torch.Tensor     ) and x.shape == column[0].shape and x.dtype == column[0].dtype
                     for x in column     ):
                      return torch.stack(a__     )
            return column



       def        __lowercase    (				self      ,    a__     )			->   Union[str, Any]:






            '''simple docstring'''





            import torch

            if isinstance(a__      ,    (str, bytes, type(a__     ))     ):
                 return value
            elif isinstance(a__      ,    (np.character, np.ndarray)     ) and np.issubdtype(value.dtype      ,    np.character     ):
                 return value.tolist()

            __snake_case      :Optional[int] 							=    {}

            if isinstance(a__      ,    (np.number, np.ndarray)     ) and np.issubdtype(value.dtype      ,    np.integer     ):
                 __snake_case      :List[Any] 							=    {"""dtype""": torch.intaa}
            elif isinstance(a__      ,    (np.number, np.ndarray)     ) and np.issubdtype(value.dtype      ,    np.floating     ):
                 __snake_case      :Tuple 							=    {"""dtype""": torch.floataa}
            elif config.PIL_AVAILABLE and "PIL" in sys.modules:
                 import PIL.Image

                 if isinstance(a__      ,    PIL.Image.Image     ):
                      __snake_case      :Union[str, Any] 							=    np.asarray(a__     )
            return torch.tensor(a__      ,    **{**default_dtype, **self.torch_tensor_kwargs}     )



       def        __lowercase    (				self      ,    a__     )			->   str:






            '''simple docstring'''





            import torch

            # support for torch, tf, jax etc.
            if hasattr(a__      ,    """__array__"""     ) and not isinstance(a__      ,    torch.Tensor     ):
                 __snake_case      :Optional[Any] 							=    data_struct.__array__()
            # support for nested types like struct of list of struct
            if isinstance(a__      ,    np.ndarray     ):
                 if data_struct.dtype == object:  # torch tensors cannot be instantied from an array of objects
                      return self._consolidate([self.recursive_tensorize(a__     ) for substruct in data_struct]     )
            elif isinstance(a__      ,    (list, tuple)     ):
                 return self._consolidate([self.recursive_tensorize(a__     ) for substruct in data_struct]     )
            return self._tensorize(a__     )



       def        __lowercase    (				self      ,    a__     )			->   Union[str, Any]:






            '''simple docstring'''





            return map_nested(self._recursive_tensorize      ,    a__      ,    map_list=a__     )



       def        __lowercase    (				self      ,    a__     )			->   Mapping:






            '''simple docstring'''





            __snake_case      :Tuple 							=    self.numpy_arrow_extractor().extract_row(a__     )
            __snake_case      :Any 							=    self.python_features_decoder.decode_row(a__     )
            return self.recursive_tensorize(a__     )



       def        __lowercase    (				self      ,    a__     )			->   "torch.Tensor":






            '''simple docstring'''





            __snake_case      :List[str] 							=    self.numpy_arrow_extractor().extract_column(a__     )
            __snake_case      :List[Any] 							=    self.python_features_decoder.decode_column(a__      ,    pa_table.column_names[0]     )
            __snake_case      :Optional[Any] 							=    self.recursive_tensorize(a__     )
            __snake_case      :Any 							=    self._consolidate(a__     )
            return column



       def        __lowercase    (				self      ,    a__     )			->   Mapping:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    self.numpy_arrow_extractor().extract_batch(a__     )
            __snake_case      :Tuple 							=    self.python_features_decoder.decode_batch(a__     )
            __snake_case      :Tuple 							=    self.recursive_tensorize(a__     )
            for column_name in batch:
                 __snake_case      :str 							=    self._consolidate(batch[column_name]     )
            return batch


