



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 					_snake_case      :




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




				def 				A__   (    self       :    Any    ):
							lowercase__       	=					ids_tensor([self.batch_size, self.seq_length],       self.vocab_size    )

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

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

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

							lowercase__       	=					self.get_config()

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




				def 				A__   (    self       :    Tuple    ):
							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 				A__   (    self       :    Dict,       __lowercase       :    Dict,       __lowercase       :    str,       __lowercase       :    int,       __lowercase       :    Optional[int],       __lowercase       :    str,       __lowercase       :    List[Any],       __lowercase       :    List[str]    ):
							lowercase__       	=					AlbertModel(config=snake_case__    )
							model.to(snake_case__    )
							model.eval()
							lowercase__       	=					model(snake_case__,       attention_mask=snake_case__,       token_type_ids=snake_case__    )
							lowercase__       	=					model(snake_case__,       token_type_ids=snake_case__    )
							lowercase__       	=					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 				A__   (    self       :    Union[str, Any],       __lowercase       :    Union[str, Any],       __lowercase       :    Union[str, Any],       __lowercase       :    str,       __lowercase       :    List[str],       __lowercase       :    str,       __lowercase       :    List[Any],       __lowercase       :    Optional[Any]    ):
							lowercase__       	=					AlbertForPreTraining(config=snake_case__    )
							model.to(snake_case__    )
							model.eval()
							lowercase__       	=					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 				A__   (    self       :    List[str],       __lowercase       :    List[str],       __lowercase       :    List[str],       __lowercase       :    Dict,       __lowercase       :    List[Any],       __lowercase       :    Tuple,       __lowercase       :    Dict,       __lowercase       :    Optional[Any]    ):
							lowercase__       	=					AlbertForMaskedLM(config=snake_case__    )
							model.to(snake_case__    )
							model.eval()
							lowercase__       	=					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 				A__   (    self       :    Any,       __lowercase       :    Optional[Any],       __lowercase       :    Tuple,       __lowercase       :    str,       __lowercase       :    Optional[Any],       __lowercase       :    List[Any],       __lowercase       :    Union[str, Any],       __lowercase       :    Optional[Any]    ):
							lowercase__       	=					AlbertForQuestionAnswering(config=snake_case__    )
							model.to(snake_case__    )
							model.eval()
							lowercase__       	=					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 				A__   (    self       :    str,       __lowercase       :    List[Any],       __lowercase       :    Optional[int],       __lowercase       :    List[str],       __lowercase       :    Union[str, Any],       __lowercase       :    Dict,       __lowercase       :    List[str],       __lowercase       :    Optional[Any]    ):
							lowercase__       	=					self.num_labels
							lowercase__       	=					AlbertForSequenceClassification(snake_case__    )
							model.to(snake_case__    )
							model.eval()
							lowercase__       	=					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 				A__   (    self       :    List[str],       __lowercase       :    Dict,       __lowercase       :    int,       __lowercase       :    List[str],       __lowercase       :    Dict,       __lowercase       :    Dict,       __lowercase       :    Dict,       __lowercase       :    Tuple    ):
							lowercase__       	=					self.num_labels
							lowercase__       	=					AlbertForTokenClassification(config=snake_case__    )
							model.to(snake_case__    )
							model.eval()
							lowercase__       	=					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 				A__   (    self       :    List[str],       __lowercase       :    Union[str, Any],       __lowercase       :    List[str],       __lowercase       :    List[Any],       __lowercase       :    int,       __lowercase       :    Union[str, Any],       __lowercase       :    List[Any],       __lowercase       :    Optional[int]    ):
							lowercase__       	=					self.num_choices
							lowercase__       	=					AlbertForMultipleChoice(config=snake_case__    )
							model.to(snake_case__    )
							model.eval()
							lowercase__       	=					input_ids.unsqueeze(1    ).expand(-1,       self.num_choices,       -1    ).contiguous()
							lowercase__       	=					token_type_ids.unsqueeze(1    ).expand(-1,       self.num_choices,       -1    ).contiguous()
							lowercase__       	=					input_mask.unsqueeze(1    ).expand(-1,       self.num_choices,       -1    ).contiguous()
							lowercase__       	=					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 				A__   (    self       :    Union[str, Any]    ):
							lowercase__       	=					self.prepare_config_and_inputs()
							(
							    (
							    lowercase__				
							)	,							(
							    lowercase__				
							)	,							(
							    lowercase__				
							)	,							(
							    lowercase__				
							)	,							(
							    lowercase__				
							)	,							(
							    lowercase__				
							)	,							(
							    lowercase__				
							)	,											
							)       	=					config_and_inputs
							lowercase__       	=					{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
							return config, inputs_dict
@require_torch
class 					_snake_case      (       SCREAMING_SNAKE_CASE__     ,						SCREAMING_SNAKE_CASE__     ,						unittest.TestCase):
				UpperCamelCase__					:       Any  		=(
				    (
				        AlbertModel,
				        AlbertForPreTraining,
				        AlbertForMaskedLM,
				        AlbertForMultipleChoice,
				        AlbertForSequenceClassification,
				        AlbertForTokenClassification,
				        AlbertForQuestionAnswering,
				    )
				    if is_torch_available()
				    else ()
				)
				UpperCamelCase__					:       int  		=(
				    {
				        """feature-extraction""": AlbertModel,
				        """fill-mask""": AlbertForMaskedLM,
				        """question-answering""": AlbertForQuestionAnswering,
				        """text-classification""": AlbertForSequenceClassification,
				        """token-classification""": AlbertForTokenClassification,
				        """zero-shot""": AlbertForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				UpperCamelCase__					:       Union[str, Any]  		=True




				def 				A__   (    self       :    str,       __lowercase       :    Any,       __lowercase       :    List[Any],       __lowercase       :    Union[str, Any]=False    ):
							lowercase__       	=					super()._prepare_for_class(snake_case__,       snake_case__,       return_labels=snake_case__    )

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




				def 				A__   (    self       :    Tuple    ):
							lowercase__       	=					AlbertModelTester(self    )
							lowercase__       	=					ConfigTester(self,       config_class=snake_case__,       hidden_size=37    )




				def 				A__   (    self       :    Optional[Any]    ):
							self.config_tester.run_common_tests()




				def 				A__   (    self       :    Optional[Any]    ):
							lowercase__       	=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*snake_case__    )




				def 				A__   (    self       :    Optional[Any]    ):
							lowercase__       	=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_pretraining(*snake_case__    )




				def 				A__   (    self       :    Optional[Any]    ):
							lowercase__       	=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*snake_case__    )




				def 				A__   (    self       :    Any    ):
							lowercase__       	=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*snake_case__    )




				def 				A__   (    self       :    Dict    ):
							lowercase__       	=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*snake_case__    )




				def 				A__   (    self       :    str    ):
							lowercase__       	=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*snake_case__    )




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




				@slow
				def 				A__   (    self       :    Optional[int]    ):

							for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowercase__       	=					AlbertModel.from_pretrained(snake_case__    )
										self.assertIsNotNone(snake_case__    )
