



"""simple docstring"""

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    __lowercase     (					UpperCamelCase__):







						"""simple docstring"""







						_A		:							Optional[Any]							       =		42
						_A		:							int							       =		jnp.floataa
						_A		:							int							       =		True







						def 						__UpperCamelCase							(self       ):
											super().setup()
											snake_case_					:		List[Any] 					=					nn.Dense(5   ,		dtype=self.dtype       )







						def __call__(self   ,		*lowercase__   ,		**lowercase__       ):
											snake_case_					:		Any 					=					super().__call__(*__A   ,		**__A       )
											snake_case_					:		Optional[Any] 					=					self.cls(outputs[2]       )
											return outputs[:2] + (cls_out,)







class    __lowercase     (					UpperCamelCase__):







						"""simple docstring"""







						_A		:							List[Any]							       =		FlaxBigBirdForNaturalQuestionsModule







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  List[str]				,    SCREAMING_SNAKE_CASE__ :  List[str]				,    SCREAMING_SNAKE_CASE__ :  Tuple				,    SCREAMING_SNAKE_CASE__ :  List[str]				,    SCREAMING_SNAKE_CASE__ :  Optional[int]				,    SCREAMING_SNAKE_CASE__ :  Optional[Any]			):



					"""simple docstring"""





					def cross_entropy(SCREAMING_SNAKE_CASE__ :  List[Any]				,    SCREAMING_SNAKE_CASE__ :  Tuple				,    SCREAMING_SNAKE_CASE__ :  Union[str, Any]=None			):
										snake_case_					:		Optional[Any] 					=					logits.shape[-1]
										snake_case_					:		str 					=					(labels[..., None] == jnp.arange(UpperCamelCase__			)[None]).astype("""f4"""			)
										snake_case_					:		Optional[Any] 					=					jax.nn.log_softmax(UpperCamelCase__				,    axis=-1			)
										snake_case_					:		Any 					=					-jnp.sum(labels * logits				,    axis=-1			)
										if reduction is not None:
															snake_case_					:		List[Any] 					=					reduction(UpperCamelCase__			)
										return loss

					snake_case_					:		Optional[Any] 					=					partial(UpperCamelCase__				,    reduction=jnp.mean			)
					snake_case_					:		str 					=					cross_entropy(UpperCamelCase__				,    UpperCamelCase__			)
					snake_case_					:		Union[str, Any] 					=					cross_entropy(UpperCamelCase__				,    UpperCamelCase__			)
					snake_case_					:		Union[str, Any] 					=					cross_entropy(UpperCamelCase__				,    UpperCamelCase__			)
					return (start_loss + end_loss + pooled_loss) / 3







@dataclass
class    __lowercase     :







						"""simple docstring"""







						_A		:							int							       =		"""google/bigbird-roberta-base"""
						_A		:							int							       =		3000
						_A		:							Optional[int]							       =		10500

						_A		:							Optional[Any]							       =		128
						_A		:							List[Any]							       =		3

						_A		:							List[str]							       =		1
						_A		:							Any							       =		5

						# tx_args
						_A		:							Optional[Any]							       =		3e-5
						_A		:							Union[str, Any]							       =		0.0
						_A		:							int							       =		20000
						_A		:							List[Any]							       =		0.0095

						_A		:							Union[str, Any]							       =		"""bigbird-roberta-natural-questions"""
						_A		:							Optional[int]							       =		"""training-expt"""
						_A		:							Dict							       =		"""data/nq-training.jsonl"""
						_A		:							Optional[int]							       =		"""data/nq-validation.jsonl"""







						def 						__UpperCamelCase							(self       ):
											os.makedirs(self.base_dir   ,		exist_ok=__A       )
											snake_case_					:		Optional[int] 					=					os.path.join(self.base_dir   ,		self.save_dir       )
											snake_case_					:		str 					=					self.batch_size_per_device * jax.device_count()







