




import os
import tempfile
import unittest

from transformers import DistilBertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, 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 (
      DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      DistilBertForMaskedLM,
      DistilBertForMultipleChoice,
      DistilBertForQuestionAnswering,
      DistilBertForSequenceClassification,
      DistilBertForTokenClassification,
      DistilBertModel,
  )







class 						lowercase							(   __a		):

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

     def 					a__	(  self				)						->							Optional[Any]:
           _A							:     Tuple			=						ids_tensor([self.batch_size, self.seq_length]		,     self.vocab_size				)

           _A							:     Optional[int]			=						None
           if self.use_input_mask:
                 _A							:     Dict			=						random_attention_mask([self.batch_size, self.seq_length]				)

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

           _A							:     Any			=						self.get_config()

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

     def 					a__	(  self				)						->							str:
           return DistilBertConfig(
               vocab_size=self.vocab_size		,     dim=self.hidden_size		,     n_layers=self.num_hidden_layers		,     n_heads=self.num_attention_heads		,     hidden_dim=self.intermediate_size		,     hidden_act=self.hidden_act		,     dropout=self.hidden_dropout_prob		,     attention_dropout=self.attention_probs_dropout_prob		,     max_position_embeddings=self.max_position_embeddings		,     initializer_range=self.initializer_range		,     )

     def 					a__	(  self		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a				)						->							Optional[int]:
           _A							:     Any			=						DistilBertModel(config=_a				)
           model.to(_a				)
           model.eval()
           _A							:     Optional[Any]			=						model(_a		,     _a				)
           _A							:     Optional[int]			=						model(_a				)
           self.parent.assertEqual(result.last_hidden_state.shape		,     (self.batch_size, self.seq_length, self.hidden_size)				)

     def 					a__	(  self		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a				)						->							Any:
           _A							:     Union[str, Any]			=						DistilBertForMaskedLM(config=_a				)
           model.to(_a				)
           model.eval()
           _A							:     Optional[int]			=						model(_a		,     attention_mask=_a		,     labels=_a				)
           self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.seq_length, self.vocab_size)				)

     def 					a__	(  self		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a				)						->							str:
           _A							:     Dict			=						DistilBertForQuestionAnswering(config=_a				)
           model.to(_a				)
           model.eval()
           _A							:     Any			=						model(
               _a		,     attention_mask=_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		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a				)						->							str:
           _A							:     Optional[Any]			=						self.num_labels
           _A							:     Tuple			=						DistilBertForSequenceClassification(_a				)
           model.to(_a				)
           model.eval()
           _A							:     Tuple			=						model(_a		,     attention_mask=_a		,     labels=_a				)
           self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.num_labels)				)

     def 					a__	(  self		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a				)						->							List[str]:
           _A							:     Dict			=						self.num_labels
           _A							:     Tuple			=						DistilBertForTokenClassification(config=_a				)
           model.to(_a				)
           model.eval()

           _A							:     Optional[int]			=						model(_a		,     attention_mask=_a		,     labels=_a				)
           self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.seq_length, self.num_labels)				)

     def 					a__	(  self		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a				)						->							str:
           _A							:     Optional[Any]			=						self.num_choices
           _A							:     Optional[Any]			=						DistilBertForMultipleChoice(config=_a				)
           model.to(_a				)
           model.eval()
           _A							:     Any			=						input_ids.unsqueeze(1				).expand(-1		,     self.num_choices		,     -1				).contiguous()
           _A							:     Optional[int]			=						input_mask.unsqueeze(1				).expand(-1		,     self.num_choices		,     -1				).contiguous()
           _A							:     Dict			=						model(
               _a		,     attention_mask=_a		,     labels=_a		,     )
           self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.num_choices)				)



     def 					a__	(  self				)						->							Tuple:
           _A							:     Optional[int]			=						self.prepare_config_and_inputs()
           ((_A)			,   (_A)			,   (_A)			,   (_A)			,   (_A)			,   (_A))							:     Union[str, Any]			=						config_and_inputs
           _A							:     Tuple			=						{"""input_ids""": input_ids, """attention_mask""": input_mask}
           return config, inputs_dict







