






import unittest

import numpy as np

from transformers import RobertaConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
							from transformers.models.roberta.modeling_flax_roberta import (
							    FlaxRobertaForCausalLM,
							    FlaxRobertaForMaskedLM,
							    FlaxRobertaForMultipleChoice,
							    FlaxRobertaForQuestionAnswering,
							    FlaxRobertaForSequenceClassification,
							    FlaxRobertaForTokenClassification,
							    FlaxRobertaModel,
							)
class       _UpperCAmelCase					(       unittest.TestCase	):

				"""simple docstring"""




				def __init__(     self						:    Dict,  lowerCamelCase						:    Any,  lowerCamelCase						:    Dict=13,  lowerCamelCase						:    Union[str, Any]=7,  lowerCamelCase						:    Optional[int]=True,  lowerCamelCase						:    str=True,  lowerCamelCase						:    Tuple=True,  lowerCamelCase						:    List[str]=True,  lowerCamelCase						:    int=99,  lowerCamelCase						:    Union[str, Any]=32,  lowerCamelCase						:    List[Any]=5,  lowerCamelCase						:    List[str]=4,  lowerCamelCase						:    Optional[int]=37,  lowerCamelCase						:    int="gelu",  lowerCamelCase						:    List[Any]=0.1,  lowerCamelCase						:    Any=0.1,  lowerCamelCase						:    List[str]=512,  lowerCamelCase						:    Optional[int]=16,  lowerCamelCase						:    List[str]=2,  lowerCamelCase						:    Union[str, Any]=0.02,  lowerCamelCase						:    Dict=4,  ):

								'''simple docstring'''
								lowercase__    =			parent
								lowercase__    =			batch_size
								lowercase__    =			seq_length
								lowercase__    =			is_training
								lowercase__    =			use_attention_mask
								lowercase__    =			use_token_type_ids
								lowercase__    =			use_labels
								lowercase__    =			vocab_size
								lowercase__    =			hidden_size
								lowercase__    =			num_hidden_layers
								lowercase__    =			num_attention_heads
								lowercase__    =			intermediate_size
								lowercase__    =			hidden_act
								lowercase__    =			hidden_dropout_prob
								lowercase__    =			attention_probs_dropout_prob
								lowercase__    =			max_position_embeddings
								lowercase__    =			type_vocab_size
								lowercase__    =			type_sequence_label_size
								lowercase__    =			initializer_range
								lowercase__    =			num_choices




				def 				lowercase__   (     self						:    str  ):

								'''simple docstring'''
								lowercase__    =			ids_tensor([self.batch_size, self.seq_length],  self.vocab_size  )

								lowercase__    =			None
								if self.use_attention_mask:
												lowercase__    =			random_attention_mask([self.batch_size, self.seq_length]  )

								lowercase__    =			None
								if self.use_token_type_ids:
												lowercase__    =			ids_tensor([self.batch_size, self.seq_length],  self.type_vocab_size  )

								lowercase__    =			RobertaConfig(
								    vocab_size=self.vocab_size,  hidden_size=self.hidden_size,  num_hidden_layers=self.num_hidden_layers,  num_attention_heads=self.num_attention_heads,  intermediate_size=self.intermediate_size,  hidden_act=self.hidden_act,  hidden_dropout_prob=self.hidden_dropout_prob,  attention_probs_dropout_prob=self.attention_probs_dropout_prob,  max_position_embeddings=self.max_position_embeddings,  type_vocab_size=self.type_vocab_size,  is_decoder=lowerCamelCase,  initializer_range=self.initializer_range,  )

								return config, input_ids, token_type_ids, attention_mask




				def 				lowercase__   (     self						:    str  ):

								'''simple docstring'''
								lowercase__    =			self.prepare_config_and_inputs()
								lowercase__ ,						lowercase__ ,						lowercase__ ,						lowercase__    =			config_and_inputs
								lowercase__    =			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
								return config, inputs_dict




				def 				lowercase__   (     self						:    List[Any]  ):

								'''simple docstring'''
								lowercase__    =			self.prepare_config_and_inputs()
								lowercase__ ,						lowercase__ ,						lowercase__ ,						lowercase__    =			config_and_inputs

								lowercase__    =			True
								lowercase__    =			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
								lowercase__    =			ids_tensor([self.batch_size, self.seq_length],  vocab_size=2  )

								return (
								    config,
								    input_ids,
								    token_type_ids,
								    encoder_hidden_states,
								    encoder_attention_mask,
								)



@require_flax
class       _UpperCAmelCase					(       A__		,unittest.TestCase	):

				"""simple docstring"""




				lowercase__											=       True

				lowercase__											=       (
				    (
				        FlaxRobertaModel,
				        FlaxRobertaForCausalLM,
				        FlaxRobertaForMaskedLM,
				        FlaxRobertaForSequenceClassification,
				        FlaxRobertaForTokenClassification,
				        FlaxRobertaForMultipleChoice,
				        FlaxRobertaForQuestionAnswering,
				    )
				    if is_flax_available()
				    else ()
				)




				def 				lowercase__   (     self						:    int  ):

								'''simple docstring'''
								lowercase__    =			FlaxRobertaModelTester(self  )




				@slow
				def 				lowercase__   (     self						:    Optional[int]  ):

								'''simple docstring'''
								for model_class_name in self.all_model_classes:
												lowercase__    =			model_class_name.from_pretrained('''roberta-base''',  from_pt=lowerCamelCase  )
												lowercase__    =			model(np.ones((1, 1)  )  )
												self.assertIsNotNone(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_				):






				'''simple docstring'''
				lowercase__    =			[]
				lowercase__    =			[]
				lowercase__    =			[]

				for rt in rc.restypes:
								lowercase__    =			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]				)
								lowercase__    =			{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				)

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

				lowercase__    =			residx_atomaa_mask
				lowercase__    =			residx_atomaa_to_atomaa.long()

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

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

				lowercase__    =			restype_atomaa_mask[protein_aatype]
				lowercase__    =			residx_atomaa_mask

				return protein






def   a			(						lowerCamelCase_				):






				'''simple docstring'''
				lowercase__    =			tree_map(lambda lowerCamelCase_				: torch.tensor(lowerCamelCase_		,   device=batch['''aatype'''].device				)		,   lowerCamelCase_		,   np.ndarray				)
				lowercase__    =			tensor_tree_map(lambda lowerCamelCase_				: np.array(lowerCamelCase_				)		,   make_atomaa_masks(lowerCamelCase_				)				)
				return out




