





"""simple docstring"""







import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
		from transformers.models.esm.modeling_esm import (
		    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    EsmEmbeddings,
		    create_position_ids_from_input_ids,
		)






class       SCREAMING_SNAKE_CASE__					:


						def __init__(							self				:			str						,						lowerCAmelCase				:			List[Any]						,						lowerCAmelCase				:			List[str]=13						,						lowerCAmelCase				:			Dict=7						,						lowerCAmelCase				:			Any=False						,						lowerCAmelCase				:			List[Any]=True						,						lowerCAmelCase				:			List[str]=False						,						lowerCAmelCase				:			Any=True						,						lowerCAmelCase				:			Optional[Any]=33						,						lowerCAmelCase				:			Any=32						,						lowerCAmelCase				:			Tuple=5						,						lowerCAmelCase				:			Optional[Any]=4						,						lowerCAmelCase				:			str=37						,						lowerCAmelCase				:			Dict="gelu"						,						lowerCAmelCase				:			Any=0.1						,						lowerCAmelCase				:			Tuple=0.1						,						lowerCAmelCase				:			Optional[int]=512						,						lowerCAmelCase				:			Tuple=16						,						lowerCAmelCase				:			str=2						,						lowerCAmelCase				:			List[str]=0.02						,						lowerCAmelCase				:			Any=3						,						lowerCAmelCase				:			Union[str, Any]=4						,						lowerCAmelCase				:			int=None						,						):
										lowerCAmelCase  =				parent
										lowerCAmelCase  =				batch_size
										lowerCAmelCase  =				seq_length
										lowerCAmelCase  =				is_training
										lowerCAmelCase  =				use_input_mask
										lowerCAmelCase  =				use_token_type_ids
										lowerCAmelCase  =				use_labels
										lowerCAmelCase  =				vocab_size
										lowerCAmelCase  =				hidden_size
										lowerCAmelCase  =				num_hidden_layers
										lowerCAmelCase  =				num_attention_heads
										lowerCAmelCase  =				intermediate_size
										lowerCAmelCase  =				hidden_act
										lowerCAmelCase  =				hidden_dropout_prob
										lowerCAmelCase  =				attention_probs_dropout_prob
										lowerCAmelCase  =				max_position_embeddings
										lowerCAmelCase  =				type_vocab_size
										lowerCAmelCase  =				type_sequence_label_size
										lowerCAmelCase  =				initializer_range
										lowerCAmelCase  =				num_labels
										lowerCAmelCase  =				num_choices
										lowerCAmelCase  =				scope


						def   __lowercase (							self				:			Optional[Any]					):
										lowerCAmelCase  =				ids_tensor([self.batch_size, self.seq_length]						,						self.vocab_size					)

										lowerCAmelCase  =				None
										if self.use_input_mask:
														lowerCAmelCase  =				random_attention_mask([self.batch_size, self.seq_length]					)

										lowerCAmelCase  =				None
										lowerCAmelCase  =				None
										lowerCAmelCase  =				None
										if self.use_labels:
														lowerCAmelCase  =				ids_tensor([self.batch_size]						,						self.type_sequence_label_size					)
														lowerCAmelCase  =				ids_tensor([self.batch_size, self.seq_length]						,						self.num_labels					)
														lowerCAmelCase  =				ids_tensor([self.batch_size]						,						self.num_choices					)

										lowerCAmelCase  =				self.get_config()

										return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels


						def   __lowercase (							self				:			Dict					):
										return EsmConfig(
										    vocab_size=self.vocab_size						,						hidden_size=self.hidden_size						,						pad_token_id=1						,						num_hidden_layers=self.num_hidden_layers						,						num_attention_heads=self.num_attention_heads						,						intermediate_size=self.intermediate_size						,						hidden_act=self.hidden_act						,						hidden_dropout_prob=self.hidden_dropout_prob						,						attention_probs_dropout_prob=self.attention_probs_dropout_prob						,						max_position_embeddings=self.max_position_embeddings						,						type_vocab_size=self.type_vocab_size						,						initializer_range=self.initializer_range						,						)


						def   __lowercase (							self				:			Any						,						lowerCAmelCase				:			Dict						,						lowerCAmelCase				:			List[str]						,						lowerCAmelCase				:			Any						,						lowerCAmelCase				:			Optional[Any]						,						lowerCAmelCase				:			List[Any]						,						lowerCAmelCase				:			Tuple					):
										lowerCAmelCase  =				EsmModel(config=lowerCAmelCase					)
										model.to(lowerCAmelCase					)
										model.eval()
										lowerCAmelCase  =				model(lowerCAmelCase						,						attention_mask=lowerCAmelCase					)
										lowerCAmelCase  =				model(lowerCAmelCase					)
										lowerCAmelCase  =				model(lowerCAmelCase					)

										self.parent.assertEqual(result.last_hidden_state.shape						,						(self.batch_size, self.seq_length, self.hidden_size)					)
										self.parent.assertEqual(result.pooler_output.shape						,						(self.batch_size, self.hidden_size)					)


						def   __lowercase (							self				:			Dict						,						lowerCAmelCase				:			List[Any]						,						lowerCAmelCase				:			Optional[Any]						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			Union[str, Any]						,						lowerCAmelCase				:			Optional[int]						,						lowerCAmelCase				:			int					):
										lowerCAmelCase  =				EsmForMaskedLM(config=lowerCAmelCase					)
										model.to(lowerCAmelCase					)
										model.eval()
										lowerCAmelCase  =				model(lowerCAmelCase						,						attention_mask=lowerCAmelCase						,						labels=lowerCAmelCase					)
										self.parent.assertEqual(result.logits.shape						,						(self.batch_size, self.seq_length, self.vocab_size)					)


						def   __lowercase (							self				:			List[str]						,						lowerCAmelCase				:			Optional[int]						,						lowerCAmelCase				:			Optional[int]						,						lowerCAmelCase				:			List[Any]						,						lowerCAmelCase				:			List[Any]						,						lowerCAmelCase				:			Dict						,						lowerCAmelCase				:			Union[str, Any]					):
										lowerCAmelCase  =				self.num_labels
										lowerCAmelCase  =				EsmForTokenClassification(config=lowerCAmelCase					)
										model.to(lowerCAmelCase					)
										model.eval()
										lowerCAmelCase  =				model(lowerCAmelCase						,						attention_mask=lowerCAmelCase						,						labels=lowerCAmelCase					)
										self.parent.assertEqual(result.logits.shape						,						(self.batch_size, self.seq_length, self.num_labels)					)




						def   __lowercase (							self				:			Tuple					):
										lowerCAmelCase  =				self.prepare_config_and_inputs()
										(
										    (
										    lowerCAmelCase						
										) ,						(
										    lowerCAmelCase						
										) ,						(
										    lowerCAmelCase						
										) ,						(
										    lowerCAmelCase						
										) ,						(
										    lowerCAmelCase						
										) ,						(
										    lowerCAmelCase						
										) ,												
										)  =				config_and_inputs
										lowerCAmelCase  =				{"""input_ids""": input_ids, """attention_mask""": input_mask}
										return config, inputs_dict






