



'''simple docstring'''

import os
import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from huggingface_hub.file_download import http_get
from requests.exceptions import HTTPError

from transformers import (
    AlbertTokenizer,
    AutoTokenizer,
    BertTokenizer,
    BertTokenizerFast,
    GPTaTokenizerFast,
    is_tokenizers_available,
)
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_tokenizers
from transformers.tokenization_utils import Trie


sys.path.append(str(Path(__file__).parent.parent / 'utils'))

from test_module.custom_tokenization import CustomTokenizer  # noqa E402


if is_tokenizers_available():
				from test_module.custom_tokenization_fast import CustomTokenizerFast

class __snake_case( unittest.TestCase   ):


						'''simple docstring'''



						def 							__snake_case       (   self )      ->      Any:
							# A mock response for an HTTP head request to emulate server down
							lowerCAmelCase       =      mock.Mock()
							lowerCAmelCase       =      500
							lowerCAmelCase       =      {}
							lowerCAmelCase       =      HTTPError
							lowerCAmelCase       =      {}

							# Download this model to make sure it's in the cache.
							lowerCAmelCase       =      BertTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bert""" )

							# Under the mock environment we get a 500 error when trying to reach the tokenizer.
							with mock.patch("""requests.Session.request"""      ,			return_value=A_ ) as mock_head:
								lowerCAmelCase       =      BertTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bert""" )
								# This check we did call the fake head request
								mock_head.assert_called()


						@require_tokenizers
						def 							__snake_case       (   self )      ->      Optional[int]:
							# A mock response for an HTTP head request to emulate server down
							lowerCAmelCase       =      mock.Mock()
							lowerCAmelCase       =      500
							lowerCAmelCase       =      {}
							lowerCAmelCase       =      HTTPError
							lowerCAmelCase       =      {}

							# Download this model to make sure it's in the cache.
							lowerCAmelCase       =      GPTaTokenizerFast.from_pretrained("""gpt2""" )

							# Under the mock environment we get a 500 error when trying to reach the tokenizer.
							with mock.patch("""requests.Session.request"""      ,			return_value=A_ ) as mock_head:
								lowerCAmelCase       =      GPTaTokenizerFast.from_pretrained("""gpt2""" )
								# This check we did call the fake head request
								mock_head.assert_called()


						def 							__snake_case       (   self )      ->      int:
							# This test is for deprecated behavior and can be removed in v5
							try:
								lowerCAmelCase       =      tempfile.mktemp()
								with open(A_      ,			"""wb""" ) as f:
									http_get("""https://huggingface.co/albert-base-v1/resolve/main/spiece.model"""      ,			A_ )

								lowerCAmelCase       =      AlbertTokenizer.from_pretrained(A_ )
							finally:
								os.remove(A_ )

							# Supporting this legacy load introduced a weird bug where the tokenizer would load local files if they are in
							# the current folder and have the right name.
							if os.path.isfile("""tokenizer.json""" ):
								# We skip the test if the user has a `tokenizer.json` in this folder to avoid deleting it.
								return
							try:
								with open("""tokenizer.json"""      ,			"""wb""" ) as f:
									http_get("""https://huggingface.co/hf-internal-testing/tiny-random-bert/blob/main/tokenizer.json"""      ,			A_ )
								lowerCAmelCase       =      AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2""" )
								# The tiny random BERT has a vocab size of 1024, tiny gpt2 as a vocab size of 1000
								self.assertEqual(tokenizer.vocab_size      ,			1000 )
								# Tokenizer should depend on the remote checkpoint, not the local tokenizer.json file.

							finally:
								os.remove("""tokenizer.json""" )




						def 							__snake_case       (   self )      ->      Any:
							# This test is for deprecated behavior and can be removed in v5
							lowerCAmelCase       =      AlbertTokenizer.from_pretrained("""https://huggingface.co/albert-base-v1/resolve/main/spiece.model""" )







