
from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch





@require_torch
@require_tf
class 			_lowerCAmelCase      (			__a ):




				def   __a		(      self  )       ->   Optional[Any]:
							lowerCAmelCase_      			=			SMALL_MODEL_IDENTIFIER
							lowerCAmelCase_      			=			"pt"
							lowerCAmelCase_      			=			"tf"




				def   __a		(      self       ,       _UpperCamelCase  )       ->   Any:
							lowerCAmelCase_      			=			AutoModel.from_pretrained(self.test_model  )
							model_pt.save_pretrained(_UpperCamelCase  )




				def   __a		(      self       ,       _UpperCamelCase  )       ->   Optional[int]:
							lowerCAmelCase_      			=			TFAutoModel.from_pretrained(self.test_model       ,       from_pt=_UpperCamelCase  )
							model_tf.save_pretrained(_UpperCamelCase  )




				def   __a		(      self  )       ->   List[str]:
							lowerCAmelCase_      			=			"mock_framework"

							# Framework provided - return whatever the user provides
							lowerCAmelCase_      			=			FeaturesManager.determine_framework(self.test_model       ,       _UpperCamelCase  )
							self.assertEqual(_UpperCamelCase       ,       _UpperCamelCase  )

							# Local checkpoint and framework provided - return provided framework
							# PyTorch checkpoint
							with TemporaryDirectory() as local_pt_ckpt:
										self._setup_pt_ckpt(_UpperCamelCase  )
										lowerCAmelCase_      			=			FeaturesManager.determine_framework(_UpperCamelCase       ,       _UpperCamelCase  )
										self.assertEqual(_UpperCamelCase       ,       _UpperCamelCase  )

							# TensorFlow checkpoint
							with TemporaryDirectory() as local_tf_ckpt:
										self._setup_tf_ckpt(_UpperCamelCase  )
										lowerCAmelCase_      			=			FeaturesManager.determine_framework(_UpperCamelCase       ,       _UpperCamelCase  )
										self.assertEqual(_UpperCamelCase       ,       _UpperCamelCase  )




				def   __a		(      self  )       ->   List[str]:
							# PyTorch checkpoint
							with TemporaryDirectory() as local_pt_ckpt:
										self._setup_pt_ckpt(_UpperCamelCase  )
										lowerCAmelCase_      			=			FeaturesManager.determine_framework(_UpperCamelCase  )
										self.assertEqual(_UpperCamelCase       ,       self.framework_pt  )

							# TensorFlow checkpoint
							with TemporaryDirectory() as local_tf_ckpt:
										self._setup_tf_ckpt(_UpperCamelCase  )
										lowerCAmelCase_      			=			FeaturesManager.determine_framework(_UpperCamelCase  )
										self.assertEqual(_UpperCamelCase       ,       self.framework_tf  )

							# Invalid local checkpoint
							with TemporaryDirectory() as local_invalid_ckpt:
										with self.assertRaises(_UpperCamelCase  ):
													lowerCAmelCase_      			=			FeaturesManager.determine_framework(_UpperCamelCase  )




				def   __a		(      self  )       ->   Dict:
							lowerCAmelCase_      			=			MagicMock(return_value=_UpperCamelCase  )
							with patch("transformers.onnx.features.is_tf_available"       ,       _UpperCamelCase  ):
										lowerCAmelCase_      			=			FeaturesManager.determine_framework(self.test_model  )
										self.assertEqual(_UpperCamelCase       ,       self.framework_pt  )

							# PyTorch not in environment -> use TensorFlow
							lowerCAmelCase_      			=			MagicMock(return_value=_UpperCamelCase  )
							with patch("transformers.onnx.features.is_torch_available"       ,       _UpperCamelCase  ):
										lowerCAmelCase_      			=			FeaturesManager.determine_framework(self.test_model  )
										self.assertEqual(_UpperCamelCase       ,       self.framework_tf  )

							# Both in environment -> use PyTorch
							lowerCAmelCase_      			=			MagicMock(return_value=_UpperCamelCase  )
							lowerCAmelCase_      			=			MagicMock(return_value=_UpperCamelCase  )
							with patch("transformers.onnx.features.is_tf_available"       ,       _UpperCamelCase  ), patch(
							    "transformers.onnx.features.is_torch_available"       ,       _UpperCamelCase  ):
										lowerCAmelCase_      			=			FeaturesManager.determine_framework(self.test_model  )
										self.assertEqual(_UpperCamelCase       ,       self.framework_pt  )

							# Both not in environment -> raise error
							lowerCAmelCase_      			=			MagicMock(return_value=_UpperCamelCase  )
							lowerCAmelCase_      			=			MagicMock(return_value=_UpperCamelCase  )
							with patch("transformers.onnx.features.is_tf_available"       ,       _UpperCamelCase  ), patch(
							    "transformers.onnx.features.is_torch_available"       ,       _UpperCamelCase  ):
										with self.assertRaises(_UpperCamelCase  ):
													lowerCAmelCase_      			=			FeaturesManager.determine_framework(self.test_model  )





from ...processing_utils import ProcessorMixin





