






"""simple docstring"""

from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowercase__          = {
    """configuration_whisper""": ["""WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP""", """WhisperConfig""", """WhisperOnnxConfig"""],
    """feature_extraction_whisper""": ["""WhisperFeatureExtractor"""],
    """processing_whisper""": ["""WhisperProcessor"""],
    """tokenization_whisper""": ["""WhisperTokenizer"""],
}

try:
					if not is_tokenizers_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase__          = ["""WhisperTokenizerFast"""]

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase__          = [
					    """WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """WhisperForConditionalGeneration""",
					    """WhisperModel""",
					    """WhisperPreTrainedModel""",
					    """WhisperForAudioClassification""",
					]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase__          = [
					    """TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """TFWhisperForConditionalGeneration""",
					    """TFWhisperModel""",
					    """TFWhisperPreTrainedModel""",
					]

try:
					if not is_flax_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase__          = [
					    """FlaxWhisperForConditionalGeneration""",
					    """FlaxWhisperModel""",
					    """FlaxWhisperPreTrainedModel""",
					    """FlaxWhisperForAudioClassification""",
					]


if TYPE_CHECKING:
					from .configuration_whisper import WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP, WhisperConfig, WhisperOnnxConfig
					from .feature_extraction_whisper import WhisperFeatureExtractor
					from .processing_whisper import WhisperProcessor
					from .tokenization_whisper import WhisperTokenizer

					try:
										if not is_tokenizers_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_whisper_fast import WhisperTokenizerFast

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_whisper import (
										    WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    WhisperForAudioClassification,
										    WhisperForConditionalGeneration,
										    WhisperModel,
										    WhisperPreTrainedModel,
										)

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_whisper import (
										    TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    TFWhisperForConditionalGeneration,
										    TFWhisperModel,
										    TFWhisperPreTrainedModel,
										)

					try:
										if not is_flax_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_flax_whisper import (
										    FlaxWhisperForAudioClassification,
										    FlaxWhisperForConditionalGeneration,
										    FlaxWhisperModel,
										    FlaxWhisperPreTrainedModel,
										)

else:
					import sys

					lowercase__          = _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)








"""simple docstring"""

# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
					import jax
					import jaxlib

lowercase__          = get_logger()

lowercase__          = None


