



"""simple docstring"""






import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis_float32 (there's also the fix_lavis branch)
# also note: to convert Vicuna checkpoints, we had to include /home/niels/python_projects/checkpoints/FastChat/vicuna-7b in lavis/configs/models/blip2/blip2_instruct_vicuna7b.yaml
# same for Vicuna-13b
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipImageProcessor,
    InstructBlipConfig,
    InstructBlipForConditionalGeneration,
    InstructBlipProcessor,
    InstructBlipQFormerConfig,
    InstructBlipVisionConfig,
    LlamaConfig,
    LlamaTokenizerFast,
    TaConfig,
    TaTokenizerFast,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD






def 			snake_case_    ( ):


       '''simple docstring'''





       _lowerCamelCase		:							Any            =					'''https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'''
       _lowerCamelCase		:							str            =					Image.open(requests.get(A_,							stream=A_		).raw		).convert('''RGB'''		)

       return image






def 			snake_case_    ( A_   :						Union[str, Any]		):


       '''simple docstring'''





       _lowerCamelCase		:							Optional[Any]            =					[]
       # fmt: off

       # vision encoder
       rename_keys.append(('''visual_encoder.cls_token''', '''vision_model.embeddings.class_embedding''')		)
       rename_keys.append(('''visual_encoder.pos_embed''', '''vision_model.embeddings.position_embedding''')		)
       rename_keys.append(('''visual_encoder.patch_embed.proj.weight''', '''vision_model.embeddings.patch_embedding.weight''')		)
       rename_keys.append(('''visual_encoder.patch_embed.proj.bias''', '''vision_model.embeddings.patch_embedding.bias''')		)
       rename_keys.append(('''ln_vision.weight''', '''vision_model.post_layernorm.weight''')		)
       rename_keys.append(('''ln_vision.bias''', '''vision_model.post_layernorm.bias''')		)

       for i in range(config.vision_config.num_hidden_layers		):
              rename_keys.append((F'''visual_encoder.blocks.{i}.norm1.weight''', F'''vision_model.encoder.layers.{i}.layer_norm1.weight''')		)
              rename_keys.append((F'''visual_encoder.blocks.{i}.norm1.bias''', F'''vision_model.encoder.layers.{i}.layer_norm1.bias''')		)
              rename_keys.append((F'''visual_encoder.blocks.{i}.norm2.weight''', F'''vision_model.encoder.layers.{i}.layer_norm2.weight''')		)
              rename_keys.append((F'''visual_encoder.blocks.{i}.norm2.bias''', F'''vision_model.encoder.layers.{i}.layer_norm2.bias''')		)
              rename_keys.append((F'''visual_encoder.blocks.{i}.attn.qkv.weight''', F'''vision_model.encoder.layers.{i}.self_attn.qkv.weight''')		)
              rename_keys.append((F'''visual_encoder.blocks.{i}.attn.proj.weight''', F'''vision_model.encoder.layers.{i}.self_attn.projection.weight''',)		)
              rename_keys.append((F'''visual_encoder.blocks.{i}.attn.proj.bias''', F'''vision_model.encoder.layers.{i}.self_attn.projection.bias''')		)
              rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc1.weight''', F'''vision_model.encoder.layers.{i}.mlp.fc1.weight''')		)
              rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc1.bias''', F'''vision_model.encoder.layers.{i}.mlp.fc1.bias''')		)
              rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc2.weight''', F'''vision_model.encoder.layers.{i}.mlp.fc2.weight''')		)
              rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc2.bias''', F'''vision_model.encoder.layers.{i}.mlp.fc2.bias''')		)

       # QFormer
       rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.weight''', '''qformer.embeddings.layernorm.weight''')		)
       rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.bias''', '''qformer.embeddings.layernorm.bias''')		)

       # fmt: on
       return rename_keys






def 			snake_case_    ( A_   :						str,							A_   :						Optional[Any],							A_   :						Dict		):


       '''simple docstring'''





       _lowerCamelCase		:							List[Any]            =					dct.pop(A_		)
       _lowerCamelCase		:							Dict            =					val






def 			snake_case_    ( A_   :						List[Any],							A_   :						List[Any]		):


       '''simple docstring'''





       for i in range(config.vision_config.num_hidden_layers		):
              # read in original q and v biases
              _lowerCamelCase		:							Tuple            =					state_dict.pop(F'''visual_encoder.blocks.{i}.attn.q_bias'''		)
              _lowerCamelCase		:							Any            =					state_dict.pop(F'''visual_encoder.blocks.{i}.attn.v_bias'''		)

              # next, set bias in the state dict
              _lowerCamelCase		:							Union[str, Any]            =					torch.cat((q_bias, torch.zeros_like(A_,							requires_grad=A_		), v_bias)		)
              _lowerCamelCase		:							Any            =					qkv_bias






