






import io
import os
import unicodedata
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

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


UpperCAmelCase__							  =						logging.get_logger(__name__)

UpperCAmelCase__							  =						"▁"

UpperCAmelCase__							  =						{"vocab_file": "vocab.txt", "sentencepiece_model_ckpt": "sentencepiece.bpe.model"}

UpperCAmelCase__							  =						{
    "sentencepiece_model_file": "sentencepiece.bpe.model",
    "vocab_file": "vocab.txt",
}

UpperCAmelCase__							  =						{
    "vocab_file": {
        "ernie-m-base": "https://huggingface.co/susnato/ernie-m-base_pytorch/blob/main/vocab.txt",
        "ernie-m-large": "https://huggingface.co/susnato/ernie-m-base_pytorch/blob/main/vocab.txt",
    },
    "sentencepiece_model_file": {
        "ernie-m-base": "https://huggingface.co/susnato/ernie-m-base_pytorch/blob/main/sentencepiece.bpe.model",
        "ernie-m-large": "https://huggingface.co/susnato/ernie-m-base_pytorch/blob/main/sentencepiece.bpe.model",
    },
}

UpperCAmelCase__							  =						{
    "ernie-m-base": 514,
    "ernie-m-large": 514,
}

UpperCAmelCase__							  =						{
    "ernie-m-base": {"do_lower_case": False},
    "ernie-m-large": {"do_lower_case": False},
}






