


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 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Union[str, Any]      ,					SCREAMING_SNAKE_CASE__      : Any    ):
     __UpperCamelCase   							=BertAbsConfig(
         temp_dir='.'      ,					finetune_bert=SCREAMING_SNAKE_CASE__      ,					large=SCREAMING_SNAKE_CASE__      ,					share_emb=SCREAMING_SNAKE_CASE__      ,					use_bert_emb=SCREAMING_SNAKE_CASE__      ,					encoder='bert'      ,					max_pos=5_12      ,					enc_layers=6      ,					enc_hidden_size=5_12      ,					enc_heads=8      ,					enc_ff_size=5_12      ,					enc_dropout=0.2      ,					dec_layers=6      ,					dec_hidden_size=7_68      ,					dec_heads=8      ,					dec_ff_size=20_48      ,					dec_dropout=0.2      ,					)
     __UpperCamelCase   							=torch.load(SCREAMING_SNAKE_CASE__      ,					lambda SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    : storage    )
     __UpperCamelCase   							=AbsSummarizer(SCREAMING_SNAKE_CASE__      ,					torch.device('cpu'    )      ,					SCREAMING_SNAKE_CASE__    )
     original.eval()

     __UpperCamelCase   							=BertAbsSummarizer(SCREAMING_SNAKE_CASE__      ,					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_12 - len(SCREAMING_SNAKE_CASE__    ))    )
     __UpperCamelCase   							=torch.tensor(SCREAMING_SNAKE_CASE__    ).unsqueeze(0    )
     __UpperCamelCase   							=tokenizer.encode('This is sample 3 éàalj\'-.'    )
     decoder_input_ids.extend([tokenizer.pad_token_id] * (5_12 - len(SCREAMING_SNAKE_CASE__    ))    )
     __UpperCamelCase   							=torch.tensor(SCREAMING_SNAKE_CASE__    ).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(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    )[0]
     __UpperCamelCase   							=original.generator(SCREAMING_SNAKE_CASE__    )

     __UpperCamelCase   							=new_model(
         SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    )[0]
     __UpperCamelCase   							=new_model.generator(SCREAMING_SNAKE_CASE__    )

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

     __UpperCamelCase   							=torch.allclose(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					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,
   )









import warnings
from contextlib import contextmanager

from ....processing_utils import ProcessorMixin






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    UpperCAmelCase__						:		Optional[int]              =  "MCTCTFeatureExtractor"
    UpperCAmelCase__						:		str              =  "AutoTokenizer"






    def __init__(      self ,       A_ ,       A_   )  ->	Dict:
         super().__init__(A_ ,       A_   )
         __UpperCamelCase   							=self.feature_extractor
         __UpperCamelCase   							=False






    def __call__(      self ,       *A_ ,       **A_   )  ->	Optional[Any]:
         # For backward compatibility
         if self._in_target_context_manager:
              return self.current_processor(*A_ ,       **A_   )

         if "raw_speech" in kwargs:
              warnings.warn('Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'   )
              __UpperCamelCase   							=kwargs.pop('raw_speech'   )
         else:
              __UpperCamelCase   							=kwargs.pop('audio' ,       A_   )
         __UpperCamelCase   							=kwargs.pop('sampling_rate' ,       A_   )
         __UpperCamelCase   							=kwargs.pop('text' ,       A_   )
         if len(A_   ) > 0:
              __UpperCamelCase   							=args[0]
              __UpperCamelCase   							=args[1:]

         if audio is None and text is None:
              raise ValueError('You need to specify either an `audio` or `text` input to process.'   )

         if audio is not None:
              __UpperCamelCase   							=self.feature_extractor(A_ ,       *A_ ,       sampling_rate=A_ ,       **A_   )
         if text is not None:
              __UpperCamelCase   							=self.tokenizer(A_ ,       **A_   )

         if text is None:
              return inputs
         elif audio is None:
              return encodings
         else:
              __UpperCamelCase   							=encodings['input_ids']
              return inputs






    def _a (      self ,       *A_ ,       **A_   )  ->	str:
         return self.tokenizer.batch_decode(*A_ ,       **A_   )






    def _a (      self ,       *A_ ,       **A_   )  ->	List[Any]:
         # For backward compatibility
         if self._in_target_context_manager:
              return self.current_processor.pad(*A_ ,       **A_   )

         __UpperCamelCase   							=kwargs.pop('input_features' ,       A_   )
         __UpperCamelCase   							=kwargs.pop('labels' ,       A_   )
         if len(A_   ) > 0:
              __UpperCamelCase   							=args[0]
              __UpperCamelCase   							=args[1:]

         if input_features is not None:
              __UpperCamelCase   							=self.feature_extractor.pad(A_ ,       *A_ ,       **A_   )
         if labels is not None:
              __UpperCamelCase   							=self.tokenizer.pad(A_ ,       **A_   )

         if labels is None:
              return input_features
         elif input_features is None:
              return labels
         else:
              __UpperCamelCase   							=labels['input_ids']
              return input_features






    def _a (      self ,       *A_ ,       **A_   )  ->	Optional[int]:
         return self.tokenizer.decode(*A_ ,       **A_   )






    @contextmanager
    def _a (      self   )  ->	str:
         warnings.warn(
             '`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '
             'labels by using the argument `text` of the regular `__call__` method (either in the same call as '
             'your audio inputs, or in a separate call.'   )
         __UpperCamelCase   							=True
         __UpperCamelCase   							=self.tokenizer
         yield
         __UpperCamelCase   							=self.feature_extractor
         __UpperCamelCase   							=False