class 		__lowerCamelCase				( TensorFormatter[Mapping, """jax.Array""", Mapping]			):





		'''simple docstring'''






		def __init__(   self							:	Optional[Any]						,			a_							:	Union[str, Any]=None						,			a_							:	Optional[Any]=None						,			**a_							:	Tuple	):
				super().__init__(features=a_	)
				import jax
				from jaxlib.xla_client import Device

				if isinstance(a_						,			a_	):
						raise ValueError(
						    f'''Expected {device} to be a `str` not {type(a_	)}, as `jaxlib.xla_extension.Device` '''
						    "is not serializable neither with `pickle` nor with `dill`. Instead you can surround "
						    "the device with `str()` to get its string identifier that will be internally mapped "
						    "to the actual `jaxlib.xla_extension.Device`."	)
				lowerCAmelCase_   :							List[Any]									=			device if isinstance(a_						,			a_	) else str(jax.devices()[0]	)
				# using global variable since `jaxlib.xla_extension.Device` is not serializable neither
				# with `pickle` nor with `dill`, so we need to use a global variable instead
				global DEVICE_MAPPING
				if DEVICE_MAPPING is None:
						lowerCAmelCase_   :							Dict									=			self._map_devices_to_str()
				if self.device not in list(DEVICE_MAPPING.keys()	):
						logger.warning(
						    f'''Device with string identifier {self.device} not listed among the available '''
						    f'''devices: {list(DEVICE_MAPPING.keys()	)}, so falling back to the default '''
						    f'''device: {str(jax.devices()[0]	)}.'''	)
						lowerCAmelCase_   :							Tuple									=			str(jax.devices()[0]	)
				lowerCAmelCase_   :							Dict									=			jnp_array_kwargs





		@staticmethod
		def    lowerCamelCase       (   ):
				import jax

				return {str(a_	): device for device in jax.devices()}





		def    lowerCamelCase       (   self							:	Tuple						,			a_							:	Dict	):
				import jax
				import jax.numpy as jnp

				if isinstance(a_						,			a_	) and column:
						if all(
						    isinstance(a_						,			jax.Array	) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column	):
								return jnp.stack(a_						,			axis=0	)
				return column





		def    lowerCamelCase       (   self							:	Tuple						,			a_							:	Tuple	):
				import jax
				import jax.numpy as jnp

				if isinstance(a_						,			(str, bytes, type(a_	))	):
						return value
				elif isinstance(a_						,			(np.character, np.ndarray)	) and np.issubdtype(value.dtype						,			np.character	):
						return value.tolist()

				lowerCAmelCase_   :							Any									=			{}

				if isinstance(a_						,			(np.number, np.ndarray)	) and np.issubdtype(value.dtype						,			np.integer	):
						# the default int precision depends on the jax config
						# see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
						if jax.config.jax_enable_xaa:
								lowerCAmelCase_   :							List[Any]									=			{"dtype": jnp.intaa}
						else:
								lowerCAmelCase_   :							Tuple									=			{"dtype": jnp.intaa}
				elif isinstance(a_						,			(np.number, np.ndarray)	) and np.issubdtype(value.dtype						,			np.floating	):
						lowerCAmelCase_   :							Optional[int]									=			{"dtype": jnp.floataa}
				elif config.PIL_AVAILABLE and "PIL" in sys.modules:
						import PIL.Image

						if isinstance(a_						,			PIL.Image.Image	):
								lowerCAmelCase_   :							List[str]									=			np.asarray(a_	)

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
				global DEVICE_MAPPING
				if DEVICE_MAPPING is None:
						lowerCAmelCase_   :							Dict									=			self._map_devices_to_str()

				with jax.default_device(DEVICE_MAPPING[self.device]	):
						# calling jnp.array on a np.ndarray does copy the data
						# see https://github.com/google/jax/issues/4486
						return jnp.array(a_						,			**{**default_dtype, **self.jnp_array_kwargs}	)





		def    lowerCamelCase       (   self							:	Any						,			a_							:	List[str]	):
				import jax

				# support for torch, tf, jax etc.
				if config.TORCH_AVAILABLE and "torch" in sys.modules:
						import torch

						if isinstance(a_						,			torch.Tensor	):
								return self._tensorize(data_struct.detach().cpu().numpy()[()]	)
				if hasattr(a_						,			"__array__"	) and not isinstance(a_						,			jax.Array	):
						lowerCAmelCase_   :							str									=			data_struct.__array__()
				# support for nested types like struct of list of struct
				if isinstance(a_						,			np.ndarray	):
						if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
								return self._consolidate([self.recursive_tensorize(a_	) for substruct in data_struct]	)
				elif isinstance(a_						,			(list, tuple)	):
						return self._consolidate([self.recursive_tensorize(a_	) for substruct in data_struct]	)
				return self._tensorize(a_	)





		def    lowerCamelCase       (   self							:	Dict						,			a_							:	dict	):
				return map_nested(self._recursive_tensorize						,			a_						,			map_list=a_	)





		def    lowerCamelCase       (   self							:	List[str]						,			a_							:	pa.Table	):
				lowerCAmelCase_   :							Tuple									=			self.numpy_arrow_extractor().extract_row(a_	)
				lowerCAmelCase_   :							Union[str, Any]									=			self.python_features_decoder.decode_row(a_	)
				return self.recursive_tensorize(a_	)





		def    lowerCamelCase       (   self							:	Any						,			a_							:	pa.Table	):
				lowerCAmelCase_   :							Dict									=			self.numpy_arrow_extractor().extract_column(a_	)
				lowerCAmelCase_   :							Optional[int]									=			self.python_features_decoder.decode_column(a_						,			pa_table.column_names[0]	)
				lowerCAmelCase_   :							List[str]									=			self.recursive_tensorize(a_	)
				lowerCAmelCase_   :							Optional[Any]									=			self._consolidate(a_	)
				return column





		def    lowerCamelCase       (   self							:	Tuple						,			a_							:	pa.Table	):
				lowerCAmelCase_   :							Tuple									=			self.numpy_arrow_extractor().extract_batch(a_	)
				lowerCAmelCase_   :							Tuple									=			self.python_features_decoder.decode_batch(a_	)
				lowerCAmelCase_   :							List[str]									=			self.recursive_tensorize(a_	)
				for column_name in batch:
						lowerCAmelCase_   :							Tuple									=			self._consolidate(batch[column_name]	)
				return batch
