


"""simple docstring"""




import json
import os
import tempfile
from unittest.mock import patch

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment


def     lowercase			()   ->						Optional[int]:
						SCREAMING_SNAKE_CASE													=					torch.nn.Linear(2 ,			4  )
						SCREAMING_SNAKE_CASE													=					torch.optim.AdamW(model.parameters() ,			lr=1.0  )
						SCREAMING_SNAKE_CASE													=					torch.optim.lr_scheduler.OneCycleLR(SCREAMING_SNAKE_CASE_ ,			max_lr=0.01 ,			steps_per_epoch=2 ,			epochs=1  )
						SCREAMING_SNAKE_CASE													=					DataLoader(TensorDataset(torch.tensor([1, 2, 3]  )  )  )
						SCREAMING_SNAKE_CASE													=					DataLoader(TensorDataset(torch.tensor([4, 5, 6]  )  )  )

						return model, optimizer, scheduler, train_dl, valid_dl


def     lowercase			(SCREAMING_SNAKE_CASE_		:				Optional[Any]  )   ->						int:
						return (model.weight.abs().sum() + model.bias.abs().sum()).item()


def     lowercase			(SCREAMING_SNAKE_CASE_		:				str  )   ->						List[str]:
						SCREAMING_SNAKE_CASE													=					torch.nn.Linear(*tuple(model.weight.T.shape  )  ).state_dict()
						model.load_state_dict(SCREAMING_SNAKE_CASE_  )