@require_torch
class 					_snake_case      (       unittest.TestCase):




				@slow
				def 				A__   (    self       :    Optional[int]    ):
							lowercase__       	=					AlbertModel.from_pretrained("albert-base-v2"    )
							lowercase__       	=					torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]    )
							lowercase__       	=					torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
							with torch.no_grad():
										lowercase__       	=					model(snake_case__,       attention_mask=snake_case__    )[0]
							lowercase__       	=					torch.Size((1, 11, 768)    )
							self.assertEqual(output.shape,       snake_case__    )
							lowercase__       	=					torch.tensor(
							    [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]    )

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










import argparse
import os
from pathlib import Path

import fairseq
import torch
from packaging import version
from torch import nn

from transformers import (
    BartConfig,
    BartForConditionalGeneration,
    BartForSequenceClassification,
    BartModel,
    BartTokenizer,
)
from transformers.utils import logging


SCREAMING_SNAKE_CASE:     int									= ['''bart.large''', '''bart.large.mnli''', '''bart.large.cnn''', '''bart_xsum/model.pt''']
SCREAMING_SNAKE_CASE:     Tuple									= {'''bart.large''': BartModel, '''bart.large.mnli''': BartForSequenceClassification}
if version.parse(fairseq.__version__) < version.parse('''0.9.0'''):
						raise Exception('''requires fairseq >= 0.9.0''')


logging.set_verbosity_info()
SCREAMING_SNAKE_CASE:     int									= logging.get_logger(__name__)

SCREAMING_SNAKE_CASE:     Any									= ''' Hello world! cécé herlolip'''

SCREAMING_SNAKE_CASE:     Dict									= [
    ('''model.classification_heads.mnli.dense.weight''', '''classification_head.dense.weight'''),
    ('''model.classification_heads.mnli.dense.bias''', '''classification_head.dense.bias'''),
    ('''model.classification_heads.mnli.out_proj.weight''', '''classification_head.out_proj.weight'''),
    ('''model.classification_heads.mnli.out_proj.bias''', '''classification_head.out_proj.bias'''),
]


def 					_a							(					lowerCAmelCase )->	Any:
		SCREAMING_SNAKE_CASE_					      =							[
		    'encoder.version',
		    'decoder.version',
		    'model.encoder.version',
		    'model.decoder.version',
		    '_float_tensor',
		]
		for k in ignore_keys:
				state_dict.pop(lowerCAmelCase    ,						lowerCAmelCase )


def 					_a							(					lowerCAmelCase    ,						lowerCAmelCase    ,						lowerCAmelCase )->	str:
		SCREAMING_SNAKE_CASE_					      =							dct.pop(lowerCAmelCase )
		SCREAMING_SNAKE_CASE_					      =							val


def 					_a							(					lowerCAmelCase )->	List[str]:
		SCREAMING_SNAKE_CASE_					      =							torch.load(lowerCAmelCase    ,						map_location='cpu' )
		SCREAMING_SNAKE_CASE_					      =							torch.hub.load('pytorch/fairseq'    ,						'bart.large.cnn' ).eval()
		hub_interface.model.load_state_dict(sd['model'] )
		return hub_interface


