


'''simple docstring'''



import argparse
import collections
import json
import os
import re
import string
import sys

import numpy as np


A__		:					Optional[int]												=  re.compile(R"""\b(a|an|the)\b""", re.UNICODE)

A__		:					Tuple												=  None



def 			UpperCAmelCase__					(	)     ->		Optional[Any]:
     __lowerCamelCase   :	Optional[int]								=		argparse.ArgumentParser('Official evaluation script for SQuAD version 2.0.'							)
     parser.add_argument('data_file'      ,    metavar='data.json'      ,    help='Input data JSON file.'							)
     parser.add_argument('pred_file'      ,    metavar='pred.json'      ,    help='Model predictions.'							)
     parser.add_argument(
         '--out-file'      ,    '-o'      ,    metavar='eval.json'      ,    help='Write accuracy metrics to file (default is stdout).'							)
     parser.add_argument(
         '--na-prob-file'      ,    '-n'      ,    metavar='na_prob.json'      ,    help='Model estimates of probability of no answer.'							)
     parser.add_argument(
         '--na-prob-thresh'      ,    '-t'      ,    type=UpperCAmelCase_      ,    default=1.0      ,    help='Predict "" if no-answer probability exceeds this (default = 1.0).'      ,    )
     parser.add_argument(
         '--out-image-dir'      ,    '-p'      ,    metavar='out_images'      ,    default=UpperCAmelCase_      ,    help='Save precision-recall curves to directory.'							)
     parser.add_argument('--verbose'      ,    '-v'      ,    action='store_true'							)
     if len(sys.argv							) == 1:
          parser.print_help()
          sys.exit(1							)
     return parser.parse_args()



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				List[Any]							)     ->		Optional[Any]:
     __lowerCamelCase   :	int								=		{}
     for article in dataset:
          for p in article["paragraphs"]:
               for qa in p["qas"]:
                    __lowerCamelCase   :	Union[str, Any]								=		bool(qa['answers']['text']							)
     return qid_to_has_ans



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				List[str]							)     ->		Optional[int]:

     def remove_articles(UpperCAmelCase_    :				Optional[int]							):
          return ARTICLES_REGEX.sub(' '      ,    UpperCAmelCase_							)

     def white_space_fix(UpperCAmelCase_    :				Dict							):
          return " ".join(text.split()							)

     def remove_punc(UpperCAmelCase_    :				Optional[int]							):
          __lowerCamelCase   :	Dict								=		set(string.punctuation							)
          return "".join(ch for ch in text if ch not in exclude							)

     def lower(UpperCAmelCase_    :				Union[str, Any]							):
          return text.lower()

     return white_space_fix(remove_articles(remove_punc(lower(UpperCAmelCase_							)							)							)							)



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				List[str]							)     ->		Optional[int]:
     if not s:
          return []
     return normalize_answer(UpperCAmelCase_							).split()



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Dict      ,    UpperCAmelCase_    :				Tuple							)     ->		int:
     return int(normalize_answer(UpperCAmelCase_							) == normalize_answer(UpperCAmelCase_							)							)



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				List[Any]      ,    UpperCAmelCase_    :				int							)     ->		Optional[int]:
     __lowerCamelCase   :	List[Any]								=		get_tokens(UpperCAmelCase_							)
     __lowerCamelCase   :	Union[str, Any]								=		get_tokens(UpperCAmelCase_							)
     __lowerCamelCase   :	List[str]								=		collections.Counter(UpperCAmelCase_							) & collections.Counter(UpperCAmelCase_							)
     __lowerCamelCase   :	List[str]								=		sum(common.values()							)
     if len(UpperCAmelCase_							) == 0 or len(UpperCAmelCase_							) == 0:
          # If either is no-answer, then F1 is 1 if they agree, 0 otherwise
          return int(gold_toks == pred_toks							)
     if num_same == 0:
          return 0
     __lowerCamelCase   :	str								=		1.0 * num_same / len(UpperCAmelCase_							)
     __lowerCamelCase   :	Dict								=		1.0 * num_same / len(UpperCAmelCase_							)
     __lowerCamelCase   :	int								=		(2 * precision * recall) / (precision + recall)
     return fa



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Optional[int]      ,    UpperCAmelCase_    :				Union[str, Any]							)     ->		Tuple:
     __lowerCamelCase   :	Optional[int]								=		{}
     __lowerCamelCase   :	List[Any]								=		{}
     for article in dataset:
          for p in article["paragraphs"]:
               for qa in p["qas"]:
                    __lowerCamelCase   :	Union[str, Any]								=		qa['id']
                    __lowerCamelCase   :	List[str]								=		[t for t in qa['answers']['text'] if normalize_answer(UpperCAmelCase_							)]
                    if not gold_answers:
                         # For unanswerable questions, only correct answer is empty string
                         __lowerCamelCase   :	Optional[int]								=		['']
                    if qid not in preds:
                         print(F'Missing prediction for {qid}'							)
                         continue
                    __lowerCamelCase   :	str								=		preds[qid]
                    # Take max over all gold answers
                    __lowerCamelCase   :	Optional[Any]								=		max(compute_exact(UpperCAmelCase_      ,    UpperCAmelCase_							) for a in gold_answers							)
                    __lowerCamelCase   :	str								=		max(compute_fa(UpperCAmelCase_      ,    UpperCAmelCase_							) for a in gold_answers							)
     return exact_scores, fa_scores



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Dict      ,    UpperCAmelCase_    :				List[str]      ,    UpperCAmelCase_    :				Union[str, Any]      ,    UpperCAmelCase_    :				Dict							)     ->		Tuple:
     __lowerCamelCase   :	str								=		{}
     for qid, s in scores.items():
          __lowerCamelCase   :	Union[str, Any]								=		na_probs[qid] > na_prob_thresh
          if pred_na:
               __lowerCamelCase   :	int								=		float(not qid_to_has_ans[qid]							)
          else:
               __lowerCamelCase   :	List[str]								=		s
     return new_scores



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				List[Any]      ,    UpperCAmelCase_    :				Dict      ,    UpperCAmelCase_    :				str=None							)     ->		int:
     if not qid_list:
          __lowerCamelCase   :	List[str]								=		len(UpperCAmelCase_							)
          return collections.OrderedDict(
              [
                  ('exact', 100.0 * sum(exact_scores.values()							) / total),
                  ('f1', 100.0 * sum(fa_scores.values()							) / total),
                  ('total', total),
              ]							)
     else:
          __lowerCamelCase   :	List[Any]								=		len(UpperCAmelCase_							)
          return collections.OrderedDict(
              [
                  ('exact', 100.0 * sum(exact_scores[k] for k in qid_list							) / total),
                  ('f1', 100.0 * sum(fa_scores[k] for k in qid_list							) / total),
                  ('total', total),
              ]							)



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Optional[int]      ,    UpperCAmelCase_    :				List[Any]      ,    UpperCAmelCase_    :				Any							)     ->		Union[str, Any]:
     for k in new_eval:
          __lowerCamelCase   :	int								=		new_eval[k]



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Optional[Any]      ,    UpperCAmelCase_    :				int      ,    UpperCAmelCase_    :				Any      ,    UpperCAmelCase_    :				int							)     ->		Optional[Any]:
     plt.step(UpperCAmelCase_      ,    UpperCAmelCase_      ,    color='b'      ,    alpha=0.2      ,    where='post'							)
     plt.fill_between(UpperCAmelCase_      ,    UpperCAmelCase_      ,    step='post'      ,    alpha=0.2      ,    color='b'							)
     plt.xlabel('Recall'							)
     plt.ylabel('Precision'							)
     plt.xlim([0.0, 1.05]							)
     plt.ylim([0.0, 1.05]							)
     plt.title(UpperCAmelCase_							)
     plt.savefig(UpperCAmelCase_							)
     plt.clf()



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				str      ,    UpperCAmelCase_    :				Optional[Any]      ,    UpperCAmelCase_    :				str      ,    UpperCAmelCase_    :				Any      ,    UpperCAmelCase_    :				List[Any]=None      ,    UpperCAmelCase_    :				Dict=None							)     ->		Dict:
     __lowerCamelCase   :	int								=		sorted(UpperCAmelCase_      ,    key=lambda UpperCAmelCase_							: na_probs[k]							)
     __lowerCamelCase   :	List[Any]								=		0.0
     __lowerCamelCase   :	Optional[int]								=		1.0
     __lowerCamelCase   :	List[Any]								=		0.0
     __lowerCamelCase   :	Optional[int]								=		[1.0]
     __lowerCamelCase   :	Optional[int]								=		[0.0]
     __lowerCamelCase   :	Union[str, Any]								=		0.0
     for i, qid in enumerate(UpperCAmelCase_							):
          if qid_to_has_ans[qid]:
               true_pos += scores[qid]
          __lowerCamelCase   :	List[Any]								=		true_pos / float(i + 1							)
          __lowerCamelCase   :	str								=		true_pos / float(UpperCAmelCase_							)
          if i == len(UpperCAmelCase_							) - 1 or na_probs[qid] != na_probs[qid_list[i + 1]]:
               # i.e., if we can put a threshold after this point
               avg_prec += cur_p * (cur_r - recalls[-1])
               precisions.append(UpperCAmelCase_							)
               recalls.append(UpperCAmelCase_							)
     if out_image:
          plot_pr_curve(UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_							)
     return {"ap": 100.0 * avg_prec}



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Optional[int]      ,    UpperCAmelCase_    :				Union[str, Any]      ,    UpperCAmelCase_    :				List[str]      ,    UpperCAmelCase_    :				str      ,    UpperCAmelCase_    :				str      ,    UpperCAmelCase_    :				List[str]							)     ->		Any:
     if out_image_dir and not os.path.exists(UpperCAmelCase_							):
          os.makedirs(UpperCAmelCase_							)
     __lowerCamelCase   :	str								=		sum(1 for v in qid_to_has_ans.values() if v							)
     if num_true_pos == 0:
          return
     __lowerCamelCase   :	Optional[Any]								=		make_precision_recall_eval(
         UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    out_image=os.path.join(UpperCAmelCase_      ,    'pr_exact.png'							)      ,    title='Precision-Recall curve for Exact Match score'      ,    )
     __lowerCamelCase   :	Optional[Any]								=		make_precision_recall_eval(
         UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    out_image=os.path.join(UpperCAmelCase_      ,    'pr_f1.png'							)      ,    title='Precision-Recall curve for F1 score'      ,    )
     __lowerCamelCase   :	Optional[Any]								=		{k: float(UpperCAmelCase_							) for k, v in qid_to_has_ans.items()}
     __lowerCamelCase   :	Tuple								=		make_precision_recall_eval(
         UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    out_image=os.path.join(UpperCAmelCase_      ,    'pr_oracle.png'							)      ,    title='Oracle Precision-Recall curve (binary task of HasAns vs. NoAns)'      ,    )
     merge_eval(UpperCAmelCase_      ,    UpperCAmelCase_      ,    'pr_exact'							)
     merge_eval(UpperCAmelCase_      ,    UpperCAmelCase_      ,    'pr_f1'							)
     merge_eval(UpperCAmelCase_      ,    UpperCAmelCase_      ,    'pr_oracle'							)



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Union[str, Any]      ,    UpperCAmelCase_    :				str      ,    UpperCAmelCase_    :				Any      ,    UpperCAmelCase_    :				List[str]							)     ->		Tuple:
     if not qid_list:
          return
     __lowerCamelCase   :	List[Any]								=		[na_probs[k] for k in qid_list]
     __lowerCamelCase   :	Union[str, Any]								=		np.ones_like(UpperCAmelCase_							) / float(len(UpperCAmelCase_							)							)
     plt.hist(UpperCAmelCase_      ,    weights=UpperCAmelCase_      ,    bins=20      ,    range=(0.0, 1.0)							)
     plt.xlabel('Model probability of no-answer'							)
     plt.ylabel('Proportion of dataset'							)
     plt.title(F'Histogram of no-answer probability: {name}'							)
     plt.savefig(os.path.join(UpperCAmelCase_      ,    F'na_prob_hist_{name}.png'							)							)
     plt.clf()



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Any      ,    UpperCAmelCase_    :				List[Any]      ,    UpperCAmelCase_    :				List[Any]      ,    UpperCAmelCase_    :				Optional[Any]							)     ->		Optional[Any]:
     __lowerCamelCase   :	str								=		sum(1 for k in qid_to_has_ans if not qid_to_has_ans[k]							)
     __lowerCamelCase   :	Tuple								=		num_no_ans
     __lowerCamelCase   :	Tuple								=		cur_score
     __lowerCamelCase   :	int								=		0.0
     __lowerCamelCase   :	List[str]								=		sorted(UpperCAmelCase_      ,    key=lambda UpperCAmelCase_							: na_probs[k]							)
     for i, qid in enumerate(UpperCAmelCase_							):
          if qid not in scores:
               continue
          if qid_to_has_ans[qid]:
               __lowerCamelCase   :	Optional[int]								=		scores[qid]
          else:
               if preds[qid]:
                    __lowerCamelCase   :	str								=		-1
               else:
                    __lowerCamelCase   :	Any								=		0
          cur_score += diff
          if cur_score > best_score:
               __lowerCamelCase   :	int								=		cur_score
               __lowerCamelCase   :	Union[str, Any]								=		na_probs[qid]
     return 100.0 * best_score / len(UpperCAmelCase_							), best_thresh



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				List[Any]      ,    UpperCAmelCase_    :				int      ,    UpperCAmelCase_    :				Tuple      ,    UpperCAmelCase_    :				str      ,    UpperCAmelCase_    :				int      ,    UpperCAmelCase_    :				Optional[Any]							)     ->		int:
     __lowerCamelCase				,       __lowerCamelCase   :	int								=		find_best_thresh(UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_							)
     __lowerCamelCase				,       __lowerCamelCase   :	List[Any]								=		find_best_thresh(UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_							)
     __lowerCamelCase   :	str								=		best_exact
     __lowerCamelCase   :	List[str]								=		exact_thresh
     __lowerCamelCase   :	Tuple								=		best_fa
     __lowerCamelCase   :	int								=		fa_thresh



