





'''simple docstring'''



import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()


class    __SCREAMING_SNAKE_CASE						(    __SCREAMING_SNAKE_CASE       ,     __SCREAMING_SNAKE_CASE       ,     unittest.TestCase      ):



     '''simple docstring'''






     UpperCAmelCase__							:    str     	= AutoencoderKL
     UpperCAmelCase__							:    Tuple     	= '''sample'''
     UpperCAmelCase__							:    Tuple     	= 1E-2

     @property
     def 						UpperCamelCase(       self						):
       _snake_case  		=					4
       _snake_case  		=					3
       _snake_case  		=					(32, 32)

       _snake_case  		=					floats_tensor((batch_size, num_channels) + sizes						).to(UpperCamelCase__						)

       return {"sample": image}

     @property
     def 						UpperCamelCase(       self						):
       return (3, 32, 32)

     @property
     def 						UpperCamelCase(       self						):
       return (3, 32, 32)

     def 						UpperCamelCase(       self						):
       _snake_case  		=					{
           '''block_out_channels''': [32, 64],
           '''in_channels''': 3,
           '''out_channels''': 3,
           '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
           '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
           '''latent_channels''': 4,
       }
       _snake_case  		=					self.dummy_input
       return init_dict, inputs_dict

     def 						UpperCamelCase(       self						):
       pass

     def 						UpperCamelCase(       self						):
       pass

     @unittest.skipIf(torch_device == "mps"    , "Gradient checkpointing skipped on MPS"						)
     def 						UpperCamelCase(       self						):
       _snake_case  		=					self.prepare_init_args_and_inputs_for_common()
       _snake_case  		=					self.model_class(**UpperCamelCase__						)
       model.to(UpperCamelCase__						)

       assert not model.is_gradient_checkpointing and model.training

       _snake_case  		=					model(**UpperCamelCase__						).sample
       # run the backwards pass on the model. For backwards pass, for simplicity purpose,
       # we won't calculate the loss and rather backprop on out.sum()
       model.zero_grad()

       _snake_case  		=					torch.randn_like(UpperCamelCase__						)
       _snake_case  		=					(out - labels).mean()
       loss.backward()

       # re-instantiate the model now enabling gradient checkpointing
       _snake_case  		=					self.model_class(**UpperCamelCase__						)
       # clone model
       model_a.load_state_dict(model.state_dict()						)
       model_a.to(UpperCamelCase__						)
       model_a.enable_gradient_checkpointing()

       assert model_a.is_gradient_checkpointing and model_a.training

       _snake_case  		=					model_a(**UpperCamelCase__						).sample
       # run the backwards pass on the model. For backwards pass, for simplicity purpose,
       # we won't calculate the loss and rather backprop on out.sum()
       model_a.zero_grad()
       _snake_case  		=					(out_a - labels).mean()
       loss_a.backward()

       # compare the output and parameters gradients
       self.assertTrue((loss - loss_a).abs() < 1e-5						)
       _snake_case  		=					dict(model.named_parameters()						)
       _snake_case  		=					dict(model_a.named_parameters()						)
       for name, param in named_params.items():
         self.assertTrue(torch_all_close(param.grad.data    , named_params_a[name].grad.data    , atol=5e-5						)						)

     def 						UpperCamelCase(       self						):
       _snake_case  		=					AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy"    , output_loading_info=UpperCamelCase__						)
       self.assertIsNotNone(UpperCamelCase__						)
       self.assertEqual(len(loading_info["missing_keys"]						)    , 0						)

       model.to(UpperCamelCase__						)
       _snake_case  		=					model(**self.dummy_input						)

       assert image is not None, "Make sure output is not None"






     def 						UpperCamelCase(       self						):
       _snake_case  		=					AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy"						)
       _snake_case  		=					model.to(UpperCamelCase__						)
       model.eval()

       if torch_device == "mps":
         _snake_case  		=					torch.manual_seed(0						)
       else:
         _snake_case  		=					torch.Generator(device=UpperCamelCase__						).manual_seed(0						)

       _snake_case  		=					torch.randn(
           1    , model.config.in_channels    , model.config.sample_size    , model.config.sample_size    , generator=torch.manual_seed(0						)    , )
       _snake_case  		=					image.to(UpperCamelCase__						)
       with torch.no_grad():
         _snake_case  		=					model(UpperCamelCase__    , sample_posterior=UpperCamelCase__    , generator=UpperCamelCase__						).sample

       _snake_case  		=					output[0, -1, -3:, -3:].flatten().cpu()

       # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
       # the expected output slices are not the same for CPU and GPU.
       if torch_device == "mps":
         _snake_case  		=					torch.tensor(
             [
                 -4.0078e-01,
                 -3.8323e-04,
                 -1.2681e-01,
                 -1.1462e-01,
                 2.0095e-01,
                 1.0893e-01,
                 -8.8247e-02,
                 -3.0361e-01,
                 -9.8644e-03,
             ]						)
       elif torch_device == "cpu":
         _snake_case  		=					torch.tensor(
             [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]						)
       else:
         _snake_case  		=					torch.tensor(
             [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485]						)

       self.assertTrue(torch_all_close(UpperCamelCase__    , UpperCamelCase__    , rtol=1e-2						)						)







