





from typing import List, Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_herbert import HerbertTokenizer


_snake_case  :   int						   =  logging.get_logger(__name__)

_snake_case  :   Dict						   =  {'vocab_file': 'vocab.json', 'merges_file': 'merges.txt', 'tokenizer_file': 'tokenizer.json'}

_snake_case  :   Any						   =  {
    'vocab_file': {
        'allegro/herbert-base-cased': 'https://huggingface.co/allegro/herbert-base-cased/resolve/main/vocab.json'
    },
    'merges_file': {
        'allegro/herbert-base-cased': 'https://huggingface.co/allegro/herbert-base-cased/resolve/main/merges.txt'
    },
}

_snake_case  :   Tuple						   =  {'allegro/herbert-base-cased': 514}
_snake_case  :   List[str]						   =  {}



class     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


  a_      =			VOCAB_FILES_NAMES
  a_      =			PRETRAINED_VOCAB_FILES_MAP
  a_      =			PRETRAINED_INIT_CONFIGURATION
  a_      =			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  a_      =			HerbertTokenizer






  def __init__(    self			:		Union[str, Any] ,  lowerCAmelCase_			:		Dict=None ,  lowerCAmelCase_			:		Tuple=None ,  lowerCAmelCase_			:		List[str]=None ,  lowerCAmelCase_			:		List[Any]="<s>" ,  lowerCAmelCase_			:		Dict="<unk>" ,  lowerCAmelCase_			:		Dict="<pad>" ,  lowerCAmelCase_			:		Optional[int]="<mask>" ,  lowerCAmelCase_			:		Any="</s>" ,  **lowerCAmelCase_			:		Optional[Any] ,  )							->						str:
     super().__init__(
         lowerCAmelCase_ ,  lowerCAmelCase_ ,  tokenizer_file=lowerCAmelCase_ ,  cls_token=lowerCAmelCase_ ,  unk_token=lowerCAmelCase_ ,  pad_token=lowerCAmelCase_ ,  mask_token=lowerCAmelCase_ ,  sep_token=lowerCAmelCase_ ,  **lowerCAmelCase_ ,  )






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

     return cls + token_ids_a + sep + token_ids_a + sep






  def     lowercase		(    self			:		Tuple ,  lowerCAmelCase_			:		List[int] ,  lowerCAmelCase_			:		Optional[List[int]] = None ,  lowerCAmelCase_			:		bool = False					)							->						List[int]:
     if already_has_special_tokens:
        return super().get_special_tokens_mask(
            token_ids_a=lowerCAmelCase_ ,  token_ids_a=lowerCAmelCase_ ,  already_has_special_tokens=lowerCAmelCase_					)

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






  def     lowercase		(    self			:		Any ,  lowerCAmelCase_			:		List[int] ,  lowerCAmelCase_			:		Optional[List[int]] = None					)							->						List[int]:
     __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			:		str ,  lowerCAmelCase_			:		str ,  lowerCAmelCase_			:		Optional[str] = None					)							->						Tuple[str]:
     __lowerCAmelCase							     =     self._tokenizer.model.save(lowerCAmelCase_ ,  name=lowerCAmelCase_					)
     return tuple(lowerCAmelCase_					)








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

