


'''simple docstring'''







import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch






def 			lowercase				(_A					):  # picklable for multiprocessing







       """simple docstring"""







       return x.sum()






def 			lowercase				(_A					):  # picklable for multiprocessing







       """simple docstring"""







       return i + 1
@dataclass
class 	UpperCamelCase__					:



     """simple docstring"""






     __magic_name__    			=    42
     __magic_name__    			=    42
class 	UpperCamelCase__					( SCREAMING_SNAKE_CASE_	):



     """simple docstring"""






     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[int]   			=       {}
            _lowerCAmelCase    :							Union[str, Any]   			=       []
            _lowerCAmelCase    :							str   			=       1
            _lowerCAmelCase    :							List[Any]   			=       [1, 2]
            _lowerCAmelCase    :							List[Any]   			=       {'a': 1, 'b': 2}
            _lowerCAmelCase    :							int   			=       {'a': [1, 2], 'b': [3, 4]}
            _lowerCAmelCase    :							int   			=       {'a': {'1': 1}, 'b': 2}
            _lowerCAmelCase    :							Optional[Any]   			=       {'a': 1, 'b': 2, 'c': 3, 'd': 4}
            _lowerCAmelCase    :							Tuple   			=       {}
            _lowerCAmelCase    :							int   			=       []
            _lowerCAmelCase    :							Union[str, Any]   			=       2
            _lowerCAmelCase    :							Optional[int]   			=       [2, 3]
            _lowerCAmelCase    :							Optional[int]   			=       {'a': 2, 'b': 3}
            _lowerCAmelCase    :							str   			=       {'a': [2, 3], 'b': [4, 5]}
            _lowerCAmelCase    :							Union[str, Any]   			=       {'a': {'1': 2}, 'b': 3}
            _lowerCAmelCase    :							Optional[int]   			=       {'a': 2, 'b': 3, 'c': 4, 'd': 5}
            self.assertEqual(map_nested(snake_case__      ,	snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__   )      ,	snake_case__   )

            _lowerCAmelCase    :							Union[str, Any]   			=       2
            self.assertEqual(map_nested(snake_case__      ,	snake_case__      ,	num_proc=snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__      ,	num_proc=snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__      ,	num_proc=snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__      ,	num_proc=snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__      ,	num_proc=snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__      ,	num_proc=snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__      ,	num_proc=snake_case__   )      ,	snake_case__   )
            self.assertEqual(map_nested(snake_case__      ,	snake_case__      ,	num_proc=snake_case__   )      ,	snake_case__   )

            _lowerCAmelCase    :							Optional[Any]   			=       {'a': np.eye(2   ), 'b': np.zeros(3   ), 'c': np.ones(2   )}
            _lowerCAmelCase    :							List[str]   			=       {'a': 2, 'b': 0, 'c': 2}
            _lowerCAmelCase    :							List[Any]   			=       {
                'a': np.eye(2   ).astype(snake_case__   ),
                'b': np.zeros(3   ).astype(snake_case__   ),
                'c': np.ones(2   ).astype(snake_case__   ),
            }
            self.assertEqual(map_nested(snake_case__      ,	snake_case__      ,	map_numpy=snake_case__   )      ,	snake_case__   )
            self.assertEqual(
                {k: v.tolist() for k, v in map_nested(snake_case__      ,	snake_case__      ,	map_numpy=snake_case__   ).items()}      ,	{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}      ,	)
            self.assertEqual(map_nested(snake_case__      ,	snake_case__      ,	map_numpy=snake_case__      ,	num_proc=snake_case__   )      ,	snake_case__   )
            self.assertEqual(
                {k: v.tolist() for k, v in map_nested(snake_case__      ,	snake_case__      ,	map_numpy=snake_case__      ,	num_proc=snake_case__   ).items()}      ,	{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}      ,	)
            with self.assertRaises(snake_case__   ):  # can't pickle a local lambda
                   map_nested(lambda snake_case__   : x + 1      ,	snake_case__      ,	num_proc=snake_case__   )







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[int]   			=       {'a': 1, 'b': 2}
            _lowerCAmelCase    :							Optional[Any]   			=       {'a': 3, 'b': 4}
            _lowerCAmelCase    :							Dict   			=       {'a': 5, 'b': 6}
            _lowerCAmelCase    :							Union[str, Any]   			=       sorted([('a', (1, 3, 5)), ('b', (2, 4, 6))]   )
            self.assertEqual(sorted(zip_dict(snake_case__      ,	snake_case__      ,	snake_case__   )   )      ,	snake_case__   )







     def a	( self   ):

            '''simple docstring'''

            class 	UpperCamelCase__					:



                 """simple docstring"""






                 __magic_name__    			=    "bar"
            _lowerCAmelCase    :							Any   			=       Foo()
            self.assertEqual(foo.my_attr      ,	'bar'   )
            with temporary_assignment(snake_case__      ,	'my_attr'      ,	'BAR'   ):
                   self.assertEqual(foo.my_attr      ,	'BAR'   )
            self.assertEqual(foo.my_attr      ,	'bar'   )






