


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	=							'0.12'  # assumed parallelism: 8


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

							@classmethod
							def 				lowerCamelCase_			(							cls:     int							)      ->       int:





									"""simple docstring"""







									lowercase__						=				TOKEN
									HfFolder.save_token(UpperCamelCase_							)

							@classmethod
							def 				lowerCamelCase_			(							cls:     Optional[Any]							)      ->       List[str]:





									"""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:     List[Any]							)      ->       Tuple:





									"""simple docstring"""







									lowercase__						=				BertConfig(
									    vocab_size=99   ,						hidden_size=32   ,						num_hidden_layers=5   ,						num_attention_heads=4   ,						intermediate_size=37							)
									lowercase__						=				FlaxBertModel(UpperCamelCase_							)
									model.push_to_hub('''test-model-flax'''   ,						use_auth_token=self._token							)

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

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

									for key in base_params.keys():
											lowercase__						=				(base_params[key] - new_params[key]).sum().item()
											self.assertLessEqual(UpperCamelCase_   ,						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(UpperCamelCase_   ,						repo_id='''test-model-flax'''   ,						push_to_hub=UpperCamelCase_   ,						use_auth_token=self._token							)

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

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

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


							def 				lowerCamelCase_			(							self:     Optional[int]							)      ->       Any:





									"""simple docstring"""







									lowercase__						=				BertConfig(
									    vocab_size=99   ,						hidden_size=32   ,						num_hidden_layers=5   ,						num_attention_heads=4   ,						intermediate_size=37							)
									lowercase__						=				FlaxBertModel(UpperCamelCase_							)
									model.push_to_hub('''valid_org/test-model-flax-org'''   ,						use_auth_token=self._token							)

									lowercase__						=				FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''							)

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

									for key in base_params.keys():
											lowercase__						=				(base_params[key] - new_params[key]).sum().item()
											self.assertLessEqual(UpperCamelCase_   ,						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(
											    UpperCamelCase_   ,						repo_id='''valid_org/test-model-flax-org'''   ,						push_to_hub=UpperCamelCase_   ,						use_auth_token=self._token							)

									lowercase__						=				FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''							)

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

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


		"""simple docstring"""




		lowercase__						=				True
		lowercase__						=				flatten_dict(modela.params					)
		lowercase__						=				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:
						lowercase__						=				False

		return models_are_equal


@require_flax
class 	_a							(      unittest.TestCase      ):

							def 				lowerCamelCase_			(							self:     int							)      ->       Optional[int]:





									"""simple docstring"""







									lowercase__						=				BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''							)
									lowercase__						=				FlaxBertModel(UpperCamelCase_							)

									lowercase__						=				'''bert'''
									with tempfile.TemporaryDirectory() as tmp_dir:
											model.save_pretrained(os.path.join(UpperCamelCase_   ,						UpperCamelCase_							)							)

											with self.assertRaises(UpperCamelCase_							):
													lowercase__						=				FlaxBertModel.from_pretrained(UpperCamelCase_							)

											lowercase__						=				FlaxBertModel.from_pretrained(UpperCamelCase_   ,						subfolder=UpperCamelCase_							)

									self.assertTrue(check_models_equal(UpperCamelCase_   ,						UpperCamelCase_							)							)

							def 				lowerCamelCase_			(							self:     str							)      ->       List[str]:





									"""simple docstring"""







									lowercase__						=				BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''							)
									lowercase__						=				FlaxBertModel(UpperCamelCase_							)

									lowercase__						=				'''bert'''
									with tempfile.TemporaryDirectory() as tmp_dir:
											model.save_pretrained(os.path.join(UpperCamelCase_   ,						UpperCamelCase_							)   ,						max_shard_size='''10KB'''							)

											with self.assertRaises(UpperCamelCase_							):
													lowercase__						=				FlaxBertModel.from_pretrained(UpperCamelCase_							)

											lowercase__						=				FlaxBertModel.from_pretrained(UpperCamelCase_   ,						subfolder=UpperCamelCase_							)

									self.assertTrue(check_models_equal(UpperCamelCase_   ,						UpperCamelCase_							)							)

							def 				lowerCamelCase_			(							self:     Any							)      ->       Union[str, Any]:





									"""simple docstring"""







									lowercase__						=				'''bert'''
									lowercase__						=				'''hf-internal-testing/tiny-random-bert-subfolder'''

									with self.assertRaises(UpperCamelCase_							):
											lowercase__						=				FlaxBertModel.from_pretrained(UpperCamelCase_							)

									lowercase__						=				FlaxBertModel.from_pretrained(UpperCamelCase_   ,						subfolder=UpperCamelCase_							)

									self.assertIsNotNone(UpperCamelCase_							)


							def 				lowerCamelCase_			(							self:     str							)      ->       int:





									"""simple docstring"""







									lowercase__						=				'''bert'''
									lowercase__						=				'''hf-internal-testing/tiny-random-bert-sharded-subfolder'''
									with self.assertRaises(UpperCamelCase_							):
											lowercase__						=				FlaxBertModel.from_pretrained(UpperCamelCase_							)

									lowercase__						=				FlaxBertModel.from_pretrained(UpperCamelCase_   ,						subfolder=UpperCamelCase_							)

									self.assertIsNotNone(UpperCamelCase_							)






import json
import logging
import os
import socket

import git
import numpy as np
import torch


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - PID: %(process)d -  %(message)s',
    datefmt='%m/%d/%Y %H:%M:%S',
    level=logging.INFO,
)
lowerCAmelCase	=							logging.getLogger(__name__)
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		lowercase__						=				git.Repo(search_parent_directories=SCREAMING_SNAKE_CASE					)
		lowercase__						=				{
		    '''repo_id''': str(SCREAMING_SNAKE_CASE					),
		    '''repo_sha''': str(repo.head.object.hexsha					),
		    '''repo_branch''': str(repo.active_branch					),
		}

		with open(os.path.join(SCREAMING_SNAKE_CASE						,    '''git_log.json'''					)						,    '''w'''					) as f:
				json.dump(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    indent=4					)
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		if params.n_gpu <= 0:
				lowercase__						=				0
				lowercase__						=				-1
				lowercase__						=				True
				lowercase__						=				False
				return

		assert torch.cuda.is_available()

		logger.info('''Initializing GPUs'''					)
		if params.n_gpu > 1:
				assert params.local_rank != -1

				lowercase__						=				int(os.environ['''WORLD_SIZE''']					)
				lowercase__						=				int(os.environ['''N_GPU_NODE''']					)
				lowercase__						=				int(os.environ['''RANK''']					)

				# number of nodes / node ID
				lowercase__						=				params.world_size // params.n_gpu_per_node
				lowercase__						=				params.global_rank // params.n_gpu_per_node
				lowercase__						=				True

				assert params.n_nodes == int(os.environ['''N_NODES''']					)
				assert params.node_id == int(os.environ['''NODE_RANK''']					)

		# local job (single GPU)
		else:
				assert params.local_rank == -1

				lowercase__						=				1
				lowercase__						=				0
				lowercase__						=				0
				lowercase__						=				0
				lowercase__						=				1
				lowercase__						=				1
				lowercase__						=				False

		# sanity checks
		assert params.n_nodes >= 1
		assert 0 <= params.node_id < params.n_nodes
		assert 0 <= params.local_rank <= params.global_rank < params.world_size
		assert params.world_size == params.n_nodes * params.n_gpu_per_node

		# define whether this is the master process / if we are in multi-node distributed mode
		lowercase__						=				params.node_id == 0 and params.local_rank == 0
		lowercase__						=				params.n_nodes > 1

		# summary
		lowercase__						=				f'--- Global rank: {params.global_rank} - '
		logger.info(PREFIX + '''Number of nodes: %i''' % params.n_nodes					)
		logger.info(PREFIX + '''Node ID        : %i''' % params.node_id					)
		logger.info(PREFIX + '''Local rank     : %i''' % params.local_rank					)
		logger.info(PREFIX + '''World size     : %i''' % params.world_size					)
		logger.info(PREFIX + '''GPUs per node  : %i''' % params.n_gpu_per_node					)
		logger.info(PREFIX + '''Master         : %s''' % str(params.is_master					)					)
		logger.info(PREFIX + '''Multi-node     : %s''' % str(params.multi_node					)					)
		logger.info(PREFIX + '''Multi-GPU      : %s''' % str(params.multi_gpu					)					)
		logger.info(PREFIX + '''Hostname       : %s''' % socket.gethostname()					)

		# set GPU device
		torch.cuda.set_device(params.local_rank					)

		# initialize multi-GPU
		if params.multi_gpu:
				logger.info('''Initializing PyTorch distributed'''					)
				torch.distributed.init_process_group(
				    init_method='''env://'''						,    backend='''nccl'''						,    )
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		np.random.seed(args.seed					)
		torch.manual_seed(args.seed					)
		if args.n_gpu > 0:
				torch.cuda.manual_seed_all(args.seed					)


