





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      A__       (    SCREAMING_SNAKE_CASE_   :      Tuple						)  ->  Any:  # picklable for multiprocessing






   """simple docstring"""







   return x.sum()
def      A__       (    SCREAMING_SNAKE_CASE_   :      Tuple						)  ->  Union[str, Any]:  # picklable for multiprocessing






   """simple docstring"""







   return i + 1




@dataclass
class 			__UpperCamelCase    :
     __A	:  int
     __A	:  str




class 			__UpperCamelCase    (							A__    ):





     def 		UpperCamelCase(     self					):
        _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(_UpperCamelCase				,							_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase					)				,							_UpperCamelCase					)

        _UpperCAmelCase           =   2
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase				,							num_proc=_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase				,							num_proc=_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase				,							num_proc=_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase				,							num_proc=_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase				,							num_proc=_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase				,							num_proc=_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase				,							num_proc=_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase				,							num_proc=_UpperCamelCase					)				,							_UpperCamelCase					)

        _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(_UpperCamelCase					),
            '''b''': np.zeros(3					).astype(_UpperCamelCase					),
            '''c''': np.ones(2					).astype(_UpperCamelCase					),
        }
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase				,							map_numpy=_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(
            {k: v.tolist() for k, v in map_nested(_UpperCamelCase				,							_UpperCamelCase				,							map_numpy=_UpperCamelCase					).items()}				,							{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}				,							)
        self.assertEqual(map_nested(_UpperCamelCase				,							_UpperCamelCase				,							map_numpy=_UpperCamelCase				,							num_proc=_UpperCamelCase					)				,							_UpperCamelCase					)
        self.assertEqual(
            {k: v.tolist() for k, v in map_nested(_UpperCamelCase				,							_UpperCamelCase				,							map_numpy=_UpperCamelCase				,							num_proc=_UpperCamelCase					).items()}				,							{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}				,							)
        with self.assertRaises(_UpperCamelCase					):  # can't pickle a local lambda
           map_nested(lambda _UpperCamelCase					: x + 1				,							_UpperCamelCase				,							num_proc=_UpperCamelCase					)





     def 		UpperCamelCase(     self					):
        _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(_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase					)					)				,							_UpperCamelCase					)





     def 		UpperCamelCase(     self					):




        class 			__UpperCamelCase    :
             __A	:  int							  =   """bar"""




        _UpperCAmelCase           =   Foo()
        self.assertEqual(foo.my_attr				,							'''bar'''					)
        with temporary_assignment(_UpperCamelCase				,							'''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      A__       (    SCREAMING_SNAKE_CASE_   :      int				,	SCREAMING_SNAKE_CASE_   :      List[Any]				,	SCREAMING_SNAKE_CASE_   :      Tuple						)  ->  Tuple:






   """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:
      _UpperCAmelCase           =   {F'''{i}''': i for i in range(SCREAMING_SNAKE_CASE_						)}
      _UpperCAmelCase           =   map_nested(lambda SCREAMING_SNAKE_CASE_						: x + 10				,	SCREAMING_SNAKE_CASE_				,	num_proc=SCREAMING_SNAKE_CASE_				,	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 			__UpperCamelCase    (							A__    ):





     @require_tf
     def 		UpperCamelCase(     self					):
        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(_UpperCamelCase					).numpy()

        with temp_seed(42				,							set_tensorflow=_UpperCamelCase					):
           _UpperCAmelCase           =   gen_random_output()
        with temp_seed(42				,							set_tensorflow=_UpperCamelCase					):
           _UpperCAmelCase           =   gen_random_output()
        _UpperCAmelCase           =   gen_random_output()

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





     @require_torch
     def 		UpperCamelCase(     self					):
        import torch

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

        with temp_seed(42				,							set_pytorch=_UpperCamelCase					):
           _UpperCAmelCase           =   gen_random_output()
        with temp_seed(42				,							set_pytorch=_UpperCamelCase					):
           _UpperCAmelCase           =   gen_random_output()
        _UpperCAmelCase           =   gen_random_output()

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





     def 		UpperCamelCase(     self					):
        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(_UpperCamelCase				,							_UpperCamelCase					)
        self.assertGreater(np.abs(outa - outa					).sum()				,							0					)




@pytest.mark.parametrize('''input_data'''				,	[{}]						)
def      A__       (    SCREAMING_SNAKE_CASE_   :      int						)  ->  Dict:






   """simple docstring"""







   _UpperCAmelCase           =   NestedDataStructure(SCREAMING_SNAKE_CASE_						).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      A__       (    SCREAMING_SNAKE_CASE_   :      int				,	SCREAMING_SNAKE_CASE_   :      Dict						)  ->  Optional[Any]:






   """simple docstring"""







   _UpperCAmelCase           =   NestedDataStructure(SCREAMING_SNAKE_CASE_						).flatten()
   assert output == expected_output
def      A__       (    )  ->  int:






   """simple docstring"""







   _UpperCAmelCase           =   A(x=1				,	y='''foobar'''						)
   _UpperCAmelCase           =   {'''x''': 1, '''y''': '''foobar'''}
   assert asdict(SCREAMING_SNAKE_CASE_						) == 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(SCREAMING_SNAKE_CASE_						) == expected_output

   with pytest.raises(SCREAMING_SNAKE_CASE_						):
      asdict([1, A(x=10				,	y='''foo'''						)]						)
def      A__       (    SCREAMING_SNAKE_CASE_   :      str						)  ->  int:






   """simple docstring"""







   return text.split()
def      A__       (    SCREAMING_SNAKE_CASE_   :      List[Any]						)  ->  Any:






   """simple docstring"""







   yield (time.time(), content)
   time.sleep(2						)
   yield (time.time(), content)
def      A__       (    )  ->  str:






   """simple docstring"""







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

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

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






import gc
import unittest

import numpy as np
import torch

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS, UNCONDITIONAL_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class 			__UpperCamelCase    (							A__	, unittest.TestCase    ):
     __A	:  Any							  =   DanceDiffusionPipeline
     __A	:  Any							  =   UNCONDITIONAL_AUDIO_GENERATION_PARAMS
     __A	:  Tuple							  =   PipelineTesterMixin.required_optional_params - {
         """callback""",
         """latents""",
         """callback_steps""",
         """output_type""",
         """num_images_per_prompt""",
     }
     __A	:  Tuple							  =   UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
     __A	:  List[str]							  =   False
     __A	:  str							  =   False





     def 		UpperCamelCase(     self					):
        torch.manual_seed(0					)
        _UpperCAmelCase           =   UNetaDModel(
            block_out_channels=(32, 32, 64)				,							extra_in_channels=16				,							sample_size=512				,							sample_rate=16000				,							in_channels=2				,							out_channels=2				,							flip_sin_to_cos=_UpperCamelCase				,							use_timestep_embedding=_UpperCamelCase				,							time_embedding_type='''fourier'''				,							mid_block_type='''UNetMidBlock1D'''				,							down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''')				,							up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''')				,							)
        _UpperCAmelCase           =   IPNDMScheduler()

        _UpperCAmelCase           =   {
            '''unet''': unet,
            '''scheduler''': scheduler,
        }
        return components





     def 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase=0					):
        if str(_UpperCamelCase					).startswith('''mps'''					):
           _UpperCAmelCase           =   torch.manual_seed(_UpperCamelCase					)
        else:
           _UpperCAmelCase           =   torch.Generator(device=_UpperCamelCase					).manual_seed(_UpperCamelCase					)
        _UpperCAmelCase           =   {
            '''batch_size''': 1,
            '''generator''': generator,
            '''num_inference_steps''': 4,
        }
        return inputs





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _UpperCAmelCase           =   self.get_dummy_components()
        _UpperCAmelCase           =   DanceDiffusionPipeline(**_UpperCamelCase					)
        _UpperCAmelCase           =   pipe.to(_UpperCamelCase					)
        pipe.set_progress_bar_config(disable=_UpperCamelCase					)

        _UpperCAmelCase           =   self.get_dummy_inputs(_UpperCamelCase					)
        _UpperCAmelCase           =   pipe(**_UpperCamelCase					)
        _UpperCAmelCase           =   output.audios

        _UpperCAmelCase           =   audio[0, -3:, -3:]

        assert audio.shape == (1, 2, components["unet"].sample_size)
        _UpperCAmelCase           =   np.array([-0.7265, 1.0000, -0.8388, 0.1175, 0.9498, -1.0000]					)
        assert np.abs(audio_slice.flatten() - expected_slice					).max() < 1e-2





     @skip_mps
     def 		UpperCamelCase(     self					):
        return super().test_save_load_local()





     @skip_mps
     def 		UpperCamelCase(     self					):
        return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3					)





     @skip_mps
     def 		UpperCamelCase(     self					):
        return super().test_save_load_optional_components()





     @skip_mps
     def 		UpperCamelCase(     self					):
        return super().test_attention_slicing_forward_pass()





     def 		UpperCamelCase(     self					):
        super().test_inference_batch_single_identical(expected_max_diff=3e-3					)




@slow
@require_torch_gpu
class 			__UpperCamelCase    (							unittest.TestCase    ):





     def 		UpperCamelCase(     self					):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   torch_device

        _UpperCAmelCase           =   DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''					)
        _UpperCAmelCase           =   pipe.to(_UpperCamelCase					)
        pipe.set_progress_bar_config(disable=_UpperCamelCase					)

        _UpperCAmelCase           =   torch.manual_seed(0					)
        _UpperCAmelCase           =   pipe(generator=_UpperCamelCase				,							num_inference_steps=100				,							audio_length_in_s=4.096					)
        _UpperCAmelCase           =   output.audios

        _UpperCAmelCase           =   audio[0, -3:, -3:]

        assert audio.shape == (1, 2, pipe.unet.sample_size)
        _UpperCAmelCase           =   np.array([-0.0192, -0.0231, -0.0318, -0.0059, 0.0002, -0.0020]					)

        assert np.abs(audio_slice.flatten() - expected_slice					).max() < 1e-2





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   torch_device

        _UpperCAmelCase           =   DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''				,							torch_dtype=torch.floataa					)
        _UpperCAmelCase           =   pipe.to(_UpperCamelCase					)
        pipe.set_progress_bar_config(disable=_UpperCamelCase					)

        _UpperCAmelCase           =   torch.manual_seed(0					)
        _UpperCAmelCase           =   pipe(generator=_UpperCamelCase				,							num_inference_steps=100				,							audio_length_in_s=4.096					)
        _UpperCAmelCase           =   output.audios

        _UpperCAmelCase           =   audio[0, -3:, -3:]

        assert audio.shape == (1, 2, pipe.unet.sample_size)
        _UpperCAmelCase           =   np.array([-0.0367, -0.0488, -0.0771, -0.0525, -0.0444, -0.0341]					)

        assert np.abs(audio_slice.flatten() - expected_slice					).max() < 1e-2