





import inspect
import os

import torch
from transformers import AutoModel
from transformers.testing_utils import mockenv_context
from transformers.trainer_utils import set_seed

import accelerate
from accelerate.accelerator import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils.testing import (
    AccelerateTestCase,
    TempDirTestCase,
    execute_subprocess_async,
    require_cuda,
    require_fsdp,
    require_multi_gpu,
    slow,
)
from accelerate.utils.constants import (
    FSDP_AUTO_WRAP_POLICY,
    FSDP_BACKWARD_PREFETCH,
    FSDP_SHARDING_STRATEGY,
    FSDP_STATE_DICT_TYPE,
)
from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin
from accelerate.utils.other import patch_environment


set_seed(42)

a__:      Dict									=		'bert-base-cased'
a__:      List[Any]									=		'fp16'
a__:      Union[str, Any]									=		'bf16'
a__:      Optional[Any]									=		[FPaa, BFaa]




@require_fsdp
@require_cuda
class     SCREAMING_SNAKE_CASE__ (						UpperCamelCase__      ):



 def 				UpperCamelCase			(							self ):
        super().setUp()

        A__											=						dict(
            ACCELERATE_USE_FSDP='''true''',MASTER_ADDR='''localhost''',MASTER_PORT='''10999''',RANK='''0''',LOCAL_RANK='''0''',WORLD_SIZE='''1''',)



 def 				UpperCamelCase			(							self ):
        from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

        for i, strategy in enumerate(__lowerCamelCase ):
               A__											=						self.dist_env.copy()
               A__											=						f"{i + 1}"
               A__											=						strategy
               with mockenv_context(**__lowerCamelCase ):
                      A__											=						FullyShardedDataParallelPlugin()
                      self.assertEqual(fsdp_plugin.sharding_strategy,ShardingStrategy(i + 1 ) )



 def 				UpperCamelCase			(							self ):
        from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

        for i, prefetch_policy in enumerate(__lowerCamelCase ):
               A__											=						self.dist_env.copy()
               A__											=						prefetch_policy
               with mockenv_context(**__lowerCamelCase ):
                      A__											=						FullyShardedDataParallelPlugin()
                      if prefetch_policy == "NO_PREFETCH":
                             self.assertIsNone(fsdp_plugin.backward_prefetch )
                      else:
                             self.assertEqual(fsdp_plugin.backward_prefetch,BackwardPrefetch(i + 1 ) )



 def 				UpperCamelCase			(							self ):
        from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

        for i, state_dict_type in enumerate(__lowerCamelCase ):
               A__											=						self.dist_env.copy()
               A__											=						state_dict_type
               with mockenv_context(**__lowerCamelCase ):
                      A__											=						FullyShardedDataParallelPlugin()
                      self.assertEqual(fsdp_plugin.state_dict_type,StateDictType(i + 1 ) )
                      if state_dict_type == "FULL_STATE_DICT":
                             self.assertTrue(fsdp_plugin.state_dict_config.offload_to_cpu )
                             self.assertTrue(fsdp_plugin.state_dict_config.ranka_only )



 def 				UpperCamelCase			(							self ):
        A__											=						AutoModel.from_pretrained(__lowerCamelCase )
        for policy in FSDP_AUTO_WRAP_POLICY:
               A__											=						self.dist_env.copy()
               A__											=						policy
               if policy == "TRANSFORMER_BASED_WRAP":
                      A__											=						'''BertLayer'''
               elif policy == "SIZE_BASED_WRAP":
                      A__											=						'''2000'''
               with mockenv_context(**__lowerCamelCase ):
                      A__											=						FullyShardedDataParallelPlugin()
                      fsdp_plugin.set_auto_wrap_policy(__lowerCamelCase )
                      if policy == "NO_WRAP":
                             self.assertIsNone(fsdp_plugin.auto_wrap_policy )
                      else:
                             self.assertIsNotNone(fsdp_plugin.auto_wrap_policy )

        A__											=						self.dist_env.copy()
        A__											=						'''TRANSFORMER_BASED_WRAP'''
        A__											=						'''T5Layer'''
        with mockenv_context(**__lowerCamelCase ):
               A__											=						FullyShardedDataParallelPlugin()
               with self.assertRaises(__lowerCamelCase ) as cm:
                      fsdp_plugin.set_auto_wrap_policy(__lowerCamelCase )
               self.assertTrue('''Could not find the transformer layer class to wrap in the model.''' in str(cm.exception ) )

        A__											=						self.dist_env.copy()
        A__											=						'''SIZE_BASED_WRAP'''
        A__											=						'''0'''
        with mockenv_context(**__lowerCamelCase ):
               A__											=						FullyShardedDataParallelPlugin()
               fsdp_plugin.set_auto_wrap_policy(__lowerCamelCase )
               self.assertIsNone(fsdp_plugin.auto_wrap_policy )



 def 				UpperCamelCase			(							self ):
        from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
        from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

        for mp_dtype in dtypes:
               A__											=						self.dist_env.copy()
               A__											=						mp_dtype
               with mockenv_context(**__lowerCamelCase ):
                      A__											=						Accelerator()
                      if mp_dtype == "fp16":
                             A__											=						torch.floataa
                      elif mp_dtype == "bf16":
                             A__											=						torch.bfloataa
                      A__											=						MixedPrecision(param_dtype=__lowerCamelCase,reduce_dtype=__lowerCamelCase,buffer_dtype=__lowerCamelCase )
                      self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy,__lowerCamelCase )
                      if mp_dtype == FPaa:
                             self.assertTrue(isinstance(accelerator.scaler,__lowerCamelCase ) )
                      elif mp_dtype == BFaa:
                             self.assertIsNone(accelerator.scaler )
                      AcceleratorState._reset_state(__lowerCamelCase )




 def 				UpperCamelCase			(							self ):
        from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

        for flag in [True, False]:
               A__											=						self.dist_env.copy()
               A__											=						str(__lowerCamelCase ).lower()
               with mockenv_context(**__lowerCamelCase ):
                      A__											=						FullyShardedDataParallelPlugin()
                      self.assertEqual(fsdp_plugin.cpu_offload,CPUOffload(offload_params=__lowerCamelCase ) )




