


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	=							get_logger(__name__)


lowerCAmelCase	=							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 	_a							:

							@add_start_docstrings(UpperCamelCase_							)
							def __call__(							self:     Dict   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     jnp.ndarray							)      ->       jnp.ndarray:





									"""simple docstring"""







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


class 	_a							:

							@add_start_docstrings(UpperCamelCase_							)
							def __call__(							self:     Tuple   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     jnp.ndarray							)      ->       jnp.ndarray:





									"""simple docstring"""







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


class 	_a							(      UpperCamelCase__      ):

							@add_start_docstrings(UpperCamelCase_							)
							def __call__(							self:     Dict   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     int   ,						**UpperCamelCase_:     str							)      ->       jnp.ndarray:





									"""simple docstring"""







									for processor in self:
											lowercase__						=				inspect.signature(processor.__call__							).parameters
											if len(UpperCamelCase_							) > 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.'							)
													lowercase__						=				processor(UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						**UpperCamelCase_							)
											else:
													lowercase__						=				processor(UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_							)
									return scores


class 	_a							(      UpperCamelCase__      ):

							def __init__(							self:     List[Any]   ,						UpperCamelCase_:     float							)      ->       Optional[Any]:





									"""simple docstring"""







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

									lowercase__						=				temperature


							def __call__(							self:     Union[str, Any]   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     int							)      ->       jnp.ndarray:





									"""simple docstring"""







									lowercase__						=				scores / self.temperature
									return scores


class 	_a							(      UpperCamelCase__      ):

							def __init__(							self:     int   ,						UpperCamelCase_:     float   ,						UpperCamelCase_:     float = -float('''Inf'''							)   ,						UpperCamelCase_:     int = 1							)      ->       Optional[Any]:





									"""simple docstring"""







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

									lowercase__						=				top_p
									lowercase__						=				filter_value
									lowercase__						=				min_tokens_to_keep


							def __call__(							self:     List[Any]   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     int							)      ->       jnp.ndarray:





									"""simple docstring"""







									lowercase__					,		lowercase__						=				lax.top_k(UpperCamelCase_   ,						scores.shape[-1]							)

									lowercase__						=				jnp.full_like(UpperCamelCase_   ,						self.filter_value							)
									lowercase__						=				jax.nn.softmax(UpperCamelCase_   ,						axis=-1							).cumsum(axis=-1							)
									lowercase__						=				cumulative_probs < self.top_p

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

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

									lowercase__						=				jnp.where(UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_							)
									lowercase__						=				jax.lax.sort_key_val(UpperCamelCase_   ,						UpperCamelCase_							)[-1]

									return next_scores


class 	_a							(      UpperCamelCase__      ):

							def __init__(							self:     Dict   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     float = -float('''Inf'''							)   ,						UpperCamelCase_:     int = 1							)      ->       Dict:





									"""simple docstring"""







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

									lowercase__						=				max(UpperCamelCase_   ,						UpperCamelCase_							)
									lowercase__						=				filter_value


							def __call__(							self:     Any   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     int							)      ->       jnp.ndarray:





									"""simple docstring"""







									lowercase__					,		lowercase__						=				scores.shape
									lowercase__						=				jnp.full(batch_size * vocab_size   ,						self.filter_value							)

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

									lowercase__						=				next_scores_flat.at[topk_indices_flat].set(UpperCamelCase_							)
									lowercase__						=				next_scores_flat.reshape(UpperCamelCase_   ,						UpperCamelCase_							)
									return next_scores


class 	_a							(      UpperCamelCase__      ):

							def __init__(							self:     Dict   ,						UpperCamelCase_:     int							)      ->       Any:





									"""simple docstring"""







									lowercase__						=				bos_token_id


							def __call__(							self:     Optional[Any]   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     int							)      ->       jnp.ndarray:





									"""simple docstring"""







									lowercase__						=				jnp.full(scores.shape   ,						-float('''inf'''							)							)

									lowercase__						=				1 - jnp.bool_(cur_len - 1							)

									lowercase__						=				jnp.where(UpperCamelCase_   ,						new_scores.at[:, self.bos_token_id].set(0							)   ,						UpperCamelCase_							)

									return scores


class 	_a							(      UpperCamelCase__      ):

							def __init__(							self:     Dict   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int							)      ->       Any:





									"""simple docstring"""







									lowercase__						=				max_length
									lowercase__						=				eos_token_id


							def __call__(							self:     List[str]   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     int							)      ->       jnp.ndarray:





									"""simple docstring"""







									lowercase__						=				jnp.full(scores.shape   ,						-float('''inf'''							)							)

									lowercase__						=				1 - jnp.bool_(cur_len - self.max_length + 1							)

									lowercase__						=				jnp.where(UpperCamelCase_   ,						new_scores.at[:, self.eos_token_id].set(0							)   ,						UpperCamelCase_							)

									return scores


