





"""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


a_ =			logging.get_logger(__name__)


a_ =			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 		__snake_case					(		SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 @add_start_docstrings(__lowerCamelCase  )
 def __call__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						**__lowerCamelCase  ):



     '''simple docstring'''




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






class 		__snake_case					(		SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase					,						__lowerCamelCase = None  ):



     '''simple docstring'''




     __A    :						Any						 =     max_length
     __A    :						Any						 =     max_position_embeddings







 @add_start_docstrings(__lowerCamelCase  )
 def __call__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						**__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						int						 =     input_ids.shape[-1]
     __A    :						Tuple						 =     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 		__snake_case					(		SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''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.'''					,						__lowerCamelCase					,						)
     __A    :						Union[str, Any]						 =     start_length
     __A    :						str						 =     max_new_tokens
     __A    :						Dict						 =     start_length + max_new_tokens







 @add_start_docstrings(__lowerCamelCase  )
 def __call__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						**__lowerCamelCase  ):



     '''simple docstring'''




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






class 		__snake_case					(		SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase					,						__lowerCamelCase = None  ):



     '''simple docstring'''




     __A    :						Optional[int]						 =     max_time
     __A    :						Tuple						 =     time.time() if initial_timestamp is None else initial_timestamp







 @add_start_docstrings(__lowerCamelCase  )
 def __call__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						**__lowerCamelCase  ):



     '''simple docstring'''




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






class 		__snake_case					(		SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 @add_start_docstrings(__lowerCamelCase  )
 def __call__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						**__lowerCamelCase  ):



     '''simple docstring'''




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







 @property
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




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







def 		__lowercase     (   snake_case_  :  StoppingCriteriaList				,snake_case_  :  int						)   ->StoppingCriteriaList:



    '''simple docstring'''







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











"""simple docstring"""



import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


a_ =			50000
a_ =			5000

a_ ,			a_ =			os.path.split(__file__)
a_ =			os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))







@get_duration
def 		__lowercase     (   snake_case_  :  datasets.Dataset				,snake_case_  :  Tuple						)   ->Tuple:



    '''simple docstring'''







    for i in range(snake_case_						):
        __A    :						int						 =     dataset[i]







@get_duration
def 		__lowercase     (   snake_case_  :  datasets.Dataset				,snake_case_  :  Optional[Any]				,snake_case_  :  int						)   ->Tuple:



    '''simple docstring'''







    for i in range(0				,len(snake_case_						)				,snake_case_						):
        __A    :						List[str]						 =     dataset[i : i + batch_size]







@get_duration
def 		__lowercase     (   snake_case_  :  datasets.Dataset				,snake_case_  :  List[Any]				,snake_case_  :  Any						)   ->int:



    '''simple docstring'''







    with dataset.formatted_as(type=snake_case_						):
        for i in range(snake_case_						):
            __A    :						Union[str, Any]						 =     dataset[i]







@get_duration
def 		__lowercase     (   snake_case_  :  datasets.Dataset				,snake_case_  :  Any				,snake_case_  :  Union[str, Any]				,snake_case_  :  Optional[int]						)   ->Union[str, Any]:



    '''simple docstring'''







    with dataset.formatted_as(type=snake_case_						):
        for i in range(0				,snake_case_				,snake_case_						):
            __A    :						Dict						 =     dataset[i : i + batch_size]







def 		__lowercase     (   )   ->Optional[int]:



    '''simple docstring'''







    __A    :						int						 =     {'''num examples''': SPEED_TEST_N_EXAMPLES}
    __A    :						Optional[int]						 =     [
        (read, {'''length''': SMALL_TEST}),
        (read, {'''length''': SPEED_TEST_N_EXAMPLES}),
        (read_batch, {'''length''': SPEED_TEST_N_EXAMPLES, '''batch_size''': 10}),
        (read_batch, {'''length''': SPEED_TEST_N_EXAMPLES, '''batch_size''': 100}),
        (read_batch, {'''length''': SPEED_TEST_N_EXAMPLES, '''batch_size''': 1000}),
        (read_formatted, {'''type''': '''numpy''', '''length''': SMALL_TEST}),
        (read_formatted, {'''type''': '''pandas''', '''length''': SMALL_TEST}),
        (read_formatted, {'''type''': '''torch''', '''length''': SMALL_TEST}),
        (read_formatted, {'''type''': '''tensorflow''', '''length''': SMALL_TEST}),
        (read_formatted_batch, {'''type''': '''numpy''', '''length''': SMALL_TEST, '''batch_size''': 10}),
        (read_formatted_batch, {'''type''': '''numpy''', '''length''': SMALL_TEST, '''batch_size''': 1000}),
    ]

    __A    :						int						 =     [
        (read, {'''length''': SMALL_TEST}),
        (read, {'''length''': SPEED_TEST_N_EXAMPLES}),
        (read_batch, {'''length''': SPEED_TEST_N_EXAMPLES, '''batch_size''': 10}),
        (read_batch, {'''length''': SPEED_TEST_N_EXAMPLES, '''batch_size''': 100}),
        (read_batch, {'''length''': SPEED_TEST_N_EXAMPLES, '''batch_size''': 1000}),
        (read_formatted, {'''type''': '''numpy''', '''length''': SMALL_TEST}),
        (read_formatted_batch, {'''type''': '''numpy''', '''length''': SMALL_TEST, '''batch_size''': 10}),
        (read_formatted_batch, {'''type''': '''numpy''', '''length''': SMALL_TEST, '''batch_size''': 1000}),
    ]
    with tempfile.TemporaryDirectory() as tmp_dir:
        print('''generating dataset'''						)
        __A    :						Any						 =     datasets.Features(
            {'''list''': datasets.Sequence(datasets.Value('''float32'''						)						), '''numbers''': datasets.Value('''float32'''						)}						)
        __A    :						List[Any]						 =     generate_example_dataset(
            os.path.join(snake_case_				,'''dataset.arrow'''						)				,snake_case_				,num_examples=snake_case_				,seq_shapes={'''list''': (100,)}				,)
        print('''first set of iterations'''						)
        for func, kwargs in functions:
            print(func.__name__				,str(snake_case_						)						)
            __A    :						Dict						 =     func(snake_case_				,**snake_case_						)

        print('''shuffling dataset'''						)
        __A    :						int						 =     dataset.shuffle()
        print('''Second set of iterations (after shuffling'''						)
        for func, kwargs in functions_shuffled:
            print('''shuffled '''				,func.__name__				,str(snake_case_						)						)
            __A    :						Optional[Any]						 =     func(
                snake_case_				,**snake_case_						)

    with open(snake_case_				,'''wb'''						) as f:
        f.write(json.dumps(snake_case_						).encode('''utf-8'''						)						)


if __name__ == "__main__":  # useful to run the profiler
       benchmark_iterating()




