






'''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     :      Tuple   					='''bert-base-cased'''
_A     :      Union[str, Any]   					='''fp16'''
_A     :      Dict   					='''bf16'''
_A     :      Any   					=[FPaa, BFaa]

@require_fsdp
@require_cuda
class _lowercase     (    _lowercase		):




							def 		lowerCamelCase_    (							self:    Any					):
														super().setUp()

														lowerCamelCase__ :     str	  =	dict(
														    ACCELERATE_USE_FSDP="""true"""      ,							MASTER_ADDR="""localhost"""      ,							MASTER_PORT="""10999"""      ,							RANK="""0"""      ,							LOCAL_RANK="""0"""      ,							WORLD_SIZE="""1"""      ,							)




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

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




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

														for i, prefetch_policy in enumerate(UpperCamelCase__					):
																					lowerCamelCase__ :     List[str]	  =	self.dist_env.copy()
																					lowerCamelCase__ :     List[Any]	  =	prefetch_policy
																					with mockenv_context(**UpperCamelCase__					):
																												lowerCamelCase__ :     str	  =	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:    Union[str, Any]					):
														from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

														for i, state_dict_type in enumerate(UpperCamelCase__					):
																					lowerCamelCase__ :     Tuple	  =	self.dist_env.copy()
																					lowerCamelCase__ :     List[str]	  =	state_dict_type
																					with mockenv_context(**UpperCamelCase__					):
																												lowerCamelCase__ :     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 		lowerCamelCase_    (							self:    Dict					):
														lowerCamelCase__ :     Optional[int]	  =	AutoModel.from_pretrained(UpperCamelCase__					)
														for policy in FSDP_AUTO_WRAP_POLICY:
																					lowerCamelCase__ :     Any	  =	self.dist_env.copy()
																					lowerCamelCase__ :     str	  =	policy
																					if policy == "TRANSFORMER_BASED_WRAP":
																												lowerCamelCase__ :     Any	  =	"""BertLayer"""
																					elif policy == "SIZE_BASED_WRAP":
																												lowerCamelCase__ :     List[str]	  =	"""2000"""
																					with mockenv_context(**UpperCamelCase__					):
																												lowerCamelCase__ :     str	  =	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					)

														lowerCamelCase__ :     str	  =	self.dist_env.copy()
														lowerCamelCase__ :     str	  =	"""TRANSFORMER_BASED_WRAP"""
														lowerCamelCase__ :     str	  =	"""T5Layer"""
														with mockenv_context(**UpperCamelCase__					):
																					lowerCamelCase__ :     Optional[int]	  =	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					)					)

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




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

														for mp_dtype in dtypes:
																					lowerCamelCase__ :     str	  =	self.dist_env.copy()
																					lowerCamelCase__ :     Optional[Any]	  =	mp_dtype
																					with mockenv_context(**UpperCamelCase__					):
																												lowerCamelCase__ :     int	  =	Accelerator()
																												if mp_dtype == "fp16":
																																			lowerCamelCase__ :     str	  =	torch.floataa
																												elif mp_dtype == "bf16":
																																			lowerCamelCase__ :     Any	  =	torch.bfloataa
																												lowerCamelCase__ :     Dict	  =	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]:
																					lowerCamelCase__ :     List[str]	  =	self.dist_env.copy()
																					lowerCamelCase__ :     Optional[int]	  =	str(UpperCamelCase__					).lower()
																					with mockenv_context(**UpperCamelCase__					):
																												lowerCamelCase__ :     Optional[Any]	  =	FullyShardedDataParallelPlugin()
																												self.assertEqual(fsdp_plugin.cpu_offload      ,							CPUOffload(offload_params=UpperCamelCase__					)					)



@require_fsdp
@require_multi_gpu
@slow
class _lowercase     (    _lowercase		):




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

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




							def 		lowerCamelCase_    (							self:    List[Any]					):
														lowerCamelCase__ :     Union[str, Any]	  =	os.path.join(self.test_scripts_folder      ,							"""test_performance.py"""					)
														lowerCamelCase__ :     Dict	  =	["""accelerate""", """launch""", """--num_processes=2""", """--num_machines=1""", """--machine_rank=0""", """--use_fsdp"""]
														for config in self.performance_configs:
																					lowerCamelCase__ :     Optional[Any]	  =	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:    List[str]					):
														lowerCamelCase__ :     Dict	  =	os.path.join(self.test_scripts_folder      ,							"""test_checkpointing.py"""					)
														lowerCamelCase__ :     str	  =	[
														    """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__					):
																					lowerCamelCase__ :     Optional[Any]	  =	cmd.copy()
																					cmd_config.append(F'''--fsdp_sharding_strategy={i+1}'''					)
																					if strategy != "FULL_SHARD":
																												continue
																					lowerCamelCase__ :     List[Any]	  =	len(UpperCamelCase__					)
																					for state_dict_type in FSDP_STATE_DICT_TYPE:
																												lowerCamelCase__ :     str	  =	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()					)

																												lowerCamelCase__ :     int	  =	cmd_config[:-1]
																												lowerCamelCase__ :     Tuple	  =	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]					):
														lowerCamelCase__ :     Union[str, Any]	  =	os.path.join(self.test_scripts_folder      ,							"""test_peak_memory_usage.py"""					)
														lowerCamelCase__ :     Optional[Any]	  =	[
														    """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():
																					lowerCamelCase__ :     Optional[int]	  =	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'''







