
"""simple docstring"""



import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


a					:	Tuple							=		logging.get_logger(__name__)

a					:	Optional[Any]							=		{
    '''microsoft/wavlm-base''': '''https://huggingface.co/microsoft/wavlm-base/resolve/main/config.json''',
    # See all WavLM models at https://huggingface.co/models?filter=wavlm
}




class 		a_						( _UpperCAmelCase	):
    a      :				int     =				'wavlm'







    def __init__(			self			:       List[str]     ,					__UpperCamelCase			:       Optional[int]=32     ,					__UpperCamelCase			:       Dict=7_68     ,					__UpperCamelCase			:       List[Any]=12     ,					__UpperCamelCase			:       List[str]=12     ,					__UpperCamelCase			:       str=30_72     ,					__UpperCamelCase			:       int="gelu"     ,					__UpperCamelCase			:       int=0.1     ,					__UpperCamelCase			:       Optional[Any]=0.1     ,					__UpperCamelCase			:       Dict=0.1     ,					__UpperCamelCase			:       Dict=0.0     ,					__UpperCamelCase			:       Tuple=0.1     ,					__UpperCamelCase			:       List[Any]=0.1     ,					__UpperCamelCase			:       List[Any]=0.0_2     ,					__UpperCamelCase			:       Optional[Any]=1e-5     ,					__UpperCamelCase			:       Optional[Any]="group"     ,					__UpperCamelCase			:       Union[str, Any]="gelu"     ,					__UpperCamelCase			:       str=(5_12, 5_12, 5_12, 5_12, 5_12, 5_12, 5_12)     ,					__UpperCamelCase			:       Union[str, Any]=(5, 2, 2, 2, 2, 2, 2)     ,					__UpperCamelCase			:       Union[str, Any]=(10, 3, 3, 3, 3, 2, 2)     ,					__UpperCamelCase			:       str=False     ,					__UpperCamelCase			:       List[Any]=1_28     ,					__UpperCamelCase			:       Union[str, Any]=16     ,					__UpperCamelCase			:       Optional[int]=3_20     ,					__UpperCamelCase			:       Any=8_00     ,					__UpperCamelCase			:       Optional[int]=False     ,					__UpperCamelCase			:       List[Any]=True     ,					__UpperCamelCase			:       Optional[int]=0.0_5     ,					__UpperCamelCase			:       Optional[Any]=10     ,					__UpperCamelCase			:       int=2     ,					__UpperCamelCase			:       str=0.0     ,					__UpperCamelCase			:       Optional[int]=10     ,					__UpperCamelCase			:       Any=3_20     ,					__UpperCamelCase			:       Any=2     ,					__UpperCamelCase			:       Tuple=0.1     ,					__UpperCamelCase			:       List[str]=1_00     ,					__UpperCamelCase			:       int=2_56     ,					__UpperCamelCase			:       Dict=2_56     ,					__UpperCamelCase			:       Tuple=0.1     ,					__UpperCamelCase			:       Optional[int]="mean"     ,					__UpperCamelCase			:       Union[str, Any]=False     ,					__UpperCamelCase			:       str=False     ,					__UpperCamelCase			:       List[str]=2_56     ,					__UpperCamelCase			:       Tuple=(5_12, 5_12, 5_12, 5_12, 15_00)     ,					__UpperCamelCase			:       Union[str, Any]=(5, 3, 3, 1, 1)     ,					__UpperCamelCase			:       Any=(1, 2, 3, 1, 1)     ,					__UpperCamelCase			:       int=5_12     ,					__UpperCamelCase			:       List[Any]=80     ,					__UpperCamelCase			:       Optional[int]=0     ,					__UpperCamelCase			:       Dict=1     ,					__UpperCamelCase			:       List[Any]=2     ,					__UpperCamelCase			:       Tuple=False     ,					__UpperCamelCase			:       Dict=3     ,					__UpperCamelCase			:       str=2     ,					__UpperCamelCase			:       Dict=3     ,					__UpperCamelCase			:       int=None     ,					**__UpperCamelCase			:       Optional[Any]     ,					)		->Union[str, Any]:






        '''simple docstring'''






        super().__init__(**__UpperCamelCase     ,					pad_token_id=__UpperCamelCase     ,					bos_token_id=__UpperCamelCase     ,					eos_token_id=__UpperCamelCase     )
        _UpperCAmelCase          =   hidden_size
        _UpperCAmelCase          =   feat_extract_norm
        _UpperCAmelCase          =   feat_extract_activation
        _UpperCAmelCase          =   list(__UpperCamelCase     )
        _UpperCAmelCase          =   list(__UpperCamelCase     )
        _UpperCAmelCase          =   list(__UpperCamelCase     )
        _UpperCAmelCase          =   conv_bias
        _UpperCAmelCase          =   num_buckets
        _UpperCAmelCase          =   max_bucket_distance
        _UpperCAmelCase          =   num_conv_pos_embeddings
        _UpperCAmelCase          =   num_conv_pos_embedding_groups
        _UpperCAmelCase          =   len(self.conv_dim     )
        _UpperCAmelCase          =   num_hidden_layers
        _UpperCAmelCase          =   intermediate_size
        _UpperCAmelCase          =   hidden_act
        _UpperCAmelCase          =   num_attention_heads
        _UpperCAmelCase          =   hidden_dropout
        _UpperCAmelCase          =   attention_dropout
        _UpperCAmelCase          =   activation_dropout
        _UpperCAmelCase          =   feat_proj_dropout
        _UpperCAmelCase          =   final_dropout
        _UpperCAmelCase          =   layerdrop
        _UpperCAmelCase          =   layer_norm_eps
        _UpperCAmelCase          =   initializer_range
        _UpperCAmelCase          =   num_ctc_classes
        _UpperCAmelCase          =   vocab_size
        _UpperCAmelCase          =   do_stable_layer_norm
        _UpperCAmelCase          =   use_weighted_layer_sum
        _UpperCAmelCase          =   classifier_proj_size

        if (
            (len(self.conv_stride     ) != self.num_feat_extract_layers)
            or (len(self.conv_kernel     ) != self.num_feat_extract_layers)
            or (len(self.conv_dim     ) != self.num_feat_extract_layers)
        ):
            raise ValueError(
                """Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =="""
                """ `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ="""
                f""" {len(self.conv_dim     )}`, `len(config.conv_stride) = {len(self.conv_stride     )}`,"""
                f""" `len(config.conv_kernel) = {len(self.conv_kernel     )}`."""     )

        # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
        _UpperCAmelCase          =   apply_spec_augment
        _UpperCAmelCase          =   mask_time_prob
        _UpperCAmelCase          =   mask_time_length
        _UpperCAmelCase          =   mask_time_min_masks
        _UpperCAmelCase          =   mask_feature_prob
        _UpperCAmelCase          =   mask_feature_length

        # parameters for pretraining with codevector quantized representations
        _UpperCAmelCase          =   num_codevectors_per_group
        _UpperCAmelCase          =   num_codevector_groups
        _UpperCAmelCase          =   contrastive_logits_temperature
        _UpperCAmelCase          =   num_negatives
        _UpperCAmelCase          =   codevector_dim
        _UpperCAmelCase          =   proj_codevector_dim
        _UpperCAmelCase          =   diversity_loss_weight

        # ctc loss
        _UpperCAmelCase          =   ctc_loss_reduction
        _UpperCAmelCase          =   ctc_zero_infinity

        # adapter
        _UpperCAmelCase          =   add_adapter
        _UpperCAmelCase          =   adapter_kernel_size
        _UpperCAmelCase          =   adapter_stride
        _UpperCAmelCase          =   num_adapter_layers
        _UpperCAmelCase          =   output_hidden_size or hidden_size

        # SequenceClassification-specific parameter. Feel free to ignore for other classes.
        _UpperCAmelCase          =   classifier_proj_size

        # XVector-specific parameters. Feel free to ignore for other classes.
        _UpperCAmelCase          =   list(__UpperCamelCase     )
        _UpperCAmelCase          =   list(__UpperCamelCase     )
        _UpperCAmelCase          =   list(__UpperCamelCase     )
        _UpperCAmelCase          =   xvector_output_dim







    @property
    def      _snake_case					(			self			:       str     )		->Tuple:






        '''simple docstring'''






        return functools.reduce(operator.mul     ,					self.conv_stride     ,					1     )

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