@pytest.mark.parametrize(
    'iterable_length, num_proc, expected_num_proc'   ,						[
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (1_6, 1_6, 1_6),
        (1_6, 1_7, 1_6),
        (1_7, 1_6, 1_6),
    ]   ,						)
def 			lowercase				(_A   ,						_A   ,						_A					):







       """simple docstring"""







       with patch('datasets.utils.py_utils._single_map_nested'					) as mock_single_map_nested, patch(
           'datasets.parallel.parallel.Pool'					) as mock_multiprocessing_pool:
              _lowerCAmelCase    :							Any   			=       {f'{i}': i for i in range(_A					)}
              _lowerCAmelCase    :							Union[str, Any]   			=       map_nested(lambda _A					: x + 1_0   ,						_A   ,						num_proc=_A   ,						parallel_min_length=1_6					)
              if expected_num_proc == 1:
                     assert mock_single_map_nested.called
                     assert not mock_multiprocessing_pool.called
              else:
                     assert not mock_single_map_nested.called
                     assert mock_multiprocessing_pool.called
                     assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc
class 	UpperCamelCase__					( SCREAMING_SNAKE_CASE_	):



     """simple docstring"""






     @require_tf
     def a	( self   ):

            '''simple docstring'''

            import tensorflow as tf
            from tensorflow.keras import layers

            _lowerCAmelCase    :							Any   			=       layers.Dense(2   )

            def gen_random_output():
                   _lowerCAmelCase    :							List[str]   			=       tf.random.uniform((1, 3)   )
                   return model(snake_case__   ).numpy()

            with temp_seed(42      ,	set_tensorflow=snake_case__   ):
                   _lowerCAmelCase    :							Any   			=       gen_random_output()
            with temp_seed(42      ,	set_tensorflow=snake_case__   ):
                   _lowerCAmelCase    :							List[str]   			=       gen_random_output()
            _lowerCAmelCase    :							Dict   			=       gen_random_output()

            np.testing.assert_equal(snake_case__      ,	snake_case__   )
            self.assertGreater(np.abs(outa - outa   ).sum()      ,	0   )







     @require_torch
     def a	( self   ):

            '''simple docstring'''

            import torch

            def gen_random_output():
                   _lowerCAmelCase    :							List[Any]   			=       torch.nn.Linear(3      ,	2   )
                   _lowerCAmelCase    :							int   			=       torch.rand(1      ,	3   )
                   return model(snake_case__   ).detach().numpy()

            with temp_seed(42      ,	set_pytorch=snake_case__   ):
                   _lowerCAmelCase    :							List[str]   			=       gen_random_output()
            with temp_seed(42      ,	set_pytorch=snake_case__   ):
                   _lowerCAmelCase    :							Optional[int]   			=       gen_random_output()
            _lowerCAmelCase    :							Optional[int]   			=       gen_random_output()

            np.testing.assert_equal(snake_case__      ,	snake_case__   )
            self.assertGreater(np.abs(outa - outa   ).sum()      ,	0   )







     def a	( self   ):

            '''simple docstring'''

            def gen_random_output():
                   return np.random.rand(1      ,	3   )

            with temp_seed(42   ):
                   _lowerCAmelCase    :							Dict   			=       gen_random_output()
            with temp_seed(42   ):
                   _lowerCAmelCase    :							int   			=       gen_random_output()
            _lowerCAmelCase    :							str   			=       gen_random_output()

            np.testing.assert_equal(snake_case__      ,	snake_case__   )
            self.assertGreater(np.abs(outa - outa   ).sum()      ,	0   )






@pytest.mark.parametrize('input_data'   ,						[{}]					)
def 			lowercase				(_A					):







       """simple docstring"""







       _lowerCAmelCase    :							List[Any]   			=       NestedDataStructure(_A					).data
       assert output_data == input_data






