


"""simple docstring"""



import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import 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 (
     MODEL_FOR_PRETRAINING_MAPPING,
     AlbertForMaskedLM,
     AlbertForMultipleChoice,
     AlbertForPreTraining,
     AlbertForQuestionAnswering,
     AlbertForSequenceClassification,
     AlbertForTokenClassification,
     AlbertModel,
 )
 from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST
class 					__magic_name__		:






       def __init__(   self  :    Optional[Any]  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    Dict=1_3  ,     snake_case__  :    Tuple=7  ,     snake_case__  :    List[Any]=True  ,     snake_case__  :    Any=True  ,     snake_case__  :    str=True  ,     snake_case__  :    int=True  ,     snake_case__  :    Tuple=9_9  ,     snake_case__  :    Union[str, Any]=1_6  ,     snake_case__  :    str=3_6  ,     snake_case__  :    List[str]=6  ,     snake_case__  :    List[Any]=6  ,     snake_case__  :    Dict=6  ,     snake_case__  :    str=3_7  ,     snake_case__  :    Dict="gelu"  ,     snake_case__  :    Optional[int]=0.1  ,     snake_case__  :    Any=0.1  ,     snake_case__  :    List[str]=5_1_2  ,     snake_case__  :    Optional[Any]=1_6  ,     snake_case__  :    Tuple=2  ,     snake_case__  :    Tuple=0.02  ,     snake_case__  :    List[str]=3  ,     snake_case__  :    Optional[Any]=4  ,     snake_case__  :    Any=None  ,     ):

              '''simple docstring'''






              lowercase		:Optional[int]       							=				parent
              lowercase		:int       							=				batch_size
              lowercase		:Optional[int]       							=				seq_length
              lowercase		:str       							=				is_training
              lowercase		:Optional[Any]       							=				use_input_mask
              lowercase		:List[str]       							=				use_token_type_ids
              lowercase		:Optional[int]       							=				use_labels
              lowercase		:Tuple       							=				vocab_size
              lowercase		:List[str]       							=				embedding_size
              lowercase		:Tuple       							=				hidden_size
              lowercase		:List[Any]       							=				num_hidden_layers
              lowercase		:Dict       							=				num_hidden_groups
              lowercase		:List[Any]       							=				num_attention_heads
              lowercase		:str       							=				intermediate_size
              lowercase		:List[Any]       							=				hidden_act
              lowercase		:Optional[int]       							=				hidden_dropout_prob
              lowercase		:Any       							=				attention_probs_dropout_prob
              lowercase		:Optional[int]       							=				max_position_embeddings
              lowercase		:Union[str, Any]       							=				type_vocab_size
              lowercase		:Union[str, Any]       							=				type_sequence_label_size
              lowercase		:Tuple       							=				initializer_range
              lowercase		:List[Any]       							=				num_labels
              lowercase		:Union[str, Any]       							=				num_choices
              lowercase		:Dict       							=				scope






       def  __snake_case						(   self  :    Any				):

              '''simple docstring'''






              lowercase		:List[Any]       							=				ids_tensor([self.batch_size, self.seq_length]  ,     self.vocab_size				)

              lowercase		:Dict       							=				None
              if self.use_input_mask:
                     lowercase		:Any       							=				random_attention_mask([self.batch_size, self.seq_length]				)

              lowercase		:Optional[Any]       							=				None
              if self.use_token_type_ids:
                     lowercase		:Any       							=				ids_tensor([self.batch_size, self.seq_length]  ,     self.type_vocab_size				)

              lowercase		:List[str]       							=				None
              lowercase		:Tuple       							=				None
              lowercase		:Union[str, Any]       							=				None
              if self.use_labels:
                     lowercase		:Dict       							=				ids_tensor([self.batch_size]  ,     self.type_sequence_label_size				)
                     lowercase		:Tuple       							=				ids_tensor([self.batch_size, self.seq_length]  ,     self.num_labels				)
                     lowercase		:str       							=				ids_tensor([self.batch_size]  ,     self.num_choices				)

              lowercase		:Optional[Any]       							=				self.get_config()

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






       def  __snake_case						(   self  :    Optional[int]				):

              '''simple docstring'''






              return AlbertConfig(
                  vocab_size=self.vocab_size  ,     hidden_size=self.hidden_size  ,     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  ,     num_hidden_groups=self.num_hidden_groups  ,     )






       def  __snake_case						(   self  :    int  ,     snake_case__  :    Optional[Any]  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    int  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    str  ,     snake_case__  :    Dict  ,     snake_case__  :    Optional[int]				):

              '''simple docstring'''






              lowercase		:str       							=				AlbertModel(config=snake_case__				)
              model.to(snake_case__				)
              model.eval()
              lowercase		:Dict       							=				model(snake_case__  ,     attention_mask=snake_case__  ,     token_type_ids=snake_case__				)
              lowercase		:Tuple       							=				model(snake_case__  ,     token_type_ids=snake_case__				)
              lowercase		:Union[str, Any]       							=				model(snake_case__				)
              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  __snake_case						(   self  :    Optional[Any]  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    int  ,     snake_case__  :    Tuple  ,     snake_case__  :    Optional[Any]  ,     snake_case__  :    Union[str, Any]  ,     snake_case__  :    Optional[int]				):

              '''simple docstring'''






              lowercase		:Union[str, Any]       							=				AlbertForPreTraining(config=snake_case__				)
              model.to(snake_case__				)
              model.eval()
              lowercase		:Optional[Any]       							=				model(
                  snake_case__  ,     attention_mask=snake_case__  ,     token_type_ids=snake_case__  ,     labels=snake_case__  ,     sentence_order_label=snake_case__  ,     )
              self.parent.assertEqual(result.prediction_logits.shape  ,     (self.batch_size, self.seq_length, self.vocab_size)				)
              self.parent.assertEqual(result.sop_logits.shape  ,     (self.batch_size, config.num_labels)				)






       def  __snake_case						(   self  :    Tuple  ,     snake_case__  :    Union[str, Any]  ,     snake_case__  :    Optional[Any]  ,     snake_case__  :    int  ,     snake_case__  :    Optional[Any]  ,     snake_case__  :    int  ,     snake_case__  :    Union[str, Any]  ,     snake_case__  :    List[Any]				):

              '''simple docstring'''






              lowercase		:Optional[int]       							=				AlbertForMaskedLM(config=snake_case__				)
              model.to(snake_case__				)
              model.eval()
              lowercase		:Union[str, Any]       							=				model(snake_case__  ,     attention_mask=snake_case__  ,     token_type_ids=snake_case__  ,     labels=snake_case__				)
              self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.seq_length, self.vocab_size)				)






       def  __snake_case						(   self  :    Optional[int]  ,     snake_case__  :    Dict  ,     snake_case__  :    Dict  ,     snake_case__  :    List[Any]  ,     snake_case__  :    int  ,     snake_case__  :    List[str]  ,     snake_case__  :    Union[str, Any]  ,     snake_case__  :    Optional[Any]				):

              '''simple docstring'''






              lowercase		:Union[str, Any]       							=				AlbertForQuestionAnswering(config=snake_case__				)
              model.to(snake_case__				)
              model.eval()
              lowercase		:int       							=				model(
                  snake_case__  ,     attention_mask=snake_case__  ,     token_type_ids=snake_case__  ,     start_positions=snake_case__  ,     end_positions=snake_case__  ,     )
              self.parent.assertEqual(result.start_logits.shape  ,     (self.batch_size, self.seq_length)				)
              self.parent.assertEqual(result.end_logits.shape  ,     (self.batch_size, self.seq_length)				)






       def  __snake_case						(   self  :    Optional[int]  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    int  ,     snake_case__  :    List[str]  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    Dict  ,     snake_case__  :    str  ,     snake_case__  :    Optional[Any]				):

              '''simple docstring'''






              lowercase		:Union[str, Any]       							=				self.num_labels
              lowercase		:Optional[Any]       							=				AlbertForSequenceClassification(snake_case__				)
              model.to(snake_case__				)
              model.eval()
              lowercase		:List[str]       							=				model(snake_case__  ,     attention_mask=snake_case__  ,     token_type_ids=snake_case__  ,     labels=snake_case__				)
              self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.num_labels)				)






       def  __snake_case						(   self  :    str  ,     snake_case__  :    List[Any]  ,     snake_case__  :    Dict  ,     snake_case__  :    List[str]  ,     snake_case__  :    Tuple  ,     snake_case__  :    Tuple  ,     snake_case__  :    Union[str, Any]  ,     snake_case__  :    List[Any]				):

              '''simple docstring'''






              lowercase		:int       							=				self.num_labels
              lowercase		:Union[str, Any]       							=				AlbertForTokenClassification(config=snake_case__				)
              model.to(snake_case__				)
              model.eval()
              lowercase		:Optional[Any]       							=				model(snake_case__  ,     attention_mask=snake_case__  ,     token_type_ids=snake_case__  ,     labels=snake_case__				)
              self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.seq_length, self.num_labels)				)






       def  __snake_case						(   self  :    Optional[Any]  ,     snake_case__  :    Optional[Any]  ,     snake_case__  :    Tuple  ,     snake_case__  :    List[str]  ,     snake_case__  :    Optional[Any]  ,     snake_case__  :    Tuple  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    List[str]				):

              '''simple docstring'''






              lowercase		:Optional[Any]       							=				self.num_choices
              lowercase		:Dict       							=				AlbertForMultipleChoice(config=snake_case__				)
              model.to(snake_case__				)
              model.eval()
              lowercase		:Dict       							=				input_ids.unsqueeze(1				).expand(-1  ,     self.num_choices  ,     -1				).contiguous()
              lowercase		:Optional[int]       							=				token_type_ids.unsqueeze(1				).expand(-1  ,     self.num_choices  ,     -1				).contiguous()
              lowercase		:Union[str, Any]       							=				input_mask.unsqueeze(1				).expand(-1  ,     self.num_choices  ,     -1				).contiguous()
              lowercase		:int       							=				model(
                  snake_case__  ,     attention_mask=snake_case__  ,     token_type_ids=snake_case__  ,     labels=snake_case__  ,     )
              self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.num_choices)				)







       def  __snake_case						(   self  :    Union[str, Any]				):

              '''simple docstring'''






              lowercase		:str       							=				self.prepare_config_and_inputs()
              (
                  (
                  lowercase
              )							,  (
                  lowercase
              )							,  (
                  lowercase
              )							,  (
                  lowercase
              )							,  (
                  lowercase
              )							,  (
                  lowercase
              )							,  (
                  lowercase
              )							,  
              )		:Optional[int]       							=				config_and_inputs
              lowercase		:List[str]       							=				{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
              return config, inputs_dict







@require_torch
class 					__magic_name__		(			__UpperCAmelCase			,						__UpperCAmelCase			,						unittest.TestCase     ):
       __A			:      List[Any]        =   (
           (
               AlbertModel,
               AlbertForPreTraining,
               AlbertForMaskedLM,
               AlbertForMultipleChoice,
               AlbertForSequenceClassification,
               AlbertForTokenClassification,
               AlbertForQuestionAnswering,
           )
           if is_torch_available()
           else ()
       )
       __A			:      Optional[Any]        =   (
           {
               "feature-extraction": AlbertModel,
               "fill-mask": AlbertForMaskedLM,
               "question-answering": AlbertForQuestionAnswering,
               "text-classification": AlbertForSequenceClassification,
               "token-classification": AlbertForTokenClassification,
               "zero-shot": AlbertForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       __A			:      List[str]        =   True






       def  __snake_case						(   self  :    Tuple  ,     snake_case__  :    List[str]  ,     snake_case__  :    List[Any]  ,     snake_case__  :    Any=False				):

              '''simple docstring'''






              lowercase		:Any       							=				super()._prepare_for_class(snake_case__  ,     snake_case__  ,     return_labels=snake_case__				)

              if return_labels:
                     if model_class in get_values(snake_case__				):
                            lowercase		:Any       							=				torch.zeros(
                                (self.model_tester.batch_size, self.model_tester.seq_length)  ,     dtype=torch.long  ,     device=snake_case__				)
                            lowercase		:int       							=				torch.zeros(
                                self.model_tester.batch_size  ,     dtype=torch.long  ,     device=snake_case__				)
              return inputs_dict






       def  __snake_case						(   self  :    Any				):

              '''simple docstring'''






              lowercase		:Dict       							=				AlbertModelTester(self				)
              lowercase		:str       							=				ConfigTester(self  ,     config_class=snake_case__  ,     hidden_size=3_7				)






       def  __snake_case						(   self  :    Any				):

              '''simple docstring'''






              self.config_tester.run_common_tests()






       def  __snake_case						(   self  :    Union[str, Any]				):

              '''simple docstring'''






              lowercase		:Optional[int]       							=				self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_model(*snake_case__				)






       def  __snake_case						(   self  :    List[Any]				):

              '''simple docstring'''






              lowercase		:Tuple       							=				self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_for_pretraining(*snake_case__				)






       def  __snake_case						(   self  :    Tuple				):

              '''simple docstring'''






              lowercase		:Any       							=				self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_for_masked_lm(*snake_case__				)






       def  __snake_case						(   self  :    Optional[int]				):

              '''simple docstring'''






              lowercase		:List[str]       							=				self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_for_multiple_choice(*snake_case__				)






       def  __snake_case						(   self  :    List[str]				):

              '''simple docstring'''






              lowercase		:int       							=				self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_for_question_answering(*snake_case__				)






       def  __snake_case						(   self  :    str				):

              '''simple docstring'''






              lowercase		:List[str]       							=				self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_for_sequence_classification(*snake_case__				)






       def  __snake_case						(   self  :    Dict				):

              '''simple docstring'''






              lowercase		:Optional[Any]       							=				self.model_tester.prepare_config_and_inputs()
              for type in ["absolute", "relative_key", "relative_key_query"]:
                     lowercase		:Optional[int]       							=				type
                     self.model_tester.create_and_check_model(*snake_case__				)







       @slow
       def  __snake_case						(   self  :    Tuple				):

              '''simple docstring'''






              for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                     lowercase		:Union[str, Any]       							=				AlbertModel.from_pretrained(snake_case__				)
                     self.assertIsNotNone(snake_case__				)







@require_torch
class 					__magic_name__		(			unittest.TestCase     ):






       @slow
       def  __snake_case						(   self  :    Union[str, Any]				):

              '''simple docstring'''






              lowercase		:Dict       							=				AlbertModel.from_pretrained('''albert-base-v2'''				)
              lowercase		:Union[str, Any]       							=				torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]				)
              lowercase		:Dict       							=				torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]				)
              with torch.no_grad():
                     lowercase		:List[str]       							=				model(snake_case__  ,     attention_mask=snake_case__				)[0]
              lowercase		:Union[str, Any]       							=				torch.Size((1, 1_1, 7_6_8)				)
              self.assertEqual(output.shape  ,     snake_case__				)
              lowercase		:int       							=				torch.tensor(
                  [[[-0.65_13, 1.50_35, -0.27_66], [-0.65_15, 1.50_46, -0.27_80], [-0.65_12, 1.50_49, -0.27_84]]]				)

              self.assertTrue(torch.allclose(output[:, 1:4, 1:4]  ,     snake_case__  ,     atol=1e-4				)				)