class 						lowercase_   (      lowercase  ):






      '''simple docstring'''




      __snake_case	    =    ["input_ids"]

      __snake_case	    =    VOCAB_FILES_NAMES
      __snake_case	    =    PRETRAINED_INIT_CONFIGURATION
      __snake_case	    =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __snake_case	    =    PRETRAINED_VOCAB_FILES_MAP
      __snake_case	    =    RESOURCE_FILES_NAMES







      def __init__( self :		Union[str, Any]					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		Optional[int]=None					,   __UpperCAmelCase :		Any=False					,   __UpperCAmelCase :		str="utf8"					,   __UpperCAmelCase :		int="[UNK]"					,   __UpperCAmelCase :		Union[str, Any]="[SEP]"					,   __UpperCAmelCase :		List[Any]="[PAD]"					,   __UpperCAmelCase :		List[str]="[CLS]"					,   __UpperCAmelCase :		List[str]="[MASK]"					,   __UpperCAmelCase :		Optional[Dict[str, Any]] = None					,   **__UpperCAmelCase :		Optional[int]					,   )       ->None:




       """simple docstring"""


       a						    = {} if sp_model_kwargs is None else sp_model_kwargs
       super().__init__(
           do_lower_case=__UpperCAmelCase					,   unk_token=__UpperCAmelCase					,   sep_token=__UpperCAmelCase					,   pad_token=__UpperCAmelCase					,   cls_token=__UpperCAmelCase					,   mask_token=__UpperCAmelCase					,   vocab_file=__UpperCAmelCase					,   encoding=__UpperCAmelCase					,   sp_model_kwargs=self.sp_model_kwargs					,   **__UpperCAmelCase					,   )
       a						    = do_lower_case
       a						    = sentencepiece_model_ckpt
       a						    = spm.SentencePieceProcessor(**self.sp_model_kwargs    )
       self.sp_model.Load(__UpperCAmelCase    )

       # to mimic paddlenlp.transformers.ernie_m.tokenizer.ErnieMTokenizer functioning
       if vocab_file is not None:
        a						    = self.load_vocab(filepath=__UpperCAmelCase    )
       else:
        a						    = {self.sp_model.id_to_piece(__UpperCAmelCase    ): id for id in range(self.sp_model.get_piece_size()    )}
       a						    = {v: k for k, v in self.vocab.items()}







      def  __lowerCAmelCase   ( self :		str					,   __UpperCAmelCase :		str    )       ->Any:




       """simple docstring"""


       if text is None:
        return None

       a						    = self.tokenize(__UpperCAmelCase    )
       a    ,	a						    = '''''', []

       for i, ch in enumerate(__UpperCAmelCase    ):
        if ch in self.SP_CHAR_MAPPING:
         a						    = self.SP_CHAR_MAPPING.get(__UpperCAmelCase    )
        else:
         a						    = unicodedata.normalize('''NFKC'''					,   __UpperCAmelCase    )
        if self.is_whitespace(__UpperCAmelCase    ):
         continue
        normalized_text += ch
        char_mapping.extend([i] * len(__UpperCAmelCase    )    )

       a    ,	a    ,	a						    = normalized_text, [], 0

       if self.do_lower_case:
        a						    = text.lower()

       for token in split_tokens:
        if token[:1] == "▁":
         a						    = token[1:]
        a						    = text[offset:].index(__UpperCAmelCase    ) + offset
        a						    = start + len(__UpperCAmelCase    )

        token_mapping.append((char_mapping[start], char_mapping[end - 1] + 1)    )
        a						    = end
       return token_mapping







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




       """simple docstring"""


       return len(self.vocab    )







      def  __lowerCAmelCase   ( self :		str    )       ->Dict:




       """simple docstring"""


       return dict(self.vocab					,   **self.added_tokens_encoder    )







      def __getstate__( self :		Tuple    )       ->Dict:




       """simple docstring"""


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







      def __setstate__( self :		Any					,   __UpperCAmelCase :		Any    )       ->List[str]:




       """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.sentencepiece_model_ckpt    )







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




       """simple docstring"""


       return "".join((self.SP_CHAR_MAPPING.get(__UpperCAmelCase					,   __UpperCAmelCase    ) for c in text)    )







      def  __lowerCAmelCase   ( self :		Tuple					,   __UpperCAmelCase :		Optional[int]					,   __UpperCAmelCase :		Dict=False					,   __UpperCAmelCase :		Optional[int]=64					,   __UpperCAmelCase :		Union[str, Any]=0.1    )       ->List[str]:




       """simple docstring"""


       if self.sp_model_kwargs.get('''enable_sampling'''    ) is True:
        a						    = True
       if self.sp_model_kwargs.get('''alpha'''    ) is not None:
        a						    = self.sp_model_kwargs.get('''alpha'''    )
       if self.sp_model_kwargs.get('''nbest_size'''    ) is not None:
        a						    = self.sp_model_kwargs.get('''nbest_size'''    )

       if not enable_sampling:
        a						    = self.sp_model.EncodeAsPieces(__UpperCAmelCase    )
       else:
        a						    = self.sp_model.SampleEncodeAsPieces(__UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase    )
       a						    = []
       for pi, piece in enumerate(__UpperCAmelCase    ):
        if piece == SPIECE_UNDERLINE:
         if not pieces[pi + 1].startswith(__UpperCAmelCase    ) and pi != 0:
          new_pieces.append(__UpperCAmelCase    )
          continue
         else:
          continue
        a						    = 0
        for i, chunk in enumerate(__UpperCAmelCase    ):
         if chunk == SPIECE_UNDERLINE:
          continue
         if self.is_ch_char(__UpperCAmelCase    ) or self.is_punct(__UpperCAmelCase    ):
          if i > lst_i and piece[lst_i:i] != SPIECE_UNDERLINE:
           new_pieces.append(piece[lst_i:i]    )
          new_pieces.append(__UpperCAmelCase    )
          a						    = i + 1
         elif chunk.isdigit() and i > 0 and not piece[i - 1].isdigit():
          if i > lst_i and piece[lst_i:i] != SPIECE_UNDERLINE:
           new_pieces.append(piece[lst_i:i]    )
          a						    = i
         elif not chunk.isdigit() and i > 0 and piece[i - 1].isdigit():
          if i > lst_i and piece[lst_i:i] != SPIECE_UNDERLINE:
           new_pieces.append(piece[lst_i:i]    )
          a						    = i
        if len(__UpperCAmelCase    ) > lst_i:
         new_pieces.append(piece[lst_i:]    )
       return new_pieces







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




       """simple docstring"""


       a						    = ''''''.join(__UpperCAmelCase    ).replace(__UpperCAmelCase					,   ''' '''    ).strip()
       return out_string







      def  __lowerCAmelCase   ( self :		Optional[Any]					,   __UpperCAmelCase :		Optional[Any]    )       ->Tuple:




       """simple docstring"""


       a						    = self.convert_ids_to_tokens(__UpperCAmelCase    )
       a						    = ''''''.join(__UpperCAmelCase    ).replace(__UpperCAmelCase					,   ''' '''    ).strip()
       return out_string







      def  __lowerCAmelCase   ( self :		int					,   __UpperCAmelCase :		int    )       ->int:




       """simple docstring"""


       return self.vocab.get(__UpperCAmelCase					,   self.vocab.get(self.unk_token    )    )







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




       """simple docstring"""


       return self.reverse_vocab.get(__UpperCAmelCase					,   self.unk_token    )







      def  __lowerCAmelCase   ( self :		Union[str, Any]					,   __UpperCAmelCase :		int					,   __UpperCAmelCase :		Any=None    )       ->List[Any]:




       """simple docstring"""


       if token_ids_a is None:
        return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
       a						    = [self.cls_token_id]
       a						    = [self.sep_token_id]
       return _cls + token_ids_a + _sep + _sep + token_ids_a + _sep







      def  __lowerCAmelCase   ( self :		Any					,   __UpperCAmelCase :		int					,   __UpperCAmelCase :		int=None    )       ->Optional[int]:




       """simple docstring"""


       if offset_mapping_a is None:
        return [(0, 0)] + offset_mapping_a + [(0, 0)]

       return [(0, 0)] + offset_mapping_a + [(0, 0), (0, 0)] + offset_mapping_a + [(0, 0)]







      def  __lowerCAmelCase   ( self :		Any					,   __UpperCAmelCase :		Dict					,   __UpperCAmelCase :		Optional[Any]=None					,   __UpperCAmelCase :		int=False    )       ->List[Any]:




       """simple docstring"""


       if already_has_special_tokens:
        if token_ids_a is not None:
         raise ValueError(
             '''You should not supply a second sequence if the provided sequence of '''
             '''ids is already formatted with special tokens for the model.'''    )
        return [1 if x in [self.sep_token_id, self.cls_token_id] else 0 for x in token_ids_a]

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







      def  __lowerCAmelCase   ( self :		Any					,   __UpperCAmelCase :		List[int]					,   __UpperCAmelCase :		Optional[List[int]] = None    )       ->List[int]:




       """simple docstring"""


       if token_ids_a is None:
        # [CLS] X [SEP]
        return (len(__UpperCAmelCase    ) + 2) * [0]

       # [CLS] A [SEP] [SEP] B [SEP]
       return [0] * (len(__UpperCAmelCase    ) + 1) + [1] * (len(__UpperCAmelCase    ) + 3)







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




       """simple docstring"""


       if "\u4e00" <= char <= "\u9fff":
        return True
       return False







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




       """simple docstring"""


       if ("a" <= char <= "z") or ("A" <= char <= "Z"):
        return True
       return False







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




       """simple docstring"""


       if char in ",;:.?!~，；：。？！《》【】":
        return True
       return False







      def  __lowerCAmelCase   ( self :		Union[str, Any]					,   __UpperCAmelCase :		str    )       ->Tuple:




       """simple docstring"""


       if char == " " or char == "\t" or char == "\n" or char == "\r":
        return True
       if len(__UpperCAmelCase    ) == 1:
        a						    = unicodedata.category(__UpperCAmelCase    )
        if cat == "Zs":
         return True
       return False







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




       """simple docstring"""


       a						    = {}
       with io.open(__UpperCAmelCase					,   '''r'''					,   encoding='''utf-8'''    ) as f:
        for index, line in enumerate(__UpperCAmelCase    ):
         a						    = line.rstrip('''\n'''    )
         a						    = int(__UpperCAmelCase    )

       return token_to_idx







      def  __lowerCAmelCase   ( self :		Dict					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		Optional[str] = None    )       ->Tuple[str]:




       """simple docstring"""


       a						    = 0
       if os.path.isdir(__UpperCAmelCase    ):
        a						    = os.path.join(
            __UpperCAmelCase					,   (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']    )
       else:
        a						    = (filename_prefix + '''-''' if filename_prefix else '''''') + save_directory
       with open(__UpperCAmelCase					,   '''w'''					,   encoding='''utf-8'''    ) as writer:
        for token, token_index in sorted(self.vocab.items()					,   key=lambda __UpperCAmelCase    : kv[1]    ):
         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!'''    )
          a						    = token_index
         writer.write(token + '''\n'''    )
         index += 1

       a						    = os.path.join(__UpperCAmelCase					,   '''sentencepiece.bpe.model'''    )
       with open(__UpperCAmelCase					,   '''wb'''    ) as fi:
        a						    = self.sp_model.serialized_model_proto()
        fi.write(__UpperCAmelCase    )

       return (vocab_file,)










from typing import Dict

import numpy as np
import torch

from . import residue_constants as rc
from .tensor_utils import tensor_tree_map, tree_map
def 		_a       ( a					:Dict[str, torch.Tensor]					)				->     Dict[str, torch.Tensor]:
 a						    = []
 a						    = []
 a						    = []

 for rt in rc.restypes:
  a						    = rc.restype_name_to_atomaa_names[rc.restype_atoa[rt]]
  restype_atomaa_to_atomaa_list.append([(rc.atom_order[name] if name else 0) for name in atom_names]					)
  a						    = {name: i for i, name in enumerate(a					)}
  restype_atomaa_to_atomaa_list.append(
      [(atom_name_to_idxaa[name] if name in atom_name_to_idxaa else 0) for name in rc.atom_types]					)

  restype_atomaa_mask_list.append([(1.0 if name else 0.0) for name in atom_names]					)

 # Add dummy mapping for restype 'UNK'
 restype_atomaa_to_atomaa_list.append([0] * 14					)
 restype_atomaa_to_atomaa_list.append([0] * 37					)
 restype_atomaa_mask_list.append([0.0] * 14					)

 a						    = torch.tensor(
     a      ,				dtype=torch.intaa      ,				device=protein['''aatype'''].device      ,				)
 a						    = torch.tensor(
     a      ,				dtype=torch.intaa      ,				device=protein['''aatype'''].device      ,				)
 a						    = torch.tensor(
     a      ,				dtype=torch.floataa      ,				device=protein['''aatype'''].device      ,				)
 a						    = protein['''aatype'''].to(torch.long					)

 # create the mapping for (residx, atom14) --> atom37, i.e. an array
 # with shape (num_res, 14) containing the atom37 indices for this protein
 a						    = restype_atomaa_to_atomaa[protein_aatype]
 a						    = restype_atomaa_mask[protein_aatype]

 a						    = residx_atomaa_mask
 a						    = residx_atomaa_to_atomaa.long()

 # create the gather indices for mapping back
 a						    = restype_atomaa_to_atomaa[protein_aatype]
 a						    = residx_atomaa_to_atomaa.long()

 # create the corresponding mask
 a						    = torch.zeros([21, 37]      ,				dtype=torch.floataa      ,				device=protein['''aatype'''].device					)
 for restype, restype_letter in enumerate(rc.restypes					):
  a						    = rc.restype_atoa[restype_letter]
  a						    = rc.residue_atoms[restype_name]
  for atom_name in atom_names:
   a						    = rc.atom_order[atom_name]
   a						    = 1

 a						    = restype_atomaa_mask[protein_aatype]
 a						    = residx_atomaa_mask

 return protein

def 		_a       ( a					:Dict[str, torch.Tensor]					)				->     Dict[str, np.ndarray]:
 a						    = tree_map(lambda a					: torch.tensor(a      ,				device=batch['''aatype'''].device					)      ,				a      ,				np.ndarray					)
 a						    = tensor_tree_map(lambda a					: np.array(a					)      ,				make_atomaa_masks(a					)					)
 return out


