
import json
import os
from dataclasses import dataclass
from functools import partial
from typing import Callable

import flax.linen as nn
import jax
import jax.numpy as jnp
import joblib
import optax
import wandb
from flax import jax_utils, struct, traverse_util
from flax.serialization import from_bytes, to_bytes
from flax.training import train_state
from flax.training.common_utils import shard
from tqdm.auto import tqdm

from transformers import BigBirdConfig, FlaxBigBirdForQuestionAnswering
from transformers.models.big_bird.modeling_flax_big_bird import FlaxBigBirdForQuestionAnsweringModule



class 		UpperCamelCase__  (       _UpperCAmelCase    ):
   _SCREAMING_SNAKE_CASE      :							BigBirdConfig
   _SCREAMING_SNAKE_CASE      :							jnp.dtype	       =      jnp.floataa
   _SCREAMING_SNAKE_CASE      :							bool	       =      True
   def       lowerCAmelCase			(self :					List[Any]   ):
    super().setup()
    __a :  str						=  nn.Dense(5					,				dtype=self.dtype   )







   def __call__(self :					Tuple					,				*snake_case_ :					Optional[Any]					,				**snake_case_ :					Any   ):
    __a :  Optional[int]						=  super().__call__(*SCREAMING_SNAKE_CASE_					,				**SCREAMING_SNAKE_CASE_   )
    __a :  str						=  self.cls(outputs[2]   )
    return outputs[:2] + (cls_out,)





class 		UpperCamelCase__  (       _UpperCAmelCase    ):
   _SCREAMING_SNAKE_CASE      :							Optional[int]	       =      FlaxBigBirdForNaturalQuestionsModule





def    __UpperCamelCase     (				lowerCAmelCase__							:				Any  ,   lowerCAmelCase__							:				List[Any]  ,   lowerCAmelCase__							:				List[Any]  ,   lowerCAmelCase__							:				Dict  ,   lowerCAmelCase__							:				Any  ,   lowerCAmelCase__							:				List[Any]  ):
 def cross_entropy(lowerCAmelCase__							:				int  ,   lowerCAmelCase__							:				str  ,   lowerCAmelCase__							:				List[str]=None  ):
  __a :  Union[str, Any]						=  logits.shape[-1]
  __a :  Any						=  (labels[..., None] == jnp.arange(UpperCAmelCase_  )[None]).astype('''f4'''  )
  __a :  Dict						=  jax.nn.log_softmax(UpperCAmelCase_  ,   axis=-1  )
  __a :  Optional[int]						=  -jnp.sum(labels * logits  ,   axis=-1  )
  if reduction is not None:
   __a :  str						=  reduction(UpperCAmelCase_  )
  return loss

 __a :  List[str]						=  partial(UpperCAmelCase_  ,   reduction=jnp.mean  )
 __a :  List[str]						=  cross_entropy(UpperCAmelCase_  ,   UpperCAmelCase_  )
 __a :  str						=  cross_entropy(UpperCAmelCase_  ,   UpperCAmelCase_  )
 __a :  Tuple						=  cross_entropy(UpperCAmelCase_  ,   UpperCAmelCase_  )
 return (start_loss + end_loss + pooled_loss) / 3



@dataclass
class 		UpperCamelCase__  :
   _SCREAMING_SNAKE_CASE      :							str	       =      "google/bigbird-roberta-base"
   _SCREAMING_SNAKE_CASE      :							int	       =      3_000
   _SCREAMING_SNAKE_CASE      :							int	       =      10_500

   _SCREAMING_SNAKE_CASE      :							int	       =      128
   _SCREAMING_SNAKE_CASE      :							int	       =      3

   _SCREAMING_SNAKE_CASE      :							int	       =      1
   _SCREAMING_SNAKE_CASE      :							int	       =      5

   # tx_args
   _SCREAMING_SNAKE_CASE      :							float	       =      3e-5
   _SCREAMING_SNAKE_CASE      :							float	       =      0.0
   _SCREAMING_SNAKE_CASE      :							int	       =      20_000
   _SCREAMING_SNAKE_CASE      :							float	       =      0.00_95

   _SCREAMING_SNAKE_CASE      :							str	       =      "bigbird-roberta-natural-questions"
   _SCREAMING_SNAKE_CASE      :							str	       =      "training-expt"
   _SCREAMING_SNAKE_CASE      :							str	       =      "data/nq-training.jsonl"
   _SCREAMING_SNAKE_CASE      :							str	       =      "data/nq-validation.jsonl"
   def       lowerCAmelCase			(self :					Optional[int]   ):
    os.makedirs(self.base_dir					,				exist_ok=SCREAMING_SNAKE_CASE_   )
    __a :  Dict						=  os.path.join(self.base_dir					,				self.save_dir   )
    __a :  Dict						=  self.batch_size_per_device * jax.device_count()