def 					_a							(					lowerCAmelCase )->	Optional[Any]:
		SCREAMING_SNAKE_CASE_      ,				SCREAMING_SNAKE_CASE_					      =							emb.weight.shape
		SCREAMING_SNAKE_CASE_					      =							nn.Linear(lowerCAmelCase    ,						lowerCAmelCase    ,						bias=lowerCAmelCase )
		SCREAMING_SNAKE_CASE_					      =							emb.weight.data
		return lin_layer


@torch.no_grad()
def 					_a							(					lowerCAmelCase    ,						lowerCAmelCase    ,						lowerCAmelCase=None )->	Tuple:
		if not os.path.exists(lowerCAmelCase ):
				SCREAMING_SNAKE_CASE_					      =							torch.hub.load('pytorch/fairseq'    ,						lowerCAmelCase ).eval()
		else:
				SCREAMING_SNAKE_CASE_					      =							load_xsum_checkpoint(lowerCAmelCase )

		bart.model.upgrade_state_dict(bart.model.state_dict() )
		if hf_checkpoint_name is None:
				SCREAMING_SNAKE_CASE_					      =							checkpoint_path.replace('.'    ,						'-' )
		SCREAMING_SNAKE_CASE_					      =							BartConfig.from_pretrained(lowerCAmelCase )
		SCREAMING_SNAKE_CASE_					      =							bart.encode(lowerCAmelCase ).unsqueeze(0 )
		SCREAMING_SNAKE_CASE_					      =							BartTokenizer.from_pretrained(lowerCAmelCase ).encode(lowerCAmelCase    ,						return_tensors='pt' ).unsqueeze(0 )
		if not torch.eq(lowerCAmelCase    ,						lowerCAmelCase ).all():
				raise ValueError(
				    F'''converted tokenizer and pretrained tokenizer returned different output: {tokens} != {tokensa}''' )

		if checkpoint_path == "bart.large.mnli":
				SCREAMING_SNAKE_CASE_					      =							bart.state_dict()
				remove_ignore_keys_(lowerCAmelCase )
				SCREAMING_SNAKE_CASE_					      =							state_dict['model.decoder.embed_tokens.weight']
				for src, dest in mnli_rename_keys:
						rename_key(lowerCAmelCase    ,						lowerCAmelCase    ,						lowerCAmelCase )
				SCREAMING_SNAKE_CASE_					      =							BartForSequenceClassification(lowerCAmelCase ).eval()
				model.load_state_dict(lowerCAmelCase )
				SCREAMING_SNAKE_CASE_					      =							bart.predict('mnli'    ,						lowerCAmelCase    ,						return_logits=lowerCAmelCase )
				SCREAMING_SNAKE_CASE_					      =							model(lowerCAmelCase )[0]  # logits
		else:  # no classification heads to worry about
				SCREAMING_SNAKE_CASE_					      =							bart.model.state_dict()
				remove_ignore_keys_(lowerCAmelCase )
				SCREAMING_SNAKE_CASE_					      =							state_dict['decoder.embed_tokens.weight']
				SCREAMING_SNAKE_CASE_					      =							bart.extract_features(lowerCAmelCase )
				if hf_checkpoint_name == "facebook/bart-large":
						SCREAMING_SNAKE_CASE_					      =							BartModel(lowerCAmelCase ).eval()
						model.load_state_dict(lowerCAmelCase )
						SCREAMING_SNAKE_CASE_					      =							model(lowerCAmelCase ).model[0]
				else:
						SCREAMING_SNAKE_CASE_					      =							BartForConditionalGeneration(lowerCAmelCase ).eval()  # an existing summarization ckpt
						model.model.load_state_dict(lowerCAmelCase )
						if hasattr(lowerCAmelCase    ,						'lm_head' ):
								SCREAMING_SNAKE_CASE_					      =							make_linear_from_emb(model.model.shared )
						SCREAMING_SNAKE_CASE_					      =							model.model(lowerCAmelCase )[0]

    # Check results
		if fairseq_output.shape != new_model_outputs.shape:
				raise ValueError(
				    F'''`fairseq_output` shape and `new_model_output` shape are different: {fairseq_output.shape=}, {new_model_outputs.shape}''' )
		if (fairseq_output != new_model_outputs).any().item():
				raise ValueError('Some values in `fairseq_output` are different from `new_model_outputs`' )
		Path(lowerCAmelCase ).mkdir(exist_ok=lowerCAmelCase )
		model.save_pretrained(lowerCAmelCase )


if __name__ == "__main__":
						SCREAMING_SNAKE_CASE:     Dict									= 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=None, type=str, help='''Which huggingface architecture to use: bart-large-xsum'''
						)
						SCREAMING_SNAKE_CASE:     List[str]									= parser.parse_args()
						convert_bart_checkpoint(args.fairseq_path, args.pytorch_dump_folder_path, hf_checkpoint_name=args.hf_config)