






import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


__lowerCamelCase :      Optional[int]										=				get_logger(__name__)


__lowerCamelCase :      Any										=				R"""
    Args:
        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary.

            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):
            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam
            search or log softmax for each vocabulary token when using beam search
        kwargs (`Dict[str, Any]`, *optional*):
            Additional logits processor specific kwargs.

    Return:
        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.

"""



class 							SCREAMING_SNAKE_CASE__						:





					"""simple docstring"""






					@add_start_docstrings(__A	)
					def __call__( self					:  Union[str, Any]							,		__A					:  jnp.ndarray							,		__A					:  jnp.ndarray	):
						raise NotImplementedError(
						    f'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''	)



class 							SCREAMING_SNAKE_CASE__						:





					"""simple docstring"""






					@add_start_docstrings(__A	)
					def __call__( self					:  Any							,		__A					:  jnp.ndarray							,		__A					:  jnp.ndarray	):
						raise NotImplementedError(
						    f'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''	)



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					@add_start_docstrings(__A	)
					def __call__( self					:  int							,		__A					:  jnp.ndarray							,		__A					:  jnp.ndarray							,		__A					:  int							,		**__A					:  Optional[int]	):
						for processor in self:
							snake_case__			:				Optional[Any]       =  inspect.signature(processor.__call__	).parameters
							if len(__A	) > 3:
								if not all(arg in kwargs for arg in list(function_args.keys()	)[2:]	):
									raise ValueError(
									    f'''Make sure that all the required parameters: {list(function_args.keys()	)} for '''
									    f'''{processor.__class__} are passed to the logits processor.'''	)
								snake_case__			:				Optional[int]       =  processor(__A							,		__A							,		__A							,		**__A	)
							else:
								snake_case__			:				Union[str, Any]       =  processor(__A							,		__A							,		__A	)
						return scores



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					def __init__( self					:  Optional[int]							,		__A					:  float	):
						if not isinstance(__A							,		__A	) or not (temperature > 0):
							raise ValueError(f'''`temperature` has to be a strictly positive float, but is {temperature}'''	)

						snake_case__			:				Dict       =  temperature






					def __call__( self					:  Tuple							,		__A					:  jnp.ndarray							,		__A					:  jnp.ndarray							,		__A					:  int	):
						snake_case__			:				List[str]       =  scores / self.temperature
						return scores



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					def __init__( self					:  Tuple							,		__A					:  float							,		__A					:  float = -float("Inf"	)							,		__A					:  int = 1	):
						if not isinstance(__A							,		__A	) or (top_p < 0 or top_p > 1.0):
							raise ValueError(f'''`top_p` has to be a float > 0 and < 1, but is {top_p}'''	)
						if not isinstance(__A							,		__A	) or (min_tokens_to_keep < 1):
							raise ValueError(f'''`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}'''	)

						snake_case__			:				Any       =  top_p
						snake_case__			:				str       =  filter_value
						snake_case__			:				Optional[Any]       =  min_tokens_to_keep






					def __call__( self					:  Tuple							,		__A					:  jnp.ndarray							,		__A					:  jnp.ndarray							,		__A					:  int	):
						snake_case__,						snake_case__			:				Optional[Any]       =  lax.top_k(__A							,		scores.shape[-1]	)

						snake_case__			:				List[str]       =  jnp.full_like(__A							,		self.filter_value	)
						snake_case__			:				Dict       =  jax.nn.softmax(__A							,		axis=-1	).cumsum(axis=-1	)
						snake_case__			:				List[str]       =  cumulative_probs < self.top_p

						# include the token that is higher than top_p as well
						snake_case__			:				Optional[int]       =  jnp.roll(__A							,		1	)
						score_mask |= score_mask.at[:, 0].set(__A	)

						# min tokens to keep
						snake_case__			:				List[str]       =  score_mask.at[:, : self.min_tokens_to_keep].set(__A	)

						snake_case__			:				List[str]       =  jnp.where(__A							,		__A							,		__A	)
						snake_case__			:				Any       =  jax.lax.sort_key_val(__A							,		__A	)[-1]

						return next_scores



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					def __init__( self					:  str							,		__A					:  int							,		__A					:  float = -float("Inf"	)							,		__A					:  int = 1	):
						if not isinstance(__A							,		__A	) or top_k <= 0:
							raise ValueError(f'''`top_k` has to be a strictly positive integer, but is {top_k}'''	)

						snake_case__			:				int       =  max(__A							,		__A	)
						snake_case__			:				Dict       =  filter_value






					def __call__( self					:  int							,		__A					:  jnp.ndarray							,		__A					:  jnp.ndarray							,		__A					:  int	):
						snake_case__,						snake_case__			:				Tuple       =  scores.shape
						snake_case__			:				Dict       =  jnp.full(batch_size * vocab_size							,		self.filter_value	)

						snake_case__			:				List[str]       =  min(self.top_k							,		scores.shape[-1]	)  # Safety check
						snake_case__,						snake_case__			:				Any       =  lax.top_k(__A							,		__A	)
						snake_case__			:				Dict       =  jnp.broadcast_to((jnp.arange(__A	) * vocab_size)[:, None]							,		(batch_size, topk)	).flatten()
						snake_case__			:				Optional[Any]       =  topk_scores.flatten()
						snake_case__			:				List[str]       =  topk_indices.flatten() + shift

						snake_case__			:				Dict       =  next_scores_flat.at[topk_indices_flat].set(__A	)
						snake_case__			:				List[Any]       =  next_scores_flat.reshape(__A							,		__A	)
						return next_scores



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					def __init__( self					:  List[Any]							,		__A					:  int	):
						snake_case__			:				str       =  bos_token_id






					def __call__( self					:  List[Any]							,		__A					:  jnp.ndarray							,		__A					:  jnp.ndarray							,		__A					:  int	):
						snake_case__			:				Any       =  jnp.full(scores.shape							,		-float("inf"	)	)

						snake_case__			:				List[Any]       =  1 - jnp.bool_(cur_len - 1	)

						snake_case__			:				Optional[int]       =  jnp.where(__A							,		new_scores.at[:, self.bos_token_id].set(0	)							,		__A	)

						return scores



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					def __init__( self					:  str							,		__A					:  int							,		__A					:  int	):
						snake_case__			:				List[Any]       =  max_length
						snake_case__			:				List[Any]       =  eos_token_id






					def __call__( self					:  int							,		__A					:  jnp.ndarray							,		__A					:  jnp.ndarray							,		__A					:  int	):
						snake_case__			:				Union[str, Any]       =  jnp.full(scores.shape							,		-float("inf"	)	)

						snake_case__			:				Optional[int]       =  1 - jnp.bool_(cur_len - self.max_length + 1	)

						snake_case__			:				Any       =  jnp.where(__A							,		new_scores.at[:, self.eos_token_id].set(0	)							,		__A	)

						return scores



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					def __init__( self					:  Optional[int]							,		__A					:  int							,		__A					:  int	):
						if not isinstance(__A							,		__A	) or min_length < 0:
							raise ValueError(f'''`min_length` has to be a positive integer, but is {min_length}'''	)

						if not isinstance(__A							,		__A	) or eos_token_id < 0:
							raise ValueError(f'''`eos_token_id` has to be a positive integer, but is {eos_token_id}'''	)

						snake_case__			:				str       =  min_length
						snake_case__			:				Any       =  eos_token_id






					def __call__( self					:  str							,		__A					:  jnp.ndarray							,		__A					:  jnp.ndarray							,		__A					:  int	):
						# create boolean flag to decide if min length penalty should be applied
						snake_case__			:				Optional[int]       =  1 - jnp.clip(cur_len - self.min_length							,		0							,		1	)

						snake_case__			:				Tuple       =  jnp.where(__A							,		scores.at[:, self.eos_token_id].set(-float("inf"	)	)							,		__A	)

						return scores



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					def __init__( self					:  Optional[int]							,		__A					:  str							,		__A					:  Union[str, Any]	):
						snake_case__			:				Union[str, Any]       =  list(__A	)
						snake_case__			:				List[Any]       =  begin_index






					def __call__( self					:  Tuple							,		__A					:  Dict							,		__A					:  int							,		__A					:  int	):
						snake_case__			:				Union[str, Any]       =  1 - jnp.bool_(cur_len - self.begin_index	)

						snake_case__			:				Optional[int]       =  jnp.where(__A							,		scores.at[:, self.begin_suppress_tokens].set(-float("inf"	)	)							,		__A	)

						return scores



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					def __init__( self					:  Optional[int]							,		__A					:  list	):
						snake_case__			:				Optional[Any]       =  list(__A	)






					def __call__( self					:  str							,		__A					:  jnp.ndarray							,		__A					:  jnp.ndarray							,		__A					:  int	):
						snake_case__			:				Optional[int]       =  scores.at[..., self.suppress_tokens].set(-float("inf"	)	)

						return scores



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					def __init__( self					:  Tuple							,		__A					:  List[str]	):
						snake_case__			:				Optional[int]       =  dict(__A	)
						# Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
						# index of the array corresponds to the index of the token to be forced, for XLA compatibility.
						# Indexes without forced tokens will have a negative value.
						snake_case__			:				Dict       =  jnp.ones((max(force_token_map.keys()	) + 1)							,		dtype=jnp.intaa	) * -1
						for index, token in force_token_map.items():
							if token is not None:
								snake_case__			:				Any       =  force_token_array.at[index].set(__A	)
						snake_case__			:				Any       =  jnp.intaa(__A	)






					def __call__( self					:  Any							,		__A					:  jnp.ndarray							,		__A					:  jnp.ndarray							,		__A					:  int	):
						def _force_token(__A					:  List[str]	):
							snake_case__			:				Union[str, Any]       =  scores.shape[0]
							snake_case__			:				int       =  self.force_token_array[generation_idx]

							snake_case__			:				str       =  jnp.ones_like(__A							,		dtype=scores.dtype	) * -float("inf"	)
							snake_case__			:				Union[str, Any]       =  jnp.zeros((batch_size, 1)							,		dtype=scores.dtype	)
							snake_case__			:				Optional[Any]       =  lax.dynamic_update_slice(__A							,		__A							,		(0, current_token)	)
							return new_scores

						snake_case__			:				List[Any]       =  lax.cond(
						    cur_len >= self.force_token_array.shape[0]							,		lambda: scores							,		lambda: lax.cond(
						        self.force_token_array[cur_len] >= 0							,		lambda: _force_token(__A	)							,		lambda: scores							,		)							,		)
						return scores



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					def __init__( self					:  str							,		__A					:  Dict							,		__A					:  Tuple							,		__A					:  str	):
						snake_case__			:				str       =  generate_config.eos_token_id
						snake_case__			:				List[str]       =  generate_config.no_timestamps_token_id
						snake_case__			:				Optional[Any]       =  generate_config.no_timestamps_token_id + 1

						snake_case__			:				Optional[int]       =  decoder_input_length + 1

						if generate_config.is_multilingual:
							# room for language token and task token
							self.begin_index += 2
						if hasattr(__A							,		"max_initial_timestamp_index"	):
							snake_case__			:				Tuple       =  generate_config.max_initial_timestamp_index
						else:
							snake_case__			:				str       =  model_config.vocab_size
						if self.max_initial_timestamp_index is None:
							snake_case__			:				Optional[Any]       =  model_config.vocab_size






					def __call__( self					:  Tuple							,		__A					:  Dict							,		__A					:  str							,		__A					:  Optional[Any]	):
						# suppress <|notimestamps|> which is handled by without_timestamps
						snake_case__			:				Optional[int]       =  scores.at[:, self.no_timestamps_token_id].set(-float("inf"	)	)

						def handle_pairs(__A					:  List[str]							,		__A					:  int	):
							snake_case__			:				Tuple       =  jnp.where((cur_len - self.begin_index) >= 1							,		__A							,		__A	)
							snake_case__			:				Any       =  jnp.where(
							    input_ids_k[cur_len - 1] >= self.timestamp_begin							,		True and last_was_timestamp							,		__A							,		)

							snake_case__			:				List[Any]       =  jnp.where((cur_len - self.begin_index) < 2							,		__A							,		__A	)
							snake_case__			:				str       =  jnp.where(
							    input_ids_k[cur_len - 2] >= self.timestamp_begin							,		__A							,		__A							,		)

							return jnp.where(
							    __A							,		jnp.where(
							        penultimate_was_timestamp > 0							,		scores_k.at[self.timestamp_begin :].set(-float("inf"	)	)							,		scores_k.at[: self.eos_token_id].set(-float("inf"	)	)							,		)							,		__A							,		)

						snake_case__			:				Tuple       =  jax.vmap(__A	)(__A							,		__A	)

						snake_case__			:				Dict       =  jnp.where(cur_len == self.begin_index							,		__A							,		__A	)
						snake_case__			:				Tuple       =  jnp.where(
						    self.max_initial_timestamp_index is not None							,		True and apply_max_initial_timestamp							,		__A							,		)

						snake_case__			:				List[str]       =  self.timestamp_begin + self.max_initial_timestamp_index

						snake_case__			:				Optional[Any]       =  jnp.where(
						    __A							,		scores.at[:, last_allowed + 1 :].set(-float("inf"	)	)							,		__A							,		)

						# if sum of probability over timestamps is above any other token, sample timestamp
						snake_case__			:				Optional[Any]       =  jax.nn.log_softmax(__A							,		axis=-1	)

						def handle_cumulative_probs(__A					:  str							,		__A					:  Dict	):
							snake_case__			:				Dict       =  jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]							,		axis=-1	)
							snake_case__			:				List[str]       =  jnp.max(logprobs_k[: self.timestamp_begin]	)
							return jnp.where(
							    timestamp_logprob > max_text_token_logprob							,		scores_k.at[: self.timestamp_begin].set(-float("inf"	)	)							,		__A							,		)

						snake_case__			:				Union[str, Any]       =  jax.vmap(__A	)(__A							,		__A	)

						return scores