@is_staging_test
class __snake_case( unittest.TestCase   ):


						'''simple docstring'''



						UpperCAmelCase				:       Optional[Any] 	=						['''[UNK]''', '''[CLS]''', '''[SEP]''', '''[PAD]''', '''[MASK]''', '''bla''', '''blou''']


						@classmethod
						def 							__snake_case       (   cls )      ->      Optional[Any]:
							lowerCAmelCase       =      TOKEN
							HfFolder.save_token(A_ )


						@classmethod
						def 							__snake_case       (   cls )      ->      List[Any]:
							try:
								delete_repo(token=cls._token      ,			repo_id="""test-tokenizer""" )
							except HTTPError:
								pass

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

							try:
								delete_repo(token=cls._token      ,			repo_id="""test-dynamic-tokenizer""" )
							except HTTPError:
								pass


						def 							__snake_case       (   self )      ->      List[Any]:
							with tempfile.TemporaryDirectory() as tmp_dir:
								lowerCAmelCase       =      os.path.join(A_      ,			"""vocab.txt""" )
								with open(A_      ,			"""w"""      ,			encoding="""utf-8""" ) as vocab_writer:
									vocab_writer.write("""""".join([x + """\n""" for x in self.vocab_tokens] ) )
								lowerCAmelCase       =      BertTokenizer(A_ )

							tokenizer.push_to_hub("""test-tokenizer"""      ,			use_auth_token=self._token )
							lowerCAmelCase       =      BertTokenizer.from_pretrained(f'{USER}/test-tokenizer' )
							self.assertDictEqual(new_tokenizer.vocab      ,			tokenizer.vocab )

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

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

							lowerCAmelCase       =      BertTokenizer.from_pretrained(f'{USER}/test-tokenizer' )
							self.assertDictEqual(new_tokenizer.vocab      ,			tokenizer.vocab )


						def 							__snake_case       (   self )      ->      Tuple:
							with tempfile.TemporaryDirectory() as tmp_dir:
								lowerCAmelCase       =      os.path.join(A_      ,			"""vocab.txt""" )
								with open(A_      ,			"""w"""      ,			encoding="""utf-8""" ) as vocab_writer:
									vocab_writer.write("""""".join([x + """\n""" for x in self.vocab_tokens] ) )
								lowerCAmelCase       =      BertTokenizer(A_ )

							tokenizer.push_to_hub("""valid_org/test-tokenizer-org"""      ,			use_auth_token=self._token )
							lowerCAmelCase       =      BertTokenizer.from_pretrained("""valid_org/test-tokenizer-org""" )
							self.assertDictEqual(new_tokenizer.vocab      ,			tokenizer.vocab )

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

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

							lowerCAmelCase       =      BertTokenizer.from_pretrained("""valid_org/test-tokenizer-org""" )
							self.assertDictEqual(new_tokenizer.vocab      ,			tokenizer.vocab )




						@require_tokenizers
						def 							__snake_case       (   self )      ->      List[str]:
							CustomTokenizer.register_for_auto_class()
							with tempfile.TemporaryDirectory() as tmp_dir:
								lowerCAmelCase       =      os.path.join(A_      ,			"""vocab.txt""" )
								with open(A_      ,			"""w"""      ,			encoding="""utf-8""" ) as vocab_writer:
									vocab_writer.write("""""".join([x + """\n""" for x in self.vocab_tokens] ) )
								lowerCAmelCase       =      CustomTokenizer(A_ )

							# No fast custom tokenizer
							tokenizer.push_to_hub("""test-dynamic-tokenizer"""      ,			use_auth_token=self._token )

							lowerCAmelCase       =      AutoTokenizer.from_pretrained(f'{USER}/test-dynamic-tokenizer'      ,			trust_remote_code=A_ )
							# Can't make an isinstance check because the new_model.config is from the CustomTokenizer class of a dynamic module
							self.assertEqual(tokenizer.__class__.__name__      ,			"""CustomTokenizer""" )

							# Fast and slow custom tokenizer
							CustomTokenizerFast.register_for_auto_class()
							with tempfile.TemporaryDirectory() as tmp_dir:
								lowerCAmelCase       =      os.path.join(A_      ,			"""vocab.txt""" )
								with open(A_      ,			"""w"""      ,			encoding="""utf-8""" ) as vocab_writer:
									vocab_writer.write("""""".join([x + """\n""" for x in self.vocab_tokens] ) )

								lowerCAmelCase       =      BertTokenizerFast.from_pretrained(A_ )
								bert_tokenizer.save_pretrained(A_ )
								lowerCAmelCase       =      CustomTokenizerFast.from_pretrained(A_ )

							tokenizer.push_to_hub("""test-dynamic-tokenizer"""      ,			use_auth_token=self._token )

							lowerCAmelCase       =      AutoTokenizer.from_pretrained(f'{USER}/test-dynamic-tokenizer'      ,			trust_remote_code=A_ )
							# Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
							self.assertEqual(tokenizer.__class__.__name__      ,			"""CustomTokenizerFast""" )
							lowerCAmelCase       =      AutoTokenizer.from_pretrained(
							    f'{USER}/test-dynamic-tokenizer'      ,			use_fast=A_      ,			trust_remote_code=A_ )
							# Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
							self.assertEqual(tokenizer.__class__.__name__      ,			"""CustomTokenizer""" )







