

import os
from shutil import copyfile
from typing import List, Optional, Tuple

from tokenizers import processors

from ...tokenization_utils import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
     from .tokenization_nllb import NllbTokenizer
else:
     _lowerCamelCase    :       List[Any]        =							None


_lowerCamelCase    :       str        =							logging.get_logger(__name__)


_lowerCamelCase    :       str        =							{"vocab_file": "sentencepiece.bpe.model", "tokenizer_file": "tokenizer.json"}

_lowerCamelCase    :       str        =							{
    "vocab_file": {
        "facebook/nllb-200-distilled-600M": (
            "https://huggingface.co/facebook/nllb-200-distilled-600M/resolve/main/sentencepiece.bpe.model"
        ),
    },
    "tokenizer_file": {
        "facebook/nllb-200-distilled-600M": (
            "https://huggingface.co/facebook/nllb-200-distilled-600M/resolve/main/tokenizer.json"
        ),
    },
}

_lowerCamelCase    :       Any        =							{
    "facebook/nllb-large-en-ro": 1_0_2_4,
    "facebook/nllb-200-distilled-600M": 1_0_2_4,
}

# fmt: off
_lowerCamelCase    :       int        =							["ace_Arab", "ace_Latn", "acm_Arab", "acq_Arab", "aeb_Arab", "afr_Latn", "ajp_Arab", "aka_Latn", "amh_Ethi", "apc_Arab", "arb_Arab", "ars_Arab", "ary_Arab", "arz_Arab", "asm_Beng", "ast_Latn", "awa_Deva", "ayr_Latn", "azb_Arab", "azj_Latn", "bak_Cyrl", "bam_Latn", "ban_Latn", "bel_Cyrl", "bem_Latn", "ben_Beng", "bho_Deva", "bjn_Arab", "bjn_Latn", "bod_Tibt", "bos_Latn", "bug_Latn", "bul_Cyrl", "cat_Latn", "ceb_Latn", "ces_Latn", "cjk_Latn", "ckb_Arab", "crh_Latn", "cym_Latn", "dan_Latn", "deu_Latn", "dik_Latn", "dyu_Latn", "dzo_Tibt", "ell_Grek", "eng_Latn", "epo_Latn", "est_Latn", "eus_Latn", "ewe_Latn", "fao_Latn", "pes_Arab", "fij_Latn", "fin_Latn", "fon_Latn", "fra_Latn", "fur_Latn", "fuv_Latn", "gla_Latn", "gle_Latn", "glg_Latn", "grn_Latn", "guj_Gujr", "hat_Latn", "hau_Latn", "heb_Hebr", "hin_Deva", "hne_Deva", "hrv_Latn", "hun_Latn", "hye_Armn", "ibo_Latn", "ilo_Latn", "ind_Latn", "isl_Latn", "ita_Latn", "jav_Latn", "jpn_Jpan", "kab_Latn", "kac_Latn", "kam_Latn", "kan_Knda", "kas_Arab", "kas_Deva", "kat_Geor", "knc_Arab", "knc_Latn", "kaz_Cyrl", "kbp_Latn", "kea_Latn", "khm_Khmr", "kik_Latn", "kin_Latn", "kir_Cyrl", "kmb_Latn", "kon_Latn", "kor_Hang", "kmr_Latn", "lao_Laoo", "lvs_Latn", "lij_Latn", "lim_Latn", "lin_Latn", "lit_Latn", "lmo_Latn", "ltg_Latn", "ltz_Latn", "lua_Latn", "lug_Latn", "luo_Latn", "lus_Latn", "mag_Deva", "mai_Deva", "mal_Mlym", "mar_Deva", "min_Latn", "mkd_Cyrl", "plt_Latn", "mlt_Latn", "mni_Beng", "khk_Cyrl", "mos_Latn", "mri_Latn", "zsm_Latn", "mya_Mymr", "nld_Latn", "nno_Latn", "nob_Latn", "npi_Deva", "nso_Latn", "nus_Latn", "nya_Latn", "oci_Latn", "gaz_Latn", "ory_Orya", "pag_Latn", "pan_Guru", "pap_Latn", "pol_Latn", "por_Latn", "prs_Arab", "pbt_Arab", "quy_Latn", "ron_Latn", "run_Latn", "rus_Cyrl", "sag_Latn", "san_Deva", "sat_Beng", "scn_Latn", "shn_Mymr", "sin_Sinh", "slk_Latn", "slv_Latn", "smo_Latn", "sna_Latn", "snd_Arab", "som_Latn", "sot_Latn", "spa_Latn", "als_Latn", "srd_Latn", "srp_Cyrl", "ssw_Latn", "sun_Latn", "swe_Latn", "swh_Latn", "szl_Latn", "tam_Taml", "tat_Cyrl", "tel_Telu", "tgk_Cyrl", "tgl_Latn", "tha_Thai", "tir_Ethi", "taq_Latn", "taq_Tfng", "tpi_Latn", "tsn_Latn", "tso_Latn", "tuk_Latn", "tum_Latn", "tur_Latn", "twi_Latn", "tzm_Tfng", "uig_Arab", "ukr_Cyrl", "umb_Latn", "urd_Arab", "uzn_Latn", "vec_Latn", "vie_Latn", "war_Latn", "wol_Latn", "xho_Latn", "ydd_Hebr", "yor_Latn", "yue_Hant", "zho_Hans", "zho_Hant", "zul_Latn"]