@dataclass
class    __lowercase     :







						"""simple docstring"""







						_A		:							Tuple							       =		42
						_A		:							Optional[Any]							       =		4096  # no dynamic padding on TPUs







						def __call__(self   ,		lowercase__       ):
											snake_case_					:		Any 					=					self.collate_fn(__A       )
											snake_case_					:		Any 					=					jax.tree_util.tree_map(__A   ,		__A       )
											return batch







						def 						__UpperCamelCase							(self   ,		lowercase__       ):
											snake_case_					,			snake_case_					:		Tuple 					=					self.fetch_inputs(features["""input_ids"""]       )
											snake_case_					:		Tuple 					=					{
											    """input_ids""": jnp.array(__A   ,		dtype=jnp.intaa       ),
											    """attention_mask""": jnp.array(__A   ,		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 						__UpperCamelCase							(self   ,		lowercase__       ):
											snake_case_					:		Optional[int] 					=					[self._fetch_inputs(__A       ) for ids in input_ids]
											return zip(*__A       )







						def 						__UpperCamelCase							(self   ,		lowercase__       ):
											snake_case_					:		Optional[int] 					=					[1 for _ in range(len(__A       )       )]
											while len(__A       ) < self.max_length:
																input_ids.append(self.pad_id       )
																attention_mask.append(0       )
											return input_ids, attention_mask







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  List[Any]				,    SCREAMING_SNAKE_CASE__ :  Any				,    SCREAMING_SNAKE_CASE__ :  int=None			):



					"""simple docstring"""





					if seed is not None:
										snake_case_					:		Any 					=					dataset.shuffle(seed=UpperCamelCase__			)
					for i in range(len(UpperCamelCase__			) // batch_size			):
										snake_case_					:		int 					=					dataset[i * batch_size : (i + 1) * batch_size]
										yield dict(UpperCamelCase__			)







@partial(jax.pmap				,    axis_name="""batch"""			)
def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Optional[Any]				,    SCREAMING_SNAKE_CASE__ :  int				,    **SCREAMING_SNAKE_CASE__ :  Dict			):



					"""simple docstring"""





					def loss_fn(SCREAMING_SNAKE_CASE__ :  Any			):
										snake_case_					:		Union[str, Any] 					=					model_inputs.pop("""start_labels"""			)
										snake_case_					:		Dict 					=					model_inputs.pop("""end_labels"""			)
										snake_case_					:		int 					=					model_inputs.pop("""pooled_labels"""			)

										snake_case_					:		Union[str, Any] 					=					state.apply_fn(**UpperCamelCase__				,    params=UpperCamelCase__				,    dropout_rng=UpperCamelCase__				,    train=UpperCamelCase__			)
										snake_case_					,			snake_case_					,			snake_case_					:		Dict 					=					outputs

										return state.loss_fn(
										    UpperCamelCase__				,    UpperCamelCase__				,    UpperCamelCase__				,    UpperCamelCase__				,    UpperCamelCase__				,    UpperCamelCase__				,    )

					snake_case_					,			snake_case_					:		List[Any] 					=					jax.random.split(UpperCamelCase__			)
					snake_case_					:		Optional[int] 					=					jax.value_and_grad(UpperCamelCase__			)
					snake_case_					,			snake_case_					:		Any 					=					grad_fn(state.params			)
					snake_case_					:		List[Any] 					=					jax.lax.pmean({"""loss""": loss}				,    axis_name="""batch"""			)
					snake_case_					:		Optional[int] 					=					jax.lax.pmean(UpperCamelCase__				,    """batch"""			)

					snake_case_					:		List[Any] 					=					state.apply_gradients(grads=UpperCamelCase__			)
					return state, metrics, new_drp_rng







@partial(jax.pmap				,    axis_name="""batch"""			)
def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Optional[Any]				,    **SCREAMING_SNAKE_CASE__ :  Optional[Any]			):



					"""simple docstring"""





					snake_case_					:		int 					=					model_inputs.pop("""start_labels"""			)
					snake_case_					:		Optional[Any] 					=					model_inputs.pop("""end_labels"""			)
					snake_case_					:		Tuple 					=					model_inputs.pop("""pooled_labels"""			)

					snake_case_					:		Union[str, Any] 					=					state.apply_fn(**UpperCamelCase__				,    params=state.params				,    train=UpperCamelCase__			)
					snake_case_					,			snake_case_					,			snake_case_					:		List[Any] 					=					outputs

					snake_case_					:		Any 					=					state.loss_fn(UpperCamelCase__				,    UpperCamelCase__				,    UpperCamelCase__				,    UpperCamelCase__				,    UpperCamelCase__				,    UpperCamelCase__			)
					snake_case_					:		int 					=					jax.lax.pmean({"""loss""": loss}				,    axis_name="""batch"""			)
					return metrics







