




"""simple docstring"""

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         =      '''bert-base-cased'''
a         =      '''fp16'''
a         =      '''bf16'''
a         =      [FPaa, BFaa]





@require_fsdp
@require_cuda
class     lowercase_	(					__lowerCAmelCase							):



 '''simple docstring'''

 def       lowerCAmelCase_							(			self	:	int		):
       super().setUp()

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

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

       for i, strategy in enumerate(_UpperCAmelCase		):
             _A        =       self.dist_env.copy()
             _A        =       F'''{i + 1}'''
             _A        =       strategy
             with mockenv_context(**_UpperCAmelCase		):
                   _A        =       FullyShardedDataParallelPlugin()
                   self.assertEqual(fsdp_plugin.sharding_strategy   ,							ShardingStrategy(i + 1		)		)

 def       lowerCAmelCase_							(			self	:	Optional[int]		):
       from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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

 def       lowerCAmelCase_							(			self	:	List[str]		):
       from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

       for i, state_dict_type in enumerate(_UpperCAmelCase		):
             _A        =       self.dist_env.copy()
             _A        =       state_dict_type
             with mockenv_context(**_UpperCAmelCase		):
                   _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       lowerCAmelCase_							(			self	:	Union[str, Any]		):
       _A        =       AutoModel.from_pretrained(_UpperCAmelCase		)
       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(**_UpperCAmelCase		):
                   _A        =       FullyShardedDataParallelPlugin()
                   fsdp_plugin.set_auto_wrap_policy(_UpperCAmelCase		)
                   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(**_UpperCAmelCase		):
             _A        =       FullyShardedDataParallelPlugin()
             with self.assertRaises(_UpperCAmelCase		) as cm:
                   fsdp_plugin.set_auto_wrap_policy(_UpperCAmelCase		)
             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(**_UpperCAmelCase		):
             _A        =       FullyShardedDataParallelPlugin()
             fsdp_plugin.set_auto_wrap_policy(_UpperCAmelCase		)
             self.assertIsNone(fsdp_plugin.auto_wrap_policy		)

 def       lowerCAmelCase_							(			self	:	Optional[Any]		):
       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(**_UpperCAmelCase		):
                   _A        =       Accelerator()
                   if mp_dtype == "fp16":
                         _A        =       torch.floataa
                   elif mp_dtype == "bf16":
                         _A        =       torch.bfloataa
                   _A        =       MixedPrecision(param_dtype=_UpperCAmelCase   ,							reduce_dtype=_UpperCAmelCase   ,							buffer_dtype=_UpperCAmelCase		)
                   self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy   ,							_UpperCAmelCase		)
                   if mp_dtype == FPaa:
                         self.assertTrue(isinstance(accelerator.scaler   ,							_UpperCAmelCase		)		)
                   elif mp_dtype == BFaa:
                         self.assertIsNone(accelerator.scaler		)
                   AcceleratorState._reset_state(_UpperCAmelCase		)



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

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





@require_fsdp
@require_multi_gpu
@slow
class     lowercase_	(					__lowerCAmelCase							):



 '''simple docstring'''

 def       lowerCAmelCase_							(			self	:	str		):
       super().setUp()
       _A        =       0.82
       _A        =       [
           'fsdp_shard_grad_op_transformer_based_wrap',
           'fsdp_full_shard_transformer_based_wrap',
       ]
       _A        =       {
           'multi_gpu_fp16': 3_200,
           'fsdp_shard_grad_op_transformer_based_wrap_fp16': 2_000,
           'fsdp_full_shard_transformer_based_wrap_fp16': 1_900,
           # 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       lowerCAmelCase_							(			self	:	str		):
       _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(_UpperCAmelCase		):
                   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(_UpperCAmelCase   ,							env=os.environ.copy()		)

 def       lowerCAmelCase_							(			self	:	Dict		):
       _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(_UpperCAmelCase		):
             _A        =       cmd.copy()
             cmd_config.append(F'''--fsdp_sharding_strategy={i+1}'''		)
             if strategy != "FULL_SHARD":
                   continue
             _A        =       len(_UpperCAmelCase		)
             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(_UpperCAmelCase   ,							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(_UpperCAmelCase   ,							env=os.environ.copy()		)



 def       lowerCAmelCase_							(			self	:	List[str]		):
       _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(_UpperCAmelCase		):
                         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(_UpperCAmelCase   ,							env=os.environ.copy()		)







"""simple docstring"""

def 						_snake_case					( _snake_case							: str							)   -> str:






      '''simple docstring'''


      _A        =       ''
      for ch in key:
            if ch == " " or ch not in key_no_dups and ch.isalpha():
                  key_no_dups += ch
      return key_no_dups






def 						_snake_case					( _snake_case							: str							)   -> dict[str, str]:






      '''simple docstring'''


      _A        =       [chr(i + 65							) for i in range(26							)]
      # Remove duplicate characters from key
      _A        =       remove_duplicates(key.upper()							)
      _A        =       len(_snake_case							)
      # First fill cipher with key characters
      _A        =       {alphabet[i]: char for i, char in enumerate(_snake_case							)}
      # Then map remaining characters in alphabet to
      # the alphabet from the beginning
      for i in range(len(_snake_case							)    ,  26							):
            _A        =       alphabet[i - offset]
            # Ensure we are not mapping letters to letters previously mapped
            while char in key:
                  offset -= 1
                  _A        =       alphabet[i - offset]
            _A        =       char
      return cipher_alphabet






def 						_snake_case					( _snake_case							: str    ,  _snake_case							: dict[str, str]							)   -> str:






      '''simple docstring'''


      return "".join(cipher_map.get(_snake_case    ,  _snake_case							) for ch in message.upper()							)






def 						_snake_case					( _snake_case							: str    ,  _snake_case							: dict[str, str]							)   -> str:






      '''simple docstring'''


      _A        =       {v: k for k, v in cipher_map.items()}
      return "".join(rev_cipher_map.get(_snake_case    ,  _snake_case							) for ch in message.upper()							)






def 						_snake_case					( )   -> None:






      '''simple docstring'''


      _A        =       input('Enter message to encode or decode: '							).strip()
      _A        =       input('Enter keyword: '							).strip()
      _A        =       input('Encipher or decipher? E/D:'							).strip()[0].lower()
      try:
            _A        =       {'e': encipher, 'd': decipher}[option]
      except KeyError:
            raise KeyError('invalid input option'							)
      _A        =       create_cipher_map(_snake_case							)
      print(func(_snake_case    ,  _snake_case							)							)


if __name__ == "__main__":
       import doctest

       doctest.testmod()
       main()

