



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 					__snake_case     (      _UpperCAmelCase    ):  # picklable for multiprocessing
    return x.sum()





def 					__snake_case     (      _UpperCAmelCase    ):  # picklable for multiprocessing
    return i + 1

@dataclass
class    _A	:
       UpperCamelCase__  :   int
       UpperCamelCase__  :   str


class    _A	(     __UpperCAmelCase   ):



       def 		_lowerCamelCase      (		self							:					Union[str, Any]):


           '''simple docstring'''

           __a							=    {}
           __a							=    []
           __a							=    1
           __a							=    [1, 2]
           __a							=    {'''a''': 1, '''b''': 2}
           __a							=    {'''a''': [1, 2], '''b''': [3, 4]}
           __a							=    {'''a''': {'''1''': 1}, '''b''': 2}
           __a							=    {'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
           __a							=    {}
           __a							=    []
           __a							=    2
           __a							=    [2, 3]
           __a							=    {'''a''': 2, '''b''': 3}
           __a							=    {'''a''': [2, 3], '''b''': [4, 5]}
           __a							=    {'''a''': {'''1''': 2}, '''b''': 3}
           __a							=    {'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)

           __a							=    2
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE					,		num_proc=__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE					,		num_proc=__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE					,		num_proc=__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE					,		num_proc=__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE					,		num_proc=__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE					,		num_proc=__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE					,		num_proc=__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)
           self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE					,		num_proc=__SCREAMING_SNAKE_CASE)					,		__SCREAMING_SNAKE_CASE)

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



       def 		_lowerCamelCase      (		self							:					Dict):


           '''simple docstring'''

           __a							=    {'''a''': 1, '''b''': 2}
           __a							=    {'''a''': 3, '''b''': 4}
           __a							=    {'''a''': 5, '''b''': 6}
           __a							=    sorted([('''a''', (1, 3, 5)), ('''b''', (2, 4, 6))])
           self.assertEqual(sorted(zip_dict(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE))					,		__SCREAMING_SNAKE_CASE)





       def 		_lowerCamelCase      (		self							:					int):


           '''simple docstring'''

           class    _A	:
                  UpperCamelCase__  :   Dict						 =						'''bar'''


           __a							=    Foo()
           self.assertEqual(foo.my_attr					,		'''bar''')
           with temporary_assignment(__SCREAMING_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),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]       ,			)
def 					__snake_case     (      _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase    ):
    with patch('''datasets.utils.py_utils._single_map_nested'''    ) as mock_single_map_nested, patch(
        '''datasets.parallel.parallel.Pool'''    ) as mock_multiprocessing_pool:
        __a							=    {f'{i}': i for i in range(_UpperCAmelCase    )}
        __a							=    map_nested(lambda _UpperCAmelCase    : x + 10       ,			_UpperCAmelCase       ,			num_proc=_UpperCAmelCase       ,			parallel_min_length=16    )
        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    _A	(     __UpperCAmelCase   ):



       @require_tf
       def 		_lowerCamelCase      (		self							:					Union[str, Any]):


           '''simple docstring'''

           import tensorflow as tf
           from tensorflow.keras import layers

           __a							=    layers.Dense(2)

           def gen_random_output():
               __a							=    tf.random.uniform((1, 3))
               return model(__SCREAMING_SNAKE_CASE).numpy()

           with temp_seed(42					,		set_tensorflow=__SCREAMING_SNAKE_CASE):
               __a							=    gen_random_output()
           with temp_seed(42					,		set_tensorflow=__SCREAMING_SNAKE_CASE):
               __a							=    gen_random_output()
           __a							=    gen_random_output()

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



       @require_torch
       def 		_lowerCamelCase      (		self							:					Optional[Any]):


           '''simple docstring'''

           import torch

           def gen_random_output():
               __a							=    torch.nn.Linear(3					,		2)
               __a							=    torch.rand(1					,		3)
               return model(__SCREAMING_SNAKE_CASE).detach().numpy()

           with temp_seed(42					,		set_pytorch=__SCREAMING_SNAKE_CASE):
               __a							=    gen_random_output()
           with temp_seed(42					,		set_pytorch=__SCREAMING_SNAKE_CASE):
               __a							=    gen_random_output()
           __a							=    gen_random_output()

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





       def 		_lowerCamelCase      (		self							:					List[Any]):


           '''simple docstring'''

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

           with temp_seed(42):
               __a							=    gen_random_output()
           with temp_seed(42):
               __a							=    gen_random_output()
           __a							=    gen_random_output()

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





@pytest.mark.parametrize('''input_data'''       ,			[{}]    )
def 					__snake_case     (      _UpperCAmelCase    ):
    __a							=    NestedDataStructure(_UpperCAmelCase    ).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 					__snake_case     (      _UpperCAmelCase       ,			_UpperCAmelCase    ):
    __a							=    NestedDataStructure(_UpperCAmelCase    ).flatten()
    assert output == expected_output





def 					__snake_case     (      ):
    __a							=    A(x=1       ,			y='''foobar'''    )
    __a							=    {'''x''': 1, '''y''': '''foobar'''}
    assert asdict(_UpperCAmelCase    ) == expected_output

    __a							=    {'''a''': {'''b''': A(x=10       ,			y='''foo'''    )}, '''c''': [A(x=20       ,			y='''bar'''    )]}
    __a							=    {'''a''': {'''b''': {'''x''': 10, '''y''': '''foo'''}}, '''c''': [{'''x''': 20, '''y''': '''bar'''}]}
    assert asdict(_UpperCAmelCase    ) == expected_output

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