class    __lowercase     (					train_state.TrainState):







						"""simple docstring"""







						_A		:							List[Any]							       =		struct.field(pytree_node=UpperCamelCase__)







@dataclass
class    __lowercase     :







						"""simple docstring"""







						_A		:							Union[str, Any]							       =		42
						_A		:							List[str]							       =		42
						_A		:							Optional[Any]							       =		42
						_A		:							List[Any]							       =		42
						_A		:							Optional[int]							       =		42
						_A		:							int							       =		42
						_A		:							Optional[Any]							       =		None







						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__=None       ):
											snake_case_					:		Optional[int] 					=					model.params
											snake_case_					:		Any 					=					TrainState.create(
											    apply_fn=model.__call__   ,		params=__A   ,		tx=__A   ,		loss_fn=__A   ,		)
											if ckpt_dir is not None:
																snake_case_					,			snake_case_					,			snake_case_					,			snake_case_					,			snake_case_					:		Dict 					=					restore_checkpoint(__A   ,		__A       )
																snake_case_					:		Any 					=					{
																    """lr""": args.lr,
																    """init_lr""": args.init_lr,
																    """warmup_steps""": args.warmup_steps,
																    """num_train_steps""": num_train_steps,
																    """weight_decay""": args.weight_decay,
																}
																snake_case_					,			snake_case_					:		Any 					=					build_tx(**__A       )
																snake_case_					:		Optional[Any] 					=					train_state.TrainState(
																    step=__A   ,		apply_fn=model.__call__   ,		params=__A   ,		tx=__A   ,		opt_state=__A   ,		)
																snake_case_					:		List[str] 					=					args
																snake_case_					:		int 					=					data_collator
																snake_case_					:		Optional[int] 					=					lr
																snake_case_					:		Any 					=					params
											snake_case_					:		Tuple 					=					jax_utils.replicate(__A       )
											return state







						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		str 					=					self.args
											snake_case_					:		Any 					=					len(__A       ) // args.batch_size

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

																										snake_case_					:		List[str] 					=					self.evaluate(__A   ,		__A       )
																										snake_case_					:		Dict 					=					{
																										    """step""": state_step.item(),
																										    """eval_loss""": eval_loss.item(),
																										    """tr_loss""": tr_loss,
																										    """lr""": lr.item(),
																										}
																										tqdm.write(str(__A       )       )
																										self.logger.log(__A   ,		commit=__A       )

																					if i % args.save_steps == 0:
																										self.save_checkpoint(args.save_dir + f'-e{epoch}-s{i}'   ,		state=__A       )







						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Union[str, Any] 					=					get_batched_dataset(__A   ,		self.args.batch_size       )
											snake_case_					:		Tuple 					=					len(__A       ) // self.args.batch_size
											snake_case_					:		Optional[int] 					=					jnp.array(0   ,		dtype=jnp.floataa       )
											snake_case_					:		Tuple 					=					0
											for batch in tqdm(__A   ,		total=__A   ,		desc="""Evaluating ... """       ):
																snake_case_					:		List[Any] 					=					self.data_collator(__A       )
																snake_case_					:		Tuple 					=					self.val_step_fn(__A   ,		**__A       )
																running_loss += jax_utils.unreplicate(metrics["""loss"""]       )
																i += 1
											return running_loss / i







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







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  List[str]				,    SCREAMING_SNAKE_CASE__ :  str			):



					"""simple docstring"""





					print(f'RESTORING CHECKPOINT FROM {save_dir}'				,    end=""" ... """			)
					with open(os.path.join(UpperCamelCase__				,    """flax_model.msgpack"""			)				,    """rb"""			) as f:
										snake_case_					:		Any 					=					from_bytes(state.params				,    f.read()			)

					with open(os.path.join(UpperCamelCase__				,    """opt_state.msgpack"""			)				,    """rb"""			) as f:
										snake_case_					:		List[str] 					=					from_bytes(state.opt_state				,    f.read()			)

					snake_case_					:		Optional[int] 					=					joblib.load(os.path.join(UpperCamelCase__				,    """args.joblib"""			)			)
					snake_case_					:		List[str] 					=					joblib.load(os.path.join(UpperCamelCase__				,    """data_collator.joblib"""			)			)

					with open(os.path.join(UpperCamelCase__				,    """training_state.json"""			)				,    """r"""			) as f:
										snake_case_					:		Dict 					=					json.load(UpperCamelCase__			)
					snake_case_					:		Optional[Any] 					=					training_state["""step"""]

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







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  int				,    SCREAMING_SNAKE_CASE__ :  Optional[Any]				,    SCREAMING_SNAKE_CASE__ :  Optional[int]				,    SCREAMING_SNAKE_CASE__ :  Optional[int]			):



					"""simple docstring"""





					snake_case_					:		Optional[int] 					=					num_train_steps - warmup_steps
					snake_case_					:		List[Any] 					=					optax.linear_schedule(init_value=UpperCamelCase__				,    end_value=UpperCamelCase__				,    transition_steps=UpperCamelCase__			)
					snake_case_					:		Dict 					=					optax.linear_schedule(init_value=UpperCamelCase__				,    end_value=1E-7				,    transition_steps=UpperCamelCase__			)
					snake_case_					:		List[str] 					=					optax.join_schedules(schedules=[warmup_fn, decay_fn]				,    boundaries=[warmup_steps]			)
					return lr







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Optional[Any]				,    SCREAMING_SNAKE_CASE__ :  Optional[int]				,    SCREAMING_SNAKE_CASE__ :  List[str]				,    SCREAMING_SNAKE_CASE__ :  Optional[int]				,    SCREAMING_SNAKE_CASE__ :  Dict			):



					"""simple docstring"""





					def weight_decay_mask(SCREAMING_SNAKE_CASE__ :  Optional[Any]			):
										snake_case_					:		Union[str, Any] 					=					traverse_util.flatten_dict(UpperCamelCase__			)
										snake_case_					:		Union[str, Any] 					=					{k: (v[-1] != """bias""" and v[-2:] != ("""LayerNorm""", """scale""")) for k, v in params.items()}
										return traverse_util.unflatten_dict(UpperCamelCase__			)

					snake_case_					:		Tuple 					=					scheduler_fn(UpperCamelCase__				,    UpperCamelCase__				,    UpperCamelCase__				,    UpperCamelCase__			)

					snake_case_					:		Any 					=					optax.adamw(learning_rate=UpperCamelCase__				,    weight_decay=UpperCamelCase__				,    mask=UpperCamelCase__			)
					return tx, lr