from collections import UserDict
from typing import Union

import numpy as np
import requests

from ..utils import (
    add_end_docstrings,
    logging,
)
from .audio_classification import ffmpeg_read
from .base import PIPELINE_INIT_ARGS, Pipeline


_A     :      Optional[Any]   					=logging.get_logger(__name__)

@add_end_docstrings(_lowercase		)
class _lowercase     (    _lowercase		):




							def __init__(							self:    Union[str, Any]      ,							**UpperCamelCase__:    str					):
														super().__init__(**UpperCamelCase__					)

														if self.framework != "pt":
																					raise ValueError(F'''The {self.__class__} is only available in PyTorch.'''					)
														# No specific FOR_XXX available yet




							def __call__(							self:    Optional[Any]      ,							UpperCamelCase__:    Union[np.ndarray, bytes, str]      ,							**UpperCamelCase__:    int					):
														return super().__call__(UpperCamelCase__      ,							**UpperCamelCase__					)




							def 		lowerCamelCase_    (							self:    Optional[Any]      ,							**UpperCamelCase__:    int					):
														lowerCamelCase__ :     Optional[Any]	  =	{}
														if "candidate_labels" in kwargs:
																					lowerCamelCase__ :     str	  =	kwargs["""candidate_labels"""]
														if "hypothesis_template" in kwargs:
																					lowerCamelCase__ :     int	  =	kwargs["""hypothesis_template"""]

														return preprocess_params, {}, {}




							def 		lowerCamelCase_    (							self:    Optional[Any]      ,							UpperCamelCase__:    int      ,							UpperCamelCase__:    List[str]=None      ,							UpperCamelCase__:    Dict="This is a sound of {}."					):
														if isinstance(UpperCamelCase__      ,							UpperCamelCase__					):
																					if audio.startswith("""http://"""					) or audio.startswith("""https://"""					):
																												# We need to actually check for a real protocol, otherwise it's impossible to use a local file
																												# like http_huggingface_co.png
																												lowerCamelCase__ :     int	  =	requests.get(UpperCamelCase__					).content
																					else:
																												with open(UpperCamelCase__      ,							"""rb"""					) as f:
																																			lowerCamelCase__ :     Dict	  =	f.read()

														if isinstance(UpperCamelCase__      ,							UpperCamelCase__					):
																					lowerCamelCase__ :     str	  =	ffmpeg_read(UpperCamelCase__      ,							self.feature_extractor.sampling_rate					)

														if not isinstance(UpperCamelCase__      ,							np.ndarray					):
																					raise ValueError("""We expect a numpy ndarray as input"""					)
														if len(audio.shape					) != 1:
																					raise ValueError("""We expect a single channel audio input for ZeroShotAudioClassificationPipeline"""					)

														lowerCamelCase__ :     Optional[Any]	  =	self.feature_extractor(
														    [audio]      ,							sampling_rate=self.feature_extractor.sampling_rate      ,							return_tensors="""pt"""					)
														lowerCamelCase__ :     Any	  =	candidate_labels
														lowerCamelCase__ :     Any	  =	[hypothesis_template.format(UpperCamelCase__					) for x in candidate_labels]
														lowerCamelCase__ :     int	  =	self.tokenizer(UpperCamelCase__      ,							return_tensors=self.framework      ,							padding=UpperCamelCase__					)
														lowerCamelCase__ :     Tuple	  =	[text_inputs]
														return inputs




							def 		lowerCamelCase_    (							self:    Tuple      ,							UpperCamelCase__:    Union[str, Any]					):
														lowerCamelCase__ :     Any	  =	model_inputs.pop("""candidate_labels"""					)
														lowerCamelCase__ :     Union[str, Any]	  =	model_inputs.pop("""text_inputs"""					)
														if isinstance(text_inputs[0]      ,							UpperCamelCase__					):
																					lowerCamelCase__ :     Optional[int]	  =	text_inputs[0]
														else:
																					# Batching case.
																					lowerCamelCase__ :     Tuple	  =	text_inputs[0][0]

														lowerCamelCase__ :     Tuple	  =	self.model(**UpperCamelCase__      ,							**UpperCamelCase__					)

														lowerCamelCase__ :     str	  =	{
														    """candidate_labels""": candidate_labels,
														    """logits""": outputs.logits_per_audio,
														}
														return model_outputs




							def 		lowerCamelCase_    (							self:    List[str]      ,							UpperCamelCase__:    Union[str, Any]					):
														lowerCamelCase__ :     List[str]	  =	model_outputs.pop("""candidate_labels"""					)
														lowerCamelCase__ :     int	  =	model_outputs["""logits"""][0]

														if self.framework == "pt":
																					lowerCamelCase__ :     Optional[int]	  =	logits.softmax(dim=0					)
																					lowerCamelCase__ :     Any	  =	probs.tolist()
														else:
																					raise ValueError("""`tf` framework not supported."""					)

														lowerCamelCase__ :     Tuple	  =	[
														    {"""score""": score, """label""": candidate_label}
														    for score, candidate_label in sorted(zip(UpperCamelCase__      ,							UpperCamelCase__					)      ,							key=lambda UpperCamelCase__					: -x[0]					)
														]
														return result