class 					lowerCAmelCase	(					lowerCAmelCase_    ):




		'''simple docstring'''







		@require_cuda
		def 							__A						(	self				)     ->							Optional[int]:
								SCREAMING_SNAKE_CASE													=					Accelerator()
								assert PartialState._shared_state["_cpu"] is False
								assert PartialState._shared_state["device"].type == "cuda"
								with self.assertRaises(_snake_case				):
														SCREAMING_SNAKE_CASE													=					Accelerator(cpu=_snake_case				)







		def 							__A						(	self				)     ->							Dict:
								SCREAMING_SNAKE_CASE													=					Accelerator()
								SCREAMING_SNAKE_CASE													=					GradientState()
								assert state.num_steps == 1
								SCREAMING_SNAKE_CASE													=					4
								assert state.num_steps == 4

								assert state.sync_gradients is True
								SCREAMING_SNAKE_CASE													=					False
								assert state.sync_gradients is False
								GradientState._reset_state()







		def 							__A						(	self				)     ->							Optional[Any]:
								SCREAMING_SNAKE_CASE													=					Accelerator()
								SCREAMING_SNAKE_CASE													=					create_components()

								(
								    SCREAMING_SNAKE_CASE				
								)													=					accelerator.prepare(_snake_case						,       _snake_case						,       _snake_case						,       _snake_case						,       _snake_case				)

								self.assertTrue(prepared_model in accelerator._models				)
								self.assertTrue(prepared_optimizer in accelerator._optimizers				)
								self.assertTrue(prepared_scheduler in accelerator._schedulers				)
								self.assertTrue(prepared_train_dl in accelerator._dataloaders				)
								self.assertTrue(prepared_valid_dl in accelerator._dataloaders				)







		def 							__A						(	self				)     ->							int:
								SCREAMING_SNAKE_CASE													=					Accelerator()
								SCREAMING_SNAKE_CASE													=					create_components()
								accelerator.prepare(_snake_case						,       _snake_case						,       _snake_case						,       _snake_case						,       _snake_case				)
								accelerator.free_memory()

								self.assertTrue(len(accelerator._models				) == 0				)
								self.assertTrue(len(accelerator._optimizers				) == 0				)
								self.assertTrue(len(accelerator._schedulers				) == 0				)
								self.assertTrue(len(accelerator._dataloaders				) == 0				)







		def 							__A						(	self				)     ->							List[Any]:
								PartialState._reset_state()

								# Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
								def noop(*lowerCAmelCase__						,       **lowerCAmelCase__				):
														pass

								with patch('torch.cuda.set_device'						,       _snake_case				), patch_environment(ACCELERATE_TORCH_DEVICE='cuda:64'				):
														SCREAMING_SNAKE_CASE													=					Accelerator()
														self.assertEqual(str(accelerator.state.device				)						,       'cuda:64'				)







		def 							__A						(	self				)     ->							List[Any]:
								SCREAMING_SNAKE_CASE													=					Accelerator()
								SCREAMING_SNAKE_CASE													=					create_components()
								accelerator.prepare(_snake_case						,       _snake_case						,       _snake_case						,       _snake_case						,       _snake_case				)

								SCREAMING_SNAKE_CASE													=					get_signature(_snake_case				)

								with tempfile.TemporaryDirectory() as tmpdirname:
														accelerator.save_state(_snake_case				)

														# make sure random weights don't match
														load_random_weights(_snake_case				)
														self.assertTrue(abs(model_signature - get_signature(_snake_case				)				) > 1e-3				)

														# make sure loaded weights match
														accelerator.load_state(_snake_case				)
														self.assertTrue(abs(model_signature - get_signature(_snake_case				)				) < 1e-3				)







		def 							__A						(	self				)     ->							int:
								SCREAMING_SNAKE_CASE													=					Accelerator()
								SCREAMING_SNAKE_CASE													=					create_components()
								accelerator.prepare(_snake_case						,       _snake_case						,       _snake_case						,       _snake_case						,       _snake_case				)

								SCREAMING_SNAKE_CASE													=					get_signature(_snake_case				)

								# saving hook
								def save_config(lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__				):
														SCREAMING_SNAKE_CASE													=					{'class_name': models[0].__class__.__name__}

														with open(os.path.join(_snake_case						,       'data.json'				)						,       'w'				) as f:
																				json.dump(_snake_case						,       _snake_case				)

        # loading hook
								def load_config(lowerCAmelCase__						,       lowerCAmelCase__				):
														with open(os.path.join(_snake_case						,       'data.json'				)						,       'r'				) as f:
																				SCREAMING_SNAKE_CASE													=					json.load(_snake_case				)

														SCREAMING_SNAKE_CASE													=					config['class_name']

								SCREAMING_SNAKE_CASE													=					accelerator.register_save_state_pre_hook(_snake_case				)
								SCREAMING_SNAKE_CASE													=					accelerator.register_load_state_pre_hook(_snake_case				)

								with tempfile.TemporaryDirectory() as tmpdirname:
														accelerator.save_state(_snake_case				)

														# make sure random weights don't match with hooks
														load_random_weights(_snake_case				)
														self.assertTrue(abs(model_signature - get_signature(_snake_case				)				) > 1e-3				)

														# random class name to verify correct one is loaded
														SCREAMING_SNAKE_CASE													=					'random'

														# make sure loaded weights match with hooks
														accelerator.load_state(_snake_case				)
														self.assertTrue(abs(model_signature - get_signature(_snake_case				)				) < 1e-3				)

														# mode.class_name is loaded from config
														self.assertTrue(model.class_name == model.__class__.__name__				)

								# remove hooks
								save_hook.remove()
								load_hook.remove()

								with tempfile.TemporaryDirectory() as tmpdirname:
														accelerator.save_state(_snake_case				)

														# make sure random weights don't match with hooks removed
														load_random_weights(_snake_case				)
														self.assertTrue(abs(model_signature - get_signature(_snake_case				)				) > 1e-3				)

														# random class name to verify correct one is loaded
														SCREAMING_SNAKE_CASE													=					'random'

														# make sure loaded weights match with hooks removed
														accelerator.load_state(_snake_case				)
														self.assertTrue(abs(model_signature - get_signature(_snake_case				)				) < 1e-3				)

														# mode.class_name is NOT loaded from config
														self.assertTrue(model.class_name != model.__class__.__name__				)







		def 							__A						(	self				)     ->							Any:
								SCREAMING_SNAKE_CASE													=					Accelerator()
								SCREAMING_SNAKE_CASE													=					create_components()
								SCREAMING_SNAKE_CASE													=					None

								# This should work
								SCREAMING_SNAKE_CASE													=					accelerator.prepare(
								    _snake_case						,       _snake_case						,       _snake_case						,       _snake_case						,       _snake_case						,       _snake_case				)
								self.assertTrue(dummy_obj is None				)







		def 							__A						(	self				)     ->							Optional[int]:
								SCREAMING_SNAKE_CASE													=					Accelerator()
								SCREAMING_SNAKE_CASE													=					create_components()
								SCREAMING_SNAKE_CASE													=					[1, 2, 3]

								# This should work
								SCREAMING_SNAKE_CASE													=					accelerator.prepare(
								    _snake_case						,       _snake_case						,       _snake_case						,       _snake_case						,       _snake_case						,       _snake_case				)
								self.assertEqual(
								    getattr(_snake_case						,       '_is_accelerate_prepared'						,       _snake_case				)						,       _snake_case						,       'Dummy object should have `_is_accelerate_prepared` set to `True`'						,       )
								self.assertEqual(
								    getattr(_snake_case						,       '_is_accelerate_prepared'						,       _snake_case				)						,       _snake_case						,       'Model is missing `_is_accelerator_prepared` or is set to `False`'						,       )
								self.assertEqual(
								    getattr(_snake_case						,       '_is_accelerate_prepared'						,       _snake_case				)						,       _snake_case						,       'Optimizer is missing `_is_accelerator_prepared` or is set to `False`'						,       )
								self.assertEqual(
								    getattr(_snake_case						,       '_is_accelerate_prepared'						,       _snake_case				)						,       _snake_case						,       'Scheduler is missing `_is_accelerator_prepared` or is set to `False`'						,       )
								self.assertEqual(
								    getattr(_snake_case						,       '_is_accelerate_prepared'						,       _snake_case				)						,       _snake_case						,       'Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`'						,       )
								self.assertEqual(
								    getattr(_snake_case						,       '_is_accelerate_prepared'						,       _snake_case				)						,       _snake_case						,       'Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`'						,       )







		@slow
		@require_bnb
		def 							__A						(	self				)     ->							Any:
								from transformers import AutoModelForCausalLM

								SCREAMING_SNAKE_CASE													=					AutoModelForCausalLM.from_pretrained(
								    'EleutherAI/gpt-neo-125m'						,       load_in_abit=_snake_case						,       device_map={'': 0}						,       )
								SCREAMING_SNAKE_CASE													=					Accelerator()

								# This should work
								SCREAMING_SNAKE_CASE													=					accelerator.prepare(_snake_case				)







		@slow
		@require_bnb
		def 							__A						(	self				)     ->							str:
								from transformers import AutoModelForCausalLM

								SCREAMING_SNAKE_CASE													=					Accelerator()

								with init_empty_weights():
														SCREAMING_SNAKE_CASE													=					AutoModelForCausalLM.from_pretrained(
														    'EleutherAI/gpt-neo-125m'						,       )
														model.tie_weights()
														SCREAMING_SNAKE_CASE													=					infer_auto_device_map(_snake_case				)
														SCREAMING_SNAKE_CASE													=					'cpu'

								SCREAMING_SNAKE_CASE													=					AutoModelForCausalLM.from_pretrained(
								    'EleutherAI/gpt-neo-125m'						,       device_map=_snake_case						,       load_in_abit=_snake_case						,       llm_inta_enable_fpaa_cpu_offload=_snake_case				)

								# This should not work and get value error
								with self.assertRaises(_snake_case				):
														SCREAMING_SNAKE_CASE													=					accelerator.prepare(_snake_case				)







		@slow
		@require_bnb
		@require_multi_gpu
		def 							__A						(	self				)     ->							Optional[Any]:
								from transformers import AutoModelForCausalLM

								SCREAMING_SNAKE_CASE													=					{'distributed_type': DistributedType.MULTI_GPU}

								with init_empty_weights():
														SCREAMING_SNAKE_CASE													=					AutoModelForCausalLM.from_pretrained(
														    'EleutherAI/gpt-neo-125m'						,       )
														model.tie_weights()
														SCREAMING_SNAKE_CASE													=					infer_auto_device_map(_snake_case				)
														SCREAMING_SNAKE_CASE													=					1

								SCREAMING_SNAKE_CASE													=					AutoModelForCausalLM.from_pretrained(
								    'EleutherAI/gpt-neo-125m'						,       load_in_abit=_snake_case						,       device_map=_snake_case						,       )
								SCREAMING_SNAKE_CASE													=					Accelerator()

								# This should not work and get value error
								with self.assertRaises(_snake_case				):
														SCREAMING_SNAKE_CASE													=					accelerator.prepare(_snake_case				)

								PartialState._reset_state()







		@slow
		@require_bnb
		@require_multi_gpu
		def 							__A						(	self				)     ->							Tuple:
								from transformers import AutoModelForCausalLM

								with init_empty_weights():
														SCREAMING_SNAKE_CASE													=					AutoModelForCausalLM.from_pretrained(
														    'EleutherAI/gpt-neo-125m'						,       )
														SCREAMING_SNAKE_CASE													=					infer_auto_device_map(_snake_case				)
														SCREAMING_SNAKE_CASE													=					1

								SCREAMING_SNAKE_CASE													=					AutoModelForCausalLM.from_pretrained(
								    'EleutherAI/gpt-neo-125m'						,       load_in_abit=_snake_case						,       device_map=_snake_case						,       )
								SCREAMING_SNAKE_CASE													=					Accelerator()

								# This should work
								SCREAMING_SNAKE_CASE													=					accelerator.prepare(_snake_case				)







		@require_cuda
		def 							__A						(	self				)     ->							List[Any]:
								SCREAMING_SNAKE_CASE													=					torch.nn.Linear(10						,       10				)
								SCREAMING_SNAKE_CASE													=					torch.optim.SGD(model.parameters()						,       lr=0.01				)
								SCREAMING_SNAKE_CASE													=					Accelerator(cpu=_snake_case				)
								SCREAMING_SNAKE_CASE													=					accelerator.prepare(_snake_case				)







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_   :List[Any]        =      logging.get_logger(__name__)