"""simple docstring"""

import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import BatchEncoding, MarianTokenizer
from transformers.testing_utils import get_tests_dir, require_sentencepiece, slow
from transformers.utils import is_sentencepiece_available, is_tf_available, is_torch_available


if is_sentencepiece_available():
			from transformers.models.marian.tokenization_marian import VOCAB_FILES_NAMES, save_json

from ...test_tokenization_common import TokenizerTesterMixin


a_			   =					get_tests_dir('''fixtures/test_sentencepiece.model''')

a_			   =					{'''target_lang''': '''fi''', '''source_lang''': '''en'''}
a_			   =					'''>>zh<<'''
a_			   =					'''Helsinki-NLP/'''

if is_torch_available():
			a_			   =					'''pt'''
elif is_tf_available():
			a_			   =					'''tf'''
else:
			a_			   =					'''jax'''







@require_sentencepiece
class    __lowercase     (					_UpperCAmelCase					,    unittest.TestCase):







						"""simple docstring"""







						_A		:							str							       =		MarianTokenizer
						_A		:							List[str]							       =		False
						_A		:							List[str]							       =		True







						def 						__UpperCamelCase							(self       ):
											super().setUp()
											snake_case_					:		Optional[int] 					=					["""</s>""", """<unk>""", """▁This""", """▁is""", """▁a""", """▁t""", """est""", """\u0120""", """<pad>"""]
											snake_case_					:		Any 					=					dict(zip(lowercase__   ,		range(len(lowercase__       )       )       )       )
											snake_case_					:		Any 					=					Path(self.tmpdirname       )
											save_json(lowercase__   ,		save_dir / VOCAB_FILES_NAMES["""vocab"""]       )
											save_json(lowercase__   ,		save_dir / VOCAB_FILES_NAMES["""tokenizer_config_file"""]       )
											if not (save_dir / VOCAB_FILES_NAMES["source_spm"]).exists():
																copyfile(lowercase__   ,		save_dir / VOCAB_FILES_NAMES["""source_spm"""]       )
																copyfile(lowercase__   ,		save_dir / VOCAB_FILES_NAMES["""target_spm"""]       )

											snake_case_					:		Optional[Any] 					=					MarianTokenizer.from_pretrained(self.tmpdirname       )
											tokenizer.save_pretrained(self.tmpdirname       )







						def 						__UpperCamelCase							(self   ,		**lowercase__       ):
											return MarianTokenizer.from_pretrained(self.tmpdirname   ,		**lowercase__       )







						def 						__UpperCamelCase							(self   ,		lowercase__       ):
											return (
											    "This is a test",
											    "This is a test",
											)







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Union[str, Any] 					=					"""</s>"""
											snake_case_					:		Tuple 					=					0

											self.assertEqual(self.get_tokenizer()._convert_token_to_id(lowercase__       )   ,		lowercase__       )
											self.assertEqual(self.get_tokenizer()._convert_id_to_token(lowercase__       )   ,		lowercase__       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		List[str] 					=					list(self.get_tokenizer().get_vocab().keys()       )

											self.assertEqual(vocab_keys[0]   ,		"""</s>"""       )
											self.assertEqual(vocab_keys[1]   ,		"""<unk>"""       )
											self.assertEqual(vocab_keys[-1]   ,		"""<pad>"""       )
											self.assertEqual(len(lowercase__       )   ,		9       )







						def 						__UpperCamelCase							(self       ):
											self.assertEqual(self.get_tokenizer().vocab_size   ,		9       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Any 					=					MarianTokenizer.from_pretrained(f'{ORG_NAME}opus-mt-en-de'       )
											snake_case_					:		Tuple 					=					en_de_tokenizer(["""I am a small frog"""]   ,		return_tensors=lowercase__       )
											self.assertIsInstance(lowercase__   ,		lowercase__       )
											snake_case_					:		Dict 					=					[38, 1_21, 14, 6_97, 3_88_48, 0]
											self.assertListEqual(lowercase__   ,		batch.input_ids[0]       )

											snake_case_					:		Tuple 					=					tempfile.mkdtemp()
											en_de_tokenizer.save_pretrained(lowercase__       )
											snake_case_					:		str 					=					[x.name for x in Path(lowercase__       ).glob("""*"""       )]
											self.assertIn("""source.spm"""   ,		lowercase__       )
											MarianTokenizer.from_pretrained(lowercase__       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Union[str, Any] 					=					self.get_tokenizer()

											snake_case_					:		List[str] 					=					tok(
											    ["""I am a small frog""" * 10_00, """I am a small frog"""]   ,		padding=lowercase__   ,		truncation=lowercase__   ,		return_tensors=lowercase__       )
											self.assertIsInstance(lowercase__   ,		lowercase__       )
											self.assertEqual(batch.input_ids.shape   ,		(2, 5_12)       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Tuple 					=					self.get_tokenizer()
											snake_case_					:		Tuple 					=					tok(["""I am a tiny frog""", """I am a small frog"""]   ,		padding=lowercase__   ,		return_tensors=lowercase__       )
											self.assertIsInstance(lowercase__   ,		lowercase__       )
											self.assertEqual(batch_smaller.input_ids.shape   ,		(2, 10)       )







						@slow
						def 						__UpperCamelCase							(self       ):
											# fmt: off
											snake_case_					:		str 					=					{"""input_ids""": [[4_34_95, 4_62, 20, 4_21_64, 13_69, 52, 4_64, 1_32, 17_03, 4_92, 13, 74_91, 3_89_99, 6, 8, 4_64, 1_32, 17_03, 4_92, 13, 46_69, 3_78_67, 13, 75_25, 27, 15_93, 9_88, 13, 3_39_72, 70_29, 6, 20, 82_51, 3_83, 2, 2_70, 58_66, 37_88, 2, 23_53, 82_51, 1_23_38, 2, 1_39_58, 3_87, 2, 36_29, 69_53, 1_88, 29_00, 2, 1_39_58, 80_11, 1_15_01, 23, 84_60, 40_73, 3_40_09, 20, 4_35, 1_14_39, 27, 8, 84_60, 40_73, 60_04, 20, 99_88, 3_75, 27, 33, 2_66, 19_45, 10_76, 13_50, 3_78_67, 32_88, 5, 5_77, 10_76, 43_74, 8, 50_82, 5, 2_64_53, 2_57, 5_56, 4_03, 2, 2_42, 1_32, 3_83, 3_16, 4_92, 8, 1_07_67, 6, 3_16, 3_04, 42_39, 3, 0], [1_48, 1_57_22, 19, 18_39, 12, 13_50, 13, 2_23_27, 50_82, 54_18, 4_75_67, 3_59_38, 59, 3_18, 1_95_52, 1_08, 21_83, 54, 1_49_76, 48_35, 32, 5_47, 11_14, 8, 3_15, 24_17, 5, 92, 1_90_88, 3, 0, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00], [36, 63_95, 1_25_70, 3_91_47, 1_15_97, 6, 2_66, 4, 4_54_05, 72_96, 3, 0, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00]], """attention_mask""": [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
											# fmt: on

											self.tokenizer_integration_test_util(
											    expected_encoding=lowercase__   ,		model_name="""Helsinki-NLP/opus-mt-en-de"""   ,		revision="""1a8c2263da11e68e50938f97e10cd57820bd504c"""   ,		decode_kwargs={"""use_source_tokenizer""": True}   ,		)







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Any 					=					MarianTokenizer.from_pretrained("""hf-internal-testing/test-marian-two-vocabs"""       )

											snake_case_					:		Dict 					=					"""Tämä on testi"""
											snake_case_					:		List[Any] 					=					"""This is a test"""

											snake_case_					:		Optional[int] 					=					[76, 7, 20_47, 2]
											snake_case_					:		List[str] 					=					[69, 12, 11, 9_40, 2]

											snake_case_					:		Any 					=					tokenizer(lowercase__       ).input_ids
											self.assertListEqual(lowercase__   ,		lowercase__       )

											snake_case_					:		str 					=					tokenizer(text_target=lowercase__       ).input_ids
											self.assertListEqual(lowercase__   ,		lowercase__       )

											snake_case_					:		int 					=					tokenizer.decode(lowercase__   ,		skip_special_tokens=lowercase__       )
											self.assertEqual(lowercase__   ,		lowercase__       )




