
"""simple docstring"""







# Note: if you intend to run this script make sure you look under scripts/fsmt/
# to locate the appropriate script to do the work correctly. There is a set of scripts to:
# - download and prepare data and run the conversion script
# - perform eval to get the best hparam into the config
# - generate model_cards - useful if you have multiple models from the same paper

import argparse
import json
import os
import re
from collections import OrderedDict
from os.path import basename, dirname

import fairseq
import torch
from fairseq import hub_utils
from fairseq.data.dictionary import Dictionary

from transformers import FSMTConfig, FSMTForConditionalGeneration
from transformers.models.fsmt.tokenization_fsmt import VOCAB_FILES_NAMES
from transformers.tokenization_utils_base import TOKENIZER_CONFIG_FILE
from transformers.utils import WEIGHTS_NAME, logging


logging.set_verbosity_warning()

lowercase_         =       2

# based on the results of a search on a range of `num_beams`, `length_penalty` and `early_stopping`
# values against wmt19 test data to obtain the best BLEU scores, we will use the following defaults:
#
# * `num_beams`: 5 (higher scores better, but requires more memory/is slower, can be adjusted by users)
# * `early_stopping`: `False` consistently scored better
# * `length_penalty` varied, so will assign the best one depending on the model
lowercase_         =       {
    # fairseq:
    "wmt19-ru-en": {"length_penalty": 1.1},
    "wmt19-en-ru": {"length_penalty": 1.15},
    "wmt19-en-de": {"length_penalty": 1.0},
    "wmt19-de-en": {"length_penalty": 1.1},
    # allenai:
    "wmt16-en-de-dist-12-1": {"length_penalty": 0.6},
    "wmt16-en-de-dist-6-1": {"length_penalty": 0.6},
    "wmt16-en-de-12-1": {"length_penalty": 0.8},
    "wmt19-de-en-6-6-base": {"length_penalty": 0.6},
    "wmt19-de-en-6-6-big": {"length_penalty": 0.6},
}

# this remaps the different models to their organization names
lowercase_         =       {}
for m in ["wmt19-ru-en", "wmt19-en-ru", "wmt19-en-de", "wmt19-de-en"]:
       lowercase_         =       "facebook"
for m in [
    "wmt16-en-de-dist-12-1",
    "wmt16-en-de-dist-6-1",
    "wmt16-en-de-12-1",
    "wmt19-de-en-6-6-base",
    "wmt19-de-en-6-6-big",
]:
       lowercase_         =       "allenai"




def       A_      (       lowercase   )  ->				Union[str, Any]:


     """simple docstring"""







     UpperCAmelCase_    :     Optional[int]            =	dict((re.sub(R"""@@$"""    ,			""""""    ,			lowercase   ), v) if k.endswith("""@@"""   ) else (re.sub(R"""$"""    ,			"""</w>"""    ,			lowercase   ), v) for k, v in d.items()   )
     UpperCAmelCase_    :     str            =	"""<s> <pad> </s> <unk>""".split()
     # restore the special tokens
     for k in keep_keys:
          del da[f'''{k}</w>''']
          UpperCAmelCase_    :     int            =	d[k]  # restore
     return da




