
'''simple docstring'''



import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
						from transformers.models.esm.modeling_esm import (
						    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
						    EsmEmbeddings,
						    create_position_ids_from_input_ids,
						)




class     UpperCAmelCase_	:






						'''simple docstring'''






						def __init__(   self       ,  _lowercase       ,  _lowercase=13       ,  _lowercase=7       ,  _lowercase=False       ,  _lowercase=True       ,  _lowercase=False       ,  _lowercase=True       ,  _lowercase=33       ,  _lowercase=32       ,  _lowercase=5       ,  _lowercase=4       ,  _lowercase=37       ,  _lowercase="gelu"       ,  _lowercase=0.1       ,  _lowercase=0.1       ,  _lowercase=512       ,  _lowercase=16       ,  _lowercase=2       ,  _lowercase=0.02       ,  _lowercase=3       ,  _lowercase=4       ,  _lowercase=None       ,  ):

												"""simple docstring"""



												_lowerCAmelCase    		=					parent
												_lowerCAmelCase    		=					batch_size
												_lowerCAmelCase    		=					seq_length
												_lowerCAmelCase    		=					is_training
												_lowerCAmelCase    		=					use_input_mask
												_lowerCAmelCase    		=					use_token_type_ids
												_lowerCAmelCase    		=					use_labels
												_lowerCAmelCase    		=					vocab_size
												_lowerCAmelCase    		=					hidden_size
												_lowerCAmelCase    		=					num_hidden_layers
												_lowerCAmelCase    		=					num_attention_heads
												_lowerCAmelCase    		=					intermediate_size
												_lowerCAmelCase    		=					hidden_act
												_lowerCAmelCase    		=					hidden_dropout_prob
												_lowerCAmelCase    		=					attention_probs_dropout_prob
												_lowerCAmelCase    		=					max_position_embeddings
												_lowerCAmelCase    		=					type_vocab_size
												_lowerCAmelCase    		=					type_sequence_label_size
												_lowerCAmelCase    		=					initializer_range
												_lowerCAmelCase    		=					num_labels
												_lowerCAmelCase    		=					num_choices
												_lowerCAmelCase    		=					scope




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					ids_tensor([self.batch_size, self.seq_length]       ,  self.vocab_size )

												_lowerCAmelCase    		=					None
												if self.use_input_mask:
																		_lowerCAmelCase    		=					random_attention_mask([self.batch_size, self.seq_length] )

												_lowerCAmelCase    		=					None
												_lowerCAmelCase    		=					None
												_lowerCAmelCase    		=					None
												if self.use_labels:
																		_lowerCAmelCase    		=					ids_tensor([self.batch_size]       ,  self.type_sequence_label_size )
																		_lowerCAmelCase    		=					ids_tensor([self.batch_size, self.seq_length]       ,  self.num_labels )
																		_lowerCAmelCase    		=					ids_tensor([self.batch_size]       ,  self.num_choices )

												_lowerCAmelCase    		=					self.get_config()

												return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels




						def    _lowercase      (   self ):

												"""simple docstring"""



												return EsmConfig(
												    vocab_size=self.vocab_size       ,  hidden_size=self.hidden_size       ,  pad_token_id=1       ,  num_hidden_layers=self.num_hidden_layers       ,  num_attention_heads=self.num_attention_heads       ,  intermediate_size=self.intermediate_size       ,  hidden_act=self.hidden_act       ,  hidden_dropout_prob=self.hidden_dropout_prob       ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,  max_position_embeddings=self.max_position_embeddings       ,  type_vocab_size=self.type_vocab_size       ,  initializer_range=self.initializer_range       ,  )




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase ):

												"""simple docstring"""



												_lowerCAmelCase    		=					EsmModel(config=_lowercase )
												model.to(_lowercase )
												model.eval()
												_lowerCAmelCase    		=					model(_lowercase       ,  attention_mask=_lowercase )
												_lowerCAmelCase    		=					model(_lowercase )
												_lowerCAmelCase    		=					model(_lowercase )

												self.parent.assertEqual(result.last_hidden_state.shape       ,  (self.batch_size, self.seq_length, self.hidden_size) )
												self.parent.assertEqual(result.pooler_output.shape       ,  (self.batch_size, self.hidden_size) )




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase ):

												"""simple docstring"""



												_lowerCAmelCase    		=					EsmForMaskedLM(config=_lowercase )
												model.to(_lowercase )
												model.eval()
												_lowerCAmelCase    		=					model(_lowercase       ,  attention_mask=_lowercase       ,  labels=_lowercase )
												self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.seq_length, self.vocab_size) )




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.num_labels
												_lowerCAmelCase    		=					EsmForTokenClassification(config=_lowercase )
												model.to(_lowercase )
												model.eval()
												_lowerCAmelCase    		=					model(_lowercase       ,  attention_mask=_lowercase       ,  labels=_lowercase )
												self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.seq_length, self.num_labels) )




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.prepare_config_and_inputs()
												(
												    (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       
												)    		=					config_and_inputs
												_lowerCAmelCase    		=					{"""input_ids""": input_ids, """attention_mask""": input_mask}
												return config, inputs_dict






@require_torch
class     UpperCAmelCase_	(						_SCREAMING_SNAKE_CASE			,				_SCREAMING_SNAKE_CASE			,				unittest.TestCase					):






						'''simple docstring'''






						_lowercase						: List[str]										=  False

						_lowercase						: Tuple										=  (
						    (
						        EsmForMaskedLM,
						        EsmModel,
						        EsmForSequenceClassification,
						        EsmForTokenClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						_lowercase						: Tuple										=  ()
						_lowercase						: Optional[int]										=  (
						    {
						        '''feature-extraction''': EsmModel,
						        '''fill-mask''': EsmForMaskedLM,
						        '''text-classification''': EsmForSequenceClassification,
						        '''token-classification''': EsmForTokenClassification,
						        '''zero-shot''': EsmForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						_lowercase						: int										=  True




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					EsmModelTester(self )
												_lowerCAmelCase    		=					ConfigTester(self       ,  config_class=_lowercase       ,  hidden_size=37 )




						def    _lowercase      (   self ):

												"""simple docstring"""



												self.config_tester.run_common_tests()




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_lowercase )




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.model_tester.prepare_config_and_inputs()
												for type in ["absolute", "relative_key", "relative_key_query"]:
																		_lowerCAmelCase    		=					type
																		self.model_tester.create_and_check_model(*_lowercase )




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_masked_lm(*_lowercase )




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_token_classification(*_lowercase )




						@slow
						def    _lowercase      (   self ):

												"""simple docstring"""



												for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		_lowerCAmelCase    		=					EsmModel.from_pretrained(_lowercase )
																		self.assertIsNotNone(_lowercase )




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.model_tester.prepare_config_and_inputs()[0]
												_lowerCAmelCase    		=					EsmEmbeddings(config=_lowercase )

												_lowerCAmelCase    		=					torch.as_tensor([[12, 31, 13, model.padding_idx]] )
												_lowerCAmelCase    		=					torch.as_tensor(
												    [
												        [
												            0 + model.padding_idx + 1,
												            1 + model.padding_idx + 1,
												            2 + model.padding_idx + 1,
												            model.padding_idx,
												        ]
												    ] )
												_lowerCAmelCase    		=					create_position_ids_from_input_ids(_lowercase       ,  model.padding_idx )
												self.assertEqual(position_ids.shape       ,  expected_positions.shape )
												self.assertTrue(torch.all(torch.eq(_lowercase       ,  _lowercase ) ) )




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.model_tester.prepare_config_and_inputs()[0]
												_lowerCAmelCase    		=					EsmEmbeddings(config=_lowercase )

												_lowerCAmelCase    		=					torch.empty(2       ,  4       ,  30 )
												_lowerCAmelCase    		=					[
												    0 + embeddings.padding_idx + 1,
												    1 + embeddings.padding_idx + 1,
												    2 + embeddings.padding_idx + 1,
												    3 + embeddings.padding_idx + 1,
												]
												_lowerCAmelCase    		=					torch.as_tensor([expected_single_positions, expected_single_positions] )
												_lowerCAmelCase    		=					embeddings.create_position_ids_from_inputs_embeds(_lowercase )
												self.assertEqual(position_ids.shape       ,  expected_positions.shape )
												self.assertTrue(torch.all(torch.eq(_lowercase       ,  _lowercase ) ) )




						@unittest.skip("""Esm does not support embedding resizing""" )
						def    _lowercase      (   self ):

												"""simple docstring"""



												pass




						@unittest.skip("""Esm does not support embedding resizing""" )
						def    _lowercase      (   self ):

												"""simple docstring"""



												pass




						@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests.""" )
						def    _lowercase      (   self ):

												"""simple docstring"""



												pass






@require_torch
class     UpperCAmelCase_	(						_SCREAMING_SNAKE_CASE					):






						'''simple docstring'''






						@slow
						def    _lowercase      (   self ):

												"""simple docstring"""



												with torch.no_grad():
																		_lowerCAmelCase    		=					EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D""" )
																		model.eval()
																		_lowerCAmelCase    		=					torch.tensor([[0, 1, 2, 3, 4, 5]] )
																		_lowerCAmelCase    		=					model(_lowercase )[0]

																		_lowerCAmelCase    		=					33

																		_lowerCAmelCase    		=					torch.Size((1, 6, vocab_size) )
																		self.assertEqual(output.shape       ,  _lowercase )

																		_lowerCAmelCase    		=					torch.tensor(
																		    [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]] )
																		self.assertTrue(torch.allclose(output[:, :3, :3]       ,  _lowercase       ,  atol=1e-4 ) )




						@slow
						def    _lowercase      (   self ):

												"""simple docstring"""



												with torch.no_grad():
																		_lowerCAmelCase    		=					EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D""" )
																		model.eval()

																		_lowerCAmelCase    		=					torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]] )
																		_lowerCAmelCase    		=					model(_lowercase )[0]
																		# compare the actual values for a slice.
																		_lowerCAmelCase    		=					torch.tensor(
																		    [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]] )
																		self.assertTrue(torch.allclose(output[:, :3, :3]       ,  _lowercase       ,  atol=1e-4 ) )


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


_lowercase          =						logging.get_logger(__name__)




@add_end_docstrings(_SCREAMING_SNAKE_CASE					)
class     UpperCAmelCase_	(						_SCREAMING_SNAKE_CASE					):






						'''simple docstring'''






						def __init__(   self       ,  **_lowercase ):

												"""simple docstring"""



												super().__init__(**_lowercase )

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




						def __call__(   self       ,  _lowercase       ,  **_lowercase ):

												"""simple docstring"""



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




						def    _lowercase      (   self       ,  **_lowercase ):

												"""simple docstring"""



												_lowerCAmelCase    		=					{}
												if "candidate_labels" in kwargs:
																		_lowerCAmelCase    		=					kwargs["""candidate_labels"""]
												if "hypothesis_template" in kwargs:
																		_lowerCAmelCase    		=					kwargs["""hypothesis_template"""]

												return preprocess_params, {}, {}




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase=None       ,  _lowercase="This is a sound of {}." ):

												"""simple docstring"""



												if isinstance(_lowercase       ,  _lowercase ):
																		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    		=					requests.get(_lowercase ).content
																		else:
																								with open(_lowercase       ,  """rb""" ) as f:
																														_lowerCAmelCase    		=					f.read()

												if isinstance(_lowercase       ,  _lowercase ):
																		_lowerCAmelCase    		=					ffmpeg_read(_lowercase       ,  self.feature_extractor.sampling_rate )

												if not isinstance(_lowercase       ,  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    		=					self.feature_extractor(
												    [audio]       ,  sampling_rate=self.feature_extractor.sampling_rate       ,  return_tensors="""pt""" )
												_lowerCAmelCase    		=					candidate_labels
												_lowerCAmelCase    		=					[hypothesis_template.format(_lowercase ) for x in candidate_labels]
												_lowerCAmelCase    		=					self.tokenizer(_lowercase       ,  return_tensors=self.framework       ,  padding=_lowercase )
												_lowerCAmelCase    		=					[text_inputs]
												return inputs




						def    _lowercase      (   self       ,  _lowercase ):

												"""simple docstring"""



												_lowerCAmelCase    		=					model_inputs.pop("""candidate_labels""" )
												_lowerCAmelCase    		=					model_inputs.pop("""text_inputs""" )
												if isinstance(text_inputs[0]       ,  _lowercase ):
																		_lowerCAmelCase    		=					text_inputs[0]
												else:
																		# Batching case.
																		_lowerCAmelCase    		=					text_inputs[0][0]

												_lowerCAmelCase    		=					self.model(**_lowercase       ,  **_lowercase )

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




						def    _lowercase      (   self       ,  _lowercase ):

												"""simple docstring"""



												_lowerCAmelCase    		=					model_outputs.pop("""candidate_labels""" )
												_lowerCAmelCase    		=					model_outputs["""logits"""][0]

												if self.framework == "pt":
																		_lowerCAmelCase    		=					logits.softmax(dim=0 )
																		_lowerCAmelCase    		=					probs.tolist()
												else:
																		raise ValueError("""`tf` framework not supported.""" )

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