"""simple docstring"""



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


_UpperCamelCase								=		get_logger(__name__)


_UpperCamelCase								=		R"\n    Args:\n        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary.\n\n            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n            [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):\n            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam\n            search or log softmax for each vocabulary token when using beam search\n        kwargs (`Dict[str, Any]`, *optional*):\n            Additional logits processor specific kwargs.\n\n    Return:\n        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.\n\n"
class 						__UpperCAmelCase    :







			'''simple docstring'''





			@add_start_docstrings(snake_case_							)
			def __call__(			self   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										raise NotImplementedError(
										    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''							)





class 						__UpperCAmelCase    :







			'''simple docstring'''





			@add_start_docstrings(snake_case_							)
			def __call__(			self   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										raise NotImplementedError(
										    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''							)





class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





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







										'''simple docstring'''


										for processor in self:
																	A__       :     Tuple								=			inspect.signature(processor.__call__							).parameters
																	if len(snake_case_							) > 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.'''							)
																								A__       :     List[str]								=			processor(snake_case_   ,      snake_case_   ,      snake_case_   ,      **snake_case_							)
																	else:
																								A__       :     Union[str, Any]								=			processor(snake_case_   ,      snake_case_   ,      snake_case_							)
										return scores





class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





			def __init__(			self   ,      snake_case_							):







										'''simple docstring'''


										if not isinstance(snake_case_   ,      snake_case_							) or not (temperature > 0):
																	raise ValueError(F'''`temperature` has to be a strictly positive float, but is {temperature}'''							)

										A__       :     str								=			temperature





			def __call__(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     List[Any]								=			scores / self.temperature
										return scores





class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





			def __init__(			self   ,      snake_case_   ,      snake_case_ = -float("""Inf"""							)   ,      snake_case_ = 1							):







										'''simple docstring'''


										if not isinstance(snake_case_   ,      snake_case_							) 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(snake_case_   ,      snake_case_							) 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}'''							)

										A__       :     List[Any]								=			top_p
										A__       :     Optional[Any]								=			filter_value
										A__       :     Dict								=			min_tokens_to_keep





			def __call__(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__				,							A__       :     List[str]								=			lax.top_k(snake_case_   ,      scores.shape[-1]							)

										A__       :     Tuple								=			jnp.full_like(snake_case_   ,      self.filter_value							)
										A__       :     Dict								=			jax.nn.softmax(snake_case_   ,      axis=-1							).cumsum(axis=-1							)
										A__       :     Dict								=			cumulative_probs < self.top_p

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

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

										A__       :     Union[str, Any]								=			jnp.where(snake_case_   ,      snake_case_   ,      snake_case_							)
										A__       :     Any								=			jax.lax.sort_key_val(snake_case_   ,      snake_case_							)[-1]

										return next_scores





class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





			def __init__(			self   ,      snake_case_   ,      snake_case_ = -float("""Inf"""							)   ,      snake_case_ = 1							):







										'''simple docstring'''


										if not isinstance(snake_case_   ,      snake_case_							) or top_k <= 0:
																	raise ValueError(F'''`top_k` has to be a strictly positive integer, but is {top_k}'''							)

										A__       :     List[Any]								=			max(snake_case_   ,      snake_case_							)
										A__       :     List[Any]								=			filter_value





			def __call__(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__				,							A__       :     Union[str, Any]								=			scores.shape
										A__       :     Any								=			jnp.full(batch_size * vocab_size   ,      self.filter_value							)

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

										A__       :     Tuple								=			next_scores_flat.at[topk_indices_flat].set(snake_case_							)
										A__       :     Dict								=			next_scores_flat.reshape(snake_case_   ,      snake_case_							)
										return next_scores





class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





			def __init__(			self   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Optional[int]								=			bos_token_id





			def __call__(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Any								=			jnp.full(scores.shape   ,      -float("""inf"""							)							)

										A__       :     Dict								=			1 - jnp.bool_(cur_len - 1							)

										A__       :     int								=			jnp.where(snake_case_   ,      new_scores.at[:, self.bos_token_id].set(0							)   ,      snake_case_							)

										return scores





class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





			def __init__(			self   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     List[str]								=			max_length
										A__       :     Optional[Any]								=			eos_token_id





			def __call__(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Optional[Any]								=			jnp.full(scores.shape   ,      -float("""inf"""							)							)

										A__       :     List[str]								=			1 - jnp.bool_(cur_len - self.max_length + 1							)

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

										return scores





class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





			def __init__(			self   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


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

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

										A__       :     Any								=			min_length
										A__       :     Optional[int]								=			eos_token_id





			def __call__(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     List[str]								=			1 - jnp.clip(cur_len - self.min_length   ,      0   ,      1							)

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

										return scores





class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





			def __init__(			self   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Optional[int]								=			list(snake_case_							)
										A__       :     List[Any]								=			begin_index





			def __call__(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Optional[int]								=			1 - jnp.bool_(cur_len - self.begin_index							)

										A__       :     List[Any]								=			jnp.where(snake_case_   ,      scores.at[:, self.begin_suppress_tokens].set(-float("""inf"""							)							)   ,      snake_case_							)

										return scores





class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





			def __init__(			self   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Any								=			list(snake_case_							)





			def __call__(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     int								=			scores.at[..., self.suppress_tokens].set(-float("""inf"""							)							)

										return scores





class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





			def __init__(			self   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Dict								=			dict(snake_case_							)
										# 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.
										A__       :     str								=			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:
																								A__       :     Tuple								=			force_token_array.at[index].set(snake_case_							)
										A__       :     Optional[Any]								=			jnp.intaa(snake_case_							)





			def __call__(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										def _force_token(snake_case_							):
																	A__       :     Optional[Any]								=			scores.shape[0]
																	A__       :     List[str]								=			self.force_token_array[generation_idx]

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

										A__       :     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(snake_case_							)   ,      lambda: scores   ,      )   ,      )
										return scores





class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





			def __init__(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Optional[int]								=			generate_config.eos_token_id
										A__       :     List[Any]								=			generate_config.no_timestamps_token_id
										A__       :     Dict								=			generate_config.no_timestamps_token_id + 1

										A__       :     Tuple								=			decoder_input_length + 1

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





			def __call__(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Dict								=			scores.at[:, self.no_timestamps_token_id].set(-float("""inf"""							)							)

										def handle_pairs(snake_case_   ,      snake_case_							):
																	A__       :     Optional[Any]								=			jnp.where((cur_len - self.begin_index) >= 1   ,      snake_case_   ,      snake_case_							)
																	A__       :     int								=			jnp.where(
																	    input_ids_k[cur_len - 1] >= self.timestamp_begin   ,      True and last_was_timestamp   ,      snake_case_   ,      )

																	A__       :     int								=			jnp.where((cur_len - self.begin_index) < 2   ,      snake_case_   ,      snake_case_							)
																	A__       :     str								=			jnp.where(
																	    input_ids_k[cur_len - 2] >= self.timestamp_begin   ,      snake_case_   ,      snake_case_   ,      )

																	return jnp.where(
																	    snake_case_   ,      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"""							)							)   ,      )   ,      snake_case_   ,      )

										A__       :     Dict								=			jax.vmap(snake_case_							)(snake_case_   ,      snake_case_							)

										A__       :     Tuple								=			jnp.where(cur_len == self.begin_index   ,      snake_case_   ,      snake_case_							)
										A__       :     Any								=			jnp.where(
										    self.max_initial_timestamp_index is not None   ,      True and apply_max_initial_timestamp   ,      snake_case_   ,      )

										A__       :     Optional[int]								=			self.timestamp_begin + self.max_initial_timestamp_index

										A__       :     List[str]								=			jnp.where(
										    snake_case_   ,      scores.at[:, last_allowed + 1 :].set(-float("""inf"""							)							)   ,      snake_case_   ,      )

										# if sum of probability over timestamps is above any other token, sample timestamp
										A__       :     Tuple								=			jax.nn.log_softmax(snake_case_   ,      axis=-1							)

										def handle_cumulative_probs(snake_case_   ,      snake_case_							):
																	A__       :     List[Any]								=			jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]   ,      axis=-1							)
																	A__       :     int								=			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"""							)							)   ,      snake_case_   ,      )

										A__       :     Optional[int]								=			jax.vmap(snake_case_							)(snake_case_   ,      snake_case_							)

										return scores






"""simple docstring"""



from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast
from ...utils import logging


_UpperCamelCase								=		logging.get_logger(__name__)

_UpperCamelCase								=		{
    "EleutherAI/gpt-neo-1.3B": "https://huggingface.co/EleutherAI/gpt-neo-1.3B/resolve/main/config.json",
    # See all GPTNeo models at https://huggingface.co/models?filter=gpt_neo
}
class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''
			_UpperCamelCase  :						Tuple					     =				'gpt_neo'
			_UpperCamelCase  :						Optional[Any]					     =				['past_key_values']
			_UpperCamelCase  :						int					     =				{'num_attention_heads': 'num_heads', 'num_hidden_layers': 'num_layers'}





			def __init__(			self   ,      snake_case_=50_257   ,      snake_case_=2_048   ,      snake_case_=2_048   ,      snake_case_=24   ,      snake_case_=[[["global", "local"], 12]]   ,      snake_case_=16   ,      snake_case_=None   ,      snake_case_=256   ,      snake_case_="gelu_new"   ,      snake_case_=0.0   ,      snake_case_=0.0   ,      snake_case_=0.0   ,      snake_case_=0.1   ,      snake_case_=1E-5   ,      snake_case_=0.02   ,      snake_case_=True   ,      snake_case_=50_256   ,      snake_case_=50_256   ,      **snake_case_   ,      ):







										'''simple docstring'''


										A__       :     int								=			vocab_size
										A__       :     Optional[Any]								=			max_position_embeddings
										A__       :     int								=			hidden_size
										A__       :     List[Any]								=			num_layers
										A__       :     Any								=			num_heads
										A__       :     List[str]								=			intermediate_size
										A__       :     Dict								=			window_size
										A__       :     Optional[int]								=			activation_function
										A__       :     Optional[int]								=			resid_dropout
										A__       :     List[str]								=			embed_dropout
										A__       :     str								=			attention_dropout
										A__       :     List[str]								=			classifier_dropout
										A__       :     str								=			layer_norm_epsilon
										A__       :     str								=			initializer_range
										A__       :     Any								=			use_cache

										A__       :     List[str]								=			bos_token_id
										A__       :     Any								=			eos_token_id

										A__       :     Optional[Any]								=			attention_types
										A__       :     List[Any]								=			self.expand_attention_types_params(snake_case_							)

										if len(self.attention_layers							) != self.num_layers:
																	raise ValueError(
																	    """Configuration for convolutional module is incorrect. """
																	    """It is required that `len(config.attention_layers)` == `config.num_layers` """
																	    F'''but is `len(config.attention_layers) = {len(self.attention_layers							)}`, '''
																	    F'''`config.num_layers = {self.num_layers}`. '''
																	    """`config.attention_layers` is prepared using `config.attention_types`. """
																	    """Please verify the value of `config.attention_types` argument."""							)

										super().__init__(bos_token_id=snake_case_   ,      eos_token_id=snake_case_   ,      **snake_case_							)





			@staticmethod
			def 		lowerCamelCase					(			snake_case_							):







										'''simple docstring'''


										A__       :     Optional[Any]								=			[]
										for item in attention_types:
																	for _ in range(item[1]							):
																								attentions.extend(item[0]							)
										return attentions





def  _A(		lowerCAmelCase			,  lowerCAmelCase			,  lowerCAmelCase			,  lowerCAmelCase							):
							import torch

							A__       :     Union[str, Any]								=			input.size()
							A__       :     str								=			len(lowerCAmelCase							)
							A__       :     Union[str, Any]								=			shape[dimension]

							A__       :     Dict								=			torch.arange(0			,  lowerCAmelCase			,  lowerCAmelCase							)
							A__       :     Dict								=			torch.div(sizedim - size			,  lowerCAmelCase			,  rounding_mode="""floor"""							) + 1
							A__       :     str								=			torch.arange(lowerCAmelCase							) + low_indices[:min_length][:, None]

							A__       :     str								=			[slice(lowerCAmelCase							)] * rank
							A__       :     Optional[Any]								=			indices
							A__       :     Tuple								=			input[s]

							A__       :     List[str]								=			list(range(0			,  rank + 1							)							)
							perm.append(perm.pop(dimension + 1							)							)

							return sliced.permute(lowerCAmelCase							)

def  _A(		lowerCAmelCase			,  lowerCAmelCase							):
							import torch

							A__       :     Optional[Any]								=			torch.arange(1			,  lowerCAmelCase							)
							A__       :     int								=			torch.remainder(lowerCAmelCase			,  lowerCAmelCase							)
							A__       :     str								=			remainders == 0
							A__       :     Tuple								=			candidates[divisor_indices]
							A__       :     Any								=			torch.max(lowerCAmelCase							)
							return largest_divisor, torch.div(lowerCAmelCase			,  lowerCAmelCase			,  rounding_mode="""floor"""							)



class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





			@property
			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Any								=			OrderedDict({"""input_ids""": {0: """batch""", 1: """sequence"""}}							)
										if self.use_past:
																	self.fill_with_past_key_values_(snake_case_   ,      direction="""inputs"""							)
																	A__       :     str								=			{0: """batch""", 1: """past_sequence + sequence"""}
										else:
																	A__       :     Optional[Any]								=			{0: """batch""", 1: """sequence"""}

										return common_inputs





			@property
			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										return self._config.num_heads





			def 		lowerCamelCase					(			self   ,      snake_case_   ,      snake_case_ = -1   ,      snake_case_ = -1   ,      snake_case_ = False   ,      snake_case_ = None   ,      ):







										'''simple docstring'''


										A__       :     Optional[int]								=			super(snake_case_   ,      self							).generate_dummy_inputs(
										    snake_case_   ,      batch_size=snake_case_   ,      seq_length=snake_case_   ,      is_pair=snake_case_   ,      framework=snake_case_							)

										# We need to order the input in the way they appears in the forward()
										A__       :     List[Any]								=			OrderedDict({"""input_ids""": common_inputs["""input_ids"""]}							)

										# Need to add the past_keys
										if self.use_past:
																	if not is_torch_available():
																								raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""							)
																	else:
																								import torch

																								A__				,							A__       :     Dict								=			common_inputs["""input_ids"""].shape
																								# Not using the same length for past_key_values
																								A__       :     Any								=			seqlen + 2
																								A__       :     Any								=			(
																								    batch,
																								    self.num_attention_heads,
																								    past_key_values_length,
																								    self._config.hidden_size // self.num_attention_heads,
																								)
																								A__       :     List[Any]								=			[
																								    (torch.zeros(snake_case_							), torch.zeros(snake_case_							)) for _ in range(self.num_layers							)
																								]

										A__       :     Tuple								=			common_inputs["""attention_mask"""]
										if self.use_past:
																	A__       :     List[Any]								=			ordered_inputs["""attention_mask"""].dtype
																	A__       :     Optional[Any]								=			torch.cat(
																	    [ordered_inputs["""attention_mask"""], torch.ones(snake_case_   ,      snake_case_   ,      dtype=snake_case_							)]   ,      dim=1							)

										return ordered_inputs





			@property
			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										return 13





