



import argparse
import os
import re

import tensorflow as tf
import torch

from transformers import BertConfig, BertModel
from transformers.utils import logging


logging.set_verbosity_info()
__snake_case						:Any						=					logging.get_logger(__name__)





def 					__snake_case     (      _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase    ):
    __a							=    os.path.abspath(_UpperCAmelCase    )
    logger.info(f'Converting TensorFlow checkpoint from {tf_path}'    )
    # Load weights from TF model
    __a							=    tf.train.list_variables(_UpperCAmelCase    )
    __a							=    []
    __a							=    []
    __a							=    []
    for full_name, shape in init_vars:
        # logger.info(f"Loading TF weight {name} with shape {shape}")
        __a							=    full_name.split('''/'''    )
        if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]:
            logger.info(f'Skipping non-model layer {full_name}'    )
            continue
        if "optimizer" in full_name:
            logger.info(f'Skipping optimization layer {full_name}'    )
            continue
        if name[0] == "model":
            # ignore initial 'model'
            __a							=    name[1:]
        # figure out how many levels deep the name is
        __a							=    0
        for _name in name:
            if _name.startswith('''layer_with_weights'''    ):
                depth += 1
            else:
                break
        layer_depth.append(_UpperCAmelCase    )
        # read data
        __a							=    tf.train.load_variable(_UpperCAmelCase       ,			_UpperCAmelCase    )
        names.append('''/'''.join(_UpperCAmelCase    )    )
        arrays.append(_UpperCAmelCase    )
    logger.info(f'Read a total of {len(_UpperCAmelCase    ):,} layers'    )

    # Sanity check
    if len(set(_UpperCAmelCase    )    ) != 1:
        raise ValueError(f'Found layer names with different depths (layer depth {list(set(_UpperCAmelCase    )    )})'    )
    __a							=    list(set(_UpperCAmelCase    )    )[0]
    if layer_depth != 1:
        raise ValueError(
            '''The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP'''
            ''' heads.'''    )

    # convert layers
    logger.info('''Converting weights...'''    )
    for full_name, array in zip(_UpperCAmelCase       ,			_UpperCAmelCase    ):
        __a							=    full_name.split('''/'''    )
        __a							=    model
        __a							=    []
        for i, m_name in enumerate(_UpperCAmelCase    ):
            if m_name == ".ATTRIBUTES":
                # variable names end with .ATTRIBUTES/VARIABLE_VALUE
                break
            if m_name.startswith('''layer_with_weights'''    ):
                __a							=    int(m_name.split('''-'''    )[-1]    )
                if layer_num <= 2:
                    # embedding layers
                    # layer_num 0: word_embeddings
                    # layer_num 1: position_embeddings
                    # layer_num 2: token_type_embeddings
                    continue
                elif layer_num == 3:
                    # embedding LayerNorm
                    trace.extend(['''embeddings''', '''LayerNorm''']    )
                    __a							=    getattr(_UpperCAmelCase       ,			'''embeddings'''    )
                    __a							=    getattr(_UpperCAmelCase       ,			'''LayerNorm'''    )
                elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
                    # encoder layers
                    trace.extend(['''encoder''', '''layer''', str(layer_num - 4    )]    )
                    __a							=    getattr(_UpperCAmelCase       ,			'''encoder'''    )
                    __a							=    getattr(_UpperCAmelCase       ,			'''layer'''    )
                    __a							=    pointer[layer_num - 4]
                elif layer_num == config.num_hidden_layers + 4:
                    # pooler layer
                    trace.extend(['''pooler''', '''dense''']    )
                    __a							=    getattr(_UpperCAmelCase       ,			'''pooler'''    )
                    __a							=    getattr(_UpperCAmelCase       ,			'''dense'''    )
            elif m_name == "embeddings":
                trace.append('''embeddings'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''embeddings'''    )
                if layer_num == 0:
                    trace.append('''word_embeddings'''    )
                    __a							=    getattr(_UpperCAmelCase       ,			'''word_embeddings'''    )
                elif layer_num == 1:
                    trace.append('''position_embeddings'''    )
                    __a							=    getattr(_UpperCAmelCase       ,			'''position_embeddings'''    )
                elif layer_num == 2:
                    trace.append('''token_type_embeddings'''    )
                    __a							=    getattr(_UpperCAmelCase       ,			'''token_type_embeddings'''    )
                else:
                    raise ValueError(f'Unknown embedding layer with name {full_name}'    )
                trace.append('''weight'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''weight'''    )
            elif m_name == "_attention_layer":
                # self-attention layer
                trace.extend(['''attention''', '''self''']    )
                __a							=    getattr(_UpperCAmelCase       ,			'''attention'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''self'''    )
            elif m_name == "_attention_layer_norm":
                # output attention norm
                trace.extend(['''attention''', '''output''', '''LayerNorm''']    )
                __a							=    getattr(_UpperCAmelCase       ,			'''attention'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''output'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''LayerNorm'''    )
            elif m_name == "_attention_output_dense":
                # output attention dense
                trace.extend(['''attention''', '''output''', '''dense''']    )
                __a							=    getattr(_UpperCAmelCase       ,			'''attention'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''output'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''dense'''    )
            elif m_name == "_output_dense":
                # output dense
                trace.extend(['''output''', '''dense''']    )
                __a							=    getattr(_UpperCAmelCase       ,			'''output'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''dense'''    )
            elif m_name == "_output_layer_norm":
                # output dense
                trace.extend(['''output''', '''LayerNorm''']    )
                __a							=    getattr(_UpperCAmelCase       ,			'''output'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''LayerNorm'''    )
            elif m_name == "_key_dense":
                # attention key
                trace.append('''key'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''key'''    )
            elif m_name == "_query_dense":
                # attention query
                trace.append('''query'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''query'''    )
            elif m_name == "_value_dense":
                # attention value
                trace.append('''value'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''value'''    )
            elif m_name == "_intermediate_dense":
                # attention intermediate dense
                trace.extend(['''intermediate''', '''dense''']    )
                __a							=    getattr(_UpperCAmelCase       ,			'''intermediate'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''dense'''    )
            elif m_name == "_output_layer_norm":
                # output layer norm
                trace.append('''output'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''output'''    )
            # weights & biases
            elif m_name in ["bias", "beta"]:
                trace.append('''bias'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''bias'''    )
            elif m_name in ["kernel", "gamma"]:
                trace.append('''weight'''    )
                __a							=    getattr(_UpperCAmelCase       ,			'''weight'''    )
            else:
                logger.warning(f'Ignored {m_name}'    )
        # for certain layers reshape is necessary
        __a							=    '''.'''.join(_UpperCAmelCase    )
        if re.match(R'''(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)'''       ,			_UpperCAmelCase    ) or re.match(
            R'''(\S+)\.attention\.output\.dense\.weight'''       ,			_UpperCAmelCase    ):
            __a							=    array.reshape(pointer.data.shape    )
        if "kernel" in full_name:
            __a							=    array.transpose()
        if pointer.shape == array.shape:
            __a							=    torch.from_numpy(_UpperCAmelCase    )
        else:
            raise ValueError(
                f'Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:'
                f' {array.shape}'    )
        logger.info(f'Successfully set variable {full_name} to PyTorch layer {trace}'    )
    return model





