




import io
import os
import unicodedata
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


UpperCamelCase__         =      logging.get_logger(__name__)

UpperCamelCase__         =      """▁"""

UpperCamelCase__         =      {"""vocab_file""": """vocab.txt""", """sentencepiece_model_ckpt""": """sentencepiece.bpe.model"""}

UpperCamelCase__         =      {
    """sentencepiece_model_file""": """sentencepiece.bpe.model""",
    """vocab_file""": """vocab.txt""",
}

UpperCamelCase__         =      {
    """vocab_file""": {
        """ernie-m-base""": """https://huggingface.co/susnato/ernie-m-base_pytorch/blob/main/vocab.txt""",
        """ernie-m-large""": """https://huggingface.co/susnato/ernie-m-base_pytorch/blob/main/vocab.txt""",
    },
    """sentencepiece_model_file""": {
        """ernie-m-base""": """https://huggingface.co/susnato/ernie-m-base_pytorch/blob/main/sentencepiece.bpe.model""",
        """ernie-m-large""": """https://huggingface.co/susnato/ernie-m-base_pytorch/blob/main/sentencepiece.bpe.model""",
    },
}

UpperCamelCase__         =      {
    """ernie-m-base""": 514,
    """ernie-m-large""": 514,
}

UpperCamelCase__         =      {
    """ernie-m-base""": {"""do_lower_case""": False},
    """ernie-m-large""": {"""do_lower_case""": False},
}
class       a__   ( snake_case__				):
						_a							:   List[str] =   ["input_ids"]

						_a							:   List[Any] =   VOCAB_FILES_NAMES
						_a							:   Union[str, Any] =   PRETRAINED_INIT_CONFIGURATION
						_a							:   Optional[Any] =   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
						_a							:   Tuple =   PRETRAINED_VOCAB_FILES_MAP
						_a							:   Any =   RESOURCE_FILES_NAMES






						def __init__(       self		,      _A		,      _A=None		,      _A=False		,      _A="utf8"		,      _A="[UNK]"		,      _A="[SEP]"		,      _A="[PAD]"		,      _A="[CLS]"		,      _A="[MASK]"		,      _A = None		,      **_A		,      ):




											"""simple docstring"""



											__lowerCAmelCase           =	{} if sp_model_kwargs is None else sp_model_kwargs
											super().__init__(
											    do_lower_case=_A		,      unk_token=_A		,      sep_token=_A		,      pad_token=_A		,      cls_token=_A		,      mask_token=_A		,      vocab_file=_A		,      encoding=_A		,      sp_model_kwargs=self.sp_model_kwargs		,      **_A		,      )
											__lowerCAmelCase           =	do_lower_case
											__lowerCAmelCase           =	sentencepiece_model_ckpt
											__lowerCAmelCase           =	spm.SentencePieceProcessor(**self.sp_model_kwargs      )
											self.sp_model.Load(_A      )

											# to mimic paddlenlp.transformers.ernie_m.tokenizer.ErnieMTokenizer functioning
											if vocab_file is not None:
																__lowerCAmelCase           =	self.load_vocab(filepath=_A      )
											else:
																__lowerCAmelCase           =	{self.sp_model.id_to_piece(_A      ): id for id in range(self.sp_model.get_piece_size()      )}
											__lowerCAmelCase           =	{v: k for k, v in self.vocab.items()}






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											if text is None:
																return None

											__lowerCAmelCase           =	self.tokenize(_A      )
											__lowerCAmelCase       ,							__lowerCAmelCase           =	"", []

											for i, ch in enumerate(_A      ):
																if ch in self.SP_CHAR_MAPPING:
																					__lowerCAmelCase           =	self.SP_CHAR_MAPPING.get(_A      )
																else:
																					__lowerCAmelCase           =	unicodedata.normalize("NFKC"		,      _A      )
																if self.is_whitespace(_A      ):
																					continue
																normalized_text += ch
																char_mapping.extend([i] * len(_A      )      )

											__lowerCAmelCase       ,							__lowerCAmelCase       ,							__lowerCAmelCase           =	normalized_text, [], 0

											if self.do_lower_case:
																__lowerCAmelCase           =	text.lower()

											for token in split_tokens:
																if token[:1] == "▁":
																					__lowerCAmelCase           =	token[1:]
																__lowerCAmelCase           =	text[offset:].index(_A      ) + offset
																__lowerCAmelCase           =	start + len(_A      )

																token_mapping.append((char_mapping[start], char_mapping[end - 1] + 1)      )
																__lowerCAmelCase           =	end
											return token_mapping






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return len(self.vocab      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return dict(self.vocab		,      **self.added_tokens_encoder      )






						def __getstate__(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.__dict__.copy()
											__lowerCAmelCase           =	None
											return state






						def __setstate__(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	d

											# for backward compatibility
											if not hasattr(self		,      "sp_model_kwargs"      ):
																__lowerCAmelCase           =	{}

											__lowerCAmelCase           =	spm.SentencePieceProcessor(**self.sp_model_kwargs      )
											self.sp_model.Load(self.sentencepiece_model_ckpt      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											return "".join((self.SP_CHAR_MAPPING.get(_A		,      _A      ) for c in text)      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A=False		,      _A=6_4		,      _A=0.1      ):




											"""simple docstring"""



											if self.sp_model_kwargs.get("enable_sampling"      ) is True:
																__lowerCAmelCase           =	True
											if self.sp_model_kwargs.get("alpha"      ) is not None:
																__lowerCAmelCase           =	self.sp_model_kwargs.get("alpha"      )
											if self.sp_model_kwargs.get("nbest_size"      ) is not None:
																__lowerCAmelCase           =	self.sp_model_kwargs.get("nbest_size"      )

											if not enable_sampling:
																__lowerCAmelCase           =	self.sp_model.EncodeAsPieces(_A      )
											else:
																__lowerCAmelCase           =	self.sp_model.SampleEncodeAsPieces(_A		,      _A		,      _A      )
											__lowerCAmelCase           =	[]
											for pi, piece in enumerate(_A      ):
																if piece == SPIECE_UNDERLINE:
																					if not pieces[pi + 1].startswith(_A      ) and pi != 0:
																										new_pieces.append(_A      )
																										continue
																					else:
																										continue
																__lowerCAmelCase           =	0
																for i, chunk in enumerate(_A      ):
																					if chunk == SPIECE_UNDERLINE:
																										continue
																					if self.is_ch_char(_A      ) or self.is_punct(_A      ):
																										if i > lst_i and piece[lst_i:i] != SPIECE_UNDERLINE:
																															new_pieces.append(piece[lst_i:i]      )
																										new_pieces.append(_A      )
																										__lowerCAmelCase           =	i + 1
																					elif chunk.isdigit() and i > 0 and not piece[i - 1].isdigit():
																										if i > lst_i and piece[lst_i:i] != SPIECE_UNDERLINE:
																															new_pieces.append(piece[lst_i:i]      )
																										__lowerCAmelCase           =	i
																					elif not chunk.isdigit() and i > 0 and piece[i - 1].isdigit():
																										if i > lst_i and piece[lst_i:i] != SPIECE_UNDERLINE:
																															new_pieces.append(piece[lst_i:i]      )
																										__lowerCAmelCase           =	i
																if len(_A      ) > lst_i:
																					new_pieces.append(piece[lst_i:]      )
											return new_pieces






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	"".join(_A      ).replace(_A		,      " "      ).strip()
											return out_string






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.convert_ids_to_tokens(_A      )
											__lowerCAmelCase           =	"".join(_A      ).replace(_A		,      " "      ).strip()
											return out_string






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											return self.vocab.get(_A		,      self.vocab.get(self.unk_token      )      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											return self.reverse_vocab.get(_A		,      self.unk_token      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A=None      ):




											"""simple docstring"""



											if token_ids_a is None:
																return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
											__lowerCAmelCase           =	[self.cls_token_id]
											__lowerCAmelCase           =	[self.sep_token_id]
											return _cls + token_ids_a + _sep + _sep + token_ids_a + _sep






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A=None      ):




											"""simple docstring"""



											if offset_mapping_a is None:
																return [(0, 0)] + offset_mapping_a + [(0, 0)]

											return [(0, 0)] + offset_mapping_a + [(0, 0), (0, 0)] + offset_mapping_a + [(0, 0)]






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A=None		,      _A=False      ):




											"""simple docstring"""



											if already_has_special_tokens:
																if token_ids_a is not None:
																					raise ValueError(
																					    "You should not supply a second sequence if the provided sequence of "
																					    "ids is already formatted with special tokens for the model."      )
																return [1 if x in [self.sep_token_id, self.cls_token_id] else 0 for x in token_ids_a]

											if token_ids_a is not None:
																return [1] + ([0] * len(_A      )) + [1, 1] + ([0] * len(_A      )) + [1]
											return [1] + ([0] * len(_A      )) + [1]






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A = None      ):




											"""simple docstring"""



											if token_ids_a is None:
																# [CLS] X [SEP]
																return (len(_A      ) + 2) * [0]

											# [CLS] A [SEP] [SEP] B [SEP]
											return [0] * (len(_A      ) + 1) + [1] * (len(_A      ) + 3)






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											if "\u4e00" <= char <= "\u9fff":
																return True
											return False






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											if ("a" <= char <= "z") or ("A" <= char <= "Z"):
																return True
											return False






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											if char in ",;:.?!~，；：。？！《》【】":
																return True
											return False






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											if char == " " or char == "\t" or char == "\n" or char == "\r":
																return True
											if len(_A      ) == 1:
																__lowerCAmelCase           =	unicodedata.category(_A      )
																if cat == "Zs":
																					return True
											return False






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	{}
											with io.open(_A		,      "r"		,      encoding="utf-8"      ) as f:
																for index, line in enumerate(_A      ):
																					__lowerCAmelCase           =	line.rstrip("\n"      )
																					__lowerCAmelCase           =	int(_A      )

											return token_to_idx






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A = None      ):




											"""simple docstring"""



											__lowerCAmelCase           =	0
											if os.path.isdir(_A      ):
																__lowerCAmelCase           =	os.path.join(
																    _A		,      (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]      )
											else:
																__lowerCAmelCase           =	(filename_prefix + "-" if filename_prefix else "") + save_directory
											with open(_A		,      "w"		,      encoding="utf-8"      ) as writer:
																for token, token_index in sorted(self.vocab.items()		,      key=lambda _A      : kv[1]      ):
																					if index != token_index:
																										logger.warning(
																										    f"""Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."""
																										    " Please check that the vocabulary is not corrupted!"      )
																										__lowerCAmelCase           =	token_index
																					writer.write(token + "\n"      )
																					index += 1

											__lowerCAmelCase           =	os.path.join(_A		,      "sentencepiece.bpe.model"      )
											with open(_A		,      "wb"      ) as fi:
																__lowerCAmelCase           =	self.sp_model.serialized_model_proto()
																fi.write(_A      )

											return (vocab_file,)










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)