class 	_a							(      UpperCamelCase__      ):

							def __init__(							self:     Union[str, Any]   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int							)      ->       Any:





									"""simple docstring"""







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

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

									lowercase__						=				min_length
									lowercase__						=				eos_token_id


							def __call__(							self:     List[str]   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     int							)      ->       jnp.ndarray:





									"""simple docstring"""







									lowercase__						=				1 - jnp.clip(cur_len - self.min_length   ,						0   ,						1							)

									lowercase__						=				jnp.where(UpperCamelCase_   ,						scores.at[:, self.eos_token_id].set(-float('''inf'''							)							)   ,						UpperCamelCase_							)

									return scores


class 	_a							(      UpperCamelCase__      ):

							def __init__(							self:     List[Any]   ,						UpperCamelCase_:     Optional[int]   ,						UpperCamelCase_:     Union[str, Any]							)      ->       Optional[Any]:





									"""simple docstring"""







									lowercase__						=				list(UpperCamelCase_							)
									lowercase__						=				begin_index


							def __call__(							self:     Tuple   ,						UpperCamelCase_:     List[str]   ,						UpperCamelCase_:     Optional[int]   ,						UpperCamelCase_:     int							)      ->       Tuple:





									"""simple docstring"""







									lowercase__						=				1 - jnp.bool_(cur_len - self.begin_index							)

									lowercase__						=				jnp.where(UpperCamelCase_   ,						scores.at[:, self.begin_suppress_tokens].set(-float('''inf'''							)							)   ,						UpperCamelCase_							)

									return scores


class 	_a							(      UpperCamelCase__      ):

							def __init__(							self:     Optional[Any]   ,						UpperCamelCase_:     list							)      ->       int:





									"""simple docstring"""







									lowercase__						=				list(UpperCamelCase_							)


							def __call__(							self:     int   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     int							)      ->       jnp.ndarray:





									"""simple docstring"""







									lowercase__						=				scores.at[..., self.suppress_tokens].set(-float('''inf'''							)							)

									return scores


class 	_a							(      UpperCamelCase__      ):

							def __init__(							self:     str   ,						UpperCamelCase_:     Any							)      ->       Optional[Any]:





									"""simple docstring"""







									lowercase__						=				dict(UpperCamelCase_							)
									# 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.
									lowercase__						=				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:
													lowercase__						=				force_token_array.at[index].set(UpperCamelCase_							)
									lowercase__						=				jnp.intaa(UpperCamelCase_							)


							def __call__(							self:     str   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     jnp.ndarray   ,						UpperCamelCase_:     int							)      ->       jnp.ndarray:





									"""simple docstring"""







									def _force_token(UpperCamelCase_:     List[str]							):
											lowercase__						=				scores.shape[0]
											lowercase__						=				self.force_token_array[generation_idx]

											lowercase__						=				jnp.ones_like(UpperCamelCase_   ,						dtype=scores.dtype							) * -float('''inf'''							)
											lowercase__						=				jnp.zeros((batch_size, 1)   ,						dtype=scores.dtype							)
											lowercase__						=				lax.dynamic_update_slice(UpperCamelCase_   ,						UpperCamelCase_   ,						(0, current_token)							)
											return new_scores

									lowercase__						=				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(UpperCamelCase_							)   ,						lambda: scores   ,						)   ,						)
									return scores


class 	_a							(      UpperCamelCase__      ):

							def __init__(							self:     Optional[int]   ,						UpperCamelCase_:     List[str]   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int							)      ->       Union[str, Any]:





									"""simple docstring"""







									lowercase__						=				generate_config.eos_token_id
									lowercase__						=				generate_config.no_timestamps_token_id
									lowercase__						=				generate_config.no_timestamps_token_id + 1

									lowercase__						=				decoder_input_length + 1

									if generate_config.is_multilingual:
											# room for language token and task token
											self.begin_index += 2
									if hasattr(UpperCamelCase_   ,						'''max_initial_timestamp_index'''							):
											lowercase__						=				generate_config.max_initial_timestamp_index
									else:
											lowercase__						=				model_config.vocab_size
									if self.max_initial_timestamp_index is None:
											lowercase__						=				model_config.vocab_size


							def __call__(							self:     Any   ,						UpperCamelCase_:     Optional[int]   ,						UpperCamelCase_:     Optional[Any]   ,						UpperCamelCase_:     str							)      ->       Union[str, Any]:





									"""simple docstring"""







									lowercase__						=				scores.at[:, self.no_timestamps_token_id].set(-float('''inf'''							)							)

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

											lowercase__						=				jnp.where((cur_len - self.begin_index) < 2   ,						UpperCamelCase_   ,						UpperCamelCase_							)
											lowercase__						=				jnp.where(
											    input_ids_k[cur_len - 2] >= self.timestamp_begin   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						)

											return jnp.where(
											    UpperCamelCase_   ,						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'''							)							)   ,						)   ,						UpperCamelCase_   ,						)

									lowercase__						=				jax.vmap(UpperCamelCase_							)(UpperCamelCase_   ,						UpperCamelCase_							)

									lowercase__						=				jnp.where(cur_len == self.begin_index   ,						UpperCamelCase_   ,						UpperCamelCase_							)
									lowercase__						=				jnp.where(
									    self.max_initial_timestamp_index is not None   ,						True and apply_max_initial_timestamp   ,						UpperCamelCase_   ,						)

									lowercase__						=				self.timestamp_begin + self.max_initial_timestamp_index

									lowercase__						=				jnp.where(
									    UpperCamelCase_   ,						scores.at[:, last_allowed + 1 :].set(-float('''inf'''							)							)   ,						UpperCamelCase_   ,						)

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

									def handle_cumulative_probs(UpperCamelCase_:     Optional[Any]   ,						UpperCamelCase_:     Tuple							):
											lowercase__						=				jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]   ,						axis=-1							)
											lowercase__						=				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'''							)							)   ,						UpperCamelCase_   ,						)

									lowercase__						=				jax.vmap(UpperCamelCase_							)(UpperCamelCase_   ,						UpperCamelCase_							)

									return scores






