




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

_lowerCamelCase			:	Dict							 =     get_logger()

_lowerCamelCase			:	Optional[dict]							 =     None


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

    '''simple docstring'''






    def __init__(    self   :			Tuple      ,				UpperCAmelCase__   :			Optional[int]=None      ,				UpperCAmelCase__   :			Any=None      ,				**UpperCAmelCase__   :			str)		->Dict:





           '''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`.''')
           A__   				=					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:
                  A__   				=					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])}.""")
                  A__   				=					str(jax.devices()[0])
           A__   				=					jnp_array_kwargs


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





           '''simple docstring'''

           import jax

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


    def  SCREAMING_SNAKE_CASE							(    self   :			List[str]      ,				UpperCAmelCase__   :			int)		->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  SCREAMING_SNAKE_CASE							(    self   :			Union[str, Any]      ,				UpperCAmelCase__   :			List[Any])		->List[Any]:





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

           A__   				=					{}

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

                  if isinstance(UpperCAmelCase__      ,				PIL.Image.Image):
                         A__   				=					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:
                  A__   				=					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  SCREAMING_SNAKE_CASE							(    self   :			Optional[Any]      ,				UpperCAmelCase__   :			Any)		->Union[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):
                  A__   				=					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  SCREAMING_SNAKE_CASE							(    self   :			Dict      ,				UpperCAmelCase__   :			dict)		->Optional[int]:





           '''simple docstring'''

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


    def  SCREAMING_SNAKE_CASE							(    self   :			Tuple      ,				UpperCAmelCase__   :			pa.Table)		->Mapping:





           '''simple docstring'''

           A__   				=					self.numpy_arrow_extractor().extract_row(UpperCAmelCase__)
           A__   				=					self.python_features_decoder.decode_row(UpperCAmelCase__)
           return self.recursive_tensorize(UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			Union[str, Any]      ,				UpperCAmelCase__   :			pa.Table)		->"jax.Array":





           '''simple docstring'''

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





    def  SCREAMING_SNAKE_CASE							(    self   :			str      ,				UpperCAmelCase__   :			pa.Table)		->Mapping:





           '''simple docstring'''

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








from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


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

_lowerCamelCase			:	Dict							 =     {
    """xlm-mlm-en-2048""": """https://huggingface.co/xlm-mlm-en-2048/resolve/main/config.json""",
    """xlm-mlm-ende-1024""": """https://huggingface.co/xlm-mlm-ende-1024/resolve/main/config.json""",
    """xlm-mlm-enfr-1024""": """https://huggingface.co/xlm-mlm-enfr-1024/resolve/main/config.json""",
    """xlm-mlm-enro-1024""": """https://huggingface.co/xlm-mlm-enro-1024/resolve/main/config.json""",
    """xlm-mlm-tlm-xnli15-1024""": """https://huggingface.co/xlm-mlm-tlm-xnli15-1024/resolve/main/config.json""",
    """xlm-mlm-xnli15-1024""": """https://huggingface.co/xlm-mlm-xnli15-1024/resolve/main/config.json""",
    """xlm-clm-enfr-1024""": """https://huggingface.co/xlm-clm-enfr-1024/resolve/main/config.json""",
    """xlm-clm-ende-1024""": """https://huggingface.co/xlm-clm-ende-1024/resolve/main/config.json""",
    """xlm-mlm-17-1280""": """https://huggingface.co/xlm-mlm-17-1280/resolve/main/config.json""",
    """xlm-mlm-100-1280""": """https://huggingface.co/xlm-mlm-100-1280/resolve/main/config.json""",
}


