


import collections
import os
from typing import List, Optional, Tuple

from transformers.utils import is_jieba_available, requires_backends


if is_jieba_available():
 import jieba

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


__lowerCAmelCase			=							logging.get_logger(__name__)

__lowerCAmelCase			=							{"""vocab_file""": """vocab.txt"""}

__lowerCAmelCase			=							{
    """vocab_file""": {
        """openbmb/cpm-ant-10b""": """https://huggingface.co/openbmb/cpm-ant-10b/blob/main/vocab.txt""",
    },
}

__lowerCAmelCase			=							{
    """openbmb/cpm-ant-10b""": 10_24,
}




def        snake_case_					(  snake_case					)  ->    List[str]:
       lowercase__:       List[str]        =       collections.OrderedDict()
       with open(_snake_case      ,							'r'      ,							encoding='utf-8'					) as reader:
              lowercase__:       Optional[Any]        =       reader.readlines()
       for index, token in enumerate(_snake_case					):
              lowercase__:       Union[str, Any]        =       token.rstrip('\n'					)
              lowercase__:       Tuple        =       index
       return vocab





class   __a      (   __snake_case				):
     def __init__(       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__="<unk>"    ,			lowerCAmelCase__=200			) ->    str:




            '''simple docstring'''
            lowercase__:       Dict        =       vocab
            lowercase__:       Dict        =       unk_token
            lowercase__:       Optional[Any]        =       max_input_chars_per_word
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__			) ->    Dict:




            '''simple docstring'''
            lowercase__:       Any        =       list(a_			)
            if len(a_			) > self.max_input_chars_per_word:
                   return [self.unk_token]

            lowercase__:       Union[str, Any]        =       0
            lowercase__:       Union[str, Any]        =       []
            while start < len(a_			):
                   lowercase__:       int        =       len(a_			)
                   lowercase__:       List[Any]        =       None
                   while start < end:
                          lowercase__:       List[Any]        =       ''''''.join(chars[start:end]			)
                          if substr in self.vocab:
                                 lowercase__:       int        =       substr
                                 break
                          end -= 1
                   if cur_substr is None:
                          sub_tokens.append(self.unk_token			)
                          start += 1
                   else:
                          sub_tokens.append(a_			)
                          lowercase__:       Tuple        =       end

            return sub_tokens





