'''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)

__SCREAMING_SNAKE_CASE       :     Any     =  """bert-base-cased"""
__SCREAMING_SNAKE_CASE       :     Any     =  """fp16"""
__SCREAMING_SNAKE_CASE       :     Optional[Any]     =  """bf16"""
__SCREAMING_SNAKE_CASE       :     Dict     =  [FPaa, BFaa]

@require_fsdp
@require_cuda
class 	lowerCamelCase_	(snake_case__    ):
      '''simple docstring'''

      def      _A	(       self						:				Union[str, Any]       ):
             super().setUp()

             _UpperCAmelCase     :  Optional[int]							      =    dict(
                 ACCELERATE_USE_FSDP="true"						,					MASTER_ADDR="localhost"						,					MASTER_PORT="10999"						,					RANK="0"						,					LOCAL_RANK="0"						,					WORLD_SIZE="1"						,					)




      def      _A	(       self						:				Optional[Any]       ):
             from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

             for i, strategy in enumerate(A       ):
                    _UpperCAmelCase     :  Union[str, Any]							      =    self.dist_env.copy()
                    _UpperCAmelCase     :  Any							      =    F"""{i + 1}"""
                    _UpperCAmelCase     :  str							      =    strategy
                    with mockenv_context(**A       ):
                           _UpperCAmelCase     :  Any							      =    FullyShardedDataParallelPlugin()
                           self.assertEqual(fsdp_plugin.sharding_strategy						,					ShardingStrategy(i + 1       )       )




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

             for i, prefetch_policy in enumerate(A       ):
                    _UpperCAmelCase     :  Union[str, Any]							      =    self.dist_env.copy()
                    _UpperCAmelCase     :  Optional[Any]							      =    prefetch_policy
                    with mockenv_context(**A       ):
                           _UpperCAmelCase     :  Union[str, Any]							      =    FullyShardedDataParallelPlugin()
                           if prefetch_policy == "NO_PREFETCH":
                                  self.assertIsNone(fsdp_plugin.backward_prefetch       )
                           else:
                                  self.assertEqual(fsdp_plugin.backward_prefetch						,					BackwardPrefetch(i + 1       )       )




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

             for i, state_dict_type in enumerate(A       ):
                    _UpperCAmelCase     :  Optional[int]							      =    self.dist_env.copy()
                    _UpperCAmelCase     :  Tuple							      =    state_dict_type
                    with mockenv_context(**A       ):
                           _UpperCAmelCase     :  List[str]							      =    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      _A	(       self						:				List[str]       ):
             _UpperCAmelCase     :  List[str]							      =    AutoModel.from_pretrained(A       )
             for policy in FSDP_AUTO_WRAP_POLICY:
                    _UpperCAmelCase     :  Any							      =    self.dist_env.copy()
                    _UpperCAmelCase     :  Optional[int]							      =    policy
                    if policy == "TRANSFORMER_BASED_WRAP":
                           _UpperCAmelCase     :  Optional[Any]							      =    "BertLayer"
                    elif policy == "SIZE_BASED_WRAP":
                           _UpperCAmelCase     :  Any							      =    "2000"
                    with mockenv_context(**A       ):
                           _UpperCAmelCase     :  List[Any]							      =    FullyShardedDataParallelPlugin()
                           fsdp_plugin.set_auto_wrap_policy(A       )
                           if policy == "NO_WRAP":
                                  self.assertIsNone(fsdp_plugin.auto_wrap_policy       )
                           else:
                                  self.assertIsNotNone(fsdp_plugin.auto_wrap_policy       )

             _UpperCAmelCase     :  Any							      =    self.dist_env.copy()
             _UpperCAmelCase     :  Tuple							      =    "TRANSFORMER_BASED_WRAP"
             _UpperCAmelCase     :  Tuple							      =    "T5Layer"
             with mockenv_context(**A       ):
                    _UpperCAmelCase     :  Dict							      =    FullyShardedDataParallelPlugin()
                    with self.assertRaises(A       ) as cm:
                           fsdp_plugin.set_auto_wrap_policy(A       )
                    self.assertTrue("Could not find the transformer layer class to wrap in the model." in str(cm.exception       )       )

             _UpperCAmelCase     :  str							      =    self.dist_env.copy()
             _UpperCAmelCase     :  Optional[int]							      =    "SIZE_BASED_WRAP"
             _UpperCAmelCase     :  List[str]							      =    "0"
             with mockenv_context(**A       ):
                    _UpperCAmelCase     :  List[Any]							      =    FullyShardedDataParallelPlugin()
                    fsdp_plugin.set_auto_wrap_policy(A       )
                    self.assertIsNone(fsdp_plugin.auto_wrap_policy       )




      def      _A	(       self						:				List[str]       ):
             from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
             from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

             for mp_dtype in dtypes:
                    _UpperCAmelCase     :  str							      =    self.dist_env.copy()
                    _UpperCAmelCase     :  str							      =    mp_dtype
                    with mockenv_context(**A       ):
                           _UpperCAmelCase     :  int							      =    Accelerator()
                           if mp_dtype == "fp16":
                                  _UpperCAmelCase     :  Union[str, Any]							      =    torch.floataa
                           elif mp_dtype == "bf16":
                                  _UpperCAmelCase     :  Union[str, Any]							      =    torch.bfloataa
                           _UpperCAmelCase     :  Optional[int]							      =    MixedPrecision(param_dtype=A						,					reduce_dtype=A						,					buffer_dtype=A       )
                           self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy						,					A       )
                           if mp_dtype == FPaa:
                                  self.assertTrue(isinstance(accelerator.scaler						,					A       )       )
                           elif mp_dtype == BFaa:
                                  self.assertIsNone(accelerator.scaler       )
                           AcceleratorState._reset_state(A       )





      def      _A	(       self						:				Optional[Any]       ):
             from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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

@require_fsdp
@require_multi_gpu
@slow
class 	lowerCamelCase_	(snake_case__    ):
      '''simple docstring'''

      def      _A	(       self						:				List[Any]       ):
             super().setUp()
             _UpperCAmelCase     :  Optional[int]							      =    0.82
             _UpperCAmelCase     :  int							      =    [
                 "fsdp_shard_grad_op_transformer_based_wrap",
                 "fsdp_full_shard_transformer_based_wrap",
             ]
             _UpperCAmelCase     :  Tuple							      =    {
                 "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
             }
             _UpperCAmelCase     :  Tuple							      =    160
             _UpperCAmelCase     :  Any							      =    160

             _UpperCAmelCase     :  List[str]							      =    inspect.getfile(accelerate.test_utils       )
             _UpperCAmelCase     :  int							      =    os.path.sep.join(mod_file.split(os.path.sep       )[:-1] + ["scripts", "external_deps"]       )




      def      _A	(       self						:				Optional[Any]       ):
             _UpperCAmelCase     :  Union[str, Any]							      =    os.path.join(self.test_scripts_folder						,					"test_performance.py"       )
             _UpperCAmelCase     :  Optional[int]							      =    ["accelerate", "launch", "--num_processes=2", "--num_machines=1", "--machine_rank=0", "--use_fsdp"]
             for config in self.performance_configs:
                    _UpperCAmelCase     :  Tuple							      =    cmd.copy()
                    for i, strategy in enumerate(A       ):
                           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(A						,					env=os.environ.copy()       )




      def      _A	(       self						:				List[Any]       ):
             _UpperCAmelCase     :  Dict							      =    os.path.join(self.test_scripts_folder						,					"test_checkpointing.py"       )
             _UpperCAmelCase     :  Union[str, Any]							      =    [
                 "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(A       ):
                    _UpperCAmelCase     :  Optional[Any]							      =    cmd.copy()
                    cmd_config.append(F"""--fsdp_sharding_strategy={i+1}"""       )
                    if strategy != "FULL_SHARD":
                           continue
                    _UpperCAmelCase     :  Optional[Any]							      =    len(A       )
                    for state_dict_type in FSDP_STATE_DICT_TYPE:
                           _UpperCAmelCase     :  Optional[Any]							      =    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(A						,					env=os.environ.copy()       )

                           _UpperCAmelCase     :  Optional[int]							      =    cmd_config[:-1]
                           _UpperCAmelCase     :  List[str]							      =    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(A						,					env=os.environ.copy()       )





      def      _A	(       self						:				List[Any]       ):
             _UpperCAmelCase     :  str							      =    os.path.join(self.test_scripts_folder						,					"test_peak_memory_usage.py"       )
             _UpperCAmelCase     :  Tuple							      =    [
                 "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():
                    _UpperCAmelCase     :  str							      =    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(A       ):
                                  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(A						,					env=os.environ.copy()       )







import itertools
import json
import linecache
import os
import pickle
import re
import socket
import string
from collections import Counter
from logging import getLogger
from pathlib import Path
from typing import Callable, Dict, Iterable, List

import git
import torch
from torch.utils.data import Dataset

from transformers import BartTokenizer, RagTokenizer, TaTokenizer


def 						lowerCamelCase__    (		_A     ,	_A     ,	_A     ,	_A     ,	_A=True     ,	_A="pt"							):







	'''simple docstring'''





	snake_case_		    =     {"add_prefix_space": True} if isinstance(_A     ,	_A							) and not line.startswith(" "							) else {}
	snake_case_		    =     padding_side
	return tokenizer(
	    [line]     ,	max_length=_A     ,	padding="max_length" if pad_to_max_length else None     ,	truncation=_A     ,	return_tensors=_A     ,	add_special_tokens=_A     ,	**_A     ,	)


def 						lowerCamelCase__    (		_A     ,	_A     ,	_A=None     ,	):







	'''simple docstring'''





	snake_case_		    =     input_ids.ne(_A							).any(dim=0							)
	if attention_mask is None:
		return input_ids[:, keep_column_mask]
	else:
		return (input_ids[:, keep_column_mask], attention_mask[:, keep_column_mask])
class 		UpperCAmelCase   (				UpperCAmelCase__			):







						'''simple docstring'''





						def __init__(						self							:     int       ,	__lowercase							:     List[Any]       ,	__lowercase							:     Tuple       ,	__lowercase							:     List[Any]       ,	__lowercase							:     str       ,	__lowercase							:     Tuple="train"       ,	__lowercase							:     List[str]=None       ,	__lowercase							:     List[Any]=None       ,	__lowercase							:     Optional[Any]=None       ,	__lowercase							:     Union[str, Any]=""       ,	):






							"""simple docstring"""






							super().__init__()
							snake_case_		    =     Path(__lowercase						).joinpath(type_path + ".source"						)
							snake_case_		    =     Path(__lowercase						).joinpath(type_path + ".target"						)
							snake_case_		    =     self.get_char_lens(self.src_file						)
							snake_case_		    =     max_source_length
							snake_case_		    =     max_target_length
							assert min(self.src_lens						) > 0, f"found empty line in {self.src_file}"
							snake_case_		    =     tokenizer
							snake_case_		    =     prefix
							if n_obs is not None:
								snake_case_		    =     self.src_lens[:n_obs]
							snake_case_		    =     src_lang
							snake_case_		    =     tgt_lang






						def __len__(						self							:     List[Any]						):






							"""simple docstring"""






							return len(self.src_lens						)






						def __getitem__(						self							:     List[Any]       ,	__lowercase							:     Dict						):






							"""simple docstring"""






							snake_case_		    =     index + 1  # linecache starts at 1
							snake_case_		    =     self.prefix + linecache.getline(str(self.src_file						)       ,	__lowercase						).rstrip("\n"						)
							snake_case_		    =     linecache.getline(str(self.tgt_file						)       ,	__lowercase						).rstrip("\n"						)
							assert source_line, f"empty source line for index {index}"
							assert tgt_line, f"empty tgt line for index {index}"

							# Need to add eos token manually for T5
							if isinstance(self.tokenizer       ,	__lowercase						):
								source_line += self.tokenizer.eos_token
								tgt_line += self.tokenizer.eos_token

							# Pad source and target to the right
							snake_case_		    =     (
							    self.tokenizer.question_encoder if isinstance(self.tokenizer       ,	__lowercase						) else self.tokenizer
							)
							snake_case_		    =     self.tokenizer.generator if isinstance(self.tokenizer       ,	__lowercase						) else self.tokenizer

							snake_case_		    =     encode_line(__lowercase       ,	__lowercase       ,	self.max_source_length       ,	"right"						)
							snake_case_		    =     encode_line(__lowercase       ,	__lowercase       ,	self.max_target_length       ,	"right"						)

							snake_case_		    =     source_inputs["input_ids"].squeeze()
							snake_case_		    =     target_inputs["input_ids"].squeeze()
							snake_case_		    =     source_inputs["attention_mask"].squeeze()
							return {
							    "input_ids": source_ids,
							    "attention_mask": src_mask,
							    "decoder_input_ids": target_ids,
							}






						@staticmethod
						def        snake_case__						(						__lowercase							:     Optional[int]						):






							"""simple docstring"""






							return [len(__lowercase						) for x in Path(__lowercase						).open().readlines()]






						def        snake_case__						(						self							:     Dict       ,	__lowercase							:     Union[str, Any]						):






							"""simple docstring"""






							snake_case_		    =     torch.stack([x["input_ids"] for x in batch]						)
							snake_case_		    =     torch.stack([x["attention_mask"] for x in batch]						)
							snake_case_		    =     torch.stack([x["decoder_input_ids"] for x in batch]						)
							snake_case_		    =     (
							    self.tokenizer.generator.pad_token_id
							    if isinstance(self.tokenizer       ,	__lowercase						)
							    else self.tokenizer.pad_token_id
							)
							snake_case_		    =     (
							    self.tokenizer.question_encoder.pad_token_id
							    if isinstance(self.tokenizer       ,	__lowercase						)
							    else self.tokenizer.pad_token_id
							)
							snake_case_		    =     trim_batch(__lowercase       ,	__lowercase						)
							snake_case_    ,      snake_case_		    =     trim_batch(__lowercase       ,	__lowercase       ,	attention_mask=__lowercase						)
							snake_case_		    =     {
							    "input_ids": source_ids,
							    "attention_mask": source_mask,
							    "decoder_input_ids": y,
							}
							return batch







lowercase__				:     str 		=						getLogger(__name__)


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	return list(itertools.chain.from_iterable(_A							)							)


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	snake_case_		    =     get_git_info()
	save_json(_A     ,	os.path.join(_A     ,	"git_log.json"							)							)


def 						lowerCamelCase__    (		_A     ,	_A     ,	_A=4     ,	**_A							):







	'''simple docstring'''





	with open(_A     ,	"w"							) as f:
		json.dump(_A     ,	_A     ,	indent=_A     ,	**_A							)


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	with open(_A							) as f:
		return json.load(_A							)


def 						lowerCamelCase__    (		):







	'''simple docstring'''





	snake_case_		    =     git.Repo(search_parent_directories=_A							)
	snake_case_		    =     {
	    "repo_id": str(_A							),
	    "repo_sha": str(repo.head.object.hexsha							),
	    "repo_branch": str(repo.active_branch							),
	    "hostname": str(socket.gethostname()							),
	}
	return repo_infos


def 						lowerCamelCase__    (		_A     ,	_A							):







	'''simple docstring'''





	return list(map(_A     ,	_A							)							)


def 						lowerCamelCase__    (		_A     ,	_A							):







	'''simple docstring'''





	with open(_A     ,	"wb"							) as f:
		return pickle.dump(_A     ,	_A							)


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	def remove_articles(_A							):
		return re.sub(R"\b(a|an|the)\b"     ,	" "     ,	_A							)

	def white_space_fix(_A							):
		return " ".join(text.split()							)

	def remove_punc(_A							):
		snake_case_		    =     set(string.punctuation							)
		return "".join(ch for ch in text if ch not in exclude							)

	def lower(_A							):
		return text.lower()

	return white_space_fix(remove_articles(remove_punc(lower(_A							)							)							)							)


def 						lowerCamelCase__    (		_A     ,	_A							):







	'''simple docstring'''





	snake_case_		    =     normalize_answer(_A							).split()
	snake_case_		    =     normalize_answer(_A							).split()
	snake_case_		    =     Counter(_A							) & Counter(_A							)
	snake_case_		    =     sum(common.values()							)
	if num_same == 0:
		return 0
	snake_case_		    =     1.0 * num_same / len(_A							)
	snake_case_		    =     1.0 * num_same / len(_A							)
	snake_case_		    =     (2 * precision * recall) / (precision + recall)
	return fa


def 						lowerCamelCase__    (		_A     ,	_A							):







	'''simple docstring'''





	return normalize_answer(_A							) == normalize_answer(_A							)


def 						lowerCamelCase__    (		_A     ,	_A							):







	'''simple docstring'''





	assert len(_A							) == len(_A							)
	snake_case_		    =     0
	for hypo, pred in zip(_A     ,	_A							):
		em += exact_match_score(_A     ,	_A							)
	if len(_A							) > 0:
		em /= len(_A							)
	return {"em": em}


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	return model_prefix.startswith("rag"							)






def 						lowerCamelCase__    (		_A     ,	_A     ,	_A							):







	'''simple docstring'''





	snake_case_		    =     {p: p for p in extra_params}
	# T5 models don't have `dropout` param, they have `dropout_rate` instead
	snake_case_		    =     "dropout_rate"
	for p in extra_params:
		if getattr(_A     ,	_A     ,	_A							):
			if not hasattr(_A     ,	_A							) and not hasattr(_A     ,	equivalent_param[p]							):
				logger.info("config doesn't have a `{}` attribute".format(_A							)							)
				delattr(_A     ,	_A							)
				continue
			snake_case_		    =     p if hasattr(_A     ,	_A							) else equivalent_param[p]
			setattr(_A     ,	_A     ,	getattr(_A     ,	_A							)							)
			delattr(_A     ,	_A							)
	return hparams, config
