





import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_convbert import ConvBertTokenizer


_snake_case  :   List[Any]						   =  logging.get_logger(__name__)

_snake_case  :   Tuple						   =  {'vocab_file': 'vocab.txt'}

_snake_case  :   Any						   =  {
    'vocab_file': {
        'YituTech/conv-bert-base': 'https://huggingface.co/YituTech/conv-bert-base/resolve/main/vocab.txt',
        'YituTech/conv-bert-medium-small': (
            'https://huggingface.co/YituTech/conv-bert-medium-small/resolve/main/vocab.txt'
        ),
        'YituTech/conv-bert-small': 'https://huggingface.co/YituTech/conv-bert-small/resolve/main/vocab.txt',
    }
}

_snake_case  :   List[str]						   =  {
    'YituTech/conv-bert-base': 512,
    'YituTech/conv-bert-medium-small': 512,
    'YituTech/conv-bert-small': 512,
}


_snake_case  :   Union[str, Any]						   =  {
    'YituTech/conv-bert-base': {'do_lower_case': True},
    'YituTech/conv-bert-medium-small': {'do_lower_case': True},
    'YituTech/conv-bert-small': {'do_lower_case': True},
}



class     _UpperCAmelCase   (   _a   ):

		"""simple docstring"""


		a_      =			VOCAB_FILES_NAMES
		a_      =			PRETRAINED_VOCAB_FILES_MAP
		a_      =			PRETRAINED_INIT_CONFIGURATION
		a_      =			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
		a_      =			ConvBertTokenizer






		def __init__(    self			:		Optional[int] ,  lowerCAmelCase_			:		str=None ,  lowerCAmelCase_			:		List[str]=None ,  lowerCAmelCase_			:		Union[str, Any]=True ,  lowerCAmelCase_			:		Tuple="[UNK]" ,  lowerCAmelCase_			:		Union[str, Any]="[SEP]" ,  lowerCAmelCase_			:		Dict="[PAD]" ,  lowerCAmelCase_			:		Tuple="[CLS]" ,  lowerCAmelCase_			:		Dict="[MASK]" ,  lowerCAmelCase_			:		Dict=True ,  lowerCAmelCase_			:		str=None ,  **lowerCAmelCase_			:		Any ,  )							->						int:
					super().__init__(
					    _a ,  tokenizer_file=_a ,  do_lower_case=_a ,  unk_token=_a ,  sep_token=_a ,  pad_token=_a ,  cls_token=_a ,  mask_token=_a ,  tokenize_chinese_chars=_a ,  strip_accents=_a ,  **_a ,  )

					__lowerCAmelCase							     =     json.loads(self.backend_tokenizer.normalizer.__getstate__()					)
					if (
					    normalizer_state.get('lowercase' ,  _a					) != do_lower_case
					    or normalizer_state.get('strip_accents' ,  _a					) != strip_accents
					    or normalizer_state.get('handle_chinese_chars' ,  _a					) != tokenize_chinese_chars
					):
								__lowerCAmelCase							     =     getattr(_a ,  normalizer_state.pop('type'					)					)
								__lowerCAmelCase							     =     do_lower_case
								__lowerCAmelCase							     =     strip_accents
								__lowerCAmelCase							     =     tokenize_chinese_chars
								__lowerCAmelCase							     =     normalizer_class(**_a					)

					__lowerCAmelCase							     =     do_lower_case






		def     lowercase		(    self			:		List[str] ,  lowerCAmelCase_			:		Tuple ,  lowerCAmelCase_			:		Union[str, Any]=None					)							->						int:
					__lowerCAmelCase							     =     [self.cls_token_id] + token_ids_a + [self.sep_token_id]

					if token_ids_a:
								output += token_ids_a + [self.sep_token_id]

					return output






		def     lowercase		(    self			:		int ,  lowerCAmelCase_			:		List[int] ,  lowerCAmelCase_			:		Optional[List[int]] = None					)							->						List[str]:
					__lowerCAmelCase							     =     [self.sep_token_id]
					__lowerCAmelCase							     =     [self.cls_token_id]
					if token_ids_a is None:
								return len(cls + token_ids_a + sep					) * [0]
					return len(cls + token_ids_a + sep					) * [0] + len(token_ids_a + sep					) * [1]






		def     lowercase		(    self			:		Dict ,  lowerCAmelCase_			:		str ,  lowerCAmelCase_			:		Optional[str] = None					)							->						Dict:
					__lowerCAmelCase							     =     self._tokenizer.model.save(_a ,  name=_a					)
					return tuple(_a					)








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_   (			lowerCAmelCase_				:      Dict[str, torch.Tensor]     ):
   __lowerCAmelCase							     =     []
   __lowerCAmelCase							     =     []
   __lowerCAmelCase							     =     []

   for rt in rc.restypes:
      __lowerCAmelCase							     =     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]     )
      __lowerCAmelCase							     =     {name: i for i, name in enumerate(lowerCAmelCase_     )}
      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     )

   __lowerCAmelCase							     =     torch.tensor(
       lowerCAmelCase_,     dtype=torch.intaa,     device=protein['aatype'].device,     )
   __lowerCAmelCase							     =     torch.tensor(
       lowerCAmelCase_,     dtype=torch.intaa,     device=protein['aatype'].device,     )
   __lowerCAmelCase							     =     torch.tensor(
       lowerCAmelCase_,     dtype=torch.floataa,     device=protein['aatype'].device,     )
   __lowerCAmelCase							     =     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
   __lowerCAmelCase							     =     restype_atomaa_to_atomaa[protein_aatype]
   __lowerCAmelCase							     =     restype_atomaa_mask[protein_aatype]

   __lowerCAmelCase							     =     residx_atomaa_mask
   __lowerCAmelCase							     =     residx_atomaa_to_atomaa.long()

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

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

   __lowerCAmelCase							     =     restype_atomaa_mask[protein_aatype]
   __lowerCAmelCase							     =     residx_atomaa_mask

   return protein




def   a_   (			lowerCAmelCase_				:      Dict[str, torch.Tensor]     ):
   __lowerCAmelCase							     =     tree_map(lambda lowerCAmelCase_     : torch.tensor(lowerCAmelCase_,     device=batch['aatype'].device     ),     lowerCAmelCase_,     np.ndarray     )
   __lowerCAmelCase							     =     tensor_tree_map(lambda lowerCAmelCase_     : np.array(lowerCAmelCase_     ),     make_atomaa_masks(lowerCAmelCase_     )     )
   return out