def       A_      (       lowercase    ,			lowercase   )  ->				str:


     """simple docstring"""







     assert os.path.exists(lowercase   )
     os.makedirs(lowercase    ,			exist_ok=lowercase   )
     print(f'''Writing results to {pytorch_dump_folder_path}'''   )

     # handle various types of models

     UpperCAmelCase_    :     Optional[int]            =	basename(lowercase   )
     UpperCAmelCase_    :     Tuple            =	dirname(lowercase   )

     UpperCAmelCase_    :     Union[str, Any]            =	fairseq.model_parallel.models.transformer.ModelParallelTransformerModel
     UpperCAmelCase_    :     Dict            =	cls.hub_models()
     UpperCAmelCase_    :     Optional[int]            =	{"""bpe""": """fastbpe""", """tokenizer""": """moses"""}
     UpperCAmelCase_    :     List[str]            =	"""."""
     # note: since the model dump is old, fairseq has upgraded its model some
     # time later, and it does a whole lot of rewrites and splits on the saved
     # weights, therefore we can't use torch.load() directly on the model file.
     # see: upgrade_state_dict(state_dict) in fairseq_model.py
     print(f'''using checkpoint {checkpoint_file}'''   )
     UpperCAmelCase_    :     Tuple            =	hub_utils.from_pretrained(
         lowercase    ,			lowercase    ,			lowercase    ,			archive_map=lowercase    ,			**lowercase   )

     UpperCAmelCase_    :     Optional[Any]            =	vars(chkpt["""args"""]["""model"""]   )

     UpperCAmelCase_    :     Optional[Any]            =	args["""source_lang"""]
     UpperCAmelCase_    :     Any            =	args["""target_lang"""]

     UpperCAmelCase_    :     Optional[int]            =	dirname(lowercase   )
     UpperCAmelCase_    :     Any            =	basename(lowercase   )

     # dicts
     UpperCAmelCase_    :     Optional[Any]            =	os.path.join(lowercase    ,			f'''dict.{src_lang}.txt'''   )
     UpperCAmelCase_    :     Tuple            =	os.path.join(lowercase    ,			f'''dict.{tgt_lang}.txt'''   )

     UpperCAmelCase_    :     Dict            =	Dictionary.load(lowercase   )
     UpperCAmelCase_    :     Optional[int]            =	rewrite_dict_keys(src_dict.indices   )
     UpperCAmelCase_    :     str            =	len(lowercase   )
     UpperCAmelCase_    :     Union[str, Any]            =	os.path.join(lowercase    ,			"""vocab-src.json"""   )
     print(f'''Generating {src_vocab_file} of {src_vocab_size} of {src_lang} records'''   )
     with open(lowercase    ,			"""w"""    ,			encoding="""utf-8"""   ) as f:
          f.write(json.dumps(lowercase    ,			ensure_ascii=lowercase    ,			indent=lowercase   )   )

     # detect whether this is a do_lower_case situation, which can be derived by checking whether we
     # have at least one uppercase letter in the source vocab
     UpperCAmelCase_    :     Optional[Any]            =	True
     for k in src_vocab.keys():
          if not k.islower():
               UpperCAmelCase_    :     List[str]            =	False
               break

     UpperCAmelCase_    :     Optional[int]            =	Dictionary.load(lowercase   )
     UpperCAmelCase_    :     Any            =	rewrite_dict_keys(tgt_dict.indices   )
     UpperCAmelCase_    :     Any            =	len(lowercase   )
     UpperCAmelCase_    :     int            =	os.path.join(lowercase    ,			"""vocab-tgt.json"""   )
     print(f'''Generating {tgt_vocab_file} of {tgt_vocab_size} of {tgt_lang} records'''   )
     with open(lowercase    ,			"""w"""    ,			encoding="""utf-8"""   ) as f:
          f.write(json.dumps(lowercase    ,			ensure_ascii=lowercase    ,			indent=lowercase   )   )

     # merges_file (bpecodes)
     UpperCAmelCase_    :     List[str]            =	os.path.join(lowercase    ,			VOCAB_FILES_NAMES["""merges_file"""]   )
     for fn in ["bpecodes", "code"]:  # older fairseq called the merges file "code"
          UpperCAmelCase_    :     Union[str, Any]            =	os.path.join(lowercase    ,			lowercase   )
          if os.path.exists(lowercase   ):
               break
     with open(lowercase    ,			encoding="""utf-8"""   ) as fin:
          UpperCAmelCase_    :     int            =	fin.read()
     UpperCAmelCase_    :     List[str]            =	re.sub(R""" \d+$"""    ,			""""""    ,			lowercase    ,			0    ,			re.M   )  # remove frequency number
     print(f'''Generating {merges_file}'''   )
     with open(lowercase    ,			"""w"""    ,			encoding="""utf-8"""   ) as fout:
          fout.write(lowercase   )

     # model config
     UpperCAmelCase_    :     str            =	os.path.join(lowercase    ,			"""config.json"""   )

     # validate bpe/tokenizer config, as currently it's hardcoded to moses+fastbpe -
     # may have to modify the tokenizer if a different type is used by a future model
     assert args["bpe"] == "fastbpe", f'''need to extend tokenizer to support bpe={args['bpe']}'''
     assert args["tokenizer"] == "moses", f'''need to extend tokenizer to support bpe={args['tokenizer']}'''

     UpperCAmelCase_    :     str            =	{
         """architectures""": ["""FSMTForConditionalGeneration"""],
         """model_type""": """fsmt""",
         """activation_dropout""": args["""activation_dropout"""],
         """activation_function""": """relu""",
         """attention_dropout""": args["""attention_dropout"""],
         """d_model""": args["""decoder_embed_dim"""],
         """dropout""": args["""dropout"""],
         """init_std""": 0.02,
         """max_position_embeddings""": args["""max_source_positions"""],
         """num_hidden_layers""": args["""encoder_layers"""],
         """src_vocab_size""": src_vocab_size,
         """tgt_vocab_size""": tgt_vocab_size,
         """langs""": [src_lang, tgt_lang],
         """encoder_attention_heads""": args["""encoder_attention_heads"""],
         """encoder_ffn_dim""": args["""encoder_ffn_embed_dim"""],
         """encoder_layerdrop""": args["""encoder_layerdrop"""],
         """encoder_layers""": args["""encoder_layers"""],
         """decoder_attention_heads""": args["""decoder_attention_heads"""],
         """decoder_ffn_dim""": args["""decoder_ffn_embed_dim"""],
         """decoder_layerdrop""": args["""decoder_layerdrop"""],
         """decoder_layers""": args["""decoder_layers"""],
         """bos_token_id""": 0,
         """pad_token_id""": 1,
         """eos_token_id""": 2,
         """is_encoder_decoder""": True,
         """scale_embedding""": not args["""no_scale_embedding"""],
         """tie_word_embeddings""": args["""share_all_embeddings"""],
     }

     # good hparam defaults to start with
     UpperCAmelCase_    :     List[Any]            =	5
     UpperCAmelCase_    :     List[str]            =	False
     if model_dir in best_score_hparams and "length_penalty" in best_score_hparams[model_dir]:
          UpperCAmelCase_    :     int            =	best_score_hparams[model_dir]["""length_penalty"""]
     else:
          UpperCAmelCase_    :     List[Any]            =	1.0

     print(f'''Generating {fsmt_model_config_file}'''   )
     with open(lowercase    ,			"""w"""    ,			encoding="""utf-8"""   ) as f:
          f.write(json.dumps(lowercase    ,			ensure_ascii=lowercase    ,			indent=lowercase   )   )

     # tokenizer config
     UpperCAmelCase_    :     Dict            =	os.path.join(lowercase    ,			lowercase   )

     UpperCAmelCase_    :     Dict            =	{
         """langs""": [src_lang, tgt_lang],
         """model_max_length""": 1024,
         """do_lower_case""": do_lower_case,
     }

     print(f'''Generating {fsmt_tokenizer_config_file}'''   )
     with open(lowercase    ,			"""w"""    ,			encoding="""utf-8"""   ) as f:
          f.write(json.dumps(lowercase    ,			ensure_ascii=lowercase    ,			indent=lowercase   )   )

     # model
     UpperCAmelCase_    :     List[str]            =	chkpt["""models"""][0]
     UpperCAmelCase_    :     Dict            =	model.state_dict()

     # rename keys to start with 'model.'
     UpperCAmelCase_    :     str            =	OrderedDict(("""model.""" + k, v) for k, v in model_state_dict.items()   )

     # remove unneeded keys
     UpperCAmelCase_    :     List[str]            =	[
         """model.model""",
         """model.encoder.version""",
         """model.decoder.version""",
         """model.encoder_embed_tokens.weight""",
         """model.decoder_embed_tokens.weight""",
         """model.encoder.embed_positions._float_tensor""",
         """model.decoder.embed_positions._float_tensor""",
     ]
     for k in ignore_keys:
          model_state_dict.pop(lowercase    ,			lowercase   )

     UpperCAmelCase_    :     int            =	FSMTConfig.from_pretrained(lowercase   )
     UpperCAmelCase_    :     Optional[int]            =	FSMTForConditionalGeneration(lowercase   )

     # check that it loads ok
     model_new.load_state_dict(lowercase    ,			strict=lowercase   )

     # save
     UpperCAmelCase_    :     Union[str, Any]            =	os.path.join(lowercase    ,			lowercase   )
     print(f'''Generating {pytorch_weights_dump_path}'''   )
     torch.save(lowercase    ,			lowercase   )

     print("""Conversion is done!"""   )
     print("""\nLast step is to upload the files to s3"""   )
     print(f'''cd {data_root}'''   )
     print(f'''transformers-cli upload {model_dir}'''   )


