






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__ (__lowercase							:Dict[str, torch.Tensor]   )		-> Dict[str, torch.Tensor]:
				_A	:				Any   =      []
				_A	:				List[Any]   =      []
				_A	:				List[Any]   =      []

				for rt in rc.restypes:
								_A	:				Any   =      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	:				Any   =      {name: i for i, name in enumerate(__lowercase   )}
								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	:				Optional[Any]   =      torch.tensor(
				    __lowercase					,  dtype=torch.intaa					,  device=protein['''aatype'''].device					,  )
				_A	:				Dict   =      torch.tensor(
				    __lowercase					,  dtype=torch.intaa					,  device=protein['''aatype'''].device					,  )
				_A	:				Union[str, Any]   =      torch.tensor(
				    __lowercase					,  dtype=torch.floataa					,  device=protein['''aatype'''].device					,  )
				_A	:				str   =      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	:				Tuple   =      restype_atomaa_to_atomaa[protein_aatype]
				_A	:				Any   =      restype_atomaa_mask[protein_aatype]

				_A	:				Tuple   =      residx_atomaa_mask
				_A	:				Any   =      residx_atomaa_to_atomaa.long()

				# create the gather indices for mapping back
				_A	:				str   =      restype_atomaa_to_atomaa[protein_aatype]
				_A	:				Union[str, Any]   =      residx_atomaa_to_atomaa.long()

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

				_A	:				int   =      restype_atomaa_mask[protein_aatype]
				_A	:				Any   =      residx_atomaa_mask

				return protein






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











from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_UpperCamelCase			:	Dict								={
    'configuration_albert': ['ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'AlbertConfig', 'AlbertOnnxConfig'],
}

try:
	if not is_sentencepiece_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_UpperCamelCase			:	List[Any]								=['AlbertTokenizer']

try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_UpperCamelCase			:	List[Any]								=['AlbertTokenizerFast']

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_UpperCamelCase			:	Optional[int]								=[
	    'ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'AlbertForMaskedLM',
	    'AlbertForMultipleChoice',
	    'AlbertForPreTraining',
	    'AlbertForQuestionAnswering',
	    'AlbertForSequenceClassification',
	    'AlbertForTokenClassification',
	    'AlbertModel',
	    'AlbertPreTrainedModel',
	    'load_tf_weights_in_albert',
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_UpperCamelCase			:	List[Any]								=[
	    'TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'TFAlbertForMaskedLM',
	    'TFAlbertForMultipleChoice',
	    'TFAlbertForPreTraining',
	    'TFAlbertForQuestionAnswering',
	    'TFAlbertForSequenceClassification',
	    'TFAlbertForTokenClassification',
	    'TFAlbertMainLayer',
	    'TFAlbertModel',
	    'TFAlbertPreTrainedModel',
	]

try:
	if not is_flax_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_UpperCamelCase			:	Any								=[
	    'FlaxAlbertForMaskedLM',
	    'FlaxAlbertForMultipleChoice',
	    'FlaxAlbertForPreTraining',
	    'FlaxAlbertForQuestionAnswering',
	    'FlaxAlbertForSequenceClassification',
	    'FlaxAlbertForTokenClassification',
	    'FlaxAlbertModel',
	    'FlaxAlbertPreTrainedModel',
	]

if TYPE_CHECKING:
	from .configuration_albert import ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, AlbertConfig, AlbertOnnxConfig

	try:
		if not is_sentencepiece_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_albert import AlbertTokenizer

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_albert_fast import AlbertTokenizerFast

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_albert import (
		    ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    AlbertForMaskedLM,
		    AlbertForMultipleChoice,
		    AlbertForPreTraining,
		    AlbertForQuestionAnswering,
		    AlbertForSequenceClassification,
		    AlbertForTokenClassification,
		    AlbertModel,
		    AlbertPreTrainedModel,
		    load_tf_weights_in_albert,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_albert import (
		    TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFAlbertForMaskedLM,
		    TFAlbertForMultipleChoice,
		    TFAlbertForPreTraining,
		    TFAlbertForQuestionAnswering,
		    TFAlbertForSequenceClassification,
		    TFAlbertForTokenClassification,
		    TFAlbertMainLayer,
		    TFAlbertModel,
		    TFAlbertPreTrainedModel,
		)

	try:
		if not is_flax_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_flax_albert import (
		    FlaxAlbertForMaskedLM,
		    FlaxAlbertForMultipleChoice,
		    FlaxAlbertForPreTraining,
		    FlaxAlbertForQuestionAnswering,
		    FlaxAlbertForSequenceClassification,
		    FlaxAlbertForTokenClassification,
		    FlaxAlbertModel,
		    FlaxAlbertPreTrainedModel,
		)
else:
	import sys

	_UpperCamelCase			:	str								=_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



