





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


lowercase_		:				List[Any]				=						logging.get_logger(__name__)


@add_end_docstrings(a_     )
class 			lowercase				(							a_     ):



						"""simple docstring"""

						def __init__(    self			:		Union[str, Any] ,				**lowerCamelCase_			:		Dict					):





										'''simple docstring'''







										super().__init__(**lowerCamelCase_					)

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







						def __call__(    self			:		List[str] ,				lowerCamelCase_			:		Union[np.ndarray, bytes, str] ,				**lowerCamelCase_			:		str					):





										'''simple docstring'''







										return super().__call__(lowerCamelCase_ ,				**lowerCamelCase_					)







						def 		__UpperCAmelCase       (    self			:		str ,				**lowerCamelCase_			:		List[str]					):





										'''simple docstring'''







										_snake_case :			str  						=   {}
										if "candidate_labels" in kwargs:
														_snake_case :			List[Any]  						=   kwargs['candidate_labels']
										if "hypothesis_template" in kwargs:
														_snake_case :			Tuple  						=   kwargs['hypothesis_template']

										return preprocess_params, {}, {}







						def 		__UpperCAmelCase       (    self			:		Optional[int] ,				lowerCamelCase_			:		Union[str, Any] ,				lowerCamelCase_			:		str=None ,				lowerCamelCase_			:		Dict="This is a sound of {}."					):





										'''simple docstring'''







										if isinstance(lowerCamelCase_ ,				lowerCamelCase_					):
														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(lowerCamelCase_					).content
														else:
																		with open(lowerCamelCase_ ,				'rb'					) as f:
																						_snake_case :			Optional[Any]  						=   f.read()

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

										if not isinstance(lowerCamelCase_ ,				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 :			str  						=   self.feature_extractor(
										    [audio] ,				sampling_rate=self.feature_extractor.sampling_rate ,				return_tensors='pt'					)
										_snake_case :			Optional[int]  						=   candidate_labels
										_snake_case :			int  						=   [hypothesis_template.format(lowerCamelCase_					) for x in candidate_labels]
										_snake_case :			List[str]  						=   self.tokenizer(lowerCamelCase_ ,				return_tensors=self.framework ,				padding=lowerCamelCase_					)
										_snake_case :			Optional[Any]  						=   [text_inputs]
										return inputs







						def 		__UpperCAmelCase       (    self			:		Union[str, Any] ,				lowerCamelCase_			:		Union[str, Any]					):





										'''simple docstring'''







										_snake_case :			Union[str, Any]  						=   model_inputs.pop('candidate_labels'					)
										_snake_case :			List[Any]  						=   model_inputs.pop('text_inputs'					)
										if isinstance(text_inputs[0] ,				lowerCamelCase_					):
														_snake_case :			Dict  						=   text_inputs[0]
										else:
														# Batching case.
														_snake_case :			Optional[int]  						=   text_inputs[0][0]

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

										_snake_case :			str  						=   {
										    'candidate_labels': candidate_labels,
										    'logits': outputs.logits_per_audio,
										}
										return model_outputs







						def 		__UpperCAmelCase       (    self			:		Union[str, Any] ,				lowerCamelCase_			:		Optional[Any]					):





										'''simple docstring'''







										_snake_case :			int  						=   model_outputs.pop('candidate_labels'					)
										_snake_case :			Optional[Any]  						=   model_outputs['logits'][0]

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

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













import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
lowercase_		:				Optional[int]				=						object()

# For specifying empty leaf dict `{}`
lowercase_		:				List[Any]				=						object()

def       A__(		__lowerCAmelCase	,			__lowerCAmelCase ):
				_snake_case :			Optional[Any]  						=   tuple((re.compile(x + '$' ) for x in qs) )
				for i in range(len(__lowerCAmelCase ) - len(__lowerCAmelCase ) + 1 ):
								_snake_case :			Tuple  						=   [x.match(__lowerCAmelCase ) for x, y in zip(__lowerCAmelCase	,			ks[i:] )]
								if matches and all(__lowerCAmelCase ):
												return True
				return False

def       A__(		__lowerCAmelCase ):
				def replace(__lowerCAmelCase	,			__lowerCAmelCase ):
								for rule, replacement in rules:
												if _match(__lowerCAmelCase	,			__lowerCAmelCase ):
																return replacement
								return val

				return replace

def       A__(		):
				return [
				    # embeddings
				    (("transformer", "wpe", "embedding"), P('mp'	,			__lowerCAmelCase )),
				    (("transformer", "wte", "embedding"), P('mp'	,			__lowerCAmelCase )),
				    # atention
				    (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(__lowerCAmelCase	,			'mp' )),
				    (("attention", "out_proj", "kernel"), P('mp'	,			__lowerCAmelCase )),
				    (("attention", "out_proj", "bias"), None),
				    # mlp
				    (("mlp", "c_fc", "kernel"), P(__lowerCAmelCase	,			'mp' )),
				    (("mlp", "c_fc", "bias"), P('mp' )),
				    (("mlp", "c_proj", "kernel"), P('mp'	,			__lowerCAmelCase )),
				    (("mlp", "c_proj", "bias"), None),
				    # layer norms
				    ((r"ln_\d+", "bias"), None),
				    ((r"\d+", r"ln_\d+", "scale"), None),
				    (("ln_f", "bias"), None),
				    (("ln_f", "scale"), None),
				]

def       A__(		__lowerCAmelCase ):
				_snake_case :			Optional[Any]  						=   _get_partition_rules()
				_snake_case :			Optional[Any]  						=   _replacement_rules(__lowerCAmelCase )
				_snake_case :			str  						=   {k: _unmatched for k in flatten_dict(__lowerCAmelCase )}
				_snake_case :			str  						=   {k: replace(__lowerCAmelCase	,			__lowerCAmelCase ) for k, v in initd.items()}
				assert _unmatched not in result.values(), "Incomplete partition spec."
				return freeze(unflatten_dict(__lowerCAmelCase ) )