if __name__ == "__main__":
       lowercase_         =       argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           "--fsmt_checkpoint_path",
           default=None,
           type=str,
           required=True,
           help=(
               "Path to the official PyTorch checkpoint file which is expected to reside in the dump dir with dicts,"
               " bpecodes, etc."
           ),
       )
       parser.add_argument(
           "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
       )
       lowercase_         =       parser.parse_args()
       convert_fsmt_checkpoint_to_pytorch(args.fsmt_checkpoint_path, args.pytorch_dump_folder_path)



"""simple docstring"""







import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
       import jax
       import jax.numpy as jnp

       from transformers.generation import (
           FlaxForcedBOSTokenLogitsProcessor,
           FlaxForcedEOSTokenLogitsProcessor,
           FlaxLogitsProcessorList,
           FlaxMinLengthLogitsProcessor,
           FlaxTemperatureLogitsWarper,
           FlaxTopKLogitsWarper,
           FlaxTopPLogitsWarper,
       )



@require_flax
class 		UpperCAmelCase_						(unittest.TestCase    ):
 """simple docstring"""




 def 	a   (			self							:			str  , a_							:			int  , a_							:			int			)->       Optional[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     List[Any]            =	jnp.ones((batch_size, length)			) / length
      return scores




 def 	a   (			self							:			str			)->       int:



      """simple docstring"""

      UpperCAmelCase_    :     List[Any]            =	None
      UpperCAmelCase_    :     int            =	20

      UpperCAmelCase_    :     Tuple            =	self._get_uniform_logits(batch_size=2  , length=a_			)

      # tweak scores to not be uniform anymore
      UpperCAmelCase_    :     Any            =	scores.at[1, 5].set((1 / length) + 0.1			)  # peak, 1st batch
      UpperCAmelCase_    :     List[Any]            =	scores.at[1, 10].set((1 / length) - 0.4			)  # valley, 1st batch

      # compute softmax
      UpperCAmelCase_    :     str            =	jax.nn.softmax(a_  , axis=-1			)

      UpperCAmelCase_    :     Union[str, Any]            =	FlaxTemperatureLogitsWarper(temperature=0.5			)
      UpperCAmelCase_    :     Any            =	FlaxTemperatureLogitsWarper(temperature=1.3			)

      UpperCAmelCase_    :     Optional[Any]            =	jax.nn.softmax(temp_dist_warper_sharper(a_  , scores.copy()  , cur_len=a_			)  , axis=-1			)
      UpperCAmelCase_    :     List[str]            =	jax.nn.softmax(temp_dist_warper_smoother(a_  , scores.copy()  , cur_len=a_			)  , axis=-1			)

      # uniform distribution stays uniform
      self.assertTrue(jnp.allclose(probs[0, :]  , warped_prob_sharp[0, :]  , atol=1E-3			)			)
      self.assertTrue(jnp.allclose(probs[0, :]  , warped_prob_smooth[0, :]  , atol=1E-3			)			)

      # sharp peaks get higher, valleys get lower
      self.assertLess(probs[1, :].max()  , warped_prob_sharp[1, :].max()			)
      self.assertGreater(probs[1, :].min()  , warped_prob_sharp[1, :].min()			)

      # smooth peaks get lower, valleys get higher
      self.assertGreater(probs[1, :].max()  , warped_prob_smooth[1, :].max()			)
      self.assertLess(probs[1, :].min()  , warped_prob_smooth[1, :].min()			)




 def 	a   (			self							:			Union[str, Any]			)->       Tuple:



      """simple docstring"""

      UpperCAmelCase_    :     Dict            =	None
      UpperCAmelCase_    :     str            =	10
      UpperCAmelCase_    :     List[Any]            =	2

      # create ramp distribution
      UpperCAmelCase_    :     Optional[Any]            =	np.broadcast_to(np.arange(a_			)[None, :]  , (batch_size, vocab_size)			).copy()
      UpperCAmelCase_    :     Any            =	ramp_logits[1:, : vocab_size // 2] + vocab_size

      UpperCAmelCase_    :     List[Any]            =	FlaxTopKLogitsWarper(3			)

      UpperCAmelCase_    :     Tuple            =	top_k_warp(a_  , a_  , cur_len=a_			)

      # check that correct tokens are filtered
      self.assertListEqual(jnp.isinf(scores[0]			).tolist()  , 7 * [True] + 3 * [False]			)
      self.assertListEqual(jnp.isinf(scores[1]			).tolist()  , 2 * [True] + 3 * [False] + 5 * [True]			)

      # check special case
      UpperCAmelCase_    :     Any            =	5
      UpperCAmelCase_    :     Optional[int]            =	FlaxTopKLogitsWarper(top_k=1  , filter_value=0.0  , min_tokens_to_keep=3			)

      UpperCAmelCase_    :     Any            =	np.broadcast_to(np.arange(a_			)[None, :]  , (batch_size, length)			).copy()
      UpperCAmelCase_    :     Tuple            =	top_k_warp_safety_check(a_  , a_  , cur_len=a_			)

      # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
      self.assertListEqual((scores == 0.0).sum(axis=-1			).tolist()  , [2, 2]			)




 def 	a   (			self							:			Dict			)->       List[str]:



      """simple docstring"""

      UpperCAmelCase_    :     Tuple            =	None
      UpperCAmelCase_    :     List[str]            =	10
      UpperCAmelCase_    :     List[Any]            =	2

      # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
      UpperCAmelCase_    :     int            =	np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]			)			)

      UpperCAmelCase_    :     str            =	FlaxTopPLogitsWarper(0.8			)
      UpperCAmelCase_    :     Union[str, Any]            =	np.exp(top_p_warp(a_  , a_  , cur_len=a_			)			)

      # dist should be filtered to keep min num values so that sum is >= top_p
      # exp (-inf) => 0
      UpperCAmelCase_    :     List[Any]            =	np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]			)
      self.assertTrue(np.allclose(a_  , a_  , atol=1E-3			)			)

      # check edge cases with negative and extreme logits
      UpperCAmelCase_    :     Optional[Any]            =	np.broadcast_to(np.arange(a_			)[None, :]  , (batch_size, vocab_size)			).copy() - (
          vocab_size // 2
      )

      # make ramp_logits more extreme
      UpperCAmelCase_    :     Tuple            =	ramp_logits[1] * 100.0

      # make sure at least 2 tokens are kept
      UpperCAmelCase_    :     Any            =	FlaxTopPLogitsWarper(0.9  , min_tokens_to_keep=2  , filter_value=0.0			)
      UpperCAmelCase_    :     Dict            =	top_p_warp(a_  , a_  , cur_len=a_			)

      # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
      self.assertListEqual((filtered_dist != 0.0).sum(axis=-1			).tolist()  , [3, 2]			)




 def 	a   (			self							:			List[str]			)->       Any:



      """simple docstring"""

      UpperCAmelCase_    :     Dict            =	20
      UpperCAmelCase_    :     List[str]            =	4
      UpperCAmelCase_    :     Dict            =	0

      UpperCAmelCase_    :     Union[str, Any]            =	FlaxMinLengthLogitsProcessor(min_length=10  , eos_token_id=a_			)

      # check that min length is applied at length 5
      UpperCAmelCase_    :     Union[str, Any]            =	ids_tensor((batch_size, 20)  , vocab_size=20			)
      UpperCAmelCase_    :     str            =	5
      UpperCAmelCase_    :     Tuple            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     List[Any]            =	min_dist_processor(a_  , a_  , cur_len=a_			)
      self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()  , 4 * [-float("""inf"""			)]			)

      # check that min length is not applied anymore at length 15
      UpperCAmelCase_    :     str            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     Optional[Any]            =	15
      UpperCAmelCase_    :     str            =	min_dist_processor(a_  , a_  , cur_len=a_			)
      self.assertFalse(jnp.isinf(a_			).any()			)




 def 	a   (			self							:			List[str]			)->       str:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	20
      UpperCAmelCase_    :     int            =	4
      UpperCAmelCase_    :     str            =	0

      UpperCAmelCase_    :     Optional[int]            =	FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a_			)

      # check that all scores are -inf except the bos_token_id score
      UpperCAmelCase_    :     List[str]            =	ids_tensor((batch_size, 1)  , vocab_size=20			)
      UpperCAmelCase_    :     List[Any]            =	1
      UpperCAmelCase_    :     int            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     int            =	logits_processor(a_  , a_  , cur_len=a_			)
      self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]			).all()			)
      self.assertListEqual(scores[:, bos_token_id].tolist()  , 4 * [0]			)  # score for bos_token_id shold be zero

      # check that bos_token_id is not forced if current length is greater than 1
      UpperCAmelCase_    :     int            =	3
      UpperCAmelCase_    :     Tuple            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     Union[str, Any]            =	logits_processor(a_  , a_  , cur_len=a_			)
      self.assertFalse(jnp.isinf(a_			).any()			)




 def 	a   (			self							:			str			)->       List[str]:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	20
      UpperCAmelCase_    :     int            =	4
      UpperCAmelCase_    :     Union[str, Any]            =	0
      UpperCAmelCase_    :     Union[str, Any]            =	5

      UpperCAmelCase_    :     Optional[Any]            =	FlaxForcedEOSTokenLogitsProcessor(max_length=a_  , eos_token_id=a_			)

      # check that all scores are -inf except the eos_token_id when max_length is reached
      UpperCAmelCase_    :     str            =	ids_tensor((batch_size, 4)  , vocab_size=20			)
      UpperCAmelCase_    :     int            =	4
      UpperCAmelCase_    :     Optional[Any]            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     Union[str, Any]            =	logits_processor(a_  , a_  , cur_len=a_			)
      self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]			).all()			)
      self.assertListEqual(scores[:, eos_token_id].tolist()  , 4 * [0]			)  # score for eos_token_id should be zero

      # check that eos_token_id is not forced if max_length is not reached
      UpperCAmelCase_    :     Any            =	3
      UpperCAmelCase_    :     Tuple            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     int            =	logits_processor(a_  , a_  , cur_len=a_			)
      self.assertFalse(jnp.isinf(a_			).any()			)




 def 	a   (			self							:			int			)->       List[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Dict            =	4
      UpperCAmelCase_    :     List[Any]            =	10
      UpperCAmelCase_    :     Tuple            =	15
      UpperCAmelCase_    :     Optional[Any]            =	2
      UpperCAmelCase_    :     List[Any]            =	1
      UpperCAmelCase_    :     Dict            =	15

      # dummy input_ids and scores
      UpperCAmelCase_    :     Dict            =	ids_tensor((batch_size, sequence_length)  , a_			)
      UpperCAmelCase_    :     Tuple            =	input_ids.copy()

      UpperCAmelCase_    :     int            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     Optional[Any]            =	scores.copy()

      # instantiate all dist processors
      UpperCAmelCase_    :     Tuple            =	FlaxTemperatureLogitsWarper(temperature=0.5			)
      UpperCAmelCase_    :     str            =	FlaxTopKLogitsWarper(3			)
      UpperCAmelCase_    :     Tuple            =	FlaxTopPLogitsWarper(0.8			)

      # instantiate all logits processors
      UpperCAmelCase_    :     Tuple            =	FlaxMinLengthLogitsProcessor(min_length=10  , eos_token_id=a_			)
      UpperCAmelCase_    :     Optional[Any]            =	FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a_			)
      UpperCAmelCase_    :     Tuple            =	FlaxForcedEOSTokenLogitsProcessor(max_length=a_  , eos_token_id=a_			)

      UpperCAmelCase_    :     Tuple            =	10

      # no processor list
      UpperCAmelCase_    :     Optional[Any]            =	temp_dist_warp(a_  , a_  , cur_len=a_			)
      UpperCAmelCase_    :     Optional[int]            =	top_k_warp(a_  , a_  , cur_len=a_			)
      UpperCAmelCase_    :     Optional[Any]            =	top_p_warp(a_  , a_  , cur_len=a_			)
      UpperCAmelCase_    :     List[Any]            =	min_dist_proc(a_  , a_  , cur_len=a_			)
      UpperCAmelCase_    :     Dict            =	bos_dist_proc(a_  , a_  , cur_len=a_			)
      UpperCAmelCase_    :     str            =	eos_dist_proc(a_  , a_  , cur_len=a_			)

      # with processor list
      UpperCAmelCase_    :     Tuple            =	FlaxLogitsProcessorList(
          [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]			)
      UpperCAmelCase_    :     Union[str, Any]            =	processor(a_  , a_  , cur_len=a_			)

      # scores should be equal
      self.assertTrue(jnp.allclose(a_  , a_  , atol=1E-3			)			)

      # input_ids should never be changed
      self.assertListEqual(input_ids.tolist()  , input_ids_comp.tolist()			)






 def 	a   (			self							:			Optional[int]			)->       Any:



      """simple docstring"""

      UpperCAmelCase_    :     Any            =	4
      UpperCAmelCase_    :     List[Any]            =	10
      UpperCAmelCase_    :     Optional[int]            =	15
      UpperCAmelCase_    :     Optional[Any]            =	2
      UpperCAmelCase_    :     Dict            =	1
      UpperCAmelCase_    :     int            =	15

      # dummy input_ids and scores
      UpperCAmelCase_    :     List[str]            =	ids_tensor((batch_size, sequence_length)  , a_			)
      UpperCAmelCase_    :     Dict            =	input_ids.copy()

      UpperCAmelCase_    :     Dict            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     Tuple            =	scores.copy()

      # instantiate all dist processors
      UpperCAmelCase_    :     Dict            =	FlaxTemperatureLogitsWarper(temperature=0.5			)
      UpperCAmelCase_    :     int            =	FlaxTopKLogitsWarper(3			)
      UpperCAmelCase_    :     int            =	FlaxTopPLogitsWarper(0.8			)

      # instantiate all logits processors
      UpperCAmelCase_    :     str            =	FlaxMinLengthLogitsProcessor(min_length=10  , eos_token_id=a_			)
      UpperCAmelCase_    :     Dict            =	FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a_			)
      UpperCAmelCase_    :     Union[str, Any]            =	FlaxForcedEOSTokenLogitsProcessor(max_length=a_  , eos_token_id=a_			)

      UpperCAmelCase_    :     Dict            =	10

      # no processor list
      def run_no_processor_list(a_							:			Any  , a_							:			List[str]  , a_							:			Tuple			):
           UpperCAmelCase_    :     List[Any]            =	temp_dist_warp(a_  , a_  , cur_len=a_			)
           UpperCAmelCase_    :     Tuple            =	top_k_warp(a_  , a_  , cur_len=a_			)
           UpperCAmelCase_    :     int            =	top_p_warp(a_  , a_  , cur_len=a_			)
           UpperCAmelCase_    :     Tuple            =	min_dist_proc(a_  , a_  , cur_len=a_			)
           UpperCAmelCase_    :     Optional[int]            =	bos_dist_proc(a_  , a_  , cur_len=a_			)
           UpperCAmelCase_    :     List[Any]            =	eos_dist_proc(a_  , a_  , cur_len=a_			)
           return scores

      # with processor list
      def run_processor_list(a_							:			List[str]  , a_							:			Optional[int]  , a_							:			Optional[int]			):
           UpperCAmelCase_    :     List[str]            =	FlaxLogitsProcessorList(
               [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]			)
           UpperCAmelCase_    :     Optional[Any]            =	processor(a_  , a_  , cur_len=a_			)
           return scores

      UpperCAmelCase_    :     str            =	jax.jit(a_			)
      UpperCAmelCase_    :     Tuple            =	jax.jit(a_			)

      UpperCAmelCase_    :     Optional[int]            =	jitted_run_no_processor_list(a_  , a_  , a_			)
      UpperCAmelCase_    :     List[Any]            =	jitted_run_processor_list(a_  , a_  , a_			)

      # scores should be equal
      self.assertTrue(jnp.allclose(a_  , a_  , atol=1E-3			)			)

      # input_ids should never be changed
      self.assertListEqual(input_ids.tolist()  , input_ids_comp.tolist()			)