@require_fsdp
@require_multi_gpu
@slow
class     SCREAMING_SNAKE_CASE__ (						UpperCamelCase__      ):



 def 				UpperCamelCase			(							self ):
        super().setUp()
        A__											=						0.82
        A__											=						[
            '''fsdp_shard_grad_op_transformer_based_wrap''',
            '''fsdp_full_shard_transformer_based_wrap''',
        ]
        A__											=						{
            '''multi_gpu_fp16''': 3200,
            '''fsdp_shard_grad_op_transformer_based_wrap_fp16''': 2000,
            '''fsdp_full_shard_transformer_based_wrap_fp16''': 1900,
            # Disabling below test as it overwhelms the RAM memory usage
            # on CI self-hosted runner leading to tests getting killed.
            # "fsdp_full_shard_cpu_offload_transformer_based_wrap_fp32": 1500,  # fp16 was leading to indefinite hang
        }
        A__											=						160
        A__											=						160

        A__											=						inspect.getfile(accelerate.test_utils )
        A__											=						os.path.sep.join(mod_file.split(os.path.sep )[:-1] + ['''scripts''', '''external_deps'''] )



 def 				UpperCamelCase			(							self ):
        A__											=						os.path.join(self.test_scripts_folder,'''test_performance.py''' )
        A__											=						['''accelerate''', '''launch''', '''--num_processes=2''', '''--num_machines=1''', '''--machine_rank=0''', '''--use_fsdp''']
        for config in self.performance_configs:
               A__											=						cmd.copy()
               for i, strategy in enumerate(__lowerCamelCase ):
                      if strategy.lower() in config:
                             cmd_config.append(f"--fsdp_sharding_strategy={i+1}" )
                             break

               if "fp32" in config:
                      cmd_config.append('''--mixed_precision=no''' )
               else:
                      cmd_config.append('''--mixed_precision=fp16''' )

               if "cpu_offload" in config:
                      cmd_config.append('''--fsdp_offload_params=True''' )

               for policy in FSDP_AUTO_WRAP_POLICY:
                      if policy.lower() in config:
                             cmd_config.append(f"--fsdp_auto_wrap_policy={policy}" )
                             break

               if policy == "TRANSFORMER_BASED_WRAP":
                      cmd_config.append('''--fsdp_transformer_layer_cls_to_wrap=BertLayer''' )
               elif policy == "SIZE_BASED_WRAP":
                      cmd_config.append('''--fsdp_min_num_params=2000''' )

               cmd_config.extend(
                   [
                       self.test_file_path,
                       f"--output_dir={self.tmpdir}",
                       f"--performance_lower_bound={self.performance_lower_bound}",
                   ] )
               with patch_environment(omp_num_threads=1 ):
                      execute_subprocess_async(__lowerCamelCase,env=os.environ.copy() )



 def 				UpperCamelCase			(							self ):
        A__											=						os.path.join(self.test_scripts_folder,'''test_checkpointing.py''' )
        A__											=						[
            '''accelerate''',
            '''launch''',
            '''--num_processes=2''',
            '''--num_machines=1''',
            '''--machine_rank=0''',
            '''--use_fsdp''',
            '''--mixed_precision=fp16''',
            '''--fsdp_transformer_layer_cls_to_wrap=BertLayer''',
        ]

        for i, strategy in enumerate(__lowerCamelCase ):
               A__											=						cmd.copy()
               cmd_config.append(f"--fsdp_sharding_strategy={i+1}" )
               if strategy != "FULL_SHARD":
                      continue
               A__											=						len(__lowerCamelCase )
               for state_dict_type in FSDP_STATE_DICT_TYPE:
                      A__											=						cmd_config[:state_dict_config_index]
                      cmd_config.append(f"--fsdp_state_dict_type={state_dict_type}" )
                      cmd_config.extend(
                          [
                              self.test_file_path,
                              f"--output_dir={self.tmpdir}",
                              '''--partial_train_epoch=1''',
                          ] )
                      with patch_environment(omp_num_threads=1 ):
                             execute_subprocess_async(__lowerCamelCase,env=os.environ.copy() )

                      A__											=						cmd_config[:-1]
                      A__											=						os.path.join(self.tmpdir,'''epoch_0''' )
                      cmd_config.extend(
                          [
                              f"--resume_from_checkpoint={resume_from_checkpoint}",
                          ] )
                      with patch_environment(omp_num_threads=1 ):
                             execute_subprocess_async(__lowerCamelCase,env=os.environ.copy() )




 def 				UpperCamelCase			(							self ):
        A__											=						os.path.join(self.test_scripts_folder,'''test_peak_memory_usage.py''' )
        A__											=						[
            '''accelerate''',
            '''launch''',
            '''--num_processes=2''',
            '''--num_machines=1''',
            '''--machine_rank=0''',
        ]
        for spec, peak_mem_upper_bound in self.peak_memory_usage_upper_bound.items():
               A__											=						cmd.copy()
               if "fp16" in spec:
                      cmd_config.extend(['''--mixed_precision=fp16'''] )
               else:
                      cmd_config.extend(['''--mixed_precision=no'''] )

               if "multi_gpu" in spec:
                      continue
               else:
                      cmd_config.extend(['''--use_fsdp'''] )
                      for i, strategy in enumerate(__lowerCamelCase ):
                             if strategy.lower() in spec:
                                    cmd_config.append(f"--fsdp_sharding_strategy={i+1}" )
                                    break

                      if "cpu_offload" in spec:
                             cmd_config.append('''--fsdp_offload_params=True''' )

                      for policy in FSDP_AUTO_WRAP_POLICY:
                             if policy.lower() in spec:
                                    cmd_config.append(f"--fsdp_auto_wrap_policy={policy}" )
                                    break

                      if policy == "TRANSFORMER_BASED_WRAP":
                             cmd_config.append('''--fsdp_transformer_layer_cls_to_wrap=BertLayer''' )
                      elif policy == "SIZE_BASED_WRAP":
                             cmd_config.append('''--fsdp_min_num_params=2000''' )

               cmd_config.extend(
                   [
                       self.test_file_path,
                       f"--output_dir={self.tmpdir}",
                       f"--peak_memory_upper_bound={peak_mem_upper_bound}",
                       f"--n_train={self.n_train}",
                       f"--n_val={self.n_val}",
                   ] )
               with patch_environment(omp_num_threads=1 ):
                      execute_subprocess_async(__lowerCamelCase,env=os.environ.copy() )