@dataclass
class 		UpperCamelCase__  :
   _SCREAMING_SNAKE_CASE      :							int
   _SCREAMING_SNAKE_CASE      :							int	       =      4_096  # no dynamic padding on TPUs
   def __call__(self :					Union[str, Any]					,				snake_case_ :					Any   ):
    __a :  Any						=  self.collate_fn(SCREAMING_SNAKE_CASE_   )
    __a :  Dict						=  jax.tree_util.tree_map(SCREAMING_SNAKE_CASE_					,				SCREAMING_SNAKE_CASE_   )
    return batch
   def       lowerCAmelCase			(self :					Tuple					,				snake_case_ :					Dict   ):
    __a :  List[str]						=  self.fetch_inputs(features['''input_ids''']   )
    __a :  Dict						=  {
        'input_ids': jnp.array(SCREAMING_SNAKE_CASE_					,				dtype=jnp.intaa   ),
        'attention_mask': jnp.array(SCREAMING_SNAKE_CASE_					,				dtype=jnp.intaa   ),
        'start_labels': jnp.array(features['''start_token''']					,				dtype=jnp.intaa   ),
        'end_labels': jnp.array(features['''end_token''']					,				dtype=jnp.intaa   ),
        'pooled_labels': jnp.array(features['''category''']					,				dtype=jnp.intaa   ),
    }
    return batch
   def       lowerCAmelCase			(self :					str					,				snake_case_ :					str   ):
    __a :  Optional[Any]						=  [self._fetch_inputs(SCREAMING_SNAKE_CASE_   ) for ids in input_ids]
    return zip(*SCREAMING_SNAKE_CASE_   )







   def       lowerCAmelCase			(self :					Optional[Any]					,				snake_case_ :					Union[str, Any]   ):
    __a :  Tuple						=  [1 for _ in range(len(SCREAMING_SNAKE_CASE_   )   )]
    while len(SCREAMING_SNAKE_CASE_   ) < self.max_length:
     input_ids.append(self.pad_id   )
     attention_mask.append(0   )
    return input_ids, attention_mask