class   __a      (   __snake_case				):
     __lowercase			:							Tuple       					=			VOCAB_FILES_NAMES
     __lowercase			:							List[str]       					=			PRETRAINED_VOCAB_FILES_MAP
     __lowercase			:							Optional[int]       					=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     __lowercase			:							Optional[int]       					=			['input_ids', 'attention_mask']
     __lowercase			:							Optional[Any]       					=			False
     def __init__(       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__="<d>"    ,			lowerCAmelCase__="</d>"    ,			lowerCAmelCase__="<s>"    ,			lowerCAmelCase__="</s>"    ,			lowerCAmelCase__="<pad>"    ,			lowerCAmelCase__="<unk>"    ,			lowerCAmelCase__="</n>"    ,			lowerCAmelCase__="</_>"    ,			lowerCAmelCase__="left"    ,			**lowerCAmelCase__    ,			) ->    Union[str, Any]:




            '''simple docstring'''
            requires_backends(self    ,			['jieba']			)
            super().__init__(
                bod_token=a_    ,			eod_token=a_    ,			bos_token=a_    ,			eos_token=a_    ,			pad_token=a_    ,			unk_token=a_    ,			line_token=a_    ,			space_token=a_    ,			padding_side=a_    ,			**a_    ,			)
            lowercase__:       Dict        =       bod_token
            lowercase__:       Dict        =       eod_token
            lowercase__:       Optional[Any]        =       load_vocab(a_			)
            lowercase__:       Dict        =       self.encoder[space_token]
            lowercase__:       Union[str, Any]        =       self.encoder[line_token]

            del self.encoder[space_token]
            del self.encoder[line_token]

            lowercase__:       Optional[Any]        =       collections.OrderedDict(sorted(self.encoder.items()    ,			key=lambda lowerCAmelCase__			: x[1]			)			)
            lowercase__:       Optional[int]        =       {v: k for k, v in self.encoder.items()}

            lowercase__:       Union[str, Any]        =       WordpieceTokenizer(vocab=self.encoder    ,			unk_token=self.unk_token			)
     @property
     def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Tuple:




            '''simple docstring'''
            return self.encoder[self.bod_token]
     @property
     def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Union[str, Any]:




            '''simple docstring'''
            return self.encoder[self.eod_token]
     @property
     def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Dict:




            '''simple docstring'''
            return self.encoder["\n"]
     @property
     def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Optional[int]:




            '''simple docstring'''
            return len(self.encoder			)
     def 		SCREAMING_SNAKE_CASE__       (       self			) ->    int:




            '''simple docstring'''
            return dict(self.encoder    ,			**self.added_tokens_encoder			)
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__			) ->    Optional[Any]:




            '''simple docstring'''
            lowercase__:       str        =       []
            for x in jieba.cut(a_    ,			cut_all=a_			):
                   output_tokens.extend(self.wordpiece_tokenizer.tokenize(a_			)			)
            return output_tokens
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			**lowerCAmelCase__			) ->    List[str]:




            '''simple docstring'''
            lowercase__:       Union[str, Any]        =       [i for i in token_ids if i >= 0]
            lowercase__:       Optional[Any]        =       [
                x for x in token_ids if x != self.pad_token_id and x != self.eos_token_id and x != self.bos_token_id
            ]
            return super()._decode(a_    ,			**a_			)
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__			) ->    Union[str, Any]:




            '''simple docstring'''
            return token in self.encoder
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__			) ->    List[str]:




            '''simple docstring'''
            return "".join(a_			)
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__			) ->    Union[str, Any]:




            '''simple docstring'''
            return self.encoder.get(a_    ,			self.encoder.get(self.unk_token			)			)
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__			) ->    int:




            '''simple docstring'''
            return self.decoder.get(a_    ,			self.unk_token			)
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__ = None			) ->    int:




            '''simple docstring'''

            if os.path.isdir(a_			):
                   lowercase__:       Any        =       os.path.join(
                       a_    ,			(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']			)
            else:
                   lowercase__:       Dict        =       (filename_prefix + '''-''' if filename_prefix else '''''') + save_directory
            lowercase__:       Any        =       0
            if " " in self.encoder:
                   lowercase__:       Any        =       self.encoder[''' ''']
                   del self.encoder[" "]
            if "\n" in self.encoder:
                   lowercase__:       List[str]        =       self.encoder['''\n''']
                   del self.encoder["\n"]
            lowercase__:       str        =       collections.OrderedDict(sorted(self.encoder.items()    ,			key=lambda lowerCAmelCase__			: x[1]			)			)
            with open(a_    ,			'w'    ,			encoding='utf-8'			) as writer:
                   for token, token_index in self.encoder.items():
                          if index != token_index:
                                 logger.warning(
                                     F'Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive.'
                                     ' Please check that the vocabulary is not corrupted!'			)
                                 lowercase__:       Tuple        =       token_index
                          writer.write(token + '\n'			)
                          index += 1
            return (vocab_file,)
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__ = None			) ->    Optional[int]:




            '''simple docstring'''

            if token_ids_a is None:
                   return [self.bos_token_id] + token_ids_a
            return [self.bos_token_id] + token_ids_a + [self.bos_token_id] + token_ids_a
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__ = None    ,			lowerCAmelCase__ = False			) ->    Tuple:




            '''simple docstring'''

            if already_has_special_tokens:
                   return super().get_special_tokens_mask(
                       token_ids_a=a_    ,			token_ids_a=a_    ,			already_has_special_tokens=a_			)

            if token_ids_a is not None:
                   return [1] + ([0] * len(a_			)) + [1] + ([0] * len(a_			))
            return [1] + ([0] * len(a_			))






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()
__lowerCAmelCase			=							logging.get_logger(__name__)