def 					__snake_case     (      _UpperCAmelCase    ):
    return text.split()





def 					__snake_case     (      _UpperCAmelCase    ):
    yield (time.time(), content)
    time.sleep(2    )
    yield (time.time(), content)





def 					__snake_case     (      ):
    with Pool(2    ) as pool:
        __a							=    list(iflatmap_unordered(_UpperCAmelCase       ,			_split_text       ,			kwargs_iterable=[{'''text''': '''hello there'''}] * 10    )    )
        assert out.count('''hello'''    ) == 10
        assert out.count('''there'''    ) == 10
        assert len(_UpperCAmelCase    ) == 20

    # check multiprocess from pathos (uses dill for pickling)
    with multiprocess.Pool(2    ) as pool:
        __a							=    list(iflatmap_unordered(_UpperCAmelCase       ,			_split_text       ,			kwargs_iterable=[{'''text''': '''hello there'''}] * 10    )    )
        assert out.count('''hello'''    ) == 10
        assert out.count('''there'''    ) == 10
        assert len(_UpperCAmelCase    ) == 20

    # check that we get items as fast as possible
    with Pool(2    ) as pool:
        __a							=    []
        for yield_time, content in iflatmap_unordered(
            _UpperCAmelCase       ,			_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(_UpperCAmelCase    )
        assert out.count('''a'''    ) == 2
        assert out.count('''b'''    ) == 2
        assert len(_UpperCAmelCase    ) == 4










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


__snake_case						:Dict						=					logging.get_logger(__name__)

__snake_case						:List[Any]						=					{
    '''tanreinama/GPTSAN-2.8B-spout_is_uniform''': (
        '''https://huggingface.co/tanreinama/GPTSAN-2.8B-spout_is_uniform/resolve/main/config.json'''
    ),
}

class    _A	(     __UpperCAmelCase   ):
       UpperCamelCase__  :   List[str]						 =						'''gptsan-japanese'''
       UpperCamelCase__  :   Dict						 =						[
           '''past_key_values''',
       ]
       UpperCamelCase__  :   Dict						 =						{
           '''hidden_size''': '''d_model''',
           '''num_attention_heads''': '''num_heads''',
           '''num_hidden_layers''': '''num_layers''',
       }



       def __init__(		self							:					List[str]					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]=36_000					,		__SCREAMING_SNAKE_CASE							:					Tuple=1_280					,		__SCREAMING_SNAKE_CASE							:					List[Any]=1_024					,		__SCREAMING_SNAKE_CASE							:					List[Any]=8_192					,		__SCREAMING_SNAKE_CASE							:					str=4_096					,		__SCREAMING_SNAKE_CASE							:					Any=128					,		__SCREAMING_SNAKE_CASE							:					int=10					,		__SCREAMING_SNAKE_CASE							:					Optional[int]=0					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]=16					,		__SCREAMING_SNAKE_CASE							:					List[Any]=16					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]=128					,		__SCREAMING_SNAKE_CASE							:					Tuple=0.0					,		__SCREAMING_SNAKE_CASE							:					List[Any]=1E-5					,		__SCREAMING_SNAKE_CASE							:					str=False					,		__SCREAMING_SNAKE_CASE							:					Optional[int]=0.0					,		__SCREAMING_SNAKE_CASE							:					List[str]="float32"					,		__SCREAMING_SNAKE_CASE							:					List[str]=False					,		__SCREAMING_SNAKE_CASE							:					Any=False					,		__SCREAMING_SNAKE_CASE							:					Any=False					,		__SCREAMING_SNAKE_CASE							:					int=0.0_02					,		__SCREAMING_SNAKE_CASE							:					Any=False					,		__SCREAMING_SNAKE_CASE							:					Tuple=True					,		__SCREAMING_SNAKE_CASE							:					int=35_998					,		__SCREAMING_SNAKE_CASE							:					Optional[int]=35_995					,		__SCREAMING_SNAKE_CASE							:					List[str]=35_999					,		**__SCREAMING_SNAKE_CASE							:					List[str]					,		):


           '''simple docstring'''

           __a							=    vocab_size
           __a							=    max_position_embeddings
           __a							=    d_model
           __a							=    d_ff
           __a							=    d_ext
           __a							=    d_spout
           __a							=    num_switch_layers
           __a							=    num_ext_layers
           __a							=    num_switch_layers + num_ext_layers
           __a							=    num_heads
           __a							=    num_experts
           __a							=    expert_capacity
           __a							=    dropout_rate
           __a							=    layer_norm_epsilon
           __a							=    router_bias
           __a							=    router_jitter_noise
           __a							=    router_dtype
           __a							=    router_ignore_padding_tokens
           __a							=    output_hidden_states
           __a							=    output_attentions
           __a							=    initializer_factor
           __a							=    output_router_logits
           __a							=    use_cache

           super().__init__(
               separator_token_id=__SCREAMING_SNAKE_CASE					,		pad_token_id=__SCREAMING_SNAKE_CASE					,		eos_token_id=__SCREAMING_SNAKE_CASE					,		**__SCREAMING_SNAKE_CASE					,		)





