




'''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)
_snake_case				:				Optional[Any]       							=		logging.getLogger(__name__)


_snake_case				:				Optional[int]       							=		'Hello world! cécé herlolip'


_snake_case				:				List[str]       							=		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_   (UpperCamelCase	:  Any			,						UpperCamelCase	:  Tuple				):





       '''simple docstring'''


       _a     					=    BertAbsConfig(
           temp_dir='''.'''			,						finetune_bert=UpperCamelCase			,						large=UpperCamelCase			,						share_emb=UpperCamelCase			,						use_bert_emb=UpperCamelCase			,						encoder='''bert'''			,						max_pos=512			,						enc_layers=6			,						enc_hidden_size=512			,						enc_heads=8			,						enc_ff_size=512			,						enc_dropout=0.2			,						dec_layers=6			,						dec_hidden_size=768			,						dec_heads=8			,						dec_ff_size=2048			,						dec_dropout=0.2			,						)
       _a     					=    torch.load(UpperCamelCase			,						lambda UpperCamelCase			,						UpperCamelCase				: storage				)
       _a     					=    AbsSummarizer(UpperCamelCase			,						torch.device('''cpu'''				)			,						UpperCamelCase				)
       original.eval()

       _a     					=    BertAbsSummarizer(UpperCamelCase			,						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'''				)
       _a     					=    BertTokenizer.from_pretrained('''bert-base-uncased'''				)

       # prepare the model inputs
       _a     					=    tokenizer.encode('''This is sample éàalj\'-.'''				)
       encoder_input_ids.extend([tokenizer.pad_token_id] * (512 - len(UpperCamelCase				))				)
       _a     					=    torch.tensor(UpperCamelCase				).unsqueeze(0				)
       _a     					=    tokenizer.encode('''This is sample 3 éàalj\'-.'''				)
       decoder_input_ids.extend([tokenizer.pad_token_id] * (512 - len(UpperCamelCase				))				)
       _a     					=    torch.tensor(UpperCamelCase				).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
       _a     					=    encoder_input_ids
       _a     					=    decoder_input_ids
       _a     					=    _a     					=    None
       _a     					=    None
       _a     					=    _a     					=    None
       _a     					=    _a     					=    None
       _a     					=    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
       _a     					=    original(UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase				)[0]
       _a     					=    original.generator(UpperCamelCase				)

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

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

       _a     					=    torch.allclose(UpperCamelCase			,						UpperCamelCase			,						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__":
    _snake_case				:				List[str]       							=		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.',
    )
    _snake_case				:				str       							=		parser.parse_args()

    convert_bertabs_checkpoints(
        args.bertabs_checkpoint_path,
        args.pytorch_dump_folder_path,
    )







'''simple docstring'''


import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch

				from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM

