
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_  :	int 				=					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     _UpperCamelCase			:






			'''simple docstring'''

			@add_start_docstrings(snake_case_ )
			def __call__(    self	:    List[Any]  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    jnp.ndarray ):
							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	:    List[Any]  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    jnp.ndarray ):
							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	:    int  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    int  ,		**snake_case_	:    Optional[int] ):
							for processor in self:
											UpperCamelCase_:							Optional[int]		=  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.''' )
															UpperCamelCase_:							Any		=  processor(snake_case_  ,		snake_case_  ,		snake_case_  ,		**snake_case_ )
											else:
															UpperCamelCase_:							List[str]		=  processor(snake_case_  ,		snake_case_  ,		snake_case_ )
							return scores




class     _UpperCamelCase			(   _A					):






			'''simple docstring'''

			def __init__(    self	:    str  ,		snake_case_	:    float ):
							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}''' )

							UpperCamelCase_:							Dict		=  temperature




			def __call__(    self	:    str  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    int ):
							UpperCamelCase_:							Union[str, Any]		=  scores / self.temperature
							return scores




class     _UpperCamelCase			(   _A					):






			'''simple docstring'''

			def __init__(    self	:    Tuple  ,		snake_case_	:    float  ,		snake_case_	:    float = -float("""Inf""" )  ,		snake_case_	:    int = 1 ):
							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}''' )

							UpperCamelCase_:							Optional[int]		=  top_p
							UpperCamelCase_:							Dict		=  filter_value
							UpperCamelCase_:							Optional[Any]		=  min_tokens_to_keep




			def __call__(    self	:    Union[str, Any]  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    int ):
							UpperCamelCase_,				UpperCamelCase_:							Any		=  lax.top_k(snake_case_  ,		scores.shape[-1] )

							UpperCamelCase_:							Any		=  jnp.full_like(snake_case_  ,		self.filter_value )
							UpperCamelCase_:							List[Any]		=  jax.nn.softmax(snake_case_  ,		axis=-1 ).cumsum(axis=-1 )
							UpperCamelCase_:							List[Any]		=  cumulative_probs < self.top_p

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

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

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

							return next_scores




class     _UpperCamelCase			(   _A					):






			'''simple docstring'''

			def __init__(    self	:    Tuple  ,		snake_case_	:    int  ,		snake_case_	:    float = -float("""Inf""" )  ,		snake_case_	:    int = 1 ):
							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}''' )

							UpperCamelCase_:							Union[str, Any]		=  max(snake_case_  ,		snake_case_ )
							UpperCamelCase_:							Union[str, Any]		=  filter_value




			def __call__(    self	:    Optional[int]  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    int ):
							UpperCamelCase_,				UpperCamelCase_:							List[str]		=  scores.shape
							UpperCamelCase_:							Union[str, Any]		=  jnp.full(batch_size * vocab_size  ,		self.filter_value )

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

							UpperCamelCase_:							Optional[int]		=  next_scores_flat.at[topk_indices_flat].set(snake_case_ )
							UpperCamelCase_:							Dict		=  next_scores_flat.reshape(snake_case_  ,		snake_case_ )
							return next_scores