def    __UpperCamelCase     (				lowerCAmelCase__							:				Any  ,   lowerCAmelCase__							:				Dict  ,   lowerCAmelCase__							:				List[str]=None  ):
 if seed is not None:
  __a :  Union[str, Any]						=  dataset.shuffle(seed=UpperCAmelCase_  )
 for i in range(len(UpperCAmelCase_  ) // batch_size  ):
  __a :  List[Any]						=  dataset[i * batch_size : (i + 1) * batch_size]
  yield dict(UpperCAmelCase_  )




@partial(jax.pmap  ,   axis_name='''batch'''  )
def    __UpperCamelCase     (				lowerCAmelCase__							:				Union[str, Any]  ,   lowerCAmelCase__							:				Optional[int]  ,   **lowerCAmelCase__							:				List[Any]  ):
 def loss_fn(lowerCAmelCase__							:				Any  ):
  __a :  Union[str, Any]						=  model_inputs.pop('''start_labels'''  )
  __a :  int						=  model_inputs.pop('''end_labels'''  )
  __a :  int						=  model_inputs.pop('''pooled_labels'''  )

  __a :  int						=  state.apply_fn(**UpperCAmelCase_  ,   params=UpperCAmelCase_  ,   dropout_rng=UpperCAmelCase_  ,   train=UpperCAmelCase_  )
  __a :  Optional[Any]						=  outputs

  return state.loss_fn(
      UpperCAmelCase_  ,   UpperCAmelCase_  ,   UpperCAmelCase_  ,   UpperCAmelCase_  ,   UpperCAmelCase_  ,   UpperCAmelCase_  ,   )

 __a :  int						=  jax.random.split(UpperCAmelCase_  )
 __a :  List[Any]						=  jax.value_and_grad(UpperCAmelCase_  )
 __a :  Union[str, Any]						=  grad_fn(state.params  )
 __a :  int						=  jax.lax.pmean({'''loss''': loss}  ,   axis_name='''batch'''  )
 __a :  Dict						=  jax.lax.pmean(UpperCAmelCase_  ,   '''batch'''  )

 __a :  Optional[int]						=  state.apply_gradients(grads=UpperCAmelCase_  )
 return state, metrics, new_drp_rng




@partial(jax.pmap  ,   axis_name='''batch'''  )
def    __UpperCamelCase     (				lowerCAmelCase__							:				List[str]  ,   **lowerCAmelCase__							:				int  ):
 __a :  Optional[int]						=  model_inputs.pop('''start_labels'''  )
 __a :  List[Any]						=  model_inputs.pop('''end_labels'''  )
 __a :  Optional[Any]						=  model_inputs.pop('''pooled_labels'''  )

 __a :  int						=  state.apply_fn(**UpperCAmelCase_  ,   params=state.params  ,   train=UpperCAmelCase_  )
 __a :  Optional[Any]						=  outputs

 __a :  List[str]						=  state.loss_fn(UpperCAmelCase_  ,   UpperCAmelCase_  ,   UpperCAmelCase_  ,   UpperCAmelCase_  ,   UpperCAmelCase_  ,   UpperCAmelCase_  )
 __a :  Optional[int]						=  jax.lax.pmean({'''loss''': loss}  ,   axis_name='''batch'''  )
 return metrics



class 		UpperCamelCase__  (       train_state.TrainState    ):
   _SCREAMING_SNAKE_CASE      :							Callable	       =      struct.field(pytree_node=_UpperCAmelCase    )





@dataclass
class 		UpperCamelCase__  :
   _SCREAMING_SNAKE_CASE      :							Args
   _SCREAMING_SNAKE_CASE      :							Callable
   _SCREAMING_SNAKE_CASE      :							Callable
   _SCREAMING_SNAKE_CASE      :							Callable
   _SCREAMING_SNAKE_CASE      :							Callable
   _SCREAMING_SNAKE_CASE      :							wandb
   _SCREAMING_SNAKE_CASE      :							Callable	       =      None
   def       lowerCAmelCase			(self :					int					,				snake_case_ :					Optional[int]					,				snake_case_ :					Optional[int]					,				snake_case_ :					List[Any]					,				snake_case_ :					Optional[int]=None   ):
    __a :  int						=  model.params
    __a :  List[Any]						=  TrainState.create(
        apply_fn=model.__call__					,				params=SCREAMING_SNAKE_CASE_					,				tx=SCREAMING_SNAKE_CASE_					,				loss_fn=SCREAMING_SNAKE_CASE_					,				)
    if ckpt_dir is not None:
     __a :  Tuple						=  restore_checkpoint(SCREAMING_SNAKE_CASE_					,				SCREAMING_SNAKE_CASE_   )
     __a :  Dict						=  {
         'lr': args.lr,
         'init_lr': args.init_lr,
         'warmup_steps': args.warmup_steps,
         'num_train_steps': num_train_steps,
         'weight_decay': args.weight_decay,
     }
     __a :  Any						=  build_tx(**SCREAMING_SNAKE_CASE_   )
     __a :  Any						=  train_state.TrainState(
         step=SCREAMING_SNAKE_CASE_					,				apply_fn=model.__call__					,				params=SCREAMING_SNAKE_CASE_					,				tx=SCREAMING_SNAKE_CASE_					,				opt_state=SCREAMING_SNAKE_CASE_					,				)
     __a :  str						=  args
     __a :  Tuple						=  data_collator
     __a :  int						=  lr
     __a :  Optional[Any]						=  params
    __a :  Optional[Any]						=  jax_utils.replicate(SCREAMING_SNAKE_CASE_   )
    return state
   def       lowerCAmelCase			(self :					str					,				snake_case_ :					Union[str, Any]					,				snake_case_ :					Optional[Any]					,				snake_case_ :					Optional[int]   ):
    __a :  List[Any]						=  self.args
    __a :  Optional[Any]						=  len(SCREAMING_SNAKE_CASE_   ) // args.batch_size

    __a :  str						=  jax.random.PRNGKey(0   )
    __a :  Optional[Any]						=  jax.random.split(SCREAMING_SNAKE_CASE_					,				jax.device_count()   )
    for epoch in range(args.max_epochs   ):
     __a :  Union[str, Any]						=  jnp.array(0					,				dtype=jnp.floataa   )
     __a :  Optional[Any]						=  get_batched_dataset(SCREAMING_SNAKE_CASE_					,				args.batch_size					,				seed=SCREAMING_SNAKE_CASE_   )
     __a :  List[Any]						=  0
     for batch in tqdm(SCREAMING_SNAKE_CASE_					,				total=SCREAMING_SNAKE_CASE_					,				desc=f"Running EPOCH-{epoch}"   ):
      __a :  List[Any]						=  self.data_collator(SCREAMING_SNAKE_CASE_   )
      __a :  str						=  self.train_step_fn(SCREAMING_SNAKE_CASE_					,				SCREAMING_SNAKE_CASE_					,				**SCREAMING_SNAKE_CASE_   )
      running_loss += jax_utils.unreplicate(metrics['''loss''']   )
      i += 1
      if i % args.logging_steps == 0:
       __a :  Optional[int]						=  jax_utils.unreplicate(state.step   )
       __a :  Tuple						=  running_loss.item() / i
       __a :  int						=  self.scheduler_fn(state_step - 1   )

       __a :  Tuple						=  self.evaluate(SCREAMING_SNAKE_CASE_					,				SCREAMING_SNAKE_CASE_   )
       __a :  Tuple						=  {
           'step': state_step.item(),
           'eval_loss': eval_loss.item(),
           'tr_loss': tr_loss,
           'lr': lr.item(),
       }
       tqdm.write(str(SCREAMING_SNAKE_CASE_   )   )
       self.logger.log(SCREAMING_SNAKE_CASE_					,				commit=SCREAMING_SNAKE_CASE_   )

      if i % args.save_steps == 0:
       self.save_checkpoint(args.save_dir + f"-e{epoch}-s{i}"					,				state=SCREAMING_SNAKE_CASE_   )
   def       lowerCAmelCase			(self :					List[Any]					,				snake_case_ :					Dict					,				snake_case_ :					str   ):
    __a :  List[Any]						=  get_batched_dataset(SCREAMING_SNAKE_CASE_					,				self.args.batch_size   )
    __a :  int						=  len(SCREAMING_SNAKE_CASE_   ) // self.args.batch_size
    __a :  Any						=  jnp.array(0					,				dtype=jnp.floataa   )
    __a :  Tuple						=  0
    for batch in tqdm(SCREAMING_SNAKE_CASE_					,				total=SCREAMING_SNAKE_CASE_					,				desc='''Evaluating ... '''   ):
     __a :  Union[str, Any]						=  self.data_collator(SCREAMING_SNAKE_CASE_   )
     __a :  Any						=  self.val_step_fn(SCREAMING_SNAKE_CASE_					,				**SCREAMING_SNAKE_CASE_   )
     running_loss += jax_utils.unreplicate(metrics['''loss''']   )
     i += 1
    return running_loss / i







   def       lowerCAmelCase			(self :					Optional[Any]					,				snake_case_ :					Tuple					,				snake_case_ :					List[Any]   ):
    __a :  Dict						=  jax_utils.unreplicate(SCREAMING_SNAKE_CASE_   )
    print(f"SAVING CHECKPOINT IN {save_dir}"					,				end=''' ... '''   )
    self.model_save_fn(SCREAMING_SNAKE_CASE_					,				params=state.params   )
    with open(os.path.join(SCREAMING_SNAKE_CASE_					,				'''opt_state.msgpack'''   )					,				'''wb'''   ) as f:
     f.write(to_bytes(state.opt_state   )   )
    joblib.dump(self.args					,				os.path.join(SCREAMING_SNAKE_CASE_					,				'''args.joblib'''   )   )
    joblib.dump(self.data_collator					,				os.path.join(SCREAMING_SNAKE_CASE_					,				'''data_collator.joblib'''   )   )
    with open(os.path.join(SCREAMING_SNAKE_CASE_					,				'''training_state.json'''   )					,				'''w'''   ) as f:
     json.dump({'''step''': state.step.item()}					,				SCREAMING_SNAKE_CASE_   )
    print('''DONE'''   )





def    __UpperCamelCase     (				lowerCAmelCase__							:				List[str]  ,   lowerCAmelCase__							:				Union[str, Any]  ):
 print(f"RESTORING CHECKPOINT FROM {save_dir}"  ,   end=''' ... '''  )
 with open(os.path.join(UpperCAmelCase_  ,   '''flax_model.msgpack'''  )  ,   '''rb'''  ) as f:
  __a :  Optional[Any]						=  from_bytes(state.params  ,   f.read()  )

 with open(os.path.join(UpperCAmelCase_  ,   '''opt_state.msgpack'''  )  ,   '''rb'''  ) as f:
  __a :  Optional[int]						=  from_bytes(state.opt_state  ,   f.read()  )

 __a :  int						=  joblib.load(os.path.join(UpperCAmelCase_  ,   '''args.joblib'''  )  )
 __a :  Optional[Any]						=  joblib.load(os.path.join(UpperCAmelCase_  ,   '''data_collator.joblib'''  )  )

 with open(os.path.join(UpperCAmelCase_  ,   '''training_state.json'''  )  ,   '''r'''  ) as f:
  __a :  Optional[Any]						=  json.load(UpperCAmelCase_  )
 __a :  Optional[Any]						=  training_state['step']

 print('''DONE'''  )
 return params, opt_state, step, args, data_collator




def    __UpperCamelCase     (				lowerCAmelCase__							:				int  ,   lowerCAmelCase__							:				Optional[int]  ,   lowerCAmelCase__							:				str  ,   lowerCAmelCase__							:				Tuple  ):
 __a :  List[Any]						=  num_train_steps - warmup_steps
 __a :  Union[str, Any]						=  optax.linear_schedule(init_value=UpperCAmelCase_  ,   end_value=UpperCAmelCase_  ,   transition_steps=UpperCAmelCase_  )
 __a :  Dict						=  optax.linear_schedule(init_value=UpperCAmelCase_  ,   end_value=1e-7  ,   transition_steps=UpperCAmelCase_  )
 __a :  Optional[int]						=  optax.join_schedules(schedules=[warmup_fn, decay_fn]  ,   boundaries=[warmup_steps]  )
 return lr




def    __UpperCamelCase     (				lowerCAmelCase__							:				List[Any]  ,   lowerCAmelCase__							:				Optional[int]  ,   lowerCAmelCase__							:				int  ,   lowerCAmelCase__							:				Union[str, Any]  ,   lowerCAmelCase__							:				Optional[int]  ):
 def weight_decay_mask(lowerCAmelCase__							:				List[Any]  ):
  __a :  List[str]						=  traverse_util.flatten_dict(UpperCAmelCase_  )
  __a :  List[Any]						=  {k: (v[-1] != 'bias' and v[-2:] != ('LayerNorm', 'scale')) for k, v in params.items()}
  return traverse_util.unflatten_dict(UpperCAmelCase_  )

 __a :  List[Any]						=  scheduler_fn(UpperCAmelCase_  ,   UpperCAmelCase_  ,   UpperCAmelCase_  ,   UpperCAmelCase_  )

 __a :  Optional[int]						=  optax.adamw(learning_rate=UpperCAmelCase_  ,   weight_decay=UpperCAmelCase_  ,   mask=UpperCAmelCase_  )
 return tx, lr




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


A__		:					List[str]												=  get_logger(__name__)


A__		:					str												=  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(SCREAMING_SNAKE_CASE_	)
    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_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(SCREAMING_SNAKE_CASE_	)
    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         raise NotImplementedError(
             f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'	)







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    @add_start_docstrings(SCREAMING_SNAKE_CASE_	)
    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     **SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         for processor in self:
              __lowerCamelCase   :	str								=		inspect.signature(processor.__call__	).parameters
              if len(SCREAMING_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.'	)
                   __lowerCamelCase   :	Tuple								=		processor(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     **SCREAMING_SNAKE_CASE_	)
              else:
                   __lowerCamelCase   :	int								=		processor(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






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

         __lowerCamelCase   :	Optional[int]								=		temperature







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	Dict								=		scores / self.temperature
         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = -float('Inf'	)			,     SCREAMING_SNAKE_CASE_ = 1	)				->     Union[str, Any]:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_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(SCREAMING_SNAKE_CASE_			,     SCREAMING_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}'	)

         __lowerCamelCase   :	str								=		top_p
         __lowerCamelCase   :	Tuple								=		filter_value
         __lowerCamelCase   :	Tuple								=		min_tokens_to_keep







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase				,       __lowerCamelCase   :	Any								=		lax.top_k(SCREAMING_SNAKE_CASE_			,     scores.shape[-1]	)

         __lowerCamelCase   :	int								=		jnp.full_like(SCREAMING_SNAKE_CASE_			,     self.filter_value	)
         __lowerCamelCase   :	Tuple								=		jax.nn.softmax(SCREAMING_SNAKE_CASE_			,     axis=-1	).cumsum(axis=-1	)
         __lowerCamelCase   :	List[str]								=		cumulative_probs < self.top_p

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

         # min tokens to keep
         __lowerCamelCase   :	Any								=		score_mask.at[:, : self.min_tokens_to_keep].set(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Any								=		jnp.where(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Tuple								=		jax.lax.sort_key_val(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)[-1]

         return next_scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






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

         __lowerCamelCase   :	List[str]								=		max(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Dict								=		filter_value







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase				,       __lowerCamelCase   :	List[Any]								=		scores.shape
         __lowerCamelCase   :	Tuple								=		jnp.full(batch_size * vocab_size			,     self.filter_value	)

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

         __lowerCamelCase   :	Tuple								=		next_scores_flat.at[topk_indices_flat].set(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		next_scores_flat.reshape(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         return next_scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     Union[str, Any]:
         __lowerCamelCase   :	Any								=		bos_token_id







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	Optional[Any]								=		jnp.full(scores.shape			,     -float('inf'	)	)

         __lowerCamelCase   :	Optional[Any]								=		1 - jnp.bool_(cur_len - 1	)

         __lowerCamelCase   :	List[Any]								=		jnp.where(SCREAMING_SNAKE_CASE_			,     new_scores.at[:, self.bos_token_id].set(0	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Dict:
         __lowerCamelCase   :	Tuple								=		max_length
         __lowerCamelCase   :	Any								=		eos_token_id







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	List[str]								=		jnp.full(scores.shape			,     -float('inf'	)	)

         __lowerCamelCase   :	Any								=		1 - jnp.bool_(cur_len - self.max_length + 1	)

         __lowerCamelCase   :	List[str]								=		jnp.where(SCREAMING_SNAKE_CASE_			,     new_scores.at[:, self.eos_token_id].set(0	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     List[str]:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or min_length < 0:
              raise ValueError(f'`min_length` has to be a positive integer, but is {min_length}'	)

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

         __lowerCamelCase   :	str								=		min_length
         __lowerCamelCase   :	Optional[int]								=		eos_token_id







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

         __lowerCamelCase   :	str								=		jnp.where(SCREAMING_SNAKE_CASE_			,     scores.at[:, self.eos_token_id].set(-float('inf'	)	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Dict:
         __lowerCamelCase   :	Union[str, Any]								=		list(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	str								=		begin_index







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         __lowerCamelCase   :	List[Any]								=		1 - jnp.bool_(cur_len - self.begin_index	)

         __lowerCamelCase   :	str								=		jnp.where(SCREAMING_SNAKE_CASE_			,     scores.at[:, self.begin_suppress_tokens].set(-float('inf'	)	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     List[Any]:
         __lowerCamelCase   :	Tuple								=		list(SCREAMING_SNAKE_CASE_	)







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	int								=		scores.at[..., self.suppress_tokens].set(-float('inf'	)	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     Union[str, Any]:
         __lowerCamelCase   :	Optional[int]								=		dict(SCREAMING_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.
         __lowerCamelCase   :	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:
                   __lowerCamelCase   :	str								=		force_token_array.at[index].set(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[Any]								=		jnp.intaa(SCREAMING_SNAKE_CASE_	)







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         def _force_token(SCREAMING_SNAKE_CASE_	):
              __lowerCamelCase   :	List[str]								=		scores.shape[0]
              __lowerCamelCase   :	Tuple								=		self.force_token_array[generation_idx]

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

         __lowerCamelCase   :	int								=		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(SCREAMING_SNAKE_CASE_	)			,     lambda: scores			,     )			,     )
         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         __lowerCamelCase   :	Any								=		generate_config.eos_token_id
         __lowerCamelCase   :	Dict								=		generate_config.no_timestamps_token_id
         __lowerCamelCase   :	Tuple								=		generate_config.no_timestamps_token_id + 1

         __lowerCamelCase   :	List[str]								=		decoder_input_length + 1

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







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         # suppress <|notimestamps|> which is handled by without_timestamps
         __lowerCamelCase   :	Union[str, Any]								=		scores.at[:, self.no_timestamps_token_id].set(-float('inf'	)	)

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

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

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

         __lowerCamelCase   :	List[Any]								=		jax.vmap(SCREAMING_SNAKE_CASE_	)(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

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

         __lowerCamelCase   :	Any								=		self.timestamp_begin + self.max_initial_timestamp_index

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

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

         def handle_cumulative_probs(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	):
              __lowerCamelCase   :	Union[str, Any]								=		jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]			,     axis=-1	)
              __lowerCamelCase   :	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'	)	)			,     SCREAMING_SNAKE_CASE_			,     )

         __lowerCamelCase   :	Dict								=		jax.vmap(SCREAMING_SNAKE_CASE_	)(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

         return scores