@require_torch
class       SCREAMING_SNAKE_CASE__					(					_a						,				_a						,				unittest.TestCase					):
						_a    						=						False

						_a    						=						(
						    (
						        EsmForMaskedLM,
						        EsmModel,
						        EsmForSequenceClassification,
						        EsmForTokenClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						_a    						=						()
						_a    						=						(
						    {
						        'feature-extraction': EsmModel,
						        'fill-mask': EsmForMaskedLM,
						        'text-classification': EsmForSequenceClassification,
						        'token-classification': EsmForTokenClassification,
						        'zero-shot': EsmForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						_a    						=						True


						def   __lowercase (							self				:			Dict					):
										lowerCAmelCase  =				EsmModelTester(self					)
										lowerCAmelCase  =				ConfigTester(self						,						config_class=lowerCAmelCase						,						hidden_size=37					)


						def   __lowercase (							self				:			Dict					):
										self.config_tester.run_common_tests()


						def   __lowercase (							self				:			Tuple					):
										lowerCAmelCase  =				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowerCAmelCase					)


						def   __lowercase (							self				:			str					):
										lowerCAmelCase  =				self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
														lowerCAmelCase  =				type
														self.model_tester.create_and_check_model(*lowerCAmelCase					)


						def   __lowercase (							self				:			Optional[int]					):
										lowerCAmelCase  =				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase					)


						def   __lowercase (							self				:			Optional[int]					):
										lowerCAmelCase  =				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase					)


						@slow
						def   __lowercase (							self				:			str					):
										for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														lowerCAmelCase  =				EsmModel.from_pretrained(lowerCAmelCase					)
														self.assertIsNotNone(lowerCAmelCase					)


						def   __lowercase (							self				:			Tuple					):
										lowerCAmelCase  =				self.model_tester.prepare_config_and_inputs()[0]
										lowerCAmelCase  =				EsmEmbeddings(config=lowerCAmelCase					)

										lowerCAmelCase  =				torch.as_tensor([[12, 31, 13, model.padding_idx]]					)
										lowerCAmelCase  =				torch.as_tensor(
										    [
										        [
										            0 + model.padding_idx + 1,
										            1 + model.padding_idx + 1,
										            2 + model.padding_idx + 1,
										            model.padding_idx,
										        ]
										    ]					)
										lowerCAmelCase  =				create_position_ids_from_input_ids(lowerCAmelCase						,						model.padding_idx					)
										self.assertEqual(position_ids.shape						,						expected_positions.shape					)
										self.assertTrue(torch.all(torch.eq(lowerCAmelCase						,						lowerCAmelCase					)					)					)


						def   __lowercase (							self				:			Dict					):
										lowerCAmelCase  =				self.model_tester.prepare_config_and_inputs()[0]
										lowerCAmelCase  =				EsmEmbeddings(config=lowerCAmelCase					)

										lowerCAmelCase  =				torch.empty(2						,						4						,						30					)
										lowerCAmelCase  =				[
										    0 + embeddings.padding_idx + 1,
										    1 + embeddings.padding_idx + 1,
										    2 + embeddings.padding_idx + 1,
										    3 + embeddings.padding_idx + 1,
										]
										lowerCAmelCase  =				torch.as_tensor([expected_single_positions, expected_single_positions]					)
										lowerCAmelCase  =				embeddings.create_position_ids_from_inputs_embeds(lowerCAmelCase					)
										self.assertEqual(position_ids.shape						,						expected_positions.shape					)
										self.assertTrue(torch.all(torch.eq(lowerCAmelCase						,						lowerCAmelCase					)					)					)


						@unittest.skip("""Esm does not support embedding resizing"""					)
						def   __lowercase (							self				:			List[str]					):
										pass


						@unittest.skip("""Esm does not support embedding resizing"""					)
						def   __lowercase (							self				:			int					):
										pass




						@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""					)
						def   __lowercase (							self				:			int					):
										pass






