



'''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 __A     (    TensorFormatter[Mapping, 'jax.Array', Mapping]		):



 '''simple docstring'''







 def __init__(self					,  A=None					,  A=None					,  **A			)       ->	List[str]:


      """simple docstring"""



      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`.'''			)
      _a           =    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:
           _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    a__    ()       ->	Dict[str, "jaxlib.xla_extension.Device"]:


      """simple docstring"""



      import jax

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



 def    a__    (self					,  A			)       ->	int:


      """simple docstring"""



      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    a__    (self					,  A			)       ->	str:


      """simple docstring"""



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

      _a           =    {}

      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:
                _a           =    {'''dtype''': jnp.intaa}
           else:
                _a           =    {'''dtype''': jnp.intaa}
      elif isinstance(A					,  (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(A					,  PIL.Image.Image			):
                _a           =    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:
           _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(A					,  **{**default_dtype, **self.jnp_array_kwargs}			)



 def    a__    (self					,  A			)       ->	Optional[int]:


      """simple docstring"""



      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			):
           _a           =    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    a__    (self					,  A			)       ->	Dict:


      """simple docstring"""



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



 def    a__    (self					,  A			)       ->	Mapping:


      """simple docstring"""



      _a           =    self.numpy_arrow_extractor().extract_row(A			)
      _a           =    self.python_features_decoder.decode_row(A			)
      return self.recursive_tensorize(A			)



 def    a__    (self					,  A			)       ->	"jax.Array":


      """simple docstring"""



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



 def    a__    (self					,  A			)       ->	Mapping:


      """simple docstring"""



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








'''simple docstring'''

import unittest

from transformers import GPTSwaTokenizer
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


__lowerCAmelCase		     =  get_tests_dir('''fixtures/test_sentencepiece_with_bytefallback.model''')

@require_sentencepiece
@require_tokenizers
class       __magic_name__		(       _UpperCamelCase     ,		unittest.TestCase  ):
							lowerCAmelCase			:       Tuple  							=   GPTSwaTokenizer
							lowerCAmelCase			:       Tuple  							=   False
							lowerCAmelCase			:       Optional[Any]  							=   True
							lowerCAmelCase			:       Any  							=   False







							def 				__lowercase     (				self   :			Any						):
									super().setUp()

									# We have a SentencePiece fixture for testing
									_a	: int								=				GPTSwaTokenizer(_UpperCAmelCase      ,eos_token='<unk>'      ,bos_token='<unk>'      ,pad_token='<unk>'						)

									tokenizer.save_pretrained(self.tmpdirname						)







							def 				__lowercase     (				self   :			Optional[Any]      ,_UpperCAmelCase   :			int						):
									_a	: Any								=				'This is a test'
									_a	: Any								=				'This is a test'
									return input_text, output_text







							def 				__lowercase     (				self   :			str						):
									_a	: Optional[Any]								=				'<s>'
									_a	: str								=				1

									self.assertEqual(self.get_tokenizer()._convert_token_to_id(_UpperCAmelCase						)      ,_UpperCAmelCase						)
									self.assertEqual(self.get_tokenizer()._convert_id_to_token(_UpperCAmelCase						)      ,_UpperCAmelCase						)







							def 				__lowercase     (				self   :			Union[str, Any]						):
									_a	: Dict								=				list(self.get_tokenizer().get_vocab().keys()						)

									self.assertEqual(vocab_keys[0]      ,'<unk>'						)
									self.assertEqual(vocab_keys[1]      ,'<s>'						)
									self.assertEqual(vocab_keys[-1]      ,'j'						)
									self.assertEqual(len(_UpperCAmelCase						)      ,2000						)







							def 				__lowercase     (				self   :			Optional[int]						):
									self.assertEqual(self.get_tokenizer().vocab_size      ,2000						)







							def 				__lowercase     (				self   :			List[Any]						):
									_a	: Union[str, Any]								=				GPTSwaTokenizer(_UpperCAmelCase						)

									_a	: str								=				tokenizer.tokenize('This is a test'						)
									self.assertListEqual(_UpperCAmelCase      ,['▁This', '▁is', '▁a', '▁t', 'est']						)

									self.assertListEqual(tokenizer.convert_tokens_to_ids(_UpperCAmelCase						)      ,[465, 287, 265, 631, 842]						)

									_a	: Any								=				tokenizer.tokenize('I was born in 92000, and this is falsé.'						)
									# fmt: off
									self.assertListEqual(
									    _UpperCAmelCase      ,['▁I', '▁was', '▁bor', 'n', '▁in', '▁', '<0x39>', '2', '0', '0', '0', ',', '▁and', '▁this', '▁is', '▁f', 'al', 's', '<0xC3>', '<0xA9>', '.']      ,)
									# fmt: on

									_a	: Optional[Any]								=				tokenizer.convert_tokens_to_ids(_UpperCAmelCase						)
									self.assertListEqual(
									    _UpperCAmelCase      ,[262, 272, 1525, 286, 271, 268, 60, 916, 633, 633, 633, 259, 266, 301, 287, 384, 367, 263, 198, 172, 260]      ,)

									_a	: List[str]								=				tokenizer.convert_ids_to_tokens(_UpperCAmelCase						)
									# fmt: off
									self.assertListEqual(
									    _UpperCAmelCase      ,['▁I', '▁was', '▁bor', 'n', '▁in', '▁', '<0x39>', '2', '0', '0', '0', ',', '▁and', '▁this', '▁is', '▁f', 'al', 's', '<0xC3>', '<0xA9>', '.']						)
									# fmt: on







							def 				__lowercase     (				self   :			str						):
									_a	: List[Any]								=				GPTSwaTokenizer(_UpperCAmelCase						)
									_a	: Optional[Any]								=				['This is a test', 'I was born in 92000, and this is falsé.']
									_a	: Dict								=				[
									    [465, 287, 265, 631, 842],
									    [262, 272, 1525, 286, 271, 268, 60, 916, 633, 633, 633, 259, 266, 301, 287, 384, 367, 263, 198, 172, 260],
									]

									# Test that encode_fast returns the same as tokenize + convert_tokens_to_ids
									for text, expected_ids in zip(_UpperCAmelCase      ,_UpperCAmelCase						):
											self.assertListEqual(tokenizer.encode_fast(_UpperCAmelCase						)      ,_UpperCAmelCase						)

									# Test that decode_fast returns the input text
									for text, token_ids in zip(_UpperCAmelCase      ,_UpperCAmelCase						):
											self.assertEqual(tokenizer.decode_fast(_UpperCAmelCase						)      ,_UpperCAmelCase						)







							@slow
							def 				__lowercase     (				self   :			Optional[Any]						):
									_a	: Optional[int]								=				[
									    '<|python|>def fibonacci(n)\n    if n < 0:\n        print(\'Incorrect input\')',
									    'Hey there, how are you doing this fine day?',
									    'This is a text with a trailing spaces followed by a dot     .',
									    'Häj sväjs lillebrör! =)',
									    'Det är inget fel på Mr. Cool',
									]

									# fmt: off
									_a	: Tuple								=				{'input_ids': [[63423, 5, 6811, 14954, 282, 816, 3821, 63466, 63425, 63462, 18, 63978, 678, 301, 1320, 63423, 63455, 63458, 18, 63982, 4246, 3940, 1901, 47789, 5547, 18994], [19630, 1100, 63446, 1342, 633, 544, 4488, 593, 5102, 2416, 63495, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1652, 428, 268, 1936, 515, 268, 58593, 22413, 9106, 546, 268, 33213, 63979, 698, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [55130, 63450, 924, 63449, 2249, 4062, 1558, 318, 63504, 21498, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [509, 377, 2827, 2559, 332, 6575, 63443, 26801, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], 'token_type_ids': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}
									# fmt: on
									self.tokenizer_integration_test_util(
									    expected_encoding=_UpperCAmelCase      ,model_name='AI-Sweden/gpt-sw3-126m'      ,sequences=_UpperCAmelCase      ,)