class  __UpperCAmelCase					(	lowerCamelCase__						):
   UpperCamelCase       					=					VOCAB_FILES_NAMES
   UpperCamelCase       					=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   UpperCamelCase       					=					PRETRAINED_VOCAB_FILES_MAP
   UpperCamelCase       					=					["""input_ids""", """attention_mask"""]
   UpperCamelCase       					=					NllbTokenizer

   UpperCamelCase       					=					[]
   UpperCamelCase       					=					[]



   def __init__(							self							:					Optional[Any],		__A							:					Tuple=None,		__A							:					int=None,		__A							:					List[Any]="<s>",		__A							:					Tuple="</s>",		__A							:					Any="</s>",		__A							:					Optional[Any]="<s>",		__A							:					Tuple="<unk>",		__A							:					str="<pad>",		__A							:					Dict="<mask>",		__A							:					Optional[Any]=None,		__A							:					List[Any]=None,		__A							:					List[Any]=None,		__A							:					str=False,		**__A							:					Tuple,		):
     # Mask token behave like a normal word, i.e. include the space before it
     UpperCAmelCase	:							int						=			AddedToken(__A,		lstrip=__A,		rstrip=__A						) if isinstance(__A,		__A						) else mask_token
     UpperCAmelCase	:							str						=			legacy_behaviour
     super().__init__(
         vocab_file=__A,		tokenizer_file=__A,		bos_token=__A,		eos_token=__A,		sep_token=__A,		cls_token=__A,		unk_token=__A,		pad_token=__A,		mask_token=__A,		src_lang=__A,		tgt_lang=__A,		additional_special_tokens=__A,		legacy_behaviour=__A,		**__A,		)

     UpperCAmelCase	:							Optional[int]						=			vocab_file
     UpperCAmelCase	:							Optional[Any]						=			False if not self.vocab_file else True

     UpperCAmelCase	:							Optional[int]						=			FAIRSEQ_LANGUAGE_CODES.copy()

     if additional_special_tokens is not None:
       # Only add those special tokens if they are not already there.
       _additional_special_tokens.extend(
           [t for t in additional_special_tokens if t not in _additional_special_tokens]						)

     self.add_special_tokens({'''additional_special_tokens''': _additional_special_tokens}						)
     UpperCAmelCase	:							List[Any]						=			{
         lang_code: self.convert_tokens_to_ids(__A						) for lang_code in FAIRSEQ_LANGUAGE_CODES
     }

     UpperCAmelCase	:							List[Any]						=			src_lang if src_lang is not None else '''eng_Latn'''
     UpperCAmelCase	:							List[Any]						=			self.convert_tokens_to_ids(self._src_lang						)
     UpperCAmelCase	:							List[str]						=			tgt_lang
     self.set_src_lang_special_tokens(self._src_lang						)



   @property
   def 	__magic_name__   (							self							:					Optional[int]						):
     return self._src_lang



   @src_lang.setter
   def 	__magic_name__   (							self							:					Union[str, Any],		__A							:					str						):
     UpperCAmelCase	:							Tuple						=			new_src_lang
     self.set_src_lang_special_tokens(self._src_lang						)



   def 	__magic_name__   (							self							:					Any,		__A							:					List[int],		__A							:					Optional[List[int]] = None						):
     if token_ids_a is None:
       return self.prefix_tokens + token_ids_a + self.suffix_tokens
     # We don't expect to process pairs, but leave the pair logic for API consistency
     return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens



   def 	__magic_name__   (							self							:					Tuple,		__A							:					List[int],		__A							:					Optional[List[int]] = None						):
     UpperCAmelCase	:							Union[str, Any]						=			[self.sep_token_id]
     UpperCAmelCase	:							Any						=			[self.cls_token_id]

     if token_ids_a is None:
       return len(cls + token_ids_a + sep						) * [0]
     return len(cls + token_ids_a + sep + sep + token_ids_a + sep						) * [0]



   def 	__magic_name__   (							self							:					str,		__A							:					Optional[int],		__A							:					str,		__A							:					Optional[str],		__A							:					Optional[str],		**__A							:					Union[str, Any]						):
     if src_lang is None or tgt_lang is None:
       raise ValueError('''Translation requires a `src_lang` and a `tgt_lang` for this model'''						)
     UpperCAmelCase	:							Optional[Any]						=			src_lang
     UpperCAmelCase	:							Optional[Any]						=			self(__A,		add_special_tokens=__A,		return_tensors=__A,		**__A						)
     UpperCAmelCase	:							Union[str, Any]						=			self.convert_tokens_to_ids(__A						)
     UpperCAmelCase	:							int						=			tgt_lang_id
     return inputs



   def 	__magic_name__   (							self							:					List[str],		__A							:					List[str],		__A							:					str = "eng_Latn",		__A							:					Optional[List[str]] = None,		__A							:					str = "fra_Latn",		**__A							:					Tuple,		):
     UpperCAmelCase	:							Any						=			src_lang
     UpperCAmelCase	:							Tuple						=			tgt_lang
     return super().prepare_seqaseq_batch(__A,		__A,		**__A						)



   def 	__magic_name__   (							self							:					List[str]						):
     return self.set_src_lang_special_tokens(self.src_lang						)



   def 	__magic_name__   (							self							:					Optional[Any]						):
     return self.set_tgt_lang_special_tokens(self.tgt_lang						)



   def 	__magic_name__   (							self							:					Union[str, Any],		__A							:					List[str]						):
     UpperCAmelCase	:							int						=			self.convert_tokens_to_ids(__A						)

     if self.legacy_behaviour:
       UpperCAmelCase	:							Tuple						=			[]
       UpperCAmelCase	:							Optional[int]						=			[self.eos_token_id, self.cur_lang_code]
     else:
       UpperCAmelCase	:							Any						=			[self.cur_lang_code]
       UpperCAmelCase	:							Optional[Any]						=			[self.eos_token_id]

     UpperCAmelCase	:							Dict						=			self.convert_ids_to_tokens(self.prefix_tokens						)
     UpperCAmelCase	:							Union[str, Any]						=			self.convert_ids_to_tokens(self.suffix_tokens						)

     UpperCAmelCase	:							int						=			processors.TemplateProcessing(
         single=prefix_tokens_str + ['''$A'''] + suffix_tokens_str,		pair=prefix_tokens_str + ['''$A''', '''$B'''] + suffix_tokens_str,		special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str,		self.prefix_tokens + self.suffix_tokens						)						),		)



   def 	__magic_name__   (							self							:					Dict,		__A							:					str						):
     UpperCAmelCase	:							Optional[int]						=			self.convert_tokens_to_ids(__A						)
     if self.legacy_behaviour:
       UpperCAmelCase	:							List[Any]						=			[]
       UpperCAmelCase	:							Any						=			[self.eos_token_id, self.cur_lang_code]
     else:
       UpperCAmelCase	:							int						=			[self.cur_lang_code]
       UpperCAmelCase	:							List[Any]						=			[self.eos_token_id]

     UpperCAmelCase	:							str						=			self.convert_ids_to_tokens(self.prefix_tokens						)
     UpperCAmelCase	:							List[Any]						=			self.convert_ids_to_tokens(self.suffix_tokens						)

     UpperCAmelCase	:							List[str]						=			processors.TemplateProcessing(
         single=prefix_tokens_str + ['''$A'''] + suffix_tokens_str,		pair=prefix_tokens_str + ['''$A''', '''$B'''] + suffix_tokens_str,		special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str,		self.prefix_tokens + self.suffix_tokens						)						),		)





   def 	__magic_name__   (							self							:					Dict,		__A							:					str,		__A							:					Optional[str] = None						):
     if not self.can_save_slow_tokenizer:
       raise ValueError(
           '''Your fast tokenizer does not have the necessary information to save the vocabulary for a slow '''
           '''tokenizer.'''						)

     if not os.path.isdir(__A						):
       logger.error(F'''Vocabulary path ({save_directory}) should be a directory.'''						)
       return
     UpperCAmelCase	:							Tuple						=			os.path.join(
         __A,		(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']						)

     if os.path.abspath(self.vocab_file						) != os.path.abspath(__A						):
       copyfile(self.vocab_file,		__A						)

     return (out_vocab_file,)








import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
     import os

     import jax.numpy as jnp
     from jax import jit

     from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
     from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

     _lowerCamelCase    :       Any        =							"0.12"  # assumed parallelism: 8


if is_torch_available():
     import torch

def       a__					(		UpperCAmelCase							:      int    ,			UpperCAmelCase							:      List[str]    ,			UpperCAmelCase							:      List[str]=None      )  ->						List[Any]:
  if rng is None:
    UpperCAmelCase	:							Dict						=			random.Random()

  UpperCAmelCase	:							Optional[Any]						=			1
  for dim in shape:
    total_dims *= dim

  UpperCAmelCase	:							List[str]						=			[]
  for _ in range(UpperCAmelCase      ):
    values.append(rng.randint(0    ,			vocab_size - 1      )      )

  UpperCAmelCase	:							List[str]						=			np.array(UpperCAmelCase    ,			dtype=jnp.intaa      ).reshape(UpperCAmelCase      )

  return output

def       a__					(		UpperCAmelCase							:      List[str]    ,			UpperCAmelCase							:      Optional[int]=None      )  ->						List[str]:
  UpperCAmelCase	:							Optional[int]						=			ids_tensor(UpperCAmelCase    ,			vocab_size=2    ,			rng=UpperCAmelCase      )
  # make sure that at least one token is attended to for each batch
  UpperCAmelCase	:							str						=			1
  return attn_mask





@require_flax
class  __UpperCAmelCase					:
   UpperCamelCase       					=					None
   UpperCamelCase       					=					()



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase			,						UpperCAmelCase	:							Tuple						=			self.model_tester.prepare_config_and_inputs_for_common()

     # cut to half length & take max batch_size 3
     UpperCAmelCase	:							Optional[int]						=			2
     UpperCAmelCase	:							Dict						=			inputs['''input_ids'''].shape[-1] // 2
     UpperCAmelCase	:							Dict						=			inputs['''input_ids'''][:max_batch_size, :sequence_length]

     UpperCAmelCase	:							Optional[int]						=			jnp.ones_like(__A						)
     UpperCAmelCase	:							Optional[int]						=			attention_mask[:max_batch_size, :sequence_length]

     # generate max 5 tokens
     UpperCAmelCase	:							Optional[Any]						=			input_ids.shape[-1] + 5
     if config.eos_token_id is not None and config.pad_token_id is None:
       # hack to allow generate for models such as GPT2 as is done in `generate()`
       UpperCAmelCase	:							Optional[Any]						=			config.eos_token_id
     return config, input_ids, attention_mask, max_length



   @is_pt_flax_cross_test
   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Any						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							Optional[Any]						=			False
     UpperCAmelCase	:							Any						=			max_length
     UpperCAmelCase	:							List[Any]						=			0

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							Union[str, Any]						=			model_class(__A						)

       UpperCAmelCase	:							List[str]						=			model_class.__name__[4:]  # Skip the "Flax" at the beginning
       UpperCAmelCase	:							List[Any]						=			getattr(__A,		__A						)
       UpperCAmelCase	:							Union[str, Any]						=			pt_model_class(__A						).eval()
       UpperCAmelCase	:							Tuple						=			load_flax_weights_in_pytorch_model(__A,		flax_model.params						)

       UpperCAmelCase	:							Dict						=			flax_model.generate(__A						).sequences
       UpperCAmelCase	:							str						=			pt_model.generate(torch.tensor(__A,		dtype=torch.long						)						)

       if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
         UpperCAmelCase	:							Any						=			flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

       self.assertListEqual(pt_generation_outputs.numpy().tolist(),		flax_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					Tuple						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Union[str, Any]						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							str						=			False
     UpperCAmelCase	:							Dict						=			max_length

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							Union[str, Any]						=			model_class(__A						)

       UpperCAmelCase	:							Optional[int]						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							List[Any]						=			jit(model.generate						)
       UpperCAmelCase	:							Optional[Any]						=			jit_generate(__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							List[Any]						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							str						=			True
     UpperCAmelCase	:							Dict						=			max_length

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							Union[str, Any]						=			model_class(__A						)

       UpperCAmelCase	:							Optional[Any]						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							str						=			jit(model.generate						)
       UpperCAmelCase	:							List[Any]						=			jit_generate(__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							int						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							Dict						=			False
     UpperCAmelCase	:							Union[str, Any]						=			max_length
     UpperCAmelCase	:							List[Any]						=			2

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							int						=			model_class(__A						)

       UpperCAmelCase	:							str						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							int						=			jit(model.generate						)
       UpperCAmelCase	:							Union[str, Any]						=			jit_generate(__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					List[str]						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Optional[Any]						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							Any						=			False
     UpperCAmelCase	:							Optional[int]						=			max_length
     UpperCAmelCase	:							Optional[int]						=			2
     UpperCAmelCase	:							str						=			2

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							int						=			model_class(__A						)

       UpperCAmelCase	:							Optional[Any]						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[0],		input_ids.shape[0] * config.num_return_sequences						)



   def 	__magic_name__   (							self							:					Any						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							str						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							str						=			True
     UpperCAmelCase	:							Union[str, Any]						=			max_length
     UpperCAmelCase	:							Union[str, Any]						=			0.8
     UpperCAmelCase	:							str						=			1_0
     UpperCAmelCase	:							Any						=			0.3
     UpperCAmelCase	:							str						=			1
     UpperCAmelCase	:							Union[str, Any]						=			8
     UpperCAmelCase	:							Optional[Any]						=			9

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							int						=			model_class(__A						)

       UpperCAmelCase	:							List[Any]						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							Optional[int]						=			jit(model.generate						)
       UpperCAmelCase	:							Any						=			jit_generate(__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					List[Any]						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Dict						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							Optional[Any]						=			max_length
     UpperCAmelCase	:							Tuple						=			1
     UpperCAmelCase	:							Optional[Any]						=			8
     UpperCAmelCase	:							Optional[int]						=			9

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							int						=			model_class(__A						)

       UpperCAmelCase	:							List[str]						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							Dict						=			jit(model.generate						)
       UpperCAmelCase	:							List[str]						=			jit_generate(__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					Optional[Any]						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Union[str, Any]						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							List[str]						=			max_length
     UpperCAmelCase	:							Union[str, Any]						=			2
     UpperCAmelCase	:							List[Any]						=			1
     UpperCAmelCase	:							List[str]						=			8
     UpperCAmelCase	:							int						=			9

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							Optional[Any]						=			model_class(__A						)

       UpperCAmelCase	:							Union[str, Any]						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							List[str]						=			jit(model.generate						)
       UpperCAmelCase	:							Tuple						=			jit_generate(__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Optional[int]						=			self._get_input_ids_and_config()

     # pad attention mask on the left
     UpperCAmelCase	:							Union[str, Any]						=			attention_mask.at[(0, 0)].set(0						)

     UpperCAmelCase	:							Tuple						=			False
     UpperCAmelCase	:							str						=			max_length

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							Any						=			model_class(__A						)

       UpperCAmelCase	:							Union[str, Any]						=			model.generate(__A,		attention_mask=__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							List[Any]						=			jit(model.generate						)
       UpperCAmelCase	:							Optional[Any]						=			jit_generate(__A,		attention_mask=__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					Optional[int]						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Optional[int]						=			self._get_input_ids_and_config()

     # pad attention mask on the left
     UpperCAmelCase	:							Union[str, Any]						=			attention_mask.at[(0, 0)].set(0						)

     UpperCAmelCase	:							Union[str, Any]						=			True
     UpperCAmelCase	:							Dict						=			max_length

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							Any						=			model_class(__A						)

       UpperCAmelCase	:							Optional[Any]						=			model.generate(__A,		attention_mask=__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							Optional[Any]						=			jit(model.generate						)
       UpperCAmelCase	:							Optional[Any]						=			jit_generate(__A,		attention_mask=__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)





   def 	__magic_name__   (							self							:					Tuple						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Union[str, Any]						=			self._get_input_ids_and_config()

     # pad attention mask on the left
     UpperCAmelCase	:							Dict						=			attention_mask.at[(0, 0)].set(0						)

     UpperCAmelCase	:							Union[str, Any]						=			2
     UpperCAmelCase	:							str						=			max_length

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							str						=			model_class(__A						)

       UpperCAmelCase	:							int						=			model.generate(__A,		attention_mask=__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							Optional[Any]						=			jit(model.generate						)
       UpperCAmelCase	:							Dict						=			jit_generate(__A,		attention_mask=__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)






@require_flax
class  __UpperCAmelCase					(	unittest.TestCase						):



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase	:							Union[str, Any]						=			AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert'''						)
     UpperCAmelCase	:							List[str]						=			FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''						)

     UpperCAmelCase	:							int						=			'''Hello world'''
     UpperCAmelCase	:							Optional[int]						=			tokenizer(__A,		return_tensors='''np'''						).input_ids

     # typos are quickly detected (the correct argument is `do_sample`)
     with self.assertRaisesRegex(__A,		'''do_samples'''						):
       model.generate(__A,		do_samples=__A						)

     # arbitrary arguments that will not be used anywhere are also not accepted
     with self.assertRaisesRegex(__A,		'''foo'''						):
       UpperCAmelCase	:							Any						=			{'''foo''': '''bar'''}
       model.generate(__A,		**__A						)