def 					__snake_case     (      _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase    ):
    # Instantiate model
    logger.info(f'Loading model based on config from {config_path}...'    )
    __a							=    BertConfig.from_json_file(_UpperCAmelCase    )
    __a							=    BertModel(_UpperCAmelCase    )

    # Load weights from checkpoint
    logger.info(f'Loading weights from checkpoint {tf_checkpoint_path}...'    )
    load_tfa_weights_in_bert(_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase    )

    # Save pytorch-model
    logger.info(f'Saving PyTorch model to {pytorch_dump_path}...'    )
    torch.save(model.state_dict()       ,			_UpperCAmelCase    )


if __name__ == "__main__":
 __snake_case						:Dict						=					argparse.ArgumentParser()
 parser.add_argument(
     '''--tf_checkpoint_path''', type=str, required=True, help='''Path to the TensorFlow 2.x checkpoint path.'''
 )
 parser.add_argument(
     '''--bert_config_file''',
     type=str,
     required=True,
     help='''The config json file corresponding to the BERT model. This specifies the model architecture.''',
 )
 parser.add_argument(
     '''--pytorch_dump_path''',
     type=str,
     required=True,
     help='''Path to the output PyTorch model (must include filename).''',
 )
 __snake_case						:List[Any]						=					parser.parse_args()
 convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)










import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


__snake_case						:Optional[Any]						=					logging.get_logger(__name__)


__snake_case						:List[Any]						=					'''▁'''

__snake_case						:List[Any]						=					{'''vocab_file''': '''spiece.model'''}

__snake_case						:Tuple						=					{
    '''vocab_file''': {
        '''google/reformer-crime-and-punishment''': (
            '''https://huggingface.co/google/reformer-crime-and-punishment/resolve/main/spiece.model'''
        )
    }
}

__snake_case						:List[Any]						=					{
    '''google/reformer-crime-and-punishment''': 52_4288,
}

