




import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, BatchEncoding, PLBartTokenizer, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
)

from ...test_tokenization_common import TokenizerTesterMixin


UpperCamelCase__         =      get_tests_dir("""fixtures/test_sentencepiece.model""")


if is_torch_available():
					from transformers.models.plbart.modeling_plbart import shift_tokens_right

UpperCamelCase__         =      50003
UpperCamelCase__         =      50002
@require_sentencepiece
@require_tokenizers
class       a__   ( snake_case__			,    unittest.TestCase				):
						_a							:   Tuple =   PLBartTokenizer
						_a							:   List[Any] =   None
						_a							:   Union[str, Any] =   False






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											super().setUp()

											# We have a SentencePiece fixture for testing
											__lowerCAmelCase           =	PLBartTokenizer(_A		,      language_codes="base"		,      keep_accents=_A      )
											tokenizer.save_pretrained(self.tmpdirname      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	PLBartTokenizer(_A		,      language_codes="base"		,      keep_accents=_A      )

											__lowerCAmelCase           =	tokenizer.tokenize("This is a test"      )
											self.assertListEqual(_A		,      ["▁This", "▁is", "▁a", "▁t", "est"]      )

											self.assertListEqual(
											    tokenizer.convert_tokens_to_ids(_A      )		,      [value + tokenizer.fairseq_offset for value in [2_8_5, 4_6, 1_0, 1_7_0, 3_8_2]]		,      )

											__lowerCAmelCase           =	tokenizer.tokenize("I was born in 92000, and this is falsé."      )
											self.assertListEqual(
											    _A		,      [
											        SPIECE_UNDERLINE + "I",
											        SPIECE_UNDERLINE + "was",
											        SPIECE_UNDERLINE + "b",
											        "or",
											        "n",
											        SPIECE_UNDERLINE + "in",
											        SPIECE_UNDERLINE + "",
											        "9",
											        "2",
											        "0",
											        "0",
											        "0",
											        ",",
											        SPIECE_UNDERLINE + "and",
											        SPIECE_UNDERLINE + "this",
											        SPIECE_UNDERLINE + "is",
											        SPIECE_UNDERLINE + "f",
											        "al",
											        "s",
											        "é",
											        ".",
											    ]		,      )
											__lowerCAmelCase           =	tokenizer.convert_tokens_to_ids(_A      )
											self.assertListEqual(
											    _A		,      [
											        value + tokenizer.fairseq_offset
											        for value in [8, 2_1, 8_4, 5_5, 2_4, 1_9, 7, 2, 6_0_2, 3_4_7, 3_4_7, 3_4_7, 3, 1_2, 6_6, 4_6, 7_2, 8_0, 6, 2, 4]
											    ]		,      )

											__lowerCAmelCase           =	tokenizer.convert_ids_to_tokens(_A      )
											self.assertListEqual(
											    _A		,      [
											        SPIECE_UNDERLINE + "I",
											        SPIECE_UNDERLINE + "was",
											        SPIECE_UNDERLINE + "b",
											        "or",
											        "n",
											        SPIECE_UNDERLINE + "in",
											        SPIECE_UNDERLINE + "",
											        "<unk>",
											        "2",
											        "0",
											        "0",
											        "0",
											        ",",
											        SPIECE_UNDERLINE + "and",
											        SPIECE_UNDERLINE + "this",
											        SPIECE_UNDERLINE + "is",
											        SPIECE_UNDERLINE + "f",
											        "al",
											        "s",
											        "<unk>",
											        ".",
											    ]		,      )

											__lowerCAmelCase           =	tokenizer.vocab_size
											__lowerCAmelCase           =	[tokenizer.convert_ids_to_tokens(_A      ) for x in range(end - 4		,      _A      )]

											self.assertListEqual(_A		,      ["__java__", "__python__", "__en_XX__", "<mask>"]      )

											__lowerCAmelCase           =	"java.lang.Exception, python.lang.Exception, javascript, php, ruby, go"
											__lowerCAmelCase           =	tokenizer(_A      ).input_ids
											self.assertEqual(
											    tokenizer.decode(_A		,      skip_special_tokens=_A		,      clean_up_tokenization_spaces=_A      )		,      _A		,      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	PLBartTokenizer(_A		,      language_codes="multi"		,      keep_accents=_A      )

											__lowerCAmelCase           =	tokenizer.tokenize("This is a test"      )
											self.assertListEqual(_A		,      ["▁This", "▁is", "▁a", "▁t", "est"]      )

											self.assertListEqual(
											    tokenizer.convert_tokens_to_ids(_A      )		,      [value + tokenizer.fairseq_offset for value in [2_8_5, 4_6, 1_0, 1_7_0, 3_8_2]]		,      )

											__lowerCAmelCase           =	tokenizer.tokenize("I was born in 92000, and this is falsé."      )
											self.assertListEqual(
											    _A		,      [
											        SPIECE_UNDERLINE + "I",
											        SPIECE_UNDERLINE + "was",
											        SPIECE_UNDERLINE + "b",
											        "or",
											        "n",
											        SPIECE_UNDERLINE + "in",
											        SPIECE_UNDERLINE + "",
											        "9",
											        "2",
											        "0",
											        "0",
											        "0",
											        ",",
											        SPIECE_UNDERLINE + "and",
											        SPIECE_UNDERLINE + "this",
											        SPIECE_UNDERLINE + "is",
											        SPIECE_UNDERLINE + "f",
											        "al",
											        "s",
											        "é",
											        ".",
											    ]		,      )
											__lowerCAmelCase           =	tokenizer.convert_tokens_to_ids(_A      )
											self.assertListEqual(
											    _A		,      [
											        value + tokenizer.fairseq_offset
											        for value in [8, 2_1, 8_4, 5_5, 2_4, 1_9, 7, 2, 6_0_2, 3_4_7, 3_4_7, 3_4_7, 3, 1_2, 6_6, 4_6, 7_2, 8_0, 6, 2, 4]
											    ]		,      )

											__lowerCAmelCase           =	tokenizer.convert_ids_to_tokens(_A      )
											self.assertListEqual(
											    _A		,      [
											        SPIECE_UNDERLINE + "I",
											        SPIECE_UNDERLINE + "was",
											        SPIECE_UNDERLINE + "b",
											        "or",
											        "n",
											        SPIECE_UNDERLINE + "in",
											        SPIECE_UNDERLINE + "",
											        "<unk>",
											        "2",
											        "0",
											        "0",
											        "0",
											        ",",
											        SPIECE_UNDERLINE + "and",
											        SPIECE_UNDERLINE + "this",
											        SPIECE_UNDERLINE + "is",
											        SPIECE_UNDERLINE + "f",
											        "al",
											        "s",
											        "<unk>",
											        ".",
											    ]		,      )
											__lowerCAmelCase           =	tokenizer.vocab_size
											__lowerCAmelCase           =	[tokenizer.convert_ids_to_tokens(_A      ) for x in range(end - 7		,      _A      )]

											self.assertListEqual(
											    _A		,      ["__java__", "__python__", "__en_XX__", "__javascript__", "__php__", "__ruby__", "__go__"]      )
											__lowerCAmelCase           =	"java.lang.Exception, python.lang.Exception, javascript, php, ruby, go"
											__lowerCAmelCase           =	tokenizer(_A      ).input_ids
											self.assertEqual(
											    tokenizer.decode(_A		,      skip_special_tokens=_A		,      clean_up_tokenization_spaces=_A      )		,      _A		,      )






@require_torch
@require_sentencepiece
@require_tokenizers
class       a__   ( unittest.TestCase				):
						_a							:   Dict =   """uclanlp/plbart-python-en_XX"""
						_a							:   int =   [
						    """def maximum(a,b,c):NEW_LINE_INDENTreturn max([a,b,c])""",
						    """def sum(a,b,c):NEW_LINE_INDENTreturn sum([a,b,c])""",
						]
						_a							:   int =   [
						    """Returns the maximum value of a b c.""",
						    """Sums the values of a b c.""",
						]
						_a							:   Tuple =   [
						    1_3_4,
						    5_4_5_2,
						    3_3_4_6_0,
						    3_3_4_4_1,
						    3_3_4_6_3,
						    3_3_4_6_5,
						    3_3_4_6_3,
						    3_3_4_4_9,
						    9_8_8,
						    2_0,
						    3_3_4_5_6,
						    1_9,
						    3_3_4_5_6,
						    7_7_1,
						    3_9,
						    4_2_5_8,
						    8_8_9,
						    3_3_1_8,
						    3_3_4_4_1,
						    3_3_4_6_3,
						    3_3_4_6_5,
						    3_3_4_6_3,
						    3_3_4_4_9,
						    2_4_7_1,
						    2,
						    PYTHON_CODE,
						]






						@classmethod
						def     __SCREAMING_SNAKE_CASE(       cls      ):




											"""simple docstring"""



											__lowerCAmelCase           =	PLBartTokenizer.from_pretrained(
											    cls.checkpoint_name		,      language_codes="base"		,      src_lang="python"		,      tgt_lang="en_XX"      )
											__lowerCAmelCase           =	1
											return cls






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["__java__"]		,      5_0_0_0_1      )
											self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["__python__"]		,      5_0_0_0_2      )
											self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["__en_XX__"]		,      5_0_0_0_3      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.tokenizer.batch_encode_plus(self.src_text      ).input_ids[0]
											self.assertListEqual(self.expected_src_tokens		,      _A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											self.assertIn(_A		,      self.tokenizer.all_special_ids      )
											__lowerCAmelCase           =	[EN_CODE, 9_0_3_7, 3_3_4_4_2, 5_7, 7_5_2, 1_5_3, 1_4, 5_6, 1_8, 9, 2]
											__lowerCAmelCase           =	self.tokenizer.decode(_A		,      skip_special_tokens=_A      )
											__lowerCAmelCase           =	self.tokenizer.decode(generated_ids[1:]		,      skip_special_tokens=_A      )
											self.assertEqual(_A		,      _A      )
											self.assertNotIn(self.tokenizer.eos_token		,      _A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	["def sum(a,b,c):NEW_LINE_INDENTreturn sum([a,b,c])" * 2_0]
											self.assertIsInstance(src_text[0]		,      _A      )
											__lowerCAmelCase           =	1_0
											__lowerCAmelCase           =	self.tokenizer(_A		,      max_length=_A		,      truncation=_A      ).input_ids[0]
											self.assertEqual(ids[-2]		,      2      )
											self.assertEqual(ids[-1]		,      _A      )
											self.assertEqual(len(_A      )		,      _A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											self.assertListEqual(self.tokenizer.convert_tokens_to_ids(["<mask>", "__java__"]      )		,      [5_0_0_0_4, 5_0_0_0_1]      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	tempfile.mkdtemp()
											__lowerCAmelCase           =	self.tokenizer.fairseq_tokens_to_ids
											self.tokenizer.save_pretrained(_A      )
											__lowerCAmelCase           =	PLBartTokenizer.from_pretrained(_A      )
											self.assertDictEqual(new_tok.fairseq_tokens_to_ids		,      _A      )






						@require_torch
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.tokenizer(self.src_text		,      text_target=self.tgt_text		,      padding=_A		,      return_tensors="pt"      )
											__lowerCAmelCase           =	shift_tokens_right(batch["labels"]		,      self.tokenizer.pad_token_id      )

											# fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
											self.assertEqual(batch.input_ids[1][-2:].tolist()		,      [2, PYTHON_CODE]      )
											self.assertEqual(batch.decoder_input_ids[1][0]		,      _A      )
											self.assertEqual(batch.decoder_input_ids[1][-1]		,      2      )
											self.assertEqual(batch.labels[1][-2:].tolist()		,      [2, EN_CODE]      )






						@require_torch
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.tokenizer(
											    self.src_text		,      text_target=self.tgt_text		,      padding=_A		,      truncation=_A		,      max_length=len(self.expected_src_tokens      )		,      return_tensors="pt"		,      )
											__lowerCAmelCase           =	shift_tokens_right(batch["labels"]		,      self.tokenizer.pad_token_id      )

											self.assertIsInstance(_A		,      _A      )

											self.assertEqual((2, 2_6)		,      batch.input_ids.shape      )
											self.assertEqual((2, 2_6)		,      batch.attention_mask.shape      )
											__lowerCAmelCase           =	batch.input_ids.tolist()[0]
											self.assertListEqual(self.expected_src_tokens		,      _A      )
											self.assertEqual(2		,      batch.decoder_input_ids[0, -1]      )  # EOS
											# Test that special tokens are reset
											self.assertEqual(self.tokenizer.prefix_tokens		,      []      )
											self.assertEqual(self.tokenizer.suffix_tokens		,      [self.tokenizer.eos_token_id, PYTHON_CODE]      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.tokenizer(self.src_text		,      padding=_A		,      truncation=_A		,      max_length=3		,      return_tensors="pt"      )
											__lowerCAmelCase           =	self.tokenizer(
											    text_target=self.tgt_text		,      padding=_A		,      truncation=_A		,      max_length=1_0		,      return_tensors="pt"      )
											__lowerCAmelCase           =	targets["input_ids"]
											__lowerCAmelCase           =	shift_tokens_right(_A		,      self.tokenizer.pad_token_id      )

											self.assertEqual(batch.input_ids.shape[1]		,      3      )
											self.assertEqual(batch.decoder_input_ids.shape[1]		,      1_0      )






						@require_torch
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.tokenizer._build_translation_inputs(
											    "A test"		,      return_tensors="pt"		,      src_lang="en_XX"		,      tgt_lang="java"      )

											self.assertEqual(
											    nested_simplify(_A      )		,      {
											        # A, test, EOS, en_XX
											        "input_ids": [[1_5_0, 2_4_2, 2, 5_0_0_0_3]],
											        "attention_mask": [[1, 1, 1, 1]],
											        # java
											        "forced_bos_token_id": 5_0_0_0_1,
											    }		,      )










import argparse
import os

import torch

from transformers import FlavaConfig, FlavaForPreTraining
from transformers.models.flava.convert_dalle_to_flava_codebook import convert_dalle_checkpoint







def      _a				(						SCREAMING_SNAKE_CASE_							:    Any					):
					# encoder.embeddings are double copied in original FLAVA
					return sum(param.float().sum() if "encoder.embeddings" not in key else 0 for key, param in state_dict.items()					)







def      _a				(						SCREAMING_SNAKE_CASE_							:    Dict				,				SCREAMING_SNAKE_CASE_							:    Dict					):
					__lowerCAmelCase           =	{}

					for key, value in state_dict.items():
										if "text_encoder.embeddings" in key or "image_encoder.embeddings" in key:
															continue

										__lowerCAmelCase           =	key.replace("heads.cmd.mim_head.cls.predictions"				,				"mmm_image_head"					)
										__lowerCAmelCase           =	key.replace("heads.cmd.mlm_head.cls.predictions"				,				"mmm_text_head"					)
										__lowerCAmelCase           =	key.replace("heads.cmd.itm_head.cls"				,				"itm_head"					)
										__lowerCAmelCase           =	key.replace("heads.cmd.itm_head.pooler"				,				"itm_head.pooler"					)
										__lowerCAmelCase           =	key.replace("heads.cmd.clip_head.logit_scale"				,				"flava.logit_scale"					)
										__lowerCAmelCase           =	key.replace("heads.fairseq_mlm.cls.predictions"				,				"mlm_head"					)
										__lowerCAmelCase           =	key.replace("heads.imagenet.mim_head.cls.predictions"				,				"mim_head"					)
										__lowerCAmelCase           =	key.replace("mm_text_projection"				,				"flava.text_to_mm_projection"					)
										__lowerCAmelCase           =	key.replace("mm_image_projection"				,				"flava.image_to_mm_projection"					)
										__lowerCAmelCase           =	key.replace("image_encoder.module"				,				"flava.image_model"					)
										__lowerCAmelCase           =	key.replace("text_encoder.module"				,				"flava.text_model"					)
										__lowerCAmelCase           =	key.replace("mm_encoder.module.encoder.cls_token"				,				"flava.multimodal_model.cls_token"					)
										__lowerCAmelCase           =	key.replace("mm_encoder.module"				,				"flava.multimodal_model"					)
										__lowerCAmelCase           =	key.replace("text_projection"				,				"flava.text_projection"					)
										__lowerCAmelCase           =	key.replace("image_projection"				,				"flava.image_projection"					)

										__lowerCAmelCase           =	value.float()

					for key, value in codebook_state_dict.items():
										__lowerCAmelCase           =	value

					return upgrade







@torch.no_grad()
def      _a				(						SCREAMING_SNAKE_CASE_							:    List[Any]				,				SCREAMING_SNAKE_CASE_							:    Optional[int]				,				SCREAMING_SNAKE_CASE_							:    List[str]				,				SCREAMING_SNAKE_CASE_							:    int=None					):
					if config_path is not None:
										__lowerCAmelCase           =	FlavaConfig.from_pretrained(SCREAMING_SNAKE_CASE_					)
					else:
										__lowerCAmelCase           =	FlavaConfig()

					__lowerCAmelCase           =	FlavaForPreTraining(SCREAMING_SNAKE_CASE_					).eval()

					__lowerCAmelCase           =	convert_dalle_checkpoint(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_				,				save_checkpoint=SCREAMING_SNAKE_CASE_					)

					if os.path.exists(SCREAMING_SNAKE_CASE_					):
										__lowerCAmelCase           =	torch.load(SCREAMING_SNAKE_CASE_				,				map_location="cpu"					)
					else:
										__lowerCAmelCase           =	torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_				,				map_location="cpu"					)

					__lowerCAmelCase           =	upgrade_state_dict(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_					)
					hf_model.load_state_dict(SCREAMING_SNAKE_CASE_					)
					__lowerCAmelCase           =	hf_model.state_dict()
					__lowerCAmelCase           =	count_parameters(SCREAMING_SNAKE_CASE_					)
					__lowerCAmelCase           =	count_parameters(SCREAMING_SNAKE_CASE_					) + count_parameters(SCREAMING_SNAKE_CASE_					)

					assert torch.allclose(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_				,				atol=1E-3					)

					hf_model.save_pretrained(SCREAMING_SNAKE_CASE_					)


if __name__ == "__main__":
					UpperCamelCase__         =      argparse.ArgumentParser()
					parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
					parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to flava checkpoint""")
					parser.add_argument("""--codebook_path""", default=None, type=str, help="""Path to flava codebook checkpoint""")
					parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
					UpperCamelCase__         =      parser.parse_args()

					convert_flava_checkpoint(args.checkpoint_path, args.codebook_path, args.pytorch_dump_folder_path, args.config_path)