class __snake_case( unittest.TestCase   ):


						'''simple docstring'''



						def 							__snake_case       (   self )      ->      Dict:
							lowerCAmelCase       =      Trie()
							trie.add("""Hello 友達""" )
							self.assertEqual(trie.data      ,			{"""H""": {"""e""": {"""l""": {"""l""": {"""o""": {""" """: {"""友""": {"""達""": {"""""": 1}}}}}}}}} )
							trie.add("""Hello""" )
							trie.data
							self.assertEqual(trie.data      ,			{"""H""": {"""e""": {"""l""": {"""l""": {"""o""": {"""""": 1, """ """: {"""友""": {"""達""": {"""""": 1}}}}}}}}} )


						def 							__snake_case       (   self )      ->      List[Any]:
							lowerCAmelCase       =      Trie()
							self.assertEqual(trie.split("""[CLS] This is a extra_id_100""" )      ,			["""[CLS] This is a extra_id_100"""] )
							trie.add("""[CLS]""" )
							trie.add("""extra_id_1""" )
							trie.add("""extra_id_100""" )
							self.assertEqual(trie.split("""[CLS] This is a extra_id_100""" )      ,			["""[CLS]""", """ This is a """, """extra_id_100"""] )


						def 							__snake_case       (   self )      ->      List[str]:
							lowerCAmelCase       =      Trie()
							trie.add("""A""" )
							self.assertEqual(trie.split("""ABC""" )      ,			["""A""", """BC"""] )
							self.assertEqual(trie.split("""BCA""" )      ,			["""BC""", """A"""] )


						def 							__snake_case       (   self )      ->      Union[str, Any]:
							lowerCAmelCase       =      Trie()
							trie.add("""TOKEN]""" )
							trie.add("""[SPECIAL_TOKEN]""" )
							self.assertEqual(trie.split("""This is something [SPECIAL_TOKEN]""" )      ,			["""This is something """, """[SPECIAL_TOKEN]"""] )


						def 							__snake_case       (   self )      ->      Tuple:
							lowerCAmelCase       =      Trie()
							trie.add("""A""" )
							trie.add("""P""" )
							trie.add("""[SPECIAL_TOKEN]""" )
							self.assertEqual(trie.split("""This is something [SPECIAL_TOKEN]""" )      ,			["""This is something """, """[SPECIAL_TOKEN]"""] )


						def 							__snake_case       (   self )      ->      int:
							lowerCAmelCase       =      Trie()
							trie.add("""AB""" )
							trie.add("""B""" )
							trie.add("""C""" )
							self.assertEqual(trie.split("""ABC""" )      ,			["""AB""", """C"""] )


						def 							__snake_case       (   self )      ->      Optional[Any]:
							lowerCAmelCase       =      Trie()
							trie.add("""ABC""" )
							trie.add("""B""" )
							trie.add("""CD""" )
							self.assertEqual(trie.split("""ABCD""" )      ,			["""ABC""", """D"""] )




						def 							__snake_case       (   self )      ->      List[str]:
							# Even if the offsets are wrong, we necessarily output correct string
							# parts.
							lowerCAmelCase       =      Trie()
							lowerCAmelCase       =      trie.cut_text("""ABC"""      ,			[0, 0, 2, 1, 2, 3] )
							self.assertEqual(A_      ,			["""AB""", """C"""] )