class 			_lowerCAmelCase      (			__a ):
				_lowercase    	='''SpeechT5FeatureExtractor'''
				_lowercase    	='''SpeechT5Tokenizer'''




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   int:
							super().__init__(_UpperCamelCase       ,       _UpperCamelCase  )




				def __call__(      self       ,       *_UpperCamelCase       ,       **_UpperCamelCase  )       ->   Optional[int]:
							lowerCAmelCase_      			=			kwargs.pop("audio"       ,       _UpperCamelCase  )
							lowerCAmelCase_      			=			kwargs.pop("text"       ,       _UpperCamelCase  )
							lowerCAmelCase_      			=			kwargs.pop("text_target"       ,       _UpperCamelCase  )
							lowerCAmelCase_      			=			kwargs.pop("audio_target"       ,       _UpperCamelCase  )
							lowerCAmelCase_      			=			kwargs.pop("sampling_rate"       ,       _UpperCamelCase  )

							if audio is not None and text is not None:
										raise ValueError(
										    "Cannot process both `audio` and `text` inputs. Did you mean `audio_target` or `text_target`?"  )
							if audio_target is not None and text_target is not None:
										raise ValueError(
										    "Cannot process both `audio_target` and `text_target` inputs. Did you mean `audio` or `text`?"  )
							if audio is None and audio_target is None and text is None and text_target is None:
										raise ValueError(
										    "You need to specify either an `audio`, `audio_target`, `text`, or `text_target` input to process."  )

							if audio is not None:
										lowerCAmelCase_      			=			self.feature_extractor(_UpperCamelCase       ,       *_UpperCamelCase       ,       sampling_rate=_UpperCamelCase       ,       **_UpperCamelCase  )
							elif text is not None:
										lowerCAmelCase_      			=			self.tokenizer(_UpperCamelCase       ,       **_UpperCamelCase  )
							else:
										lowerCAmelCase_      			=			None

							if audio_target is not None:
										lowerCAmelCase_      			=			self.feature_extractor(audio_target=_UpperCamelCase       ,       *_UpperCamelCase       ,       sampling_rate=_UpperCamelCase       ,       **_UpperCamelCase  )
										lowerCAmelCase_      			=			targets["input_values"]
							elif text_target is not None:
										lowerCAmelCase_      			=			self.tokenizer(_UpperCamelCase       ,       **_UpperCamelCase  )
										lowerCAmelCase_      			=			targets["input_ids"]
							else:
										lowerCAmelCase_      			=			None

							if inputs is None:
										return targets

							if targets is not None:
										lowerCAmelCase_      			=			labels

										lowerCAmelCase_      			=			targets.get("attention_mask"  )
										if decoder_attention_mask is not None:
													lowerCAmelCase_      			=			decoder_attention_mask

							return inputs




				def   __a		(      self       ,       *_UpperCamelCase       ,       **_UpperCamelCase  )       ->   str:
							lowerCAmelCase_      			=			kwargs.pop("input_values"       ,       _UpperCamelCase  )
							lowerCAmelCase_      			=			kwargs.pop("input_ids"       ,       _UpperCamelCase  )
							lowerCAmelCase_      			=			kwargs.pop("labels"       ,       _UpperCamelCase  )

							if input_values is not None and input_ids is not None:
										raise ValueError("Cannot process both `input_values` and `input_ids` inputs."  )
							if input_values is None and input_ids is None and labels is None:
										raise ValueError(
										    "You need to specify either an `input_values`, `input_ids`, or `labels` input to be padded."  )

							if input_values is not None:
										lowerCAmelCase_      			=			self.feature_extractor.pad(_UpperCamelCase       ,       *_UpperCamelCase       ,       **_UpperCamelCase  )
							elif input_ids is not None:
										lowerCAmelCase_      			=			self.tokenizer.pad(_UpperCamelCase       ,       **_UpperCamelCase  )
							else:
										lowerCAmelCase_      			=			None

							if labels is not None:
										if "input_ids" in labels or (isinstance(_UpperCamelCase       ,       _UpperCamelCase  ) and "input_ids" in labels[0]):
													lowerCAmelCase_      			=			self.tokenizer.pad(_UpperCamelCase       ,       **_UpperCamelCase  )
													lowerCAmelCase_      			=			targets["input_ids"]
										else:
													lowerCAmelCase_      			=			self.feature_extractor.feature_size
													lowerCAmelCase_      			=			self.feature_extractor.num_mel_bins
													lowerCAmelCase_      			=			self.feature_extractor.pad(_UpperCamelCase       ,       *_UpperCamelCase       ,       **_UpperCamelCase  )
													lowerCAmelCase_      			=			feature_size_hack
													lowerCAmelCase_      			=			targets["input_values"]
							else:
										lowerCAmelCase_      			=			None

							if inputs is None:
										return targets

							if targets is not None:
										lowerCAmelCase_      			=			labels

										lowerCAmelCase_      			=			targets.get("attention_mask"  )
										if decoder_attention_mask is not None:
													lowerCAmelCase_      			=			decoder_attention_mask

							return inputs




				def   __a		(      self       ,       *_UpperCamelCase       ,       **_UpperCamelCase  )       ->   List[Any]:
							return self.tokenizer.batch_decode(*_UpperCamelCase       ,       **_UpperCamelCase  )




				def   __a		(      self       ,       *_UpperCamelCase       ,       **_UpperCamelCase  )       ->   Optional[Any]:
							return self.tokenizer.decode(*_UpperCamelCase       ,       **_UpperCamelCase  )