@require_torch
class       A   :







							def __init__(  self	:  Union[str, Any]							,      lowerCAmelCase_	:  List[Any]							,      lowerCAmelCase_	:  Union[str, Any]=99							,      lowerCAmelCase_	:  Optional[int]=13							,      lowerCAmelCase_	:  Tuple=16							,      lowerCAmelCase_	:  Optional[int]=7							,      lowerCAmelCase_	:  Optional[Any]=True							,      lowerCAmelCase_	:  Union[str, Any]=True							,      lowerCAmelCase_	:  Any=True							,      lowerCAmelCase_	:  Union[str, Any]=False							,      lowerCAmelCase_	:  int=True							,      lowerCAmelCase_	:  Tuple=2							,      lowerCAmelCase_	:  Any=32							,      lowerCAmelCase_	:  str=4							,      lowerCAmelCase_	:  Optional[Any]=4							,      lowerCAmelCase_	:  Optional[Any]=30							,      lowerCAmelCase_	:  int=0							,      lowerCAmelCase_	:  Any=1							,      lowerCAmelCase_	:  str=2							,      lowerCAmelCase_	:  Union[str, Any]=None							,      )      ->							Any:


														"""simple docstring"""


														_a     					=    parent
														_a     					=    batch_size
														_a     					=    decoder_seq_length
														# For common tests
														_a     					=    self.decoder_seq_length
														_a     					=    is_training
														_a     					=    use_attention_mask
														_a     					=    use_labels

														_a     					=    vocab_size
														_a     					=    d_model
														_a     					=    d_model
														_a     					=    decoder_layers
														_a     					=    decoder_layers
														_a     					=    decoder_ffn_dim
														_a     					=    decoder_attention_heads
														_a     					=    decoder_attention_heads
														_a     					=    eos_token_id
														_a     					=    bos_token_id
														_a     					=    pad_token_id
														_a     					=    decoder_start_token_id
														_a     					=    use_cache
														_a     					=    max_position_embeddings

														_a     					=    None
														_a     					=    decoder_seq_length
														_a     					=    2
														_a     					=    1







							def        __lowerCAmelCase      (  self	:  str )      ->							str:


														"""simple docstring"""


														_a     					=    ids_tensor([self.batch_size, self.decoder_seq_length]							,      self.vocab_size )

														_a     					=    None
														if self.use_attention_mask:
																					_a     					=    ids_tensor([self.batch_size, self.decoder_seq_length]							,      vocab_size=2 )

														_a     					=    None
														if self.use_labels:
																					_a     					=    ids_tensor([self.batch_size, self.decoder_seq_length]							,      self.vocab_size )

														_a     					=    TrOCRConfig(
														    vocab_size=self.vocab_size							,      d_model=self.d_model							,      decoder_layers=self.decoder_layers							,      decoder_ffn_dim=self.decoder_ffn_dim							,      decoder_attention_heads=self.decoder_attention_heads							,      eos_token_id=self.eos_token_id							,      bos_token_id=self.bos_token_id							,      use_cache=self.use_cache							,      pad_token_id=self.pad_token_id							,      decoder_start_token_id=self.decoder_start_token_id							,      max_position_embeddings=self.max_position_embeddings							,      )

														return (config, input_ids, attention_mask, lm_labels)







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


														"""simple docstring"""


														_a     					=    True
														_a     					=    TrOCRDecoder(config=lowerCAmelCase_ ).to(lowerCAmelCase_ ).eval()
														_a     					=    input_ids[:2]

														input_ids[input_ids == 0] += 1
														# first forward pass
														_a     					=    model(lowerCAmelCase_							,      use_cache=lowerCAmelCase_ )
														_a     					=    model(lowerCAmelCase_ )
														_a     					=    model(lowerCAmelCase_							,      use_cache=lowerCAmelCase_ )

														self.parent.assertTrue(len(lowerCAmelCase_ ) == len(lowerCAmelCase_ ) )
														self.parent.assertTrue(len(lowerCAmelCase_ ) == len(lowerCAmelCase_ ) + 1 )

														_a     					=    outputs['''past_key_values''']

														# create hypothetical next token and extent to next_input_ids
														_a     					=    ids_tensor((2, 1)							,      config.vocab_size - 1 ) + 1

														# append to next input_ids and
														_a     					=    torch.cat([input_ids, next_tokens]							,      dim=-1 )

														_a     					=    model(lowerCAmelCase_ )['''last_hidden_state''']
														_a     					=    model(lowerCAmelCase_							,      past_key_values=lowerCAmelCase_ )['''last_hidden_state''']

														# select random slice
														_a     					=    ids_tensor((1,)							,      output_from_past.shape[-1] ).item()
														_a     					=    output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
														_a     					=    output_from_past[:, 0, random_slice_idx].detach()

														# test that outputs are equal for slice
														assert torch.allclose(lowerCAmelCase_							,      lowerCAmelCase_							,      atol=1e-3 )







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


														"""simple docstring"""


														_a     					=    self.prepare_config_and_inputs()
														_a					,							_a					,							_a					,							_a     					=    config_and_inputs

														_a     					=    {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
														return config, inputs_dict



@require_torch
class       A   (		_a					,_a					,_a					,unittest.TestCase   ):
							lowercase_          =    (TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
							lowercase_          =    (TrOCRForCausalLM,) if is_torch_available() else ()
							lowercase_          =    {'text-generation': TrOCRForCausalLM} if is_torch_available() else {}
							lowercase_          =    True
							lowercase_          =    False







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


														"""simple docstring"""


														_a     					=    TrOCRStandaloneDecoderModelTester(self							,      is_training=lowerCAmelCase_ )
														_a     					=    ConfigTester(self							,      config_class=lowerCAmelCase_ )







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


														"""simple docstring"""


														pass







							def        __lowerCAmelCase      (  self	:  Tuple )      ->							Dict:


														"""simple docstring"""


														pass







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


														"""simple docstring"""


														pass







							def        __lowerCAmelCase      (  self	:  Union[str, Any] )      ->							Optional[int]:


														"""simple docstring"""


														self.config_tester.run_common_tests()







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


														"""simple docstring"""


														_a     					=    self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_decoder_model_past(*lowerCAmelCase_ )







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


														"""simple docstring"""


														return







							@unittest.skip('''The model doesn\'t support left padding''' )  # and it's not used enough to be worth fixing :)
							def        __lowerCAmelCase      (  self	:  Union[str, Any] )      ->							Optional[int]:


														"""simple docstring"""


														pass