@require_torch
class 						lowercase							(   __a,__a,unittest.TestCase		):
     _a							 =						(
         (
             DistilBertModel,
             DistilBertForMaskedLM,
             DistilBertForMultipleChoice,
             DistilBertForQuestionAnswering,
             DistilBertForSequenceClassification,
             DistilBertForTokenClassification,
         )
         if is_torch_available()
         else None
     )
     _a							 =						(
         {
             """feature-extraction""": DistilBertModel,
             """fill-mask""": DistilBertForMaskedLM,
             """question-answering""": DistilBertForQuestionAnswering,
             """text-classification""": DistilBertForSequenceClassification,
             """token-classification""": DistilBertForTokenClassification,
             """zero-shot""": DistilBertForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     _a							 =						True
     _a							 =						True
     _a							 =						True
     _a							 =						True

     def 					a__	(  self				)						->							Union[str, Any]:
           _A							:     int			=						DistilBertModelTester(self				)
           _A							:     Optional[int]			=						ConfigTester(self		,     config_class=_a		,     dim=37				)

     def 					a__	(  self				)						->							Dict:
           self.config_tester.run_common_tests()

     def 					a__	(  self				)						->							Optional[int]:
           _A							:     Optional[Any]			=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_model(*_a				)

     def 					a__	(  self				)						->							int:
           _A							:     int			=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_masked_lm(*_a				)

     def 					a__	(  self				)						->							List[Any]:
           _A							:     str			=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_question_answering(*_a				)

     def 					a__	(  self				)						->							Dict:
           _A							:     Dict			=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_sequence_classification(*_a				)

     def 					a__	(  self				)						->							str:
           _A							:     int			=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_token_classification(*_a				)

     def 					a__	(  self				)						->							Dict:
           _A							:     Optional[Any]			=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_multiple_choice(*_a				)

     @slow
     def 					a__	(  self				)						->							str:

           for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _A							:     Optional[int]			=						DistilBertModel.from_pretrained(_a				)
                 self.assertIsNotNone(_a				)



     @slow
     @require_torch_gpu
     def 					a__	(  self				)						->							str:
           _A			,   _A							:     Any			=						self.model_tester.prepare_config_and_inputs_for_common()
           for model_class in self.all_model_classes:
                 # BertForMultipleChoice behaves incorrectly in JIT environments.
                 if model_class == DistilBertForMultipleChoice:
                       return

                 _A							:     Dict			=						True
                 _A							:     str			=						model_class(config=_a				)

                 _A							:     Dict			=						self._prepare_for_class(_a		,     _a				)
                 _A							:     Optional[Any]			=						torch.jit.trace(
                     _a		,     (inputs_dict["""input_ids"""].to("""cpu"""				), inputs_dict["""attention_mask"""].to("""cpu"""				))				)

                 with tempfile.TemporaryDirectory() as tmp:
                       torch.jit.save(_a		,     os.path.join(_a		,     """traced_model.pt"""				)				)
                       _A							:     Optional[int]			=						torch.jit.load(os.path.join(_a		,     """traced_model.pt"""				)		,     map_location=_a				)
                       loaded(inputs_dict["""input_ids"""].to(_a				)		,     inputs_dict["""attention_mask"""].to(_a				)				)







@require_torch
class 						lowercase							(   unittest.TestCase		):

     @slow
     def 					a__	(  self				)						->							Union[str, Any]:
           _A							:     Dict			=						DistilBertModel.from_pretrained("""distilbert-base-uncased"""				)
           _A							:     Optional[Any]			=						torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]				)
           _A							:     List[Any]			=						torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]				)
           with torch.no_grad():
                 _A							:     Any			=						model(_a		,     attention_mask=_a				)[0]
           _A							:     List[str]			=						torch.Size((1, 11, 768)				)
           self.assertEqual(output.shape		,     _a				)
           _A							:     Optional[int]			=						torch.tensor(
               [[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]				)

           self.assertTrue(torch.allclose(output[:, 1:4, 1:4]		,     _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,
   )