import json
import logging
import os
import socket

import git
import numpy as np
import torch


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - PID: %(process)d -  %(message)s',
    datefmt='%m/%d/%Y %H:%M:%S',
    level=logging.INFO,
)
lowerCAmelCase	=							logging.getLogger(__name__)
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		lowercase__						=				git.Repo(search_parent_directories=SCREAMING_SNAKE_CASE					)
		lowercase__						=				{
		    '''repo_id''': str(SCREAMING_SNAKE_CASE					),
		    '''repo_sha''': str(repo.head.object.hexsha					),
		    '''repo_branch''': str(repo.active_branch					),
		}

		with open(os.path.join(SCREAMING_SNAKE_CASE						,    '''git_log.json'''					)						,    '''w'''					) as f:
				json.dump(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    indent=4					)
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		if params.n_gpu <= 0:
				lowercase__						=				0
				lowercase__						=				-1
				lowercase__						=				True
				lowercase__						=				False
				return

		assert torch.cuda.is_available()

		logger.info('''Initializing GPUs'''					)
		if params.n_gpu > 1:
				assert params.local_rank != -1

				lowercase__						=				int(os.environ['''WORLD_SIZE''']					)
				lowercase__						=				int(os.environ['''N_GPU_NODE''']					)
				lowercase__						=				int(os.environ['''RANK''']					)

				# number of nodes / node ID
				lowercase__						=				params.world_size // params.n_gpu_per_node
				lowercase__						=				params.global_rank // params.n_gpu_per_node
				lowercase__						=				True

				assert params.n_nodes == int(os.environ['''N_NODES''']					)
				assert params.node_id == int(os.environ['''NODE_RANK''']					)

		# local job (single GPU)
		else:
				assert params.local_rank == -1

				lowercase__						=				1
				lowercase__						=				0
				lowercase__						=				0
				lowercase__						=				0
				lowercase__						=				1
				lowercase__						=				1
				lowercase__						=				False

		# sanity checks
		assert params.n_nodes >= 1
		assert 0 <= params.node_id < params.n_nodes
		assert 0 <= params.local_rank <= params.global_rank < params.world_size
		assert params.world_size == params.n_nodes * params.n_gpu_per_node

		# define whether this is the master process / if we are in multi-node distributed mode
		lowercase__						=				params.node_id == 0 and params.local_rank == 0
		lowercase__						=				params.n_nodes > 1

		# summary
		lowercase__						=				f'--- Global rank: {params.global_rank} - '
		logger.info(PREFIX + '''Number of nodes: %i''' % params.n_nodes					)
		logger.info(PREFIX + '''Node ID        : %i''' % params.node_id					)
		logger.info(PREFIX + '''Local rank     : %i''' % params.local_rank					)
		logger.info(PREFIX + '''World size     : %i''' % params.world_size					)
		logger.info(PREFIX + '''GPUs per node  : %i''' % params.n_gpu_per_node					)
		logger.info(PREFIX + '''Master         : %s''' % str(params.is_master					)					)
		logger.info(PREFIX + '''Multi-node     : %s''' % str(params.multi_node					)					)
		logger.info(PREFIX + '''Multi-GPU      : %s''' % str(params.multi_gpu					)					)
		logger.info(PREFIX + '''Hostname       : %s''' % socket.gethostname()					)

		# set GPU device
		torch.cuda.set_device(params.local_rank					)

		# initialize multi-GPU
		if params.multi_gpu:
				logger.info('''Initializing PyTorch distributed'''					)
				torch.distributed.init_process_group(
				    init_method='''env://'''						,    backend='''nccl'''						,    )
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		np.random.seed(args.seed					)
		torch.manual_seed(args.seed					)
		if args.n_gpu > 0:
				torch.cuda.manual_seed_all(args.seed					)