class        UpperCamelCase_					(				UpperCAmelCase__  ):

    '''simple docstring'''






    UpperCAmelCase__													=   '''xlm'''
    UpperCAmelCase__													=   {
        '''hidden_size''': '''emb_dim''',
        '''num_attention_heads''': '''n_heads''',
        '''num_hidden_layers''': '''n_layers''',
        '''n_words''': '''vocab_size''',  # For backward compatibility
    }


    def __init__(    self   :			Tuple      ,				UpperCAmelCase__   :			Optional[Any]=30_145      ,				UpperCAmelCase__   :			List[str]=2_048      ,				UpperCAmelCase__   :			str=12      ,				UpperCAmelCase__   :			List[Any]=16      ,				UpperCAmelCase__   :			Dict=0.1      ,				UpperCAmelCase__   :			Tuple=0.1      ,				UpperCAmelCase__   :			str=True      ,				UpperCAmelCase__   :			Union[str, Any]=False      ,				UpperCAmelCase__   :			str=False      ,				UpperCAmelCase__   :			Optional[int]=False      ,				UpperCAmelCase__   :			List[str]=1      ,				UpperCAmelCase__   :			List[Any]=True      ,				UpperCAmelCase__   :			Union[str, Any]=512      ,				UpperCAmelCase__   :			List[str]=2_048**-0.5      ,				UpperCAmelCase__   :			List[Any]=1e-12      ,				UpperCAmelCase__   :			Any=0.02      ,				UpperCAmelCase__   :			str=0      ,				UpperCAmelCase__   :			int=1      ,				UpperCAmelCase__   :			Union[str, Any]=2      ,				UpperCAmelCase__   :			List[Any]=3      ,				UpperCAmelCase__   :			int=5      ,				UpperCAmelCase__   :			Any=True      ,				UpperCAmelCase__   :			Tuple="first"      ,				UpperCAmelCase__   :			Optional[int]=True      ,				UpperCAmelCase__   :			List[str]=None      ,				UpperCAmelCase__   :			Tuple=True      ,				UpperCAmelCase__   :			Dict=0.1      ,				UpperCAmelCase__   :			Any=5      ,				UpperCAmelCase__   :			List[Any]=5      ,				UpperCAmelCase__   :			int=0      ,				UpperCAmelCase__   :			Dict=0      ,				UpperCAmelCase__   :			Dict=2      ,				UpperCAmelCase__   :			str=0      ,				**UpperCAmelCase__   :			List[str]      ,				)		->str:





           '''simple docstring'''

           A__   				=					vocab_size
           A__   				=					emb_dim
           A__   				=					n_layers
           A__   				=					n_heads
           A__   				=					dropout
           A__   				=					attention_dropout
           A__   				=					gelu_activation
           A__   				=					sinusoidal_embeddings
           A__   				=					causal
           A__   				=					asm
           A__   				=					n_langs
           A__   				=					use_lang_emb
           A__   				=					layer_norm_eps
           A__   				=					bos_index
           A__   				=					eos_index
           A__   				=					pad_index
           A__   				=					unk_index
           A__   				=					mask_index
           A__   				=					is_encoder
           A__   				=					max_position_embeddings
           A__   				=					embed_init_std
           A__   				=					init_std
           A__   				=					summary_type
           A__   				=					summary_use_proj
           A__   				=					summary_activation
           A__   				=					summary_proj_to_labels
           A__   				=					summary_first_dropout
           A__   				=					start_n_top
           A__   				=					end_n_top
           A__   				=					mask_token_id
           A__   				=					lang_id

           if "n_words" in kwargs:
                  A__   				=					kwargs['''n_words''']

           super().__init__(pad_token_id=UpperCAmelCase__      ,				bos_token_id=UpperCAmelCase__      ,				**UpperCAmelCase__)


class        UpperCamelCase_					(				UpperCAmelCase__  ):

    '''simple docstring'''






    @property
    def  SCREAMING_SNAKE_CASE							(    self   :			Dict)		->Mapping[str, Mapping[int, str]]:





           '''simple docstring'''

           if self.task == "multiple-choice":
                  A__   				=					{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
           else:
                  A__   				=					{0: '''batch''', 1: '''sequence'''}
           return OrderedDict(
               [
                   ('''input_ids''', dynamic_axis),
                   ('''attention_mask''', dynamic_axis),
                   ('''token_type_ids''', dynamic_axis),
               ])


