




"""simple docstring"""




import tempfile
import unittest

import numpy as np
from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_flax


if is_flax_available():
					import os

					from flax.core.frozen_dict import unfreeze
					from flax.traverse_util import flatten_dict

					from transformers import FlaxBertModel

					lowerCAmelCase				:     Union[str, Any]					=				"""0.12"""  # assumed parallelism: 8






@require_flax
@is_staging_test
class        __magic_name__				(			unittest.TestCase							):






				'''simple docstring'''





				@classmethod
				def   _lowerCAmelCase							(			cls					):





							"""simple docstring"""
							lowerCamelCase      =   TOKEN
							HfFolder.save_token(_a					)





				@classmethod
				def   _lowerCAmelCase							(			cls					):





							"""simple docstring"""
							try:
										delete_repo(token=cls._token						,					repo_id="""test-model-flax"""					)
							except HTTPError:
										pass

							try:
										delete_repo(token=cls._token						,					repo_id="""valid_org/test-model-flax-org"""					)
							except HTTPError:
										pass





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   BertConfig(
							    vocab_size=99						,					hidden_size=32						,					num_hidden_layers=5						,					num_attention_heads=4						,					intermediate_size=37					)
							lowerCamelCase      =   FlaxBertModel(_a					)
							model.push_to_hub("""test-model-flax"""						,					use_auth_token=self._token					)

							lowerCamelCase      =   FlaxBertModel.from_pretrained(f'{USER}/test-model-flax'					)

							lowerCamelCase      =   flatten_dict(unfreeze(model.params					)					)
							lowerCamelCase      =   flatten_dict(unfreeze(new_model.params					)					)

							for key in base_params.keys():
										lowerCamelCase      =   (base_params[key] - new_params[key]).sum().item()
										self.assertLessEqual(_a						,					1e-3						,					msg=f'{key} not identical'					)

							# Reset repo
							delete_repo(token=self._token						,					repo_id="""test-model-flax"""					)

							# Push to hub via save_pretrained
							with tempfile.TemporaryDirectory() as tmp_dir:
										model.save_pretrained(_a						,					repo_id="""test-model-flax"""						,					push_to_hub=_a						,					use_auth_token=self._token					)

							lowerCamelCase      =   FlaxBertModel.from_pretrained(f'{USER}/test-model-flax'					)

							lowerCamelCase      =   flatten_dict(unfreeze(model.params					)					)
							lowerCamelCase      =   flatten_dict(unfreeze(new_model.params					)					)

							for key in base_params.keys():
										lowerCamelCase      =   (base_params[key] - new_params[key]).sum().item()
										self.assertLessEqual(_a						,					1e-3						,					msg=f'{key} not identical'					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   BertConfig(
							    vocab_size=99						,					hidden_size=32						,					num_hidden_layers=5						,					num_attention_heads=4						,					intermediate_size=37					)
							lowerCamelCase      =   FlaxBertModel(_a					)
							model.push_to_hub("""valid_org/test-model-flax-org"""						,					use_auth_token=self._token					)

							lowerCamelCase      =   FlaxBertModel.from_pretrained("""valid_org/test-model-flax-org"""					)

							lowerCamelCase      =   flatten_dict(unfreeze(model.params					)					)
							lowerCamelCase      =   flatten_dict(unfreeze(new_model.params					)					)

							for key in base_params.keys():
										lowerCamelCase      =   (base_params[key] - new_params[key]).sum().item()
										self.assertLessEqual(_a						,					1e-3						,					msg=f'{key} not identical'					)

							# Reset repo
							delete_repo(token=self._token						,					repo_id="""valid_org/test-model-flax-org"""					)

							# Push to hub via save_pretrained
							with tempfile.TemporaryDirectory() as tmp_dir:
										model.save_pretrained(
										    _a						,					repo_id="""valid_org/test-model-flax-org"""						,					push_to_hub=_a						,					use_auth_token=self._token					)

							lowerCamelCase      =   FlaxBertModel.from_pretrained("""valid_org/test-model-flax-org"""					)

							lowerCamelCase      =   flatten_dict(unfreeze(model.params					)					)
							lowerCamelCase      =   flatten_dict(unfreeze(new_model.params					)					)

							for key in base_params.keys():
										lowerCamelCase      =   (base_params[key] - new_params[key]).sum().item()
										self.assertLessEqual(_a						,					1e-3						,					msg=f'{key} not identical'					)




def        a__					(							snake_case__			,					snake_case__							)		->   int:
			lowerCamelCase      =   True
			lowerCamelCase      =   flatten_dict(modela.params							)
			lowerCamelCase      =   flatten_dict(modela.params							)
			for key in flat_params_a.keys():
						if np.sum(np.abs(flat_params_a[key] - flat_params_a[key]							)							) > 1E-4:
									lowerCamelCase      =   False

			return models_are_equal






@require_flax
class        __magic_name__				(			unittest.TestCase							):






				'''simple docstring'''





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""					)
							lowerCamelCase      =   FlaxBertModel(_a					)

							lowerCamelCase      =   """bert"""
							with tempfile.TemporaryDirectory() as tmp_dir:
										model.save_pretrained(os.path.join(_a						,					_a					)					)

										with self.assertRaises(_a					):
													lowerCamelCase      =   FlaxBertModel.from_pretrained(_a					)

										lowerCamelCase      =   FlaxBertModel.from_pretrained(_a						,					subfolder=_a					)

							self.assertTrue(check_models_equal(_a						,					_a					)					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""					)
							lowerCamelCase      =   FlaxBertModel(_a					)

							lowerCamelCase      =   """bert"""
							with tempfile.TemporaryDirectory() as tmp_dir:
										model.save_pretrained(os.path.join(_a						,					_a					)						,					max_shard_size="""10KB"""					)

										with self.assertRaises(_a					):
													lowerCamelCase      =   FlaxBertModel.from_pretrained(_a					)

										lowerCamelCase      =   FlaxBertModel.from_pretrained(_a						,					subfolder=_a					)

							self.assertTrue(check_models_equal(_a						,					_a					)					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   """bert"""
							lowerCamelCase      =   """hf-internal-testing/tiny-random-bert-subfolder"""

							with self.assertRaises(_a					):
										lowerCamelCase      =   FlaxBertModel.from_pretrained(_a					)

							lowerCamelCase      =   FlaxBertModel.from_pretrained(_a						,					subfolder=_a					)

							self.assertIsNotNone(_a					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   """bert"""
							lowerCamelCase      =   """hf-internal-testing/tiny-random-bert-sharded-subfolder"""
							with self.assertRaises(_a					):
										lowerCamelCase      =   FlaxBertModel.from_pretrained(_a					)

							lowerCamelCase      =   FlaxBertModel.from_pretrained(_a						,					subfolder=_a					)

							self.assertIsNotNone(_a					)









"""simple docstring"""




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__					(							snake_case__							)		->   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(snake_case__							)}
						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(
			    snake_case__			,					dtype=torch.intaa			,					device=protein["""aatype"""].device			,					)
			lowerCamelCase      =   torch.tensor(
			    snake_case__			,					dtype=torch.intaa			,					device=protein["""aatype"""].device			,					)
			lowerCamelCase      =   torch.tensor(
			    snake_case__			,					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__					(							snake_case__							)		->   Dict[str, np.ndarray]:
			lowerCamelCase      =   tree_map(lambda snake_case__							: torch.tensor(snake_case__			,					device=batch["""aatype"""].device							)			,					snake_case__			,					np.ndarray							)
			lowerCamelCase      =   tensor_tree_map(lambda snake_case__							: np.array(snake_case__							)			,					make_atomaa_masks(snake_case__							)							)
			return out