import os

import pytest
from attr import dataclass


__lowerCamelCase :      Any										=				"""us-east-1"""  # defaults region



@dataclass
class 							SCREAMING_SNAKE_CASE__						:





					"""simple docstring"""






					a_     =		42
					a_     =		"arn:aws:iam::558105141721:role/sagemaker_execution_role"
					a_     =		{
					    "task_name": "mnli",
					    "per_device_train_batch_size": 1_6,
					    "per_device_eval_batch_size": 1_6,
					    "do_train": True,
					    "do_eval": True,
					    "do_predict": True,
					    "output_dir": "/opt/ml/model",
					    "overwrite_output_dir": True,
					    "max_steps": 5_0_0,
					    "save_steps": 5_5_0_0,
					}
					a_     =		{**hyperparameters, "max_steps": 1_0_0_0}



					@property
					def 	_lowercase					( self					:  List[Any]	):
						if self.framework == "pytorch":
							return [
							    {"Name": "train_runtime", "Regex": r"train_runtime.*=\D*(.*?)$"},
							    {"Name": "eval_accuracy", "Regex": r"eval_accuracy.*=\D*(.*?)$"},
							    {"Name": "eval_loss", "Regex": r"eval_loss.*=\D*(.*?)$"},
							]
						else:
							return [
							    {"Name": "train_runtime", "Regex": r"train_runtime.*=\D*(.*?)$"},
							    {"Name": "eval_accuracy", "Regex": r"loss.*=\D*(.*?)]?$"},
							    {"Name": "eval_loss", "Regex": r"sparse_categorical_accuracy.*=\D*(.*?)]?$"},
							]



					@property
					def 	_lowercase					( self					:  Any	):
						return f'''{self.framework}-transfromers-test'''



					@property
					def 	_lowercase					( self					:  Optional[Any]	):
						return f'''./tests/sagemaker/scripts/{self.framework}'''






					@property
					def 	_lowercase					( self					:  Union[str, Any]	):
						if self.framework == "pytorch":
							return "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training:1.7.1-transformers4.6.1-gpu-py36-cu110-ubuntu18.04"
						else:
							return "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-tensorflow-training:2.4.1-transformers4.6.1-gpu-py37-cu110-ubuntu18.04"







@pytest.fixture(scope="class"       )
def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							Optional[Any]       ):
	snake_case__			:				Optional[int]       =  SageMakerTestEnvironment(framework=request.cls.framework       )