a					:	List[Any]							=		get_logger()

a					:	Optional[dict]							=		None




class 		a_						( TensorFormatter[Mapping, 'jax.Array', Mapping]	):







    def __init__(			self			:       int     ,					__UpperCamelCase			:       List[str]=None     ,					__UpperCamelCase			:       Optional[int]=None     ,					**__UpperCamelCase			:       int     )		->Tuple:






        '''simple docstring'''






        super().__init__(features=__UpperCamelCase     )
        import jax
        from jaxlib.xla_client import Device

        if isinstance(__UpperCamelCase     ,					__UpperCamelCase     ):
            raise ValueError(
                f"""Expected {device} to be a `str` not {type(__UpperCamelCase     )}, 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`."""     )
        _UpperCAmelCase          =   device if isinstance(__UpperCamelCase     ,					__UpperCamelCase     ) 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:
            _UpperCAmelCase          =   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]     )}."""     )
            _UpperCAmelCase          =   str(jax.devices()[0]     )
        _UpperCAmelCase          =   jnp_array_kwargs







    @staticmethod
    def      _snake_case					(			)		->Dict[str, "jaxlib.xla_extension.Device"]:






        '''simple docstring'''






        import jax

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







    def      _snake_case					(			self			:       Dict     ,					__UpperCamelCase			:       Any     )		->Union[str, Any]:






        '''simple docstring'''






        import jax
        import jax.numpy as jnp

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







    def      _snake_case					(			self			:       List[str]     ,					__UpperCamelCase			:       Any     )		->Optional[int]:






        '''simple docstring'''






        import jax
        import jax.numpy as jnp

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

        _UpperCAmelCase          =   {}

        if isinstance(__UpperCamelCase     ,					(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:
                _UpperCAmelCase          =   {"""dtype""": jnp.intaa}
            else:
                _UpperCAmelCase          =   {"""dtype""": jnp.intaa}
        elif isinstance(__UpperCamelCase     ,					(np.number, np.ndarray)     ) and np.issubdtype(value.dtype     ,					np.floating     ):
            _UpperCAmelCase          =   {"""dtype""": jnp.floataa}
        elif config.PIL_AVAILABLE and "PIL" in sys.modules:
            import PIL.Image

            if isinstance(__UpperCamelCase     ,					PIL.Image.Image     ):
                _UpperCAmelCase          =   np.asarray(__UpperCamelCase     )

        # 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:
            _UpperCAmelCase          =   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(__UpperCamelCase     ,					**{**default_dtype, **self.jnp_array_kwargs}     )







    def      _snake_case					(			self			:       Union[str, Any]     ,					__UpperCamelCase			:       List[str]     )		->Any:






        '''simple docstring'''






        import jax

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

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







    def      _snake_case					(			self			:       List[str]     ,					__UpperCamelCase			:       dict     )		->int:






        '''simple docstring'''






        return map_nested(self._recursive_tensorize     ,					__UpperCamelCase     ,					map_list=__UpperCamelCase     )







    def      _snake_case					(			self			:       Dict     ,					__UpperCamelCase			:       pa.Table     )		->Mapping:






        '''simple docstring'''






        _UpperCAmelCase          =   self.numpy_arrow_extractor().extract_row(__UpperCamelCase     )
        _UpperCAmelCase          =   self.python_features_decoder.decode_row(__UpperCamelCase     )
        return self.recursive_tensorize(__UpperCamelCase     )







    def      _snake_case					(			self			:       Optional[int]     ,					__UpperCamelCase			:       pa.Table     )		->"jax.Array":






        '''simple docstring'''






        _UpperCAmelCase          =   self.numpy_arrow_extractor().extract_column(__UpperCamelCase     )
        _UpperCAmelCase          =   self.python_features_decoder.decode_column(__UpperCamelCase     ,					pa_table.column_names[0]     )
        _UpperCAmelCase          =   self.recursive_tensorize(__UpperCamelCase     )
        _UpperCAmelCase          =   self._consolidate(__UpperCamelCase     )
        return column







    def      _snake_case					(			self			:       Optional[Any]     ,					__UpperCamelCase			:       pa.Table     )		->Mapping:






        '''simple docstring'''






        _UpperCAmelCase          =   self.numpy_arrow_extractor().extract_batch(__UpperCamelCase     )
        _UpperCAmelCase          =   self.python_features_decoder.decode_batch(__UpperCamelCase     )
        _UpperCAmelCase          =   self.recursive_tensorize(__UpperCamelCase     )
        for column_name in batch:
            _UpperCAmelCase          =   self._consolidate(batch[column_name]     )
        return batch