






'''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__ (      __UpperCamelCase    )-> List[str]:  # picklable for multiprocessing
      return x.sum()







def     lowercase__ (      __UpperCamelCase    )-> List[str]:  # picklable for multiprocessing
      return i + 1






@dataclass
class 				a_						:
  lowercase     		=      42
  lowercase     		=      42






class 				a_						(			__UpperCAmelCase				):

  def   A__						(      self      )  ->      Dict:



        """simple docstring"""






        UpperCamelCase				  =					{}
        UpperCamelCase				  =					[]
        UpperCamelCase				  =					1
        UpperCamelCase				  =					[1, 2]
        UpperCamelCase				  =					{"""a""": 1, """b""": 2}
        UpperCamelCase				  =					{"""a""": [1, 2], """b""": [3, 4]}
        UpperCamelCase				  =					{"""a""": {"""1""": 1}, """b""": 2}
        UpperCamelCase				  =					{"""a""": 1, """b""": 2, """c""": 3, """d""": 4}
        UpperCamelCase				  =					{}
        UpperCamelCase				  =					[]
        UpperCamelCase				  =					2
        UpperCamelCase				  =					[2, 3]
        UpperCamelCase				  =					{"""a""": 2, """b""": 3}
        UpperCamelCase				  =					{"""a""": [2, 3], """b""": [4, 5]}
        UpperCamelCase				  =					{"""a""": {"""1""": 2}, """b""": 3}
        UpperCamelCase				  =					{"""a""": 2, """b""": 3, """c""": 4, """d""": 5}
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__      )			,		lowerCamelCase__      )

        UpperCamelCase				  =					2
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__			,		num_proc=lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__			,		num_proc=lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__			,		num_proc=lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__			,		num_proc=lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__			,		num_proc=lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__			,		num_proc=lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__			,		num_proc=lowerCamelCase__      )			,		lowerCamelCase__      )
        self.assertEqual(map_nested(lowerCamelCase__			,		lowerCamelCase__			,		num_proc=lowerCamelCase__      )			,		lowerCamelCase__      )

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

  def   A__						(      self      )  ->      Union[str, Any]:



        """simple docstring"""






        UpperCamelCase				  =					{"""a""": 1, """b""": 2}
        UpperCamelCase				  =					{"""a""": 3, """b""": 4}
        UpperCamelCase				  =					{"""a""": 5, """b""": 6}
        UpperCamelCase				  =					sorted([("""a""", (1, 3, 5)), ("""b""", (2, 4, 6))]      )
        self.assertEqual(sorted(zip_dict(lowerCamelCase__			,		lowerCamelCase__			,		lowerCamelCase__      )      )			,		lowerCamelCase__      )


  def   A__						(      self      )  ->      List[Any]:



        """simple docstring"""






        class 				a_						:
          lowercase     		=      """bar"""


        UpperCamelCase				  =					Foo()
        self.assertEqual(foo.my_attr			,		"""bar"""      )
        with temporary_assignment(lowerCamelCase__			,		"""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     lowercase__ (      __UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase    )-> Union[str, Any]:
      with patch("""datasets.utils.py_utils._single_map_nested"""    ) as mock_single_map_nested, patch(
          """datasets.parallel.parallel.Pool"""    ) as mock_multiprocessing_pool:
            UpperCamelCase				  =					{F"{i}": i for i in range(__UpperCamelCase    )}
            UpperCamelCase				  =					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   A__						(      self      )  ->      Optional[int]:



        """simple docstring"""






        import tensorflow as tf
        from tensorflow.keras import layers

        UpperCamelCase				  =					layers.Dense(2      )

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

        with temp_seed(42			,		set_tensorflow=lowerCamelCase__      ):
              UpperCamelCase				  =					gen_random_output()
        with temp_seed(42			,		set_tensorflow=lowerCamelCase__      ):
              UpperCamelCase				  =					gen_random_output()
        UpperCamelCase				  =					gen_random_output()

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

  @require_torch
  def   A__						(      self      )  ->      List[Any]:



        """simple docstring"""






        import torch

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

        with temp_seed(42			,		set_pytorch=lowerCamelCase__      ):
              UpperCamelCase				  =					gen_random_output()
        with temp_seed(42			,		set_pytorch=lowerCamelCase__      ):
              UpperCamelCase				  =					gen_random_output()
        UpperCamelCase				  =					gen_random_output()

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


  def   A__						(      self      )  ->      int:



        """simple docstring"""






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

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

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







@pytest.mark.parametrize("""input_data"""		,						[{}]    )
def     lowercase__ (      __UpperCamelCase    )-> Dict:
      UpperCamelCase				  =					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     lowercase__ (      __UpperCamelCase		,						__UpperCamelCase    )-> Dict:
      UpperCamelCase				  =					NestedDataStructure(__UpperCamelCase    ).flatten()
      assert output == expected_output







def     lowercase__ (      )-> Any:
      UpperCamelCase				  =					A(x=1		,						y="""foobar"""    )
      UpperCamelCase				  =					{"""x""": 1, """y""": """foobar"""}
      assert asdict(__UpperCamelCase    ) == expected_output

      UpperCamelCase				  =					{"""a""": {"""b""": A(x=10		,						y="""foo"""    )}, """c""": [A(x=20		,						y="""bar"""    )]}
      UpperCamelCase				  =					{"""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     lowercase__ (      __UpperCamelCase    )-> List[Any]:
      return text.split()







def     lowercase__ (      __UpperCamelCase    )-> List[Any]:
      yield (time.time(), content)
      time.sleep(2    )
      yield (time.time(), content)







def     lowercase__ (      )-> List[Any]:
      with Pool(2    ) as pool:
            UpperCamelCase				  =					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:
            UpperCamelCase				  =					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:
            UpperCamelCase				  =					[]
            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

"""simple docstring"""






import os
import sys
import unittest


lowerCAmelCase__  :  Tuple						     =	os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, 'utils'))

import get_test_info  # noqa: E402
from get_test_info import (  # noqa: E402
    get_model_to_test_mapping,
    get_model_to_tester_mapping,
    get_test_to_tester_mapping,
)


lowerCAmelCase__  :  Tuple						     =	os.path.join('tests', 'models', 'bert', 'test_modeling_bert.py')
lowerCAmelCase__  :  Optional[Any]						     =	os.path.join('tests', 'models', 'blip', 'test_modeling_blip.py')
class 				snake_case       ( unittest.TestCase				):




					"""simple docstring"""

					def 		__lowerCAmelCase						(				self   :	List[str]			):
									UpperCAmelCase__             =							get_test_to_tester_mapping(lowerCamelCase__			)
									UpperCAmelCase__             =							get_test_to_tester_mapping(lowerCamelCase__			)

									UpperCAmelCase__             =							{'BertModelTest': 'BertModelTester'}

									UpperCAmelCase__             =							{
									    'BlipModelTest': 'BlipModelTester',
									    'BlipTextImageModelTest': 'BlipTextImageModelsModelTester',
									    'BlipTextModelTest': 'BlipTextModelTester',
									    'BlipTextRetrievalModelTest': 'BlipTextRetrievalModelTester',
									    'BlipVQAModelTest': 'BlipVQAModelTester',
									    'BlipVisionModelTest': 'BlipVisionModelTester',
									}

									self.assertEqual(get_test_info.to_json(lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(get_test_info.to_json(lowerCamelCase__			)    ,lowerCamelCase__			)







					def 		__lowerCAmelCase						(				self   :	Optional[Any]			):
									UpperCAmelCase__             =							get_model_to_test_mapping(lowerCamelCase__			)
									UpperCAmelCase__             =							get_model_to_test_mapping(lowerCamelCase__			)

									UpperCAmelCase__             =							{
									    'BertForMaskedLM': ['BertModelTest'],
									    'BertForMultipleChoice': ['BertModelTest'],
									    'BertForNextSentencePrediction': ['BertModelTest'],
									    'BertForPreTraining': ['BertModelTest'],
									    'BertForQuestionAnswering': ['BertModelTest'],
									    'BertForSequenceClassification': ['BertModelTest'],
									    'BertForTokenClassification': ['BertModelTest'],
									    'BertLMHeadModel': ['BertModelTest'],
									    'BertModel': ['BertModelTest'],
									}

									UpperCAmelCase__             =							{
									    'BlipForConditionalGeneration': ['BlipTextImageModelTest'],
									    'BlipForImageTextRetrieval': ['BlipTextRetrievalModelTest'],
									    'BlipForQuestionAnswering': ['BlipVQAModelTest'],
									    'BlipModel': ['BlipModelTest'],
									    'BlipTextModel': ['BlipTextModelTest'],
									    'BlipVisionModel': ['BlipVisionModelTest'],
									}

									self.assertEqual(get_test_info.to_json(lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(get_test_info.to_json(lowerCamelCase__			)    ,lowerCamelCase__			)







					def 		__lowerCAmelCase						(				self   :	int			):
									UpperCAmelCase__             =							get_model_to_tester_mapping(lowerCamelCase__			)
									UpperCAmelCase__             =							get_model_to_tester_mapping(lowerCamelCase__			)

									UpperCAmelCase__             =							{
									    'BertForMaskedLM': ['BertModelTester'],
									    'BertForMultipleChoice': ['BertModelTester'],
									    'BertForNextSentencePrediction': ['BertModelTester'],
									    'BertForPreTraining': ['BertModelTester'],
									    'BertForQuestionAnswering': ['BertModelTester'],
									    'BertForSequenceClassification': ['BertModelTester'],
									    'BertForTokenClassification': ['BertModelTester'],
									    'BertLMHeadModel': ['BertModelTester'],
									    'BertModel': ['BertModelTester'],
									}

									UpperCAmelCase__             =							{
									    'BlipForConditionalGeneration': ['BlipTextImageModelsModelTester'],
									    'BlipForImageTextRetrieval': ['BlipTextRetrievalModelTester'],
									    'BlipForQuestionAnswering': ['BlipVQAModelTester'],
									    'BlipModel': ['BlipModelTester'],
									    'BlipTextModel': ['BlipTextModelTester'],
									    'BlipVisionModel': ['BlipVisionModelTester'],
									}

									self.assertEqual(get_test_info.to_json(lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(get_test_info.to_json(lowerCamelCase__			)    ,lowerCamelCase__			)



