

'''simple docstring'''
import unittest

from transformers import MraConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
     import torch

     from transformers import (
         MraForMaskedLM,
         MraForMultipleChoice,
         MraForQuestionAnswering,
         MraForSequenceClassification,
         MraForTokenClassification,
         MraModel,
     )
     from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST







class 							A__    :





  def __init__(		self      :	List[Any]							,					_a      :	List[str]							,					_a      :	Optional[Any]=2							,					_a      :	List[Any]=8							,					_a      :	Tuple=True							,					_a      :	Dict=True							,					_a      :	Any=True							,					_a      :	int=True							,					_a      :	List[str]=99							,					_a      :	int=16							,					_a      :	Union[str, Any]=5							,					_a      :	Optional[int]=2							,					_a      :	Optional[Any]=36							,					_a      :	List[str]="gelu"							,					_a      :	Any=0.0							,					_a      :	str=0.0							,					_a      :	Optional[Any]=512							,					_a      :	Tuple=16							,					_a      :	Optional[int]=2							,					_a      :	int=0.02							,					_a      :	int=3							,					_a      :	Optional[Any]=4							,					_a      :	Dict=None							,					)	->    str:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =parent
       _SCREAMING_SNAKE_CASE							    =batch_size
       _SCREAMING_SNAKE_CASE							    =seq_length
       _SCREAMING_SNAKE_CASE							    =is_training
       _SCREAMING_SNAKE_CASE							    =use_input_mask
       _SCREAMING_SNAKE_CASE							    =use_token_type_ids
       _SCREAMING_SNAKE_CASE							    =use_labels
       _SCREAMING_SNAKE_CASE							    =vocab_size
       _SCREAMING_SNAKE_CASE							    =hidden_size
       _SCREAMING_SNAKE_CASE							    =num_hidden_layers
       _SCREAMING_SNAKE_CASE							    =num_attention_heads
       _SCREAMING_SNAKE_CASE							    =intermediate_size
       _SCREAMING_SNAKE_CASE							    =hidden_act
       _SCREAMING_SNAKE_CASE							    =hidden_dropout_prob
       _SCREAMING_SNAKE_CASE							    =attention_probs_dropout_prob
       _SCREAMING_SNAKE_CASE							    =max_position_embeddings
       _SCREAMING_SNAKE_CASE							    =type_vocab_size
       _SCREAMING_SNAKE_CASE							    =type_sequence_label_size
       _SCREAMING_SNAKE_CASE							    =initializer_range
       _SCREAMING_SNAKE_CASE							    =num_labels
       _SCREAMING_SNAKE_CASE							    =num_choices
       _SCREAMING_SNAKE_CASE							    =scope





  def 						A  (		self      :	Any )	->    str:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =ids_tensor([self.batch_size, self.seq_length]							,					self.vocab_size )

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

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

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

       _SCREAMING_SNAKE_CASE							    =self.get_config()

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





  def 						A  (		self      :	int )	->    Tuple:





       '''simple docstring'''






       return MraConfig(
           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							,					is_decoder=_a							,					initializer_range=self.initializer_range							,					)





  def 						A  (		self      :	Dict )	->    str:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.get_config()
       _SCREAMING_SNAKE_CASE							    =300
       return config





  def 						A  (		self      :	str )	->    List[str]:





       '''simple docstring'''






       (
           (
           _SCREAMING_SNAKE_CASE
       )    ,						(
           _SCREAMING_SNAKE_CASE
       )    ,						(
           _SCREAMING_SNAKE_CASE
       )    ,						(
           _SCREAMING_SNAKE_CASE
       )    ,						(
           _SCREAMING_SNAKE_CASE
       )    ,						(
           _SCREAMING_SNAKE_CASE
       )    ,						(
           _SCREAMING_SNAKE_CASE
       )    ,						
       )							    =self.prepare_config_and_inputs()

       _SCREAMING_SNAKE_CASE							    =True
       _SCREAMING_SNAKE_CASE							    =floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
       _SCREAMING_SNAKE_CASE							    =ids_tensor([self.batch_size, self.seq_length]							,					vocab_size=2 )

       return (
           config,
           input_ids,
           token_type_ids,
           input_mask,
           sequence_labels,
           token_labels,
           choice_labels,
           encoder_hidden_states,
           encoder_attention_mask,
       )





  def 						A  (		self      :	Any							,					_a      :	List[Any]							,					_a      :	List[Any]							,					_a      :	Union[str, Any]							,					_a      :	str							,					_a      :	List[Any]							,					_a      :	Optional[Any]							,					_a      :	List[str] )	->    Optional[Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =MraModel(config=_a )
       model.to(_a )
       model.eval()
       _SCREAMING_SNAKE_CASE							    =model(_a							,					attention_mask=_a							,					token_type_ids=_a )
       _SCREAMING_SNAKE_CASE							    =model(_a							,					token_type_ids=_a )
       _SCREAMING_SNAKE_CASE							    =model(_a )
       self.parent.assertEqual(result.last_hidden_state.shape							,					(self.batch_size, self.seq_length, self.hidden_size) )





  def 						A  (		self      :	int							,					_a      :	List[Any]							,					_a      :	List[str]							,					_a      :	List[str]							,					_a      :	Dict							,					_a      :	Optional[Any]							,					_a      :	Any							,					_a      :	List[Any]							,					_a      :	Union[str, Any]							,					_a      :	List[str]							,					)	->    str:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =True
       _SCREAMING_SNAKE_CASE							    =MraModel(_a )
       model.to(_a )
       model.eval()
       _SCREAMING_SNAKE_CASE							    =model(
           _a							,					attention_mask=_a							,					token_type_ids=_a							,					encoder_hidden_states=_a							,					encoder_attention_mask=_a							,					)
       _SCREAMING_SNAKE_CASE							    =model(
           _a							,					attention_mask=_a							,					token_type_ids=_a							,					encoder_hidden_states=_a							,					)
       _SCREAMING_SNAKE_CASE							    =model(_a							,					attention_mask=_a							,					token_type_ids=_a )
       self.parent.assertEqual(result.last_hidden_state.shape							,					(self.batch_size, self.seq_length, self.hidden_size) )





  def 						A  (		self      :	List[str]							,					_a      :	Optional[Any]							,					_a      :	Dict							,					_a      :	Optional[int]							,					_a      :	str							,					_a      :	Optional[int]							,					_a      :	Optional[Any]							,					_a      :	Tuple )	->    int:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =MraForMaskedLM(config=_a )
       model.to(_a )
       model.eval()
       _SCREAMING_SNAKE_CASE							    =model(_a							,					attention_mask=_a							,					token_type_ids=_a							,					labels=_a )
       self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length, self.vocab_size) )





  def 						A  (		self      :	List[str]							,					_a      :	Optional[int]							,					_a      :	List[Any]							,					_a      :	Union[str, Any]							,					_a      :	Union[str, Any]							,					_a      :	Optional[int]							,					_a      :	Dict							,					_a      :	List[str] )	->    List[str]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =MraForQuestionAnswering(config=_a )
       model.to(_a )
       model.eval()
       _SCREAMING_SNAKE_CASE							    =model(
           _a							,					attention_mask=_a							,					token_type_ids=_a							,					start_positions=_a							,					end_positions=_a							,					)
       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      :	List[str]							,					_a      :	Optional[int]							,					_a      :	int							,					_a      :	Any							,					_a      :	str							,					_a      :	str							,					_a      :	Optional[int]							,					_a      :	Dict )	->    Optional[Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.num_labels
       _SCREAMING_SNAKE_CASE							    =MraForSequenceClassification(_a )
       model.to(_a )
       model.eval()
       _SCREAMING_SNAKE_CASE							    =model(_a							,					attention_mask=_a							,					token_type_ids=_a							,					labels=_a )
       self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.num_labels) )





  def 						A  (		self      :	Optional[Any]							,					_a      :	List[Any]							,					_a      :	List[str]							,					_a      :	Any							,					_a      :	List[Any]							,					_a      :	List[str]							,					_a      :	int							,					_a      :	Dict )	->    Dict:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.num_labels
       _SCREAMING_SNAKE_CASE							    =MraForTokenClassification(config=_a )
       model.to(_a )
       model.eval()
       _SCREAMING_SNAKE_CASE							    =model(_a							,					attention_mask=_a							,					token_type_ids=_a							,					labels=_a )
       self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length, self.num_labels) )





  def 						A  (		self      :	List[str]							,					_a      :	int							,					_a      :	Optional[int]							,					_a      :	int							,					_a      :	Dict							,					_a      :	Union[str, Any]							,					_a      :	Optional[int]							,					_a      :	Any )	->    Tuple:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.num_choices
       _SCREAMING_SNAKE_CASE							    =MraForMultipleChoice(config=_a )
       model.to(_a )
       model.eval()
       _SCREAMING_SNAKE_CASE							    =input_ids.unsqueeze(1 ).expand(-1							,					self.num_choices							,					-1 ).contiguous()
       _SCREAMING_SNAKE_CASE							    =token_type_ids.unsqueeze(1 ).expand(-1							,					self.num_choices							,					-1 ).contiguous()
       _SCREAMING_SNAKE_CASE							    =input_mask.unsqueeze(1 ).expand(-1							,					self.num_choices							,					-1 ).contiguous()
       _SCREAMING_SNAKE_CASE							    =model(
           _a							,					attention_mask=_a							,					token_type_ids=_a							,					labels=_a							,					)
       self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.num_choices) )





  def 						A  (		self      :	Tuple )	->    str:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.prepare_config_and_inputs()
       (
           (
           _SCREAMING_SNAKE_CASE
       )    ,						(
           _SCREAMING_SNAKE_CASE
       )    ,						(
           _SCREAMING_SNAKE_CASE
       )    ,						(
           _SCREAMING_SNAKE_CASE
       )    ,						(
           _SCREAMING_SNAKE_CASE
       )    ,						(
           _SCREAMING_SNAKE_CASE
       )    ,						(
           _SCREAMING_SNAKE_CASE
       )    ,						
       )							    =config_and_inputs
       _SCREAMING_SNAKE_CASE							    ={'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
       return config, inputs_dict







@require_torch
class 							A__    (       A__     ,   unittest.TestCase  ):
  A__		 =				(
      (
          MraModel,
          MraForMaskedLM,
          MraForMultipleChoice,
          MraForQuestionAnswering,
          MraForSequenceClassification,
          MraForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  A__		 =				False
  A__		 =				False
  A__		 =				False
  A__		 =				False

  A__		 =				()





  def 						A  (		self      :	Dict )	->    Dict:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =MraModelTester(self )
       _SCREAMING_SNAKE_CASE							    =ConfigTester(self							,					config_class=_a							,					hidden_size=37 )





  def 						A  (		self      :	Tuple )	->    Optional[int]:





       '''simple docstring'''






       self.config_tester.run_common_tests()





  def 						A  (		self      :	Dict )	->    Tuple:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_a )





  def 						A  (		self      :	List[str] )	->    Any:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
            _SCREAMING_SNAKE_CASE							    =type
            self.model_tester.create_and_check_model(*_a )





  def 						A  (		self      :	str )	->    Any:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*_a )





  def 						A  (		self      :	List[str] )	->    Optional[int]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*_a )





  def 						A  (		self      :	Optional[int] )	->    Any:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*_a )





  def 						A  (		self      :	Dict )	->    Union[str, Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*_a )





  def 						A  (		self      :	List[str] )	->    str:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*_a )





  @slow
  def 						A  (		self      :	int )	->    List[Any]:





       '''simple docstring'''






       for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            _SCREAMING_SNAKE_CASE							    =MraModel.from_pretrained(_a )
            self.assertIsNotNone(_a )





  @unittest.skip(reason='MRA does not output attentions' )
  def 						A  (		self      :	Optional[int] )	->    Tuple:





       '''simple docstring'''






       return