def 			snake_case_    ( A_   :						int		):


       '''simple docstring'''





       _lowerCamelCase		:							Tuple            =					3_64 if '''coco''' in model_name else 2_24
       _lowerCamelCase		:							Optional[int]            =					InstructBlipVisionConfig(image_size=A_		).to_dict()

       # make sure the models have proper bos_token_id and eos_token_id set (important for generation)
       # seems like flan-T5 models don't have bos_token_id properly set?
       if "t5-xl" in model_name:
              _lowerCamelCase		:							List[str]            =					TaConfig.from_pretrained('''google/flan-t5-xl''',							dense_act_fn='''gelu''',							bos_token_id=1		).to_dict()
       elif "t5-xxl" in model_name:
              _lowerCamelCase		:							int            =					TaConfig.from_pretrained('''google/flan-t5-xxl''',							dense_act_fn='''gelu''',							bos_token_id=1		).to_dict()
       elif "vicuna-7b" in model_name:
              _lowerCamelCase		:							int            =					LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf''',							vocab_size=3_20_01		).to_dict()
       elif "vicuna-13b" in model_name:
              _lowerCamelCase		:							Dict            =					LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf''',							vocab_size=3_20_01		).to_dict()
       else:
              raise ValueError('''Model name not supported'''		)

       # the authors add one special "[DEC]" token to the vocab of Q-Former, hence vocab size = 30522 + 1
       _lowerCamelCase		:							Any            =					InstructBlipQFormerConfig(vocab_size=3_05_23		).to_dict()
       _lowerCamelCase		:							List[Any]            =					InstructBlipConfig(vision_config=A_,							text_config=A_,							qformer_config=A_		)

       return config, image_size