"""simple docstring"""



import argparse

import torch
from torch import nn

from transformers import MBartConfig, MBartForConditionalGeneration






def       lowerCamelCase					(a_						:Dict)     ->							Dict:
       lowercase		:Tuple       							=				[
           '''encoder.version''',
           '''decoder.version''',
           '''model.encoder.version''',
           '''model.decoder.version''',
           '''_float_tensor''',
           '''decoder.output_projection.weight''',
       ]
       for k in ignore_keys:
              state_dict.pop(a_				,				a_)






def       lowerCamelCase					(a_						:Union[str, Any])     ->							str:
       lowercase							,  lowercase		:Tuple       							=				emb.weight.shape
       lowercase		:List[str]       							=				nn.Linear(a_				,				a_				,				bias=a_)
       lowercase		:List[str]       							=				emb.weight.data
       return lin_layer






def       lowerCamelCase					(a_						:int				,				a_						:Union[str, Any]="facebook/mbart-large-en-ro"				,				a_						:Union[str, Any]=False				,				a_						:List[Any]=False)     ->							List[Any]:
       lowercase		:List[Any]       							=				torch.load(a_				,				map_location='''cpu''')['''model''']
       remove_ignore_keys_(a_)
       lowercase		:Dict       							=				state_dict['''encoder.embed_tokens.weight'''].shape[0]

       lowercase		:Tuple       							=				MBartConfig.from_pretrained(a_				,				vocab_size=a_)
       if mbart_aa and finetuned:
              lowercase		:List[Any]       							=				'''relu'''

       lowercase		:Optional[int]       							=				state_dict['''decoder.embed_tokens.weight''']
       lowercase		:Union[str, Any]       							=				MBartForConditionalGeneration(a_)
       model.model.load_state_dict(a_)

       if finetuned:
              lowercase		:Dict       							=				make_linear_from_emb(model.model.shared)

       return model


if __name__ == "__main__":
 UpperCAmelCase      		=			argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''fairseq_path''', type=str, help='''bart.large, bart.large.cnn or a path to a model.pt on local filesystem.'''
 )
 parser.add_argument('''pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
 parser.add_argument(
     '''--hf_config''',
     default='''facebook/mbart-large-cc25''',
     type=str,
     help='''Which huggingface architecture to use: mbart-large''',
 )
 parser.add_argument('''--mbart_50''', action='''store_true''', help='''whether the model is mMART-50 checkpoint''')
 parser.add_argument('''--finetuned''', action='''store_true''', help='''whether the model is a fine-tuned checkpoint''')
 UpperCAmelCase      		=			parser.parse_args()
 UpperCAmelCase      		=			convert_fairseq_mbart_checkpoint_from_disk(
     args.fairseq_path, hf_config_path=args.hf_config, finetuned=args.finetuned, mbart_aa=args.mbart_aa
 )
 model.save_pretrained(args.pytorch_dump_folder_path)