class    _A	(     __UpperCAmelCase   ):
       UpperCamelCase__  :   int						 =						VOCAB_FILES_NAMES
       UpperCamelCase__  :   str						 =						PRETRAINED_VOCAB_FILES_MAP
       UpperCamelCase__  :   Dict						 =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       UpperCamelCase__  :   Optional[int]						 =						['''input_ids''', '''attention_mask''']



       def __init__(		self							:					Any					,		__SCREAMING_SNAKE_CASE							:					str					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]="</s>"					,		__SCREAMING_SNAKE_CASE							:					List[Any]="<unk>"					,		__SCREAMING_SNAKE_CASE							:					Any=[]					,		__SCREAMING_SNAKE_CASE							:					Optional[Dict[str, Any]] = None					,		**__SCREAMING_SNAKE_CASE							:					List[Any]					,		):


           '''simple docstring'''

           __a							=    {} if sp_model_kwargs is None else sp_model_kwargs

           super().__init__(
               eos_token=__SCREAMING_SNAKE_CASE					,		unk_token=__SCREAMING_SNAKE_CASE					,		additional_special_tokens=__SCREAMING_SNAKE_CASE					,		sp_model_kwargs=self.sp_model_kwargs					,		**__SCREAMING_SNAKE_CASE					,		)

           __a							=    vocab_file
           __a							=    spm.SentencePieceProcessor(**self.sp_model_kwargs)
           self.sp_model.Load(__SCREAMING_SNAKE_CASE)



       @property
       def 		_lowerCamelCase      (		self							:					Optional[Any]):


           '''simple docstring'''

           return self.sp_model.get_piece_size()



       def 		_lowerCamelCase      (		self							:					List[str]):


           '''simple docstring'''

           __a							=    {self.convert_ids_to_tokens(__SCREAMING_SNAKE_CASE): i for i in range(self.vocab_size)}
           vocab.update(self.added_tokens_encoder)
           return vocab



       def __getstate__(		self							:					Dict):


           '''simple docstring'''

           __a							=    self.__dict__.copy()
           __a							=    None
           return state



       def __setstate__(		self							:					Dict					,		__SCREAMING_SNAKE_CASE							:					List[Any]):


           '''simple docstring'''

           __a							=    d

           # for backward compatibility
           if not hasattr(self					,		'''sp_model_kwargs'''):
               __a							=    {}

           __a							=    spm.SentencePieceProcessor(**self.sp_model_kwargs)
           self.sp_model.Load(self.vocab_file)



       def 		_lowerCamelCase      (		self							:					Union[str, Any]					,		__SCREAMING_SNAKE_CASE							:					str):


           '''simple docstring'''

           return self.sp_model.encode(__SCREAMING_SNAKE_CASE					,		out_type=__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase      (		self							:					Optional[Any]					,		__SCREAMING_SNAKE_CASE							:					List[Any]):


           '''simple docstring'''

           return self.sp_model.piece_to_id(__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase      (		self							:					List[str]					,		__SCREAMING_SNAKE_CASE							:					List[Any]):


           '''simple docstring'''

           if index < self.sp_model.get_piece_size():
               __a							=    self.sp_model.IdToPiece(__SCREAMING_SNAKE_CASE)
           return token



       def 		_lowerCamelCase      (		self							:					List[Any]					,		__SCREAMING_SNAKE_CASE							:					Any):


           '''simple docstring'''

           __a							=    []
           __a							=    ''''''
           for token in tokens:
               # make sure that special tokens are not decoded using sentencepiece model
               if token in self.all_special_tokens:
                   out_string += self.sp_model.decode(__SCREAMING_SNAKE_CASE) + token
                   __a							=    []
               else:
                   current_sub_tokens.append(__SCREAMING_SNAKE_CASE)
           out_string += self.sp_model.decode(__SCREAMING_SNAKE_CASE)
           return out_string.strip()





       def 		_lowerCamelCase      (		self							:					str					,		__SCREAMING_SNAKE_CASE							:					str					,		__SCREAMING_SNAKE_CASE							:					Optional[str] = None):


           '''simple docstring'''

           if not os.path.isdir(__SCREAMING_SNAKE_CASE):
               logger.error(F'Vocabulary path ({save_directory}) should be a directory')
               return
           __a							=    os.path.join(
               __SCREAMING_SNAKE_CASE					,		(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file'''])

           if os.path.abspath(self.vocab_file) != os.path.abspath(__SCREAMING_SNAKE_CASE) and os.path.isfile(self.vocab_file):
               copyfile(self.vocab_file					,		__SCREAMING_SNAKE_CASE)
           elif not os.path.isfile(self.vocab_file):
               with open(__SCREAMING_SNAKE_CASE					,		'''wb''') as fi:
                   __a							=    self.sp_model.serialized_model_proto()
                   fi.write(__SCREAMING_SNAKE_CASE)

           return (out_vocab_file,)