@pytest.mark.parametrize(
    'data, expected_output'   ,						[
        ({}, []),
        ([], []),
        ('foo', ['foo']),
        (['foo', 'bar'], ['foo', 'bar']),
        ([['foo', 'bar']], ['foo', 'bar']),
        ([[['foo'], ['bar']]], ['foo', 'bar']),
        ([[['foo'], 'bar']], ['foo', 'bar']),
        ({'a': 1, 'b': 2}, [1, 2]),
        ({'a': [1, 2], 'b': [3, 4]}, [1, 2, 3, 4]),
        ({'a': [[1, 2]], 'b': [[3, 4]]}, [1, 2, 3, 4]),
        ({'a': [[1, 2]], 'b': [3, 4]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [[[3], [4]]]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [[3, 4]]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [3, 4]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [3, [4]]}, [1, 2, 3, 4]),
        ({'a': {'1': 1}, 'b': 2}, [1, 2]),
        ({'a': {'1': [1]}, 'b': 2}, [1, 2]),
        ({'a': {'1': [1]}, 'b': [2]}, [1, 2]),
    ]   ,						)
def 			lowercase				(_A   ,						_A					):







       """simple docstring"""







       _lowerCAmelCase    :							List[str]   			=       NestedDataStructure(_A					).flatten()
       assert output == expected_output






def 			lowercase				():







       """simple docstring"""







       _lowerCAmelCase    :							Any   			=       A(x=1   ,						y='foobar'					)
       _lowerCAmelCase    :							Dict   			=       {'x': 1, 'y': 'foobar'}
       assert asdict(_A					) == expected_output

       _lowerCAmelCase    :							Optional[int]   			=       {'a': {'b': A(x=1_0   ,						y='foo'					)}, 'c': [A(x=2_0   ,						y='bar'					)]}
       _lowerCAmelCase    :							Dict   			=       {'a': {'b': {'x': 1_0, 'y': 'foo'}}, 'c': [{'x': 2_0, 'y': 'bar'}]}
       assert asdict(_A					) == expected_output

       with pytest.raises(_A					):
              asdict([1, A(x=1_0   ,						y='foo'					)]					)






def 			lowercase				(_A					):







       """simple docstring"""







       return text.split()






def 			lowercase				(_A					):







       """simple docstring"""







       yield (time.time(), content)
       time.sleep(2					)
       yield (time.time(), content)






def 			lowercase				():







       """simple docstring"""







       with Pool(2					) as pool:
              _lowerCAmelCase    :							List[str]   			=       list(iflatmap_unordered(_A   ,						_split_text   ,						kwargs_iterable=[{'text': 'hello there'}] * 1_0					)					)
              assert out.count('hello'					) == 1_0
              assert out.count('there'					) == 1_0
              assert len(_A					) == 2_0

       # check multiprocess from pathos (uses dill for pickling)
       with multiprocess.Pool(2					) as pool:
              _lowerCAmelCase    :							List[str]   			=       list(iflatmap_unordered(_A   ,						_split_text   ,						kwargs_iterable=[{'text': 'hello there'}] * 1_0					)					)
              assert out.count('hello'					) == 1_0
              assert out.count('there'					) == 1_0
              assert len(_A					) == 2_0

       # check that we get items as fast as possible
       with Pool(2					) as pool:
              _lowerCAmelCase    :							List[Any]   			=       []
              for yield_time, content in iflatmap_unordered(
                  _A   ,						_aseconds_generator_of_aitems_with_timing   ,						kwargs_iterable=[{'content': 'a'}, {'content': 'b'}]					):
                     assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
                     out.append(_A					)
              assert out.count('a'					) == 2
              assert out.count('b'					) == 2
              assert len(_A					) == 4





'''simple docstring'''







import unittest

from transformers import AutoTokenizer, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow


if is_flax_available():
       import jax.numpy as jnp

       from transformers import FlaxXLMRobertaModel
@require_sentencepiece
@require_tokenizers
@require_flax
class 	UpperCamelCase__					( unittest.TestCase	):



     """simple docstring"""






     @slow
     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							List[str]   			=       FlaxXLMRobertaModel.from_pretrained('xlm-roberta-base'   )
            _lowerCAmelCase    :							Optional[Any]   			=       AutoTokenizer.from_pretrained('xlm-roberta-base'   )
            _lowerCAmelCase    :							Dict   			=       'The dog is cute and lives in the garden house'
            _lowerCAmelCase    :							List[str]   			=       jnp.array([tokenizer.encode(snake_case__   )]   )

            _lowerCAmelCase    :							Optional[int]   			=       (1, 12, 768)  # batch_size, sequence_length, embedding_vector_dim
            _lowerCAmelCase    :							Tuple   			=       jnp.array(
                [[-0.0101, 0.1218, -0.0803, 0.0801, 0.1327, 0.0776, -0.1215, 0.2383, 0.3338, 0.3106, 0.0300, 0.0252]]   )

            _lowerCAmelCase    :							Union[str, Any]   			=       model(snake_case__   )['last_hidden_state']
            self.assertEqual(output.shape      ,	snake_case__   )
            # compare the actual values for a slice of last dim
            self.assertTrue(jnp.allclose(output[:, :, -1]      ,	snake_case__      ,	atol=1E-3   )   )