"""simple docstring"""


import os

import torch

from ..logging import get_logger
from .constants import FSDP_PYTORCH_VERSION, MODEL_NAME, OPTIMIZER_NAME
from .versions import is_torch_version


if is_torch_version('''>=''', FSDP_PYTORCH_VERSION):
		import torch.distributed.checkpoint as dist_cp
		from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
		from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
		from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
		from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType


_lowercase				 =					get_logger(__name__)






def    _snake_case					(		snake_case__ :  Optional[Any]						,   snake_case__ :  List[Any]						,   snake_case__ :  int						,   snake_case__ :  int						,   snake_case__ :  str=0  ):
	os.makedirs(snake_case__						,   exist_ok=snake_case__  )
	with FSDP.state_dict_type(
	    snake_case__						,   fsdp_plugin.state_dict_type						,   fsdp_plugin.state_dict_config						,   fsdp_plugin.optim_state_dict_config  ):
		A     = model.state_dict()
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			A     = F'{MODEL_NAME}.bin' if model_index == 0 else F'{MODEL_NAME}_{model_index}.bin'
			A     = os.path.join(snake_case__						,   snake_case__  )
			if accelerator.process_index == 0:
				logger.info(F'Saving model to {output_model_file}'  )
				torch.save(snake_case__						,   snake_case__  )
				logger.info(F'Model saved to {output_model_file}'  )
		elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
			A     = (
			    F'{MODEL_NAME}_rank{accelerator.process_index}.bin'
			    if model_index == 0
			    else F'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
			)
			A     = os.path.join(snake_case__						,   snake_case__  )
			logger.info(F'Saving model to {output_model_file}'  )
			torch.save(snake_case__						,   snake_case__  )
			logger.info(F'Model saved to {output_model_file}'  )
		elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
			A     = os.path.join(snake_case__						,   F'{MODEL_NAME}_{model_index}'  )
			os.makedirs(snake_case__						,   exist_ok=snake_case__  )
			logger.info(F'Saving model to {ckpt_dir}'  )
			A     = {'model': state_dict}

			dist_cp.save_state_dict(
			    state_dict=snake_case__						,   storage_writer=dist_cp.FileSystemWriter(snake_case__  )						,   planner=DefaultSavePlanner()						,   )
			logger.info(F'Model saved to {ckpt_dir}'  )






def    _snake_case					(		snake_case__ :  int						,   snake_case__ :  List[str]						,   snake_case__ :  str						,   snake_case__ :  str						,   snake_case__ :  Any=0  ):
	accelerator.wait_for_everyone()
	with FSDP.state_dict_type(
	    snake_case__						,   fsdp_plugin.state_dict_type						,   fsdp_plugin.state_dict_config						,   fsdp_plugin.optim_state_dict_config  ):
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			if type(snake_case__  ) != FSDP and accelerator.process_index != 0:
				if not fsdp_plugin.sync_module_states:
					raise ValueError(
					    'Set the `sync_module_states` flag to `True` so that model states are synced across processes when '
					    'initializing FSDP object'  )
				return
			A     = F'{MODEL_NAME}.bin' if model_index == 0 else F'{MODEL_NAME}_{model_index}.bin'
			A     = os.path.join(snake_case__						,   snake_case__  )
			logger.info(F'Loading model from {input_model_file}'  )
			A     = torch.load(snake_case__  )
			logger.info(F'Model loaded from {input_model_file}'  )
		elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
			A     = (
			    F'{MODEL_NAME}_rank{accelerator.process_index}.bin'
			    if model_index == 0
			    else F'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
			)
			A     = os.path.join(snake_case__						,   snake_case__  )
			logger.info(F'Loading model from {input_model_file}'  )
			A     = torch.load(snake_case__  )
			logger.info(F'Model loaded from {input_model_file}'  )
		elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
			A     = (
			    os.path.join(snake_case__						,   F'{MODEL_NAME}_{model_index}'  )
			    if F'{MODEL_NAME}' not in input_dir
			    else input_dir
			)
			logger.info(F'Loading model from {ckpt_dir}'  )
			A     = {'model': model.state_dict()}
			dist_cp.load_state_dict(
			    state_dict=snake_case__						,   storage_reader=dist_cp.FileSystemReader(snake_case__  )						,   planner=DefaultLoadPlanner()						,   )
			A     = state_dict['model']
			logger.info(F'Model loaded from {ckpt_dir}'  )
		model.load_state_dict(snake_case__  )