@torch.no_grad()
def 			snake_case_    ( A_   :						str,							A_   :						Any=None,							A_   :						str=False		):


       '''simple docstring'''





       _lowerCamelCase		:							Dict            =					AutoTokenizer.from_pretrained('''bert-base-uncased''',							truncation_side='''left'''		)
       qformer_tokenizer.add_special_tokens({'''bos_token''': '''[DEC]'''}		)

       if "t5" in model_name:
              _lowerCamelCase		:							Dict            =					TaTokenizerFast.from_pretrained('''google/flan-t5-xl''',							truncation_side='''left'''		)
       elif "vicuna" in model_name:
              # the following was used in the original implementation:
              # tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b", use_fast=False, truncation_side="left")
              # tokenizer.add_special_tokens({"pad_token": "[PAD]"})
              # tokenizer.add_special_tokens({"bos_token": "</s>"})
              # tokenizer.add_special_tokens({"eos_token": "</s>"})
              # tokenizer.add_special_tokens({"unk_token": "</s>"})
              _lowerCamelCase		:							Optional[int]            =					LlamaTokenizerFast.from_pretrained(
                  '''huggyllama/llama-7b''',							truncation_side='''left''',							bos_token='''</s>''',							unk_token='''</s>'''		)
              tokenizer.add_special_tokens({'''pad_token''': '''[PAD]'''}		)

       _lowerCamelCase     ,       _lowerCamelCase		:							Dict            =					get_blipa_config(A_		)
       _lowerCamelCase		:							Dict            =					InstructBlipForConditionalGeneration(A_		).eval()

       _lowerCamelCase		:							Optional[int]            =					{
           '''instructblip-vicuna-7b''': ('''blip2_vicuna_instruct''', '''vicuna7b'''),
           '''instructblip-vicuna-13b''': ('''blip2_vicuna_instruct''', '''vicuna13b'''),
           '''instructblip-flan-t5-xl''': ('''blip2_t5_instruct''', '''flant5xl'''),
           '''instructblip-flan-t5-xxl''': ('''blip2_t5_instruct''', '''flant5xxl'''),
       }

       _lowerCamelCase     ,       _lowerCamelCase		:							Any            =					model_name_to_original[model_name]

       # load original model
       print('''Loading original model...'''		)
       _lowerCamelCase		:							Tuple            =					'''cuda:1''' if torch.cuda.is_available() else '''cpu'''
       _lowerCamelCase		:							Dict            =					'''cuda:2''' if torch.cuda.is_available() else '''cpu'''
       _lowerCamelCase     ,       _lowerCamelCase     ,       _lowerCamelCase		:							str            =					load_model_and_preprocess(
           name=A_,							model_type=A_,							is_eval=A_,							device=A_		)
       original_model.eval()
       print('''Done!'''		)

       # update state dict keys
       _lowerCamelCase		:							Tuple            =					original_model.state_dict()
       _lowerCamelCase		:							Tuple            =					create_rename_keys(A_		)
       for src, dest in rename_keys:
              rename_key(A_,							A_,							A_		)

       # some keys can be renamed efficiently
       for key, val in state_dict.copy().items():
              _lowerCamelCase		:							Tuple            =					state_dict.pop(A_		)
              if key.startswith('''Qformer.bert'''		):
                     _lowerCamelCase		:							str            =					key.replace('''Qformer.bert''',							'''qformer'''		)
              if "attention.self" in key:
                     _lowerCamelCase		:							List[Any]            =					key.replace('''self''',							'''attention'''		)
              if "llm_proj" in key:
                     _lowerCamelCase		:							Dict            =					key.replace('''llm_proj''',							'''language_projection'''		)
              if "t5_proj" in key:
                     _lowerCamelCase		:							List[Any]            =					key.replace('''t5_proj''',							'''language_projection'''		)
              if key.startswith('''llm_model'''		):
                     _lowerCamelCase		:							Optional[int]            =					key.replace('''llm_model''',							'''language_model'''		)
              if key.startswith('''t5'''		):
                     _lowerCamelCase		:							List[Any]            =					key.replace('''t5''',							'''language'''		)
              _lowerCamelCase		:							List[Any]            =					val

       # read in qv biases
       read_in_q_v_bias(A_,							A_		)

       # note: weights get loaded in torch.float32 by default
       hf_model.load_state_dict(A_,							strict=A_		)

       _lowerCamelCase		:							str            =					load_demo_image()
       _lowerCamelCase		:							Union[str, Any]            =					'''What is unusual about this image?'''

       # create processor
       _lowerCamelCase		:							Tuple            =					BlipImageProcessor(
           size={'''height''': image_size, '''width''': image_size},							image_mean=A_,							image_std=A_		)
       _lowerCamelCase		:							List[str]            =					InstructBlipProcessor(
           image_processor=A_,							tokenizer=A_,							qformer_tokenizer=A_,							)
       _lowerCamelCase		:							Dict            =					processor(images=A_,							text=A_,							return_tensors='''pt'''		).to(A_		)

       # make sure processor creates exact same pixel values
       _lowerCamelCase		:							Optional[Any]            =					vis_processors['''eval'''](A_		).unsqueeze(0		).to(A_		)
       _lowerCamelCase		:							List[Any]            =					inputs.pixel_values
       assert torch.allclose(original_pixel_values.to(pixel_values.device		),							A_		)

       original_model.to(A_		)
       hf_model.to(A_		)
       with torch.no_grad():
              if "vicuna" in model_name:
                     _lowerCamelCase		:							List[Any]            =					original_model({'''image''': original_pixel_values, '''text_input''': [prompt]}		).logits
                     _lowerCamelCase		:							List[Any]            =					hf_model(**A_		).logits
              else:
                     _lowerCamelCase		:							Tuple            =					original_model(
                         {'''image''': original_pixel_values, '''text_input''': [prompt], '''text_output''': ['''\n''']}		).logits
                     _lowerCamelCase		:							str            =					tokenizer('''\n''',							return_tensors='''pt'''		).input_ids.to(A_		)
                     _lowerCamelCase		:							Optional[Any]            =					label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id,							-1_00		)
                     _lowerCamelCase		:							Any            =					hf_model(**A_,							labels=A_		).logits

       print('''First values of original logits:''',							original_logits[0, :3, :3]		)
       print('''First values of HF logits:''',							logits[0, :3, :3]		)

       # assert values
       assert original_logits.shape == logits.shape
       _lowerCamelCase		:							Tuple            =					1E-4 if '''vicuna''' in model_name else 1E-5
       assert torch.allclose(original_logits.to(logits.device		),							A_,							atol=A_		)
       print('''Looks ok!'''		)

       print('''Generating with original model...'''		)
       _lowerCamelCase		:							str            =					original_model.generate({'''image''': original_pixel_values, '''prompt''': prompt},							num_beams=5		)

       # important: we need to cast the weights of the HF model to the appropriate type
       print('''Generating with HF model...'''		)
       _lowerCamelCase		:							Tuple            =					hf_model.generate(
           **A_,							do_sample=A_,							num_beams=5,							max_length=2_56,							min_length=1,							top_p=0.9,							repetition_penalty=1.5,							length_penalty=1.0,							temperature=1,							)
       if "vicuna" in model_name:
              # convert output id 0 to 2 (eos_token_id)
              # TODO add this in the generate method?
              _lowerCamelCase		:							Union[str, Any]            =					2
       print('''Original generation:''',							A_		)
       _lowerCamelCase		:							Dict            =					processor.batch_decode(A_,							skip_special_tokens=A_		)
       _lowerCamelCase		:							str            =					[text.strip() for text in output_text]
       print('''HF generation:''',							A_		)

       if pytorch_dump_folder_path is not None:
              processor.save_pretrained(A_		)
              hf_model.save_pretrained(A_		)

       if push_to_hub:
              processor.push_to_hub(F'''Salesforce/{model_name}'''		)
              hf_model.push_to_hub(F'''Salesforce/{model_name}'''		)


if __name__ == "__main__":
   lowerCAmelCase__              =							argparse.ArgumentParser()
   lowerCAmelCase__              =							[
       '''instructblip-vicuna-7b''',
       '''instructblip-vicuna-13b''',
       '''instructblip-flan-t5-xl''',
       '''instructblip-flan-t5-xxl''',
   ]
   parser.add_argument(
       '''--model_name''',
       default='''instructblip-flan-t5-xl''',
       choices=choices,
       type=str,
       help='''Path to hf config.json of model to convert''',
   )
   parser.add_argument('''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
   parser.add_argument(
       '''--push_to_hub''',
       action='''store_true''',
       help='''Whether to push the model and processor to the hub after converting''',
   )

   lowerCAmelCase__              =							parser.parse_args()

   convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




'''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,
				    }  ,	)