def 			UpperCAmelCase__					(	)     ->		List[str]:
     with open(OPTS.data_file							) as f:
          __lowerCamelCase   :	Tuple								=		json.load(UpperCAmelCase_							)
          __lowerCamelCase   :	Any								=		dataset_json['data']
     with open(OPTS.pred_file							) as f:
          __lowerCamelCase   :	Optional[Any]								=		json.load(UpperCAmelCase_							)
     if OPTS.na_prob_file:
          with open(OPTS.na_prob_file							) as f:
               __lowerCamelCase   :	Optional[int]								=		json.load(UpperCAmelCase_							)
     else:
          __lowerCamelCase   :	Union[str, Any]								=		{k: 0.0 for k in preds}
     __lowerCamelCase   :	Optional[int]								=		make_qid_to_has_ans(UpperCAmelCase_							)  # maps qid to True/False
     __lowerCamelCase   :	List[Any]								=		[k for k, v in qid_to_has_ans.items() if v]
     __lowerCamelCase   :	Tuple								=		[k for k, v in qid_to_has_ans.items() if not v]
     __lowerCamelCase				,       __lowerCamelCase   :	Union[str, Any]								=		get_raw_scores(UpperCAmelCase_      ,    UpperCAmelCase_							)
     __lowerCamelCase   :	Any								=		apply_no_ans_threshold(UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    OPTS.na_prob_thresh							)
     __lowerCamelCase   :	Tuple								=		apply_no_ans_threshold(UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    OPTS.na_prob_thresh							)
     __lowerCamelCase   :	Optional[Any]								=		make_eval_dict(UpperCAmelCase_      ,    UpperCAmelCase_							)
     if has_ans_qids:
          __lowerCamelCase   :	int								=		make_eval_dict(UpperCAmelCase_      ,    UpperCAmelCase_      ,    qid_list=UpperCAmelCase_							)
          merge_eval(UpperCAmelCase_      ,    UpperCAmelCase_      ,    'HasAns'							)
     if no_ans_qids:
          __lowerCamelCase   :	List[str]								=		make_eval_dict(UpperCAmelCase_      ,    UpperCAmelCase_      ,    qid_list=UpperCAmelCase_							)
          merge_eval(UpperCAmelCase_      ,    UpperCAmelCase_      ,    'NoAns'							)
     if OPTS.na_prob_file:
          find_all_best_thresh(UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_							)
     if OPTS.na_prob_file and OPTS.out_image_dir:
          run_precision_recall_analysis(UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_      ,    OPTS.out_image_dir							)
          histogram_na_prob(UpperCAmelCase_      ,    UpperCAmelCase_      ,    OPTS.out_image_dir      ,    'hasAns'							)
          histogram_na_prob(UpperCAmelCase_      ,    UpperCAmelCase_      ,    OPTS.out_image_dir      ,    'noAns'							)
     if OPTS.out_file:
          with open(OPTS.out_file      ,    'w'							) as f:
               json.dump(UpperCAmelCase_      ,    UpperCAmelCase_							)
     else:
          print(json.dumps(UpperCAmelCase_      ,    indent=2							)							)


if __name__ == "__main__":
      A__		:					Optional[Any]												=  parse_args()
      if OPTS.out_image_dir:
            import matplotlib

            matplotlib.use("""Agg""")
            import matplotlib.pyplot as plt
      main()










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