def    _snake_case					(		snake_case__ :  Tuple						,   snake_case__ :  Optional[int]						,   snake_case__ :  Optional[int]						,   snake_case__ :  Dict						,   snake_case__ :  List[str]						,   snake_case__ :  Any=0  ):
	os.makedirs(snake_case__						,   exist_ok=snake_case__  )
	with FSDP.state_dict_type(
	    snake_case__						,   fsdp_plugin.state_dict_type						,   fsdp_plugin.state_dict_config						,   fsdp_plugin.optim_state_dict_config  ):
		A     = FSDP.optim_state_dict(snake_case__						,   snake_case__  )
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			if accelerator.process_index == 0:
				A     = (
				    F'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else F'{OPTIMIZER_NAME}_{optimizer_index}.bin'
				)
				A     = os.path.join(snake_case__						,   snake_case__  )
				logger.info(F'Saving Optimizer state to {output_optimizer_file}'  )
				torch.save(snake_case__						,   snake_case__  )
				logger.info(F'Optimizer state saved in {output_optimizer_file}'  )
		else:
			A     = os.path.join(snake_case__						,   F'{OPTIMIZER_NAME}_{optimizer_index}'  )
			os.makedirs(snake_case__						,   exist_ok=snake_case__  )
			logger.info(F'Saving Optimizer state to {ckpt_dir}'  )
			dist_cp.save_state_dict(
			    state_dict={'optimizer': optim_state}						,   storage_writer=dist_cp.FileSystemWriter(snake_case__  )						,   planner=DefaultSavePlanner()						,   )
			logger.info(F'Optimizer state saved in {ckpt_dir}'  )






def    _snake_case					(		snake_case__ :  str						,   snake_case__ :  Optional[int]						,   snake_case__ :  List[str]						,   snake_case__ :  List[str]						,   snake_case__ :  int						,   snake_case__ :  Optional[int]=0  ):
	accelerator.wait_for_everyone()
	with FSDP.state_dict_type(
	    snake_case__						,   fsdp_plugin.state_dict_type						,   fsdp_plugin.state_dict_config						,   fsdp_plugin.optim_state_dict_config  ):
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			A     = None
			# below check should work but currently it isn't working (mostly opytorch issue),
			# in the meantime disabling it at the cost of excess memory usage
			# if accelerator.process_index == 0 or not fsdp_plugin.optim_state_dict_config.rank0_only:
			A     = (
			    F'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else F'{OPTIMIZER_NAME}_{optimizer_index}.bin'
			)
			A     = os.path.join(snake_case__						,   snake_case__  )
			logger.info(F'Loading Optimizer state from {input_optimizer_file}'  )
			A     = torch.load(snake_case__  )
			logger.info(F'Optimizer state loaded from {input_optimizer_file}'  )
		else:
			A     = (
			    os.path.join(snake_case__						,   F'{OPTIMIZER_NAME}_{optimizer_index}'  )
			    if F'{OPTIMIZER_NAME}' not in input_dir
			    else input_dir
			)
			logger.info(F'Loading Optimizer from {ckpt_dir}'  )
			A     = load_sharded_optimizer_state_dict(
			    model_state_dict=model.state_dict()						,   optimizer_key='optimizer'						,   storage_reader=dist_cp.FileSystemReader(snake_case__  )						,   )
			A     = optim_state['optimizer']
			logger.info(F'Optimizer loaded from {ckpt_dir}'  )
		A     = FSDP.optim_state_dict_to_load(snake_case__						,   snake_case__						,   snake_case__  )
		optimizer.load_state_dict(snake_case__  )