@require_torch
class       SCREAMING_SNAKE_CASE__					(					_a					):


						@slow
						def   __lowercase (							self				:			int					):
										with torch.no_grad():
														lowerCAmelCase  =				EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""					)
														model.eval()
														lowerCAmelCase  =				torch.tensor([[0, 1, 2, 3, 4, 5]]					)
														lowerCAmelCase  =				model(lowerCAmelCase					)[0]

														lowerCAmelCase  =				33

														lowerCAmelCase  =				torch.Size((1, 6, vocab_size)					)
														self.assertEqual(output.shape						,						lowerCAmelCase					)

														lowerCAmelCase  =				torch.tensor(
														    [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]					)
														self.assertTrue(torch.allclose(output[:, :3, :3]						,						lowerCAmelCase						,						atol=1e-4					)					)




						@slow
						def   __lowercase (							self				:			Optional[int]					):
										with torch.no_grad():
														lowerCAmelCase  =				EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""					)
														model.eval()

														lowerCAmelCase  =				torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]					)
														lowerCAmelCase  =				model(lowerCAmelCase					)[0]
														# compare the actual values for a slice.
														lowerCAmelCase  =				torch.tensor(
														    [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]					)
														self.assertTrue(torch.allclose(output[:, :3, :3]						,						lowerCAmelCase						,						atol=1e-4					)					)











"""simple docstring"""







import argparse
import json

import numpy
import torch

from transformers.models.xlm.tokenization_xlm import VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()





def 							lowercase					(snake_case__   : int      ,       snake_case__   : Union[str, Any] )     ->  List[Any]:


				'''simple docstring'''






				lowerCAmelCase  =				torch.load(snake_case__      ,       map_location="""cpu""" )

				lowerCAmelCase  =				chkpt["""model"""]

				# We have the base model one level deeper than the original XLM repository
				lowerCAmelCase  =				{}
				for k, v in state_dict.items():
								if "pred_layer" in k:
												lowerCAmelCase  =				v
								else:
												lowerCAmelCase  =				v

				lowerCAmelCase  =				chkpt["""params"""]
				lowerCAmelCase  =				{n: v for n, v in config.items() if not isinstance(snake_case__      ,       (torch.FloatTensor, numpy.ndarray) )}

				lowerCAmelCase  =				chkpt["""dico_word2id"""]
				lowerCAmelCase  =				{s + """</w>""" if s.find("""@@""" ) == -1 and i > 13 else s.replace("""@@"""      ,       """""" ): i for s, i in vocab.items()}

				# Save pytorch-model
				lowerCAmelCase  =				pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
				lowerCAmelCase  =				pytorch_dump_folder_path + """/""" + CONFIG_NAME
				lowerCAmelCase  =				pytorch_dump_folder_path + """/""" + VOCAB_FILES_NAMES["""vocab_file"""]

				print(f'''Save PyTorch model to {pytorch_weights_dump_path}''' )
				torch.save(snake_case__      ,       snake_case__ )

				print(f'''Save configuration file to {pytorch_config_dump_path}''' )
				with open(snake_case__      ,       """w"""      ,       encoding="""utf-8""" ) as f:
								f.write(json.dumps(snake_case__      ,       indent=2 ) + """\n""" )

				print(f'''Save vocab file to {pytorch_config_dump_path}''' )
				with open(snake_case__      ,       """w"""      ,       encoding="""utf-8""" ) as f:
								f.write(json.dumps(snake_case__      ,       indent=2 ) + """\n""" )


if __name__ == "__main__":
		a      =							argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '--xlm_checkpoint_path', default=None, type=str, required=True, help='Path the official PyTorch dump.'
		)
		parser.add_argument(
		    '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
		)
		a      =							parser.parse_args()
		convert_xlm_checkpoint_to_pytorch(args.xlm_checkpoint_path, args.pytorch_dump_folder_path)