@add_end_docstrings(lowerCAmelCase_     )
class 					snake_case__						(       lowerCAmelCase_     ):






     """simple docstring"""


     def __init__(       self  :  Optional[Any], **_snake_case  :  str					)   ->Dict:
            super().__init__(**_snake_case					)

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






     def __call__(       self  :  Union[str, Any], _snake_case  :  Union[np.ndarray, bytes, str], **_snake_case  :  Tuple					)   ->Dict:
            return super().__call__(_snake_case, **_snake_case					)






     def    lowercase_   (       self  :  Tuple, **_snake_case  :  Any					)   ->Union[str, Any]:
            snake_case__							:	str   	=	{}
            if "candidate_labels" in kwargs:
                   snake_case__							:	str   	=	kwargs['candidate_labels']
            if "hypothesis_template" in kwargs:
                   snake_case__							:	str   	=	kwargs['hypothesis_template']

            return preprocess_params, {}, {}






     def    lowercase_   (       self  :  Dict, _snake_case  :  str, _snake_case  :  Optional[int]=None, _snake_case  :  List[str]="This is a sound of {}."					)   ->int:
            if isinstance(_snake_case, _snake_case					):
                   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
                          snake_case__							:	List[Any]   	=	requests.get(_snake_case					).content
                   else:
                          with open(_snake_case, 'rb'					) as f:
                                 snake_case__							:	Union[str, Any]   	=	f.read()

            if isinstance(_snake_case, _snake_case					):
                   snake_case__							:	List[Any]   	=	ffmpeg_read(_snake_case, self.feature_extractor.sampling_rate					)

            if not isinstance(_snake_case, 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'					)

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






     def    lowercase_   (       self  :  Optional[int], _snake_case  :  Optional[Any]					)   ->int:
            snake_case__							:	Optional[int]   	=	model_inputs.pop('candidate_labels'					)
            snake_case__							:	str   	=	model_inputs.pop('text_inputs'					)
            if isinstance(text_inputs[0], _snake_case					):
                   snake_case__							:	Optional[Any]   	=	text_inputs[0]
            else:
                   # Batching case.
                   snake_case__							:	int   	=	text_inputs[0][0]

            snake_case__							:	Any   	=	self.model(**_snake_case, **_snake_case					)

            snake_case__							:	List[Any]   	=	{
                'candidate_labels': candidate_labels,
                'logits': outputs.logits_per_audio,
            }
            return model_outputs






     def    lowercase_   (       self  :  Union[str, Any], _snake_case  :  str					)   ->List[str]:
            snake_case__							:	int   	=	model_outputs.pop('candidate_labels'					)
            snake_case__							:	List[Any]   	=	model_outputs['logits'][0]

            if self.framework == "pt":
                   snake_case__							:	Tuple   	=	logits.softmax(dim=0					)
                   snake_case__							:	Union[str, Any]   	=	probs.tolist()
            else:
                   raise ValueError('`tf` framework not supported.'					)

            snake_case__							:	Union[str, Any]   	=	[
                {'score': score, 'label': candidate_label}
                for score, candidate_label in sorted(zip(_snake_case, _snake_case					), key=lambda _snake_case					: -x[0]					)
            ]
            return result