import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_distilbert import DistilBertTokenizer


a__:      Dict									=		logging.get_logger(__name__)

a__:      str									=		{'vocab_file': 'vocab.txt', 'tokenizer_file': 'tokenizer.json'}

a__:      Any									=		{
    'vocab_file': {
        'distilbert-base-uncased': 'https://huggingface.co/distilbert-base-uncased/resolve/main/vocab.txt',
        'distilbert-base-uncased-distilled-squad': (
            'https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/vocab.txt'
        ),
        'distilbert-base-cased': 'https://huggingface.co/distilbert-base-cased/resolve/main/vocab.txt',
        'distilbert-base-cased-distilled-squad': (
            'https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/vocab.txt'
        ),
        'distilbert-base-german-cased': 'https://huggingface.co/distilbert-base-german-cased/resolve/main/vocab.txt',
        'distilbert-base-multilingual-cased': (
            'https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/vocab.txt'
        ),
    },
    'tokenizer_file': {
        'distilbert-base-uncased': 'https://huggingface.co/distilbert-base-uncased/resolve/main/tokenizer.json',
        'distilbert-base-uncased-distilled-squad': (
            'https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/tokenizer.json'
        ),
        'distilbert-base-cased': 'https://huggingface.co/distilbert-base-cased/resolve/main/tokenizer.json',
        'distilbert-base-cased-distilled-squad': (
            'https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/tokenizer.json'
        ),
        'distilbert-base-german-cased': (
            'https://huggingface.co/distilbert-base-german-cased/resolve/main/tokenizer.json'
        ),
        'distilbert-base-multilingual-cased': (
            'https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/tokenizer.json'
        ),
    },
}