class     _UpperCamelCase			(   _A					):






			'''simple docstring'''

			def __init__(    self	:    Any  ,		snake_case_	:    int ):
							UpperCamelCase_:							Any		=  bos_token_id




			def __call__(    self	:    str  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    int ):
							UpperCamelCase_:							Optional[Any]		=  jnp.full(scores.shape  ,		-float("""inf""" ) )

							UpperCamelCase_:							Any		=  1 - jnp.bool_(cur_len - 1 )

							UpperCamelCase_:							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	:    List[Any]  ,		snake_case_	:    int  ,		snake_case_	:    int ):
							UpperCamelCase_:							Tuple		=  max_length
							UpperCamelCase_:							Any		=  eos_token_id




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

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

							UpperCamelCase_:							Tuple		=  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	:    Optional[int]  ,		snake_case_	:    int  ,		snake_case_	:    int ):
							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}''' )

							UpperCamelCase_:							str		=  min_length
							UpperCamelCase_:							int		=  eos_token_id




			def __call__(    self	:    int  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    int ):
							# create boolean flag to decide if min length penalty should be applied
							UpperCamelCase_:							Dict		=  1 - jnp.clip(cur_len - self.min_length  ,		0  ,		1 )

							UpperCamelCase_:							List[str]		=  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	:    Dict  ,		snake_case_	:    List[str]  ,		snake_case_	:    Any ):
							UpperCamelCase_:							Optional[int]		=  list(snake_case_ )
							UpperCamelCase_:							Tuple		=  begin_index




			def __call__(    self	:    Optional[int]  ,		snake_case_	:    Optional[int]  ,		snake_case_	:    Tuple  ,		snake_case_	:    int ):
							UpperCamelCase_:							str		=  1 - jnp.bool_(cur_len - self.begin_index )

							UpperCamelCase_:							Optional[int]		=  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	:    Union[str, Any]  ,		snake_case_	:    list ):
							UpperCamelCase_:							str		=  list(snake_case_ )




			def __call__(    self	:    List[Any]  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    int ):
							UpperCamelCase_:							Union[str, Any]		=  scores.at[..., self.suppress_tokens].set(-float("""inf""" ) )

							return scores




class     _UpperCamelCase			(   _A					):






			'''simple docstring'''

			def __init__(    self	:    str  ,		snake_case_	:    Dict ):
							UpperCamelCase_:							Union[str, Any]		=  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.
							UpperCamelCase_:							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:
															UpperCamelCase_:							Dict		=  force_token_array.at[index].set(snake_case_ )
							UpperCamelCase_:							List[Any]		=  jnp.intaa(snake_case_ )




			def __call__(    self	:    Union[str, Any]  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    jnp.ndarray  ,		snake_case_	:    int ):
							def _force_token(snake_case_	:    Tuple ):
											UpperCamelCase_:							str		=  scores.shape[0]
											UpperCamelCase_:							Tuple		=  self.force_token_array[generation_idx]

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

							UpperCamelCase_:							str		=  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	:    Dict  ,		snake_case_	:    Union[str, Any]  ,		snake_case_	:    Optional[Any]  ,		snake_case_	:    Tuple ):
							UpperCamelCase_:							Tuple		=  generate_config.eos_token_id
							UpperCamelCase_:							List[Any]		=  generate_config.no_timestamps_token_id
							UpperCamelCase_:							List[Any]		=  generate_config.no_timestamps_token_id + 1

							UpperCamelCase_:							List[Any]		=  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""" ):
											UpperCamelCase_:							List[str]		=  generate_config.max_initial_timestamp_index
							else:
											UpperCamelCase_:							str		=  model_config.vocab_size
							if self.max_initial_timestamp_index is None:
											UpperCamelCase_:							Tuple		=  model_config.vocab_size




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

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

											UpperCamelCase_:							List[str]		=  jnp.where((cur_len - self.begin_index) < 2  ,		snake_case_  ,		snake_case_ )
											UpperCamelCase_:							Tuple		=  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_  ,		)

							UpperCamelCase_:							str		=  jax.vmap(snake_case_ )(snake_case_  ,		snake_case_ )

							UpperCamelCase_:							str		=  jnp.where(cur_len == self.begin_index  ,		snake_case_  ,		snake_case_ )
							UpperCamelCase_:							List[str]		=  jnp.where(
							    self.max_initial_timestamp_index is not None  ,		True and apply_max_initial_timestamp  ,		snake_case_  ,		)

							UpperCamelCase_:							Dict		=  self.timestamp_begin + self.max_initial_timestamp_index

							UpperCamelCase_:							Any		=  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
							UpperCamelCase_:							List[str]		=  jax.nn.log_softmax(snake_case_  ,		axis=-1 )

							def handle_cumulative_probs(snake_case_	:    Tuple  ,		snake_case_	:    List[Any] ):
											UpperCamelCase_:							Union[str, Any]		=  jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]  ,		axis=-1 )
											UpperCamelCase_:							Dict		=  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_  ,		)

							UpperCamelCase_:							List[str]		=  jax.vmap(snake_case_ )(snake_case_  ,		snake_case_ )

							return scores








import argparse

import torch

from transformers import GPTaConfig, GPTaModel, load_tf_weights_in_gpta
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()







def     A__   (  lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase  )					->		Union[str, Any]:
				# Construct model
				if gpta_config_file == "":
								UpperCamelCase_:							Union[str, Any]		=  GPTaConfig()
				else:
								UpperCamelCase_:							Optional[Any]		=  GPTaConfig.from_json_file(lowerCamelCase  )
				UpperCamelCase_:							List[Any]		=  GPTaModel(lowerCamelCase  )

				# Load weights from numpy
				load_tf_weights_in_gpta(lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase  )

				# Save pytorch-model
				UpperCamelCase_:							int		=  pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
				UpperCamelCase_:							Any		=  pytorch_dump_folder_path + """/""" + CONFIG_NAME
				print(F'''Save PyTorch model to {pytorch_weights_dump_path}'''  )
				torch.save(model.state_dict()       ,      lowerCamelCase  )
				print(F'''Save configuration file to {pytorch_config_dump_path}'''  )
				with open(lowerCamelCase       ,      """w"""       ,      encoding="""utf-8"""  ) as f:
								f.write(config.to_json_string()  )


if __name__ == "__main__":
							lowerCamelCase_  :	int 				=					argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    """--gpt2_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
							)
							parser.add_argument(
							    """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
							)
							parser.add_argument(
							    """--gpt2_config_file""",
							    default="""""",
							    type=str,
							    help=(
							        """An optional config json file corresponding to the pre-trained OpenAI model. \n"""
							        """This specifies the model architecture."""
							    ),
							)
							lowerCamelCase_  :	List[str] 				=					parser.parse_args()
							convert_gpta_checkpoint_to_pytorch(args.gpta_checkpoint_path, args.gpta_config_file, args.pytorch_dump_folder_path)