def        snake_case_					(  snake_case      ,							snake_case      ,							snake_case					)  ->    Any:
       lowercase__:       Dict        =       os.path.abspath(snake_case					)
       logger.info(f'Converting TensorFlow checkpoint from {tf_path}'					)
       # Load weights from TF model
       lowercase__:       Optional[Any]        =       tf.train.list_variables(snake_case					)
       lowercase__:       List[Any]        =       []
       lowercase__:       Tuple        =       []
       lowercase__:       int        =       []
       for full_name, shape in init_vars:
              # logger.info(f"Loading TF weight {name} with shape {shape}")
              lowercase__:       Union[str, Any]        =       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'
                     lowercase__:       str        =       name[1:]
              # figure out how many levels deep the name is
              lowercase__:       Optional[Any]        =       0
              for _name in name:
                     if _name.startswith('layer_with_weights'					):
                            depth += 1
                     else:
                            break
              layer_depth.append(snake_case					)
              # read data
              lowercase__:       Optional[Any]        =       tf.train.load_variable(snake_case      ,							snake_case					)
              names.append('/'.join(snake_case					)					)
              arrays.append(snake_case					)
       logger.info(f'Read a total of {len(snake_case					):,} layers'					)

       # Sanity check
       if len(set(snake_case					)					) != 1:
              raise ValueError(f'Found layer names with different depths (layer depth {list(set(snake_case					)					)})'					)
       lowercase__:       Any        =       list(set(snake_case					)					)[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(snake_case      ,							snake_case					):
              lowercase__:       Optional[int]        =       full_name.split('/'					)
              lowercase__:       List[Any]        =       model
              lowercase__:       Any        =       []
              for i, m_name in enumerate(snake_case					):
                     if m_name == ".ATTRIBUTES":
                            # variable names end with .ATTRIBUTES/VARIABLE_VALUE
                            break
                     if m_name.startswith('layer_with_weights'					):
                            lowercase__:       Any        =       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']					)
                                   lowercase__:       str        =       getattr(snake_case      ,							'embeddings'					)
                                   lowercase__:       int        =       getattr(snake_case      ,							'LayerNorm'					)
                            elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
                                   # encoder layers
                                   trace.extend(['encoder', 'layer', str(layer_num - 4					)]					)
                                   lowercase__:       int        =       getattr(snake_case      ,							'encoder'					)
                                   lowercase__:       List[str]        =       getattr(snake_case      ,							'layer'					)
                                   lowercase__:       Union[str, Any]        =       pointer[layer_num - 4]
                            elif layer_num == config.num_hidden_layers + 4:
                                   # pooler layer
                                   trace.extend(['pooler', 'dense']					)
                                   lowercase__:       Tuple        =       getattr(snake_case      ,							'pooler'					)
                                   lowercase__:       Tuple        =       getattr(snake_case      ,							'dense'					)
                     elif m_name == "embeddings":
                            trace.append('embeddings'					)
                            lowercase__:       Union[str, Any]        =       getattr(snake_case      ,							'embeddings'					)
                            if layer_num == 0:
                                   trace.append('word_embeddings'					)
                                   lowercase__:       Union[str, Any]        =       getattr(snake_case      ,							'word_embeddings'					)
                            elif layer_num == 1:
                                   trace.append('position_embeddings'					)
                                   lowercase__:       Dict        =       getattr(snake_case      ,							'position_embeddings'					)
                            elif layer_num == 2:
                                   trace.append('token_type_embeddings'					)
                                   lowercase__:       Optional[Any]        =       getattr(snake_case      ,							'token_type_embeddings'					)
                            else:
                                   raise ValueError(f'Unknown embedding layer with name {full_name}'					)
                            trace.append('weight'					)
                            lowercase__:       List[str]        =       getattr(snake_case      ,							'weight'					)
                     elif m_name == "_attention_layer":
                            # self-attention layer
                            trace.extend(['attention', 'self']					)
                            lowercase__:       int        =       getattr(snake_case      ,							'attention'					)
                            lowercase__:       Union[str, Any]        =       getattr(snake_case      ,							'self'					)
                     elif m_name == "_attention_layer_norm":
                            # output attention norm
                            trace.extend(['attention', 'output', 'LayerNorm']					)
                            lowercase__:       str        =       getattr(snake_case      ,							'attention'					)
                            lowercase__:       Optional[Any]        =       getattr(snake_case      ,							'output'					)
                            lowercase__:       List[Any]        =       getattr(snake_case      ,							'LayerNorm'					)
                     elif m_name == "_attention_output_dense":
                            # output attention dense
                            trace.extend(['attention', 'output', 'dense']					)
                            lowercase__:       Optional[Any]        =       getattr(snake_case      ,							'attention'					)
                            lowercase__:       Optional[int]        =       getattr(snake_case      ,							'output'					)
                            lowercase__:       Optional[Any]        =       getattr(snake_case      ,							'dense'					)
                     elif m_name == "_output_dense":
                            # output dense
                            trace.extend(['output', 'dense']					)
                            lowercase__:       Union[str, Any]        =       getattr(snake_case      ,							'output'					)
                            lowercase__:       List[Any]        =       getattr(snake_case      ,							'dense'					)
                     elif m_name == "_output_layer_norm":
                            # output dense
                            trace.extend(['output', 'LayerNorm']					)
                            lowercase__:       Any        =       getattr(snake_case      ,							'output'					)
                            lowercase__:       str        =       getattr(snake_case      ,							'LayerNorm'					)
                     elif m_name == "_key_dense":
                            # attention key
                            trace.append('key'					)
                            lowercase__:       Tuple        =       getattr(snake_case      ,							'key'					)
                     elif m_name == "_query_dense":
                            # attention query
                            trace.append('query'					)
                            lowercase__:       List[str]        =       getattr(snake_case      ,							'query'					)
                     elif m_name == "_value_dense":
                            # attention value
                            trace.append('value'					)
                            lowercase__:       Optional[int]        =       getattr(snake_case      ,							'value'					)
                     elif m_name == "_intermediate_dense":
                            # attention intermediate dense
                            trace.extend(['intermediate', 'dense']					)
                            lowercase__:       Any        =       getattr(snake_case      ,							'intermediate'					)
                            lowercase__:       str        =       getattr(snake_case      ,							'dense'					)
                     elif m_name == "_output_layer_norm":
                            # output layer norm
                            trace.append('output'					)
                            lowercase__:       Union[str, Any]        =       getattr(snake_case      ,							'output'					)
                     # weights & biases
                     elif m_name in ["bias", "beta"]:
                            trace.append('bias'					)
                            lowercase__:       str        =       getattr(snake_case      ,							'bias'					)
                     elif m_name in ["kernel", "gamma"]:
                            trace.append('weight'					)
                            lowercase__:       Tuple        =       getattr(snake_case      ,							'weight'					)
                     else:
                            logger.warning(f'Ignored {m_name}'					)
        # for certain layers reshape is necessary
              lowercase__:       Any        =       '.'.join(snake_case					)
              if re.match(R'(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)'      ,							snake_case					) or re.match(
                  R'(\S+)\.attention\.output\.dense\.weight'      ,							snake_case					):
                     lowercase__:       Union[str, Any]        =       array.reshape(pointer.data.shape					)
              if "kernel" in full_name:
                     lowercase__:       str        =       array.transpose()
              if pointer.shape == array.shape:
                     lowercase__:       Union[str, Any]        =       torch.from_numpy(snake_case					)
              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_					(  snake_case      ,							snake_case      ,							snake_case					)  ->    Any:
       # Instantiate model
       logger.info(f'Loading model based on config from {config_path}...'					)
       lowercase__:       int        =       BertConfig.from_json_file(snake_case					)
       lowercase__:       Tuple        =       BertModel(snake_case					)

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

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


if __name__ == "__main__":
 __lowerCAmelCase			=							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).''',
 )
 __lowerCAmelCase			=							parser.parse_args()
 convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)