a__:      Any									=		{
    'distilbert-base-uncased': 512,
    'distilbert-base-uncased-distilled-squad': 512,
    'distilbert-base-cased': 512,
    'distilbert-base-cased-distilled-squad': 512,
    'distilbert-base-german-cased': 512,
    'distilbert-base-multilingual-cased': 512,
}


a__:      Optional[Any]									=		{
    'distilbert-base-uncased': {'do_lower_case': True},
    'distilbert-base-uncased-distilled-squad': {'do_lower_case': True},
    'distilbert-base-cased': {'do_lower_case': False},
    'distilbert-base-cased-distilled-squad': {'do_lower_case': False},
    'distilbert-base-german-cased': {'do_lower_case': False},
    'distilbert-base-multilingual-cased': {'do_lower_case': False},
}




class     SCREAMING_SNAKE_CASE__ (						UpperCamelCase__      ):
 __SCREAMING_SNAKE_CASE							       =	VOCAB_FILES_NAMES
 __SCREAMING_SNAKE_CASE							       =	PRETRAINED_VOCAB_FILES_MAP
 __SCREAMING_SNAKE_CASE							       =	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __SCREAMING_SNAKE_CASE							       =	PRETRAINED_INIT_CONFIGURATION
 __SCREAMING_SNAKE_CASE							       =	['''input_ids''', '''attention_mask''']
 __SCREAMING_SNAKE_CASE							       =	DistilBertTokenizer



 def __init__(							self,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=True,__lowerCamelCase="[UNK]",__lowerCamelCase="[SEP]",__lowerCamelCase="[PAD]",__lowerCamelCase="[CLS]",__lowerCamelCase="[MASK]",__lowerCamelCase=True,__lowerCamelCase=None,**__lowerCamelCase,):
        super().__init__(
            __lowerCamelCase,tokenizer_file=__lowerCamelCase,do_lower_case=__lowerCamelCase,unk_token=__lowerCamelCase,sep_token=__lowerCamelCase,pad_token=__lowerCamelCase,cls_token=__lowerCamelCase,mask_token=__lowerCamelCase,tokenize_chinese_chars=__lowerCamelCase,strip_accents=__lowerCamelCase,**__lowerCamelCase,)

        A__											=						json.loads(self.backend_tokenizer.normalizer.__getstate__() )
        if (
            normalizer_state.get('''lowercase''',__lowerCamelCase ) != do_lower_case
            or normalizer_state.get('''strip_accents''',__lowerCamelCase ) != strip_accents
            or normalizer_state.get('''handle_chinese_chars''',__lowerCamelCase ) != tokenize_chinese_chars
        ):
               A__											=						getattr(__lowerCamelCase,normalizer_state.pop('''type''' ) )
               A__											=						do_lower_case
               A__											=						strip_accents
               A__											=						tokenize_chinese_chars
               A__											=						normalizer_class(**__lowerCamelCase )

        A__											=						do_lower_case



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase=None ):
        A__											=						[self.cls_token_id] + token_ids_a + [self.sep_token_id]

        if token_ids_a:
               output += token_ids_a + [self.sep_token_id]

        return output



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase = None ):
        A__											=						[self.sep_token_id]
        A__											=						[self.cls_token_id]
        if token_ids_a is None:
               return len(cls + token_ids_a + sep ) * [0]
        return len(cls + token_ids_a + sep ) * [0] + len(token_ids_a + sep ) * [1]




 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase = None ):
        A__											=						self._tokenizer.model.save(__lowerCamelCase,name=__lowerCamelCase )
        return tuple(__lowerCamelCase )