@slow
class    __SCREAMING_SNAKE_CASE						(    unittest.TestCase      ):



     '''simple docstring'''






     def 						UpperCamelCase(       self    , lowerCamelCase    , lowerCamelCase						):
       return F'''gaussian_noise_s={seed}_shape={"_".join([str(UpperCamelCase__						) for s in shape]						)}.npy'''

     def 						UpperCamelCase(       self						):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()

     def 						UpperCamelCase(       self    , lowerCamelCase=0    , lowerCamelCase=(4, 3, 512, 512)    , lowerCamelCase=False						):
       _snake_case  		=					torch.floataa if fpaa else torch.floataa
       _snake_case  		=					torch.from_numpy(load_hf_numpy(self.get_file_format(UpperCamelCase__    , UpperCamelCase__						)						)						).to(UpperCamelCase__						).to(UpperCamelCase__						)
       return image

     def 						UpperCamelCase(       self    , lowerCamelCase="CompVis/stable-diffusion-v1-4"    , lowerCamelCase=False						):
       _snake_case  		=					'''fp16''' if fpaa else None
       _snake_case  		=					torch.floataa if fpaa else torch.floataa

       _snake_case  		=					AutoencoderKL.from_pretrained(
           UpperCamelCase__    , subfolder="vae"    , torch_dtype=UpperCamelCase__    , revision=UpperCamelCase__    , )
       model.to(UpperCamelCase__						).eval()

       return model

     def 						UpperCamelCase(       self    , lowerCamelCase=0						):







       if torch_device == "mps":
         return torch.manual_seed(UpperCamelCase__						)
       return torch.Generator(device=UpperCamelCase__						).manual_seed(UpperCamelCase__						)

     @parameterized.expand(
         [
             # fmt: off
             [33, [-0.1603, 0.9878, -0.0495, -0.0790, -0.2709, 0.8375, -0.2060, -0.0824], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
             [47, [-0.2376, 0.1168, 0.1332, -0.4840, -0.2508, -0.0791, -0.0493, -0.4089], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
             # fmt: on
         ]						)
     def 						UpperCamelCase(       self    , lowerCamelCase    , lowerCamelCase    , lowerCamelCase						):
       _snake_case  		=					self.get_sd_vae_model()
       _snake_case  		=					self.get_sd_image(UpperCamelCase__						)
       _snake_case  		=					self.get_generator(UpperCamelCase__						)

       with torch.no_grad():
         _snake_case  		=					model(UpperCamelCase__    , generator=UpperCamelCase__    , sample_posterior=UpperCamelCase__						).sample

       assert sample.shape == image.shape

       _snake_case  		=					sample[-1, -2:, -2:, :2].flatten().float().cpu()
       _snake_case  		=					torch.tensor(expected_slice_mps if torch_device == "mps" else expected_slice						)

       assert torch_all_close(UpperCamelCase__    , UpperCamelCase__    , atol=3e-3						)

     @parameterized.expand(
         [
             # fmt: off
             [33, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
             [47, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
             # fmt: on
         ]						)
     @require_torch_gpu
     def 						UpperCamelCase(       self    , lowerCamelCase    , lowerCamelCase						):
       _snake_case  		=					self.get_sd_vae_model(fpaa=UpperCamelCase__						)
       _snake_case  		=					self.get_sd_image(UpperCamelCase__    , fpaa=UpperCamelCase__						)
       _snake_case  		=					self.get_generator(UpperCamelCase__						)

       with torch.no_grad():
         _snake_case  		=					model(UpperCamelCase__    , generator=UpperCamelCase__    , sample_posterior=UpperCamelCase__						).sample

       assert sample.shape == image.shape

       _snake_case  		=					sample[-1, -2:, :2, -2:].flatten().float().cpu()
       _snake_case  		=					torch.tensor(UpperCamelCase__						)

       assert torch_all_close(UpperCamelCase__    , UpperCamelCase__    , atol=1e-2						)

     @parameterized.expand(
         [
             # fmt: off
             [33, [-0.1609, 0.9866, -0.0487, -0.0777, -0.2716, 0.8368, -0.2055, -0.0814], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
             [47, [-0.2377, 0.1147, 0.1333, -0.4841, -0.2506, -0.0805, -0.0491, -0.4085], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
             # fmt: on
         ]						)
     def 						UpperCamelCase(       self    , lowerCamelCase    , lowerCamelCase    , lowerCamelCase						):
       _snake_case  		=					self.get_sd_vae_model()
       _snake_case  		=					self.get_sd_image(UpperCamelCase__						)

       with torch.no_grad():
         _snake_case  		=					model(UpperCamelCase__						).sample

       assert sample.shape == image.shape

       _snake_case  		=					sample[-1, -2:, -2:, :2].flatten().float().cpu()
       _snake_case  		=					torch.tensor(expected_slice_mps if torch_device == "mps" else expected_slice						)

       assert torch_all_close(UpperCamelCase__    , UpperCamelCase__    , atol=3e-3						)

     @parameterized.expand(
         [
             # fmt: off
             [13, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
             [37, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
             # fmt: on
         ]						)
     @require_torch_gpu
     def 						UpperCamelCase(       self    , lowerCamelCase    , lowerCamelCase						):
       _snake_case  		=					self.get_sd_vae_model()
       _snake_case  		=					self.get_sd_image(UpperCamelCase__    , shape=(3, 4, 64, 64)						)

       with torch.no_grad():
         _snake_case  		=					model.decode(UpperCamelCase__						).sample

       assert list(sample.shape						) == [3, 3, 512, 512]

       _snake_case  		=					sample[-1, -2:, :2, -2:].flatten().cpu()
       _snake_case  		=					torch.tensor(UpperCamelCase__						)

       assert torch_all_close(UpperCamelCase__    , UpperCamelCase__    , atol=1e-3						)

     @parameterized.expand(
         [
             # fmt: off
             [27, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
             [16, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
             # fmt: on
         ]						)
     @require_torch_gpu
     def 						UpperCamelCase(       self    , lowerCamelCase    , lowerCamelCase						):
       _snake_case  		=					self.get_sd_vae_model(fpaa=UpperCamelCase__						)
       _snake_case  		=					self.get_sd_image(UpperCamelCase__    , shape=(3, 4, 64, 64)    , fpaa=UpperCamelCase__						)

       with torch.no_grad():
         _snake_case  		=					model.decode(UpperCamelCase__						).sample

       assert list(sample.shape						) == [3, 3, 512, 512]

       _snake_case  		=					sample[-1, -2:, :2, -2:].flatten().float().cpu()
       _snake_case  		=					torch.tensor(UpperCamelCase__						)

       assert torch_all_close(UpperCamelCase__    , UpperCamelCase__    , atol=5e-3						)

     @parameterized.expand([(13,), (16,), (27,)]						)
     @require_torch_gpu
     @unittest.skipIf(not is_xformers_available()    , reason="xformers is not required when using PyTorch 2.0."						)
     def 						UpperCamelCase(       self    , lowerCamelCase						):
       _snake_case  		=					self.get_sd_vae_model(fpaa=UpperCamelCase__						)
       _snake_case  		=					self.get_sd_image(UpperCamelCase__    , shape=(3, 4, 64, 64)    , fpaa=UpperCamelCase__						)

       with torch.no_grad():
         _snake_case  		=					model.decode(UpperCamelCase__						).sample

       model.enable_xformers_memory_efficient_attention()
       with torch.no_grad():
         _snake_case  		=					model.decode(UpperCamelCase__						).sample

       assert list(sample.shape						) == [3, 3, 512, 512]

       assert torch_all_close(UpperCamelCase__    , UpperCamelCase__    , atol=1e-1						)

     @parameterized.expand([(13,), (16,), (37,)]						)
     @require_torch_gpu
     @unittest.skipIf(not is_xformers_available()    , reason="xformers is not required when using PyTorch 2.0."						)
     def 						UpperCamelCase(       self    , lowerCamelCase						):
       _snake_case  		=					self.get_sd_vae_model()
       _snake_case  		=					self.get_sd_image(UpperCamelCase__    , shape=(3, 4, 64, 64)						)

       with torch.no_grad():
         _snake_case  		=					model.decode(UpperCamelCase__						).sample

       model.enable_xformers_memory_efficient_attention()
       with torch.no_grad():
         _snake_case  		=					model.decode(UpperCamelCase__						).sample

       assert list(sample.shape						) == [3, 3, 512, 512]

       assert torch_all_close(UpperCamelCase__    , UpperCamelCase__    , atol=1e-2						)






     @parameterized.expand(
         [
             # fmt: off
             [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
             [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
             # fmt: on
         ]						)
     def 						UpperCamelCase(       self    , lowerCamelCase    , lowerCamelCase						):
       _snake_case  		=					self.get_sd_vae_model()
       _snake_case  		=					self.get_sd_image(UpperCamelCase__						)
       _snake_case  		=					self.get_generator(UpperCamelCase__						)

       with torch.no_grad():
         _snake_case  		=					model.encode(UpperCamelCase__						).latent_dist
         _snake_case  		=					dist.sample(generator=UpperCamelCase__						)

       assert list(sample.shape						) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

       _snake_case  		=					sample[0, -1, -3:, -3:].flatten().cpu()
       _snake_case  		=					torch.tensor(UpperCamelCase__						)

       _snake_case  		=					3e-3 if torch_device != '''mps''' else 1e-2
       assert torch_all_close(UpperCamelCase__    , UpperCamelCase__    , atol=UpperCamelCase__						)




'''simple docstring'''
import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import MaMaaaTokenizer, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
)
from transformers.utils import is_sentencepiece_available


if is_sentencepiece_available():
						from transformers.models.mam_aaa.tokenization_mam_aaa import VOCAB_FILES_NAMES, save_json

from ...test_tokenization_common import TokenizerTesterMixin


if is_sentencepiece_available():
						lowercase__						       =      get_tests_dir("fixtures/test_sentencepiece.model")


if is_torch_available():
						from transformers.models.mam_aaa.modeling_mam_aaa import shift_tokens_right

lowercase__						       =      1_2_8_0_2_2
lowercase__						       =      1_2_8_0_2_8






@require_sentencepiece
class 	snake_case__       (      __SCREAMING_SNAKE_CASE    ,					unittest.TestCase					):







		"""simple docstring"""

		lowerCamelCase     							=						MaMaaaTokenizer
		lowerCamelCase     							=						False
		lowerCamelCase     							=						False
		lowerCamelCase     							=						True



		def       lowerCAmelCase				(      self				:			Dict      )						->			Optional[Any]:






				"""simple docstring"""







				super().setUp()

				snake_case  :       int     			=						['''</s>''', '''<unk>''', '''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''', '''\u0120''', '''<pad>''']
				snake_case  :       int     			=						dict(zip(UpperCamelCase__  ,	range(len(UpperCamelCase__      )      )      )      )
				snake_case  :       Optional[int]     			=						Path(self.tmpdirname      )
				save_json(UpperCamelCase__  ,	save_dir / VOCAB_FILES_NAMES['''vocab_file''']      )
				if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
						copyfile(UpperCamelCase__  ,	save_dir / VOCAB_FILES_NAMES['''spm_file''']      )

				snake_case  :       Optional[int]     			=						MaMaaaTokenizer.from_pretrained(self.tmpdirname      )
				tokenizer.save_pretrained(self.tmpdirname      )



		def       lowerCAmelCase				(      self				:			Dict  ,	**UpperCamelCase__				:			Optional[int]      )						->			Any:






				"""simple docstring"""







				return MaMaaaTokenizer.from_pretrained(self.tmpdirname  ,	**UpperCamelCase__      )



		def       lowerCAmelCase				(      self				:			Tuple  ,	UpperCamelCase__				:			Tuple      )						->			str:






				"""simple docstring"""







				return (
				    "This is a test",
				    "This is a test",
				)



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






				"""simple docstring"""







				snake_case  :       Optional[int]     			=						'''</s>'''
				snake_case  :       str     			=						0

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



		def       lowerCAmelCase				(      self				:			Dict      )						->			int:






				"""simple docstring"""







				snake_case  :       Tuple     			=						self.get_tokenizer()
				snake_case  :       Dict     			=						list(tokenizer.get_vocab().keys()      )

				self.assertEqual(vocab_keys[0]  ,	'''</s>'''      )
				self.assertEqual(vocab_keys[1]  ,	'''<unk>'''      )
				self.assertEqual(vocab_keys[-1]  ,	'''<s>'''      )
				self.assertEqual(len(UpperCamelCase__      )  ,	tokenizer.vocab_size + len(tokenizer.get_added_vocab()      )      )



		@unittest.skip('''Skip this test while all models are still to be uploaded.'''      )
		def       lowerCAmelCase				(      self				:			str      )						->			List[str]:






				"""simple docstring"""







				pass



		def       lowerCAmelCase				(      self				:			int      )						->			Union[str, Any]:






				"""simple docstring"""







				snake_case  :       Optional[Any]     			=						self.get_tokenizer()

				snake_case  :       Optional[Any]     			=						tokenizer.tokenize('''This is a test'''      )
				self.assertListEqual(UpperCamelCase__  ,	['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''']      )

				self.assertListEqual(
				    tokenizer.convert_tokens_to_ids(UpperCamelCase__      )  ,	[2, 3, 4, 5, 6]  ,	)

				snake_case  :       Any     			=						tokenizer.convert_ids_to_tokens([2, 3, 4, 5, 6]      )
				self.assertListEqual(UpperCamelCase__  ,	['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''']      )

				snake_case  :       Optional[int]     			=						tokenizer.convert_tokens_to_string(UpperCamelCase__      )
				self.assertEqual(UpperCamelCase__  ,	'''This is a test'''      )



		@slow
		def       lowerCAmelCase				(      self				:			Union[str, Any]      )						->			List[Any]:






				"""simple docstring"""







				snake_case  :       Tuple     			=						{'''input_ids''': [[12_8022, 11_0108, 397, 11, 3_8272, 2247, 12_4811, 285, 1_8105, 1586, 207, 7, 3_9534, 4428, 397, 1019, 1_8105, 1586, 207, 7, 4_1337, 1_6786, 241, 7, 2_0214, 17, 12_5690, 1_0398, 7, 4_4378, 5_8069, 6_8342, 7798, 7343, 11, 299, 3_3310, 4, 158, 3_7350, 9_4077, 4569, 299, 3_3310, 90, 4, 5_2840, 290, 4, 3_1270, 112, 299, 682, 4, 5_2840, 3_9953, 1_4079, 193, 5_2519, 9_0894, 1_7894, 12_0697, 11, 4_0445, 551, 17, 1019, 5_2519, 9_0894, 1_7756, 963, 11, 4_0445, 480, 17, 9792, 1120, 5173, 1393, 6240, 1_6786, 241, 12_0996, 28, 1245, 1393, 11_8240, 1_1123, 1019, 9_3612, 2691, 1_0618, 9_8058, 12_0409, 1928, 279, 4, 4_0683, 367, 178, 207, 1019, 103, 10_3121, 506, 6_5296, 5, 2], [12_8022, 2_1217, 367, 117, 12_5450, 128, 719, 7, 7308, 40, 9_3612, 1_2669, 1116, 1_6704, 71, 1_7785, 3699, 1_5592, 35, 144, 9584, 241, 1_1943, 713, 950, 799, 2247, 8_8427, 150, 149, 11_8813, 12_0706, 1019, 10_6906, 8_1518, 28, 1224, 2_2799, 397, 5, 2, 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], [12_8022, 1658, 12_3311, 5155, 5578, 4722, 279, 1_4947, 2366, 1120, 1197, 14, 1348, 9232, 5, 2, 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]], '''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, 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], [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, 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=UpperCamelCase__  ,	model_name='''facebook/m2m100_418M'''  ,	revision='''c168bae485c864188cf9aa0e4108b0b6934dc91e'''  ,	)







@require_torch
@require_sentencepiece
@require_tokenizers
class 	snake_case__       (      unittest.TestCase					):







		"""simple docstring"""

		lowerCamelCase     							=						"""facebook/m2m100_418M"""
		lowerCamelCase     							=						[
		    """In my opinion, there are two levels of response from the French government.""",
		    """NSA Affair Emphasizes Complete Lack of Debate on Intelligence""",
		]
		lowerCamelCase     							=						[
		    """Selon moi, il y a deux niveaux de réponse de la part du gouvernement français.""",
		    """L'affaire NSA souligne l'absence totale de débat sur le renseignement""",
		]

		# fmt: off
		lowerCamelCase     							=						[EN_CODE, 593, 1949, 115781, 4, 71586, 4234, 60633, 126233, 432, 123808, 15592, 1197, 117132, 120618, 5, 2]



		@classmethod
		def       lowerCAmelCase				(      cls				:			List[Any]      )						->			int:






				"""simple docstring"""







				snake_case  :       MaMaaaTokenizer     			=						MaMaaaTokenizer.from_pretrained(
				    cls.checkpoint_name  ,	src_lang='''en'''  ,	tgt_lang='''fr'''      )
				snake_case  :       List[str]     			=						1
				return cls



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






				"""simple docstring"""







				self.assertEqual(self.tokenizer.get_lang_id('''ar'''      )  ,	12_8006      )
				self.assertEqual(self.tokenizer.get_lang_id('''en'''      )  ,	12_8022      )
				self.assertEqual(self.tokenizer.get_lang_id('''ro'''      )  ,	12_8076      )
				self.assertEqual(self.tokenizer.get_lang_id('''mr'''      )  ,	12_8063      )



		def       lowerCAmelCase				(      self				:			Any      )						->			Dict:






				"""simple docstring"""







				snake_case  :       List[str]     			=						self.tokenizer.get_vocab()
				self.assertEqual(len(UpperCamelCase__      )  ,	self.tokenizer.vocab_size      )
				self.assertEqual(vocab['''<unk>''']  ,	3      )
				self.assertIn(self.tokenizer.get_lang_token('''en'''      )  ,	UpperCamelCase__      )



		def       lowerCAmelCase				(      self				:			Optional[Any]      )						->			List[Any]:






				"""simple docstring"""







				snake_case  :       Dict     			=						'''en'''
				snake_case  :       Dict     			=						self.tokenizer.batch_encode_plus(self.src_text      ).input_ids[0]
				self.assertListEqual(self.expected_src_tokens  ,	UpperCamelCase__      )



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






				"""simple docstring"""







				self.assertIn(UpperCamelCase__  ,	self.tokenizer.all_special_ids      )
				# fmt: off
				snake_case  :       str     			=						[FR_CODE, 5364, 82, 8642, 4, 294, 47, 8, 1_4028, 136, 3286, 9706, 6, 9_0797, 6, 14_4012, 162, 8_8128, 3_0061, 5, 2]
				# fmt: on
				snake_case  :       int     			=						self.tokenizer.decode(UpperCamelCase__  ,	skip_special_tokens=UpperCamelCase__      )
				snake_case  :       List[Any]     			=						self.tokenizer.decode(generated_ids[1:]  ,	skip_special_tokens=UpperCamelCase__      )
				self.assertEqual(UpperCamelCase__  ,	UpperCamelCase__      )
				self.assertNotIn(self.tokenizer.eos_token  ,	UpperCamelCase__      )



		def       lowerCAmelCase				(      self				:			Optional[int]      )						->			Tuple:






				"""simple docstring"""







				snake_case  :       Union[str, Any]     			=						tempfile.mkdtemp()
				snake_case  :       Union[str, Any]     			=						self.tokenizer.lang_token_to_id
				self.tokenizer.save_pretrained(UpperCamelCase__      )
				snake_case  :       Optional[int]     			=						MaMaaaTokenizer.from_pretrained(UpperCamelCase__      )
				self.assertDictEqual(new_tok.lang_token_to_id  ,	UpperCamelCase__      )



		@require_torch
		def       lowerCAmelCase				(      self				:			Optional[Any]      )						->			Union[str, Any]:






				"""simple docstring"""







				snake_case  :       Optional[Any]     			=						'''en'''
				snake_case  :       int     			=						'''fr'''

				snake_case  :       Dict     			=						self.tokenizer(self.src_text  ,	text_target=self.tgt_text  ,	padding=UpperCamelCase__  ,	return_tensors='''pt'''      )

				snake_case  :       List[str]     			=						shift_tokens_right(
				    batch['''labels''']  ,	self.tokenizer.pad_token_id  ,	self.tokenizer.eos_token_id      )

				for k in batch:
						snake_case  :       Tuple     			=						batch[k].tolist()
				# batch = {k: v.tolist() for k,v in batch.items()}
				# fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
				# batch.decoder_inputs_ids[0][0] ==
				assert batch.input_ids[1][0] == EN_CODE
				assert batch.input_ids[1][-1] == 2
				assert batch.labels[1][0] == FR_CODE
				assert batch.labels[1][-1] == 2
				assert batch.decoder_input_ids[1][:2] == [2, FR_CODE]



		@require_torch
		def       lowerCAmelCase				(      self				:			int      )						->			Union[str, Any]:






				"""simple docstring"""







				snake_case  :       List[str]     			=						'''mr'''
				self.assertListEqual(self.tokenizer.prefix_tokens  ,	[self.tokenizer.get_lang_id('''mr'''      )]      )
				self.assertListEqual(self.tokenizer.suffix_tokens  ,	[self.tokenizer.eos_token_id]      )

				snake_case  :       Optional[int]     			=						'''zh'''
				self.assertListEqual(self.tokenizer.prefix_tokens  ,	[self.tokenizer.get_lang_id('''zh'''      )]      )
				self.assertListEqual(self.tokenizer.suffix_tokens  ,	[self.tokenizer.eos_token_id]      )



		@require_torch
		def       lowerCAmelCase				(      self				:			Union[str, Any]      )						->			Dict:






				"""simple docstring"""







				snake_case  :       List[Any]     			=						'''mr'''
				self.tokenizer._switch_to_target_mode()
				self.assertListEqual(self.tokenizer.prefix_tokens  ,	[self.tokenizer.get_lang_id('''mr'''      )]      )
				self.assertListEqual(self.tokenizer.suffix_tokens  ,	[self.tokenizer.eos_token_id]      )
				self.tokenizer._switch_to_input_mode()
				self.assertListEqual(self.tokenizer.prefix_tokens  ,	[self.tokenizer.get_lang_id(self.tokenizer.src_lang      )]      )

				snake_case  :       List[str]     			=						'''zh'''
				self.tokenizer._switch_to_target_mode()
				self.assertListEqual(self.tokenizer.prefix_tokens  ,	[self.tokenizer.get_lang_id('''zh'''      )]      )
				self.assertListEqual(self.tokenizer.suffix_tokens  ,	[self.tokenizer.eos_token_id]      )
				self.tokenizer._switch_to_input_mode()
				self.assertListEqual(self.tokenizer.prefix_tokens  ,	[self.tokenizer.get_lang_id(self.tokenizer.src_lang      )]      )



		@require_torch
		def       lowerCAmelCase				(      self				:			Tuple      )						->			str:






				"""simple docstring"""







				snake_case  :       int     			=						self.tokenizer._build_translation_inputs('''A test'''  ,	return_tensors='''pt'''  ,	src_lang='''en'''  ,	tgt_lang='''ar'''      )

				self.assertEqual(
				    nested_simplify(UpperCamelCase__      )  ,	{
				        # en_XX, A, test, EOS
				        '''input_ids''': [[12_8022, 58, 4183, 2]],
				        '''attention_mask''': [[1, 1, 1, 1]],
				        # ar_AR
				        '''forced_bos_token_id''': 12_8006,
				    }  ,	)