@require_torch
class 							A__    (       unittest.TestCase  ):





  @slow
  def 						A  (		self      :	Optional[Any] )	->    Tuple:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =MraModel.from_pretrained('uw-madison/mra-base-512-4' )
       _SCREAMING_SNAKE_CASE							    =torch.arange(256 ).unsqueeze(0 )

       with torch.no_grad():
            _SCREAMING_SNAKE_CASE							    =model(_a )[0]

       _SCREAMING_SNAKE_CASE							    =torch.Size((1, 256, 768) )
       self.assertEqual(output.shape							,					_a )

       _SCREAMING_SNAKE_CASE							    =torch.tensor(
           [[[-0.01_40, 0.08_30, -0.03_81], [0.15_46, 0.14_02, 0.02_20], [0.11_62, 0.08_51, 0.01_65]]] )

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





  @slow
  def 						A  (		self      :	Dict )	->    Union[str, Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =MraForMaskedLM.from_pretrained('uw-madison/mra-base-512-4' )
       _SCREAMING_SNAKE_CASE							    =torch.arange(256 ).unsqueeze(0 )

       with torch.no_grad():
            _SCREAMING_SNAKE_CASE							    =model(_a )[0]

       _SCREAMING_SNAKE_CASE							    =5_0265

       _SCREAMING_SNAKE_CASE							    =torch.Size((1, 256, vocab_size) )
       self.assertEqual(output.shape							,					_a )

       _SCREAMING_SNAKE_CASE							    =torch.tensor(
           [[[9.25_95, -3.60_38, 11.88_19], [9.38_69, -3.26_93, 11.09_56], [11.85_24, -3.49_38, 13.12_10]]] )

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





  @slow
  def 						A  (		self      :	List[Any] )	->    str:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =MraForMaskedLM.from_pretrained('uw-madison/mra-base-4096-8-d3' )
       _SCREAMING_SNAKE_CASE							    =torch.arange(4096 ).unsqueeze(0 )

       with torch.no_grad():
            _SCREAMING_SNAKE_CASE							    =model(_a )[0]

       _SCREAMING_SNAKE_CASE							    =5_0265

       _SCREAMING_SNAKE_CASE							    =torch.Size((1, 4096, vocab_size) )
       self.assertEqual(output.shape							,					_a )

       _SCREAMING_SNAKE_CASE							    =torch.tensor(
           [[[5.47_89, -2.35_64, 7.50_64], [7.90_67, -1.33_69, 9.96_68], [9.07_12, -1.81_06, 7.03_80]]] )

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






'''simple docstring'''







import argparse
import logging
from collections import namedtuple

import torch
from model_bertabs import BertAbsSummarizer
from models.model_builder import AbsSummarizer  # The authors' implementation

from transformers import BertTokenizer


logging.basicConfig(level=logging.INFO)
A     =logging.getLogger(__name__)


A     ='Hello world! cécé herlolip'


A     =namedtuple(
    'BertAbsConfig',
    [
        'temp_dir',
        'large',
        'use_bert_emb',
        'finetune_bert',
        'encoder',
        'share_emb',
        'max_pos',
        'enc_layers',
        'enc_hidden_size',
        'enc_heads',
        'enc_ff_size',
        'enc_dropout',
        'dec_layers',
        'dec_hidden_size',
        'dec_heads',
        'dec_ff_size',
        'dec_dropout',
    ],
)





def 						snake_case_							(_a	:							List[Any]    ,   _a	:							Any     ):
    UpperCAmelCase			       =		BertAbsConfig(
        temp_dir='''.'''    ,   finetune_bert=_a    ,   large=_a    ,   share_emb=_a    ,   use_bert_emb=_a    ,   encoder='''bert'''    ,   max_pos=5_1_2    ,   enc_layers=6    ,   enc_hidden_size=5_1_2    ,   enc_heads=8    ,   enc_ff_size=5_1_2    ,   enc_dropout=0.2    ,   dec_layers=6    ,   dec_hidden_size=7_6_8    ,   dec_heads=8    ,   dec_ff_size=2_0_4_8    ,   dec_dropout=0.2    ,   )
    UpperCAmelCase			       =		torch.load(_a    ,   lambda _a    ,   _a     : storage     )
    UpperCAmelCase			       =		AbsSummarizer(_a    ,   torch.device('''cpu'''     )    ,   _a     )
    original.eval()

    UpperCAmelCase			       =		BertAbsSummarizer(_a    ,   torch.device('''cpu'''     )     )
    new_model.eval()

    # -------------------
    # Convert the weights
    # -------------------

    logging.info('''convert the model'''     )
    new_model.bert.load_state_dict(original.bert.state_dict()     )
    new_model.decoder.load_state_dict(original.decoder.state_dict()     )
    new_model.generator.load_state_dict(original.generator.state_dict()     )

    # ----------------------------------
    # Make sure the outpus are identical
    # ----------------------------------

    logging.info('''Make sure that the models\' outputs are identical'''     )
    UpperCAmelCase			       =		BertTokenizer.from_pretrained('''bert-base-uncased'''     )

    # prepare the model inputs
    UpperCAmelCase			       =		tokenizer.encode('''This is sample éàalj\'-.'''     )
    encoder_input_ids.extend([tokenizer.pad_token_id] * (5_1_2 - len(_a     ))     )
    UpperCAmelCase			       =		torch.tensor(_a     ).unsqueeze(0     )
    UpperCAmelCase			       =		tokenizer.encode('''This is sample 3 éàalj\'-.'''     )
    decoder_input_ids.extend([tokenizer.pad_token_id] * (5_1_2 - len(_a     ))     )
    UpperCAmelCase			       =		torch.tensor(_a     ).unsqueeze(0     )

    # failsafe to make sure the weights reset does not affect the
    # loaded weights.
    assert torch.max(torch.abs(original.generator[0].weight - new_model.generator[0].weight     )     ) == 0

    # forward pass
    UpperCAmelCase			       =		encoder_input_ids
    UpperCAmelCase			       =		decoder_input_ids
    UpperCAmelCase			       =		UpperCAmelCase			       =		None
    UpperCAmelCase			       =		None
    UpperCAmelCase			       =		UpperCAmelCase			       =		None
    UpperCAmelCase			       =		UpperCAmelCase			       =		None
    UpperCAmelCase			       =		None

    # The original model does not apply the geneator layer immediatly but rather in
    # the beam search (where it combines softmax + linear layer). Since we already
    # apply the softmax in our generation process we only apply the linear layer here.
    # We make sure that the outputs of the full stack are identical
    UpperCAmelCase			       =		original(_a    ,   _a    ,   _a    ,   _a    ,   _a    ,   _a    ,   _a     )[0]
    UpperCAmelCase			       =		original.generator(_a     )

    UpperCAmelCase			       =		new_model(
        _a    ,   _a    ,   _a    ,   _a    ,   _a     )[0]
    UpperCAmelCase			       =		new_model.generator(_a     )

    UpperCAmelCase			       =		torch.max(torch.abs(output_converted_model - output_original_model     )     ).item()
    print('''Maximum absolute difference beween weights: {:.2f}'''.format(_a     )     )
    UpperCAmelCase			       =		torch.max(torch.abs(output_converted_generator - output_original_generator     )     ).item()
    print('''Maximum absolute difference beween weights: {:.2f}'''.format(_a     )     )

    UpperCAmelCase			       =		torch.allclose(_a    ,   _a    ,   atol=1E-3     )
    if are_identical:
        logging.info('''all weights are equal up to 1e-3'''     )
    else:
        raise ValueError('''the weights are different. The new model is likely different from the original one.'''     )

    # The model has been saved with torch.save(model) and this is bound to the exact
    # directory structure. We save the state_dict instead.
    logging.info('''saving the model\'s state dictionary'''     )
    torch.save(
        new_model.state_dict()    ,   '''./bertabs-finetuned-cnndm-extractive-abstractive-summarization/pytorch_model.bin'''     )


if __name__ == "__main__":
   A     =argparse.ArgumentParser()
   parser.add_argument(
       '--bertabs_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_bertabs_checkpoints(
       args.bertabs_checkpoint_path,
       args.pytorch_dump_folder_path,
   )




