





import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

from ..utils.generic import ModelOutput




class     _lowerCAmelCase(	UpperCamelCase__ ):





       """simple docstring"""







       def __init__( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase							):
         UpperCamelCase_: List[Any]   				=    dataset
         UpperCamelCase_: Tuple   				=    process
         UpperCamelCase_: int   				=    params





       def __len__( self							):
         return len(self.dataset							)





       def __getitem__( self  ,      _lowerCamelCase							):
         UpperCamelCase_: Dict   				=    self.dataset[i]
         UpperCamelCase_: Dict   				=    self.process(__a  ,      **self.params							)
         return processed







class     _lowerCAmelCase(	UpperCamelCase__ ):





       """simple docstring"""







       def __init__( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase=None							):
         UpperCamelCase_: List[Any]   				=    loader
         UpperCamelCase_: Optional[Any]   				=    infer
         UpperCamelCase_: Union[str, Any]   				=    params
         if loader_batch_size == 1:
           # Let's spare some time by deactivating altogether
           UpperCamelCase_: Any   				=    None
         UpperCamelCase_: List[str]   				=    loader_batch_size

         # Internal bookkeeping
         UpperCamelCase_: List[str]   				=    None
         UpperCamelCase_: Dict   				=    None





       def __len__( self							):
         return len(self.loader							)





       def __iter__( self							):
         UpperCamelCase_: Optional[int]   				=    iter(self.loader							)
         return self





       def       _a   ( self							):
         if isinstance(self._loader_batch_data  ,      torch.Tensor							):
           # Batch data is simple tensor, just fetch the slice
           UpperCamelCase_: List[Any]   				=    self._loader_batch_data[self._loader_batch_index]
         else:
           # Batch data is assumed to be BaseModelOutput (or dict)
           UpperCamelCase_: Union[str, Any]   				=    {}
           for k, element in self._loader_batch_data.items():
             if isinstance(__a  ,      __a							):
               # Convert ModelOutput to tuple first
               UpperCamelCase_: List[Any]   				=    element.to_tuple()
               if isinstance(element[0]  ,      torch.Tensor							):
                 UpperCamelCase_: str   				=    tuple(el[self._loader_batch_index].unsqueeze(0							) for el in element							)
               elif isinstance(element[0]  ,      np.ndarray							):
                 UpperCamelCase_: Union[str, Any]   				=    tuple(np.expand_dims(el[self._loader_batch_index]  ,      0							) for el in element							)
               continue
             if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(__a  ,      __a							):
               # Those are stored as lists of tensors so need specific unbatching.
               if isinstance(element[0]  ,      torch.Tensor							):
                 UpperCamelCase_: Dict   				=    tuple(el[self._loader_batch_index].unsqueeze(0							) for el in element							)
               elif isinstance(element[0]  ,      np.ndarray							):
                 UpperCamelCase_: Dict   				=    tuple(np.expand_dims(el[self._loader_batch_index]  ,      0							) for el in element							)
               continue
             if element is None:
               # This can happen for optional data that get passed around
               UpperCamelCase_: List[Any]   				=    None
             elif isinstance(element[self._loader_batch_index]  ,      torch.Tensor							):
               # Take correct batch data, but make it looked like batch_size=1
               # For compatibility with other methods within transformers

               UpperCamelCase_: Union[str, Any]   				=    element[self._loader_batch_index].unsqueeze(0							)
             elif isinstance(element[self._loader_batch_index]  ,      np.ndarray							):
               # Take correct batch data, but make it looked like batch_size=1
               # For compatibility with other methods within transformers
               UpperCamelCase_: Optional[Any]   				=    np.expand_dims(element[self._loader_batch_index]  ,      0							)
             else:
               # This is typically a list, so no need to `unsqueeze`.
               UpperCamelCase_: Optional[Any]   				=    element[self._loader_batch_index]
            # Recreate the element by reusing the original class to make it look
            # batch_size=1
           UpperCamelCase_: Optional[int]   				=    self._loader_batch_data.__class__(__a							)
         self._loader_batch_index += 1
         return result





       def       _a   ( self							):
         if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
           # We are currently unrolling a batch so we just need to return
           # the current item within a batch
           return self.loader_batch_item()

         # We're out of items within a batch
         UpperCamelCase_: Union[str, Any]   				=    next(self.iterator							)
         UpperCamelCase_: Union[str, Any]   				=    self.infer(__a  ,      **self.params							)
         # We now have a batch of "inferred things".
         if self.loader_batch_size is not None:
           # Try to infer the size of the batch
           if isinstance(__a  ,      torch.Tensor							):
             UpperCamelCase_: Optional[int]   				=    processed
           else:
             UpperCamelCase_: str   				=    list(processed.keys()							)[0]
             UpperCamelCase_: Dict   				=    processed[key]
           if isinstance(__a  ,      __a							):
             UpperCamelCase_: List[str]   				=    len(__a							)
           else:
             UpperCamelCase_: Optional[int]   				=    first_tensor.shape[0]
           if 0 < observed_batch_size < self.loader_batch_size:
             # could be last batch so we can't unroll as many
             # elements.
             UpperCamelCase_: Tuple   				=    observed_batch_size
           # Setting internal index to unwrap the batch
           UpperCamelCase_: Optional[int]   				=    processed
           UpperCamelCase_: List[Any]   				=    0
           return self.loader_batch_item()
         else:
           # We're not unrolling batches
           return processed







class     _lowerCAmelCase(	UpperCamelCase__ ):





       """simple docstring"""







       def __init__( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase=None							):
         super().__init__(__a  ,      __a  ,      __a							)





       def __iter__( self							):
         UpperCamelCase_: str   				=    iter(self.loader							)
         UpperCamelCase_: str   				=    None
         return self





       def       _a   ( self							):
         if self.subiterator is None:
           UpperCamelCase_: Optional[int]   				=    self.infer(next(self.iterator							)  ,      **self.params							)
         try:
           # Try to return next item
           UpperCamelCase_: str   				=    next(self.subiterator							)
         except StopIteration:
           # When a preprocess iterator ends, we can start lookig at the next item
           # ChunkIterator will keep feeding until ALL elements of iterator
           # all have created their subiterator and have been iterating against.
           #
           # Another way to look at it, is we're basically flattening lists of lists
           # into a single list, but with generators
           UpperCamelCase_: Optional[Any]   				=    self.infer(next(self.iterator							)  ,      **self.params							)
           UpperCamelCase_: List[str]   				=    next(self.subiterator							)
         return processed







class     _lowerCAmelCase(	UpperCamelCase__ ):





       """simple docstring"""







       def __iter__( self							):
         UpperCamelCase_: Tuple   				=    iter(self.loader							)
         return self





       def       _a   ( self							):
         # Extremely similar to PipelineIterator in its unpacking mechanism
         # BUT, we have an extra required item which is the presence of `is_last`
         # That is because everything is flattened by `PipelineChunkIterator` we
         # need to keep track of how to regroup here in the original `process`
         # boundaries so that `process` and `postprocess` see the same data.

         # This iterator accumulates items (possibly while unbatching) until it
         # its a `is_last` and then just passes it on to the caller.
         UpperCamelCase_: List[str]   				=    False
         UpperCamelCase_: List[str]   				=    []
         if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
           while self._loader_batch_index < self.loader_batch_size:
             UpperCamelCase_: int   				=    self.loader_batch_item()
             UpperCamelCase_: str   				=    item.pop('is_last'							)
             accumulator.append(__a							)
             if is_last:
               return accumulator

         while not is_last:
           UpperCamelCase_: Tuple   				=    self.infer(next(self.iterator							)  ,      **self.params							)
           if self.loader_batch_size is not None:
             if isinstance(__a  ,      torch.Tensor							):
               UpperCamelCase_: int   				=    processed
             else:
               UpperCamelCase_: List[str]   				=    list(processed.keys()							)[0]
               UpperCamelCase_: Optional[int]   				=    processed[key]
             if isinstance(__a  ,      __a							):
               UpperCamelCase_: int   				=    len(__a							)
             else:
               UpperCamelCase_: List[str]   				=    first_tensor.shape[0]
             if 0 < observed_batch_size < self.loader_batch_size:
               # could be last batch so we can't unroll as many
               # elements.
               UpperCamelCase_: Optional[Any]   				=    observed_batch_size
             UpperCamelCase_: int   				=    processed
             UpperCamelCase_: Optional[Any]   				=    0
             while self._loader_batch_index < self.loader_batch_size:
               UpperCamelCase_: List[str]   				=    self.loader_batch_item()
               UpperCamelCase_: str   				=    item.pop('is_last'							)
               accumulator.append(__a							)
               if is_last:
                 return accumulator
           else:
             UpperCamelCase_: str   				=    processed
             UpperCamelCase_: Tuple   				=    item.pop('is_last'							)
             accumulator.append(__a							)
         return accumulator







class     _lowerCAmelCase(	UpperCamelCase__ ):





       """simple docstring"""







       def __init__( self  ,      _lowerCamelCase  ,      _lowerCamelCase							):
         UpperCamelCase_: int   				=    dataset
         UpperCamelCase_: Dict   				=    key





       def __len__( self							):
         return len(self.dataset							)





       def __getitem__( self  ,      _lowerCamelCase							):
         return self.dataset[i][self.key]







class     _lowerCAmelCase(	UpperCamelCase__ ):





       """simple docstring"""







       def __init__( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase							):
         UpperCamelCase_: Union[str, Any]   				=    dataset
         UpperCamelCase_: str   				=    keya
         UpperCamelCase_: Any   				=    keya





       def __len__( self							):
         return len(self.dataset							)





       def __getitem__( self  ,      _lowerCamelCase							):
         return {"text": self.dataset[i][self.keya], "text_pair": self.dataset[i][self.keya]}







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__  )	->    Optional[int]:  # picklable for multiprocessing
  return x.sum()

def   snake_case (UpperCAmelCase__  )	->    Any:  # picklable for multiprocessing
  return i + 1




@dataclass
class     _lowerCAmelCase:





       """simple docstring"""







       a    :		int
       a    :		str







class     _lowerCAmelCase(	UpperCAmelCase_ ):





       """simple docstring"""







       def       _a   ( self							):
         UpperCamelCase_: Optional[Any]   				=    {}
         UpperCamelCase_: List[str]   				=    []
         UpperCamelCase_: Any   				=    1
         UpperCamelCase_: Optional[int]   				=    [1, 2]
         UpperCamelCase_: List[str]   				=    {'a': 1, 'b': 2}
         UpperCamelCase_: Tuple   				=    {'a': [1, 2], 'b': [3, 4]}
         UpperCamelCase_: Optional[int]   				=    {'a': {'1': 1}, 'b': 2}
         UpperCamelCase_: Optional[Any]   				=    {'a': 1, 'b': 2, 'c': 3, 'd': 4}
         UpperCamelCase_: Tuple   				=    {}
         UpperCamelCase_: str   				=    []
         UpperCamelCase_: List[Any]   				=    2
         UpperCamelCase_: List[Any]   				=    [2, 3]
         UpperCamelCase_: Optional[Any]   				=    {'a': 2, 'b': 3}
         UpperCamelCase_: List[str]   				=    {'a': [2, 3], 'b': [4, 5]}
         UpperCamelCase_: Any   				=    {'a': {'1': 2}, 'b': 3}
         UpperCamelCase_: List[str]   				=    {'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_: Optional[int]   				=    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_: Tuple   				=    {'a': np.eye(2							), 'b': np.zeros(3							), 'c': np.ones(2							)}
         UpperCamelCase_: Tuple   				=    {'a': 2, 'b': 0, 'c': 2}
         UpperCamelCase_: str   				=    {
             '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 _lowerCamelCase							: x + 1  ,      _lowerCamelCase  ,      num_proc=_lowerCamelCase							)





       def       _a   ( self							):
         UpperCamelCase_: Optional[Any]   				=    {'a': 1, 'b': 2}
         UpperCamelCase_: Dict   				=    {'a': 3, 'b': 4}
         UpperCamelCase_: Optional[int]   				=    {'a': 5, 'b': 6}
         UpperCamelCase_: int   				=    sorted([('a', (1, 3, 5)), ('b', (2, 4, 6))]							)
         self.assertEqual(sorted(zip_dict(_lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase							)							)  ,      _lowerCamelCase							)





       def       _a   ( self							):




         class     _lowerCAmelCase:





                """simple docstring"""







                a    :		str									='''bar'''







         UpperCamelCase_: int   				=    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),
        (1_6, 1_6, 1_6),
        (1_6, 1_7, 1_6),
        (1_7, 1_6, 1_6),
    ]     ,   )
def   snake_case (UpperCAmelCase__     ,   UpperCAmelCase__     ,   UpperCAmelCase__  )	->    Dict:
  with patch('datasets.utils.py_utils._single_map_nested'  ) as mock_single_map_nested, patch(
      'datasets.parallel.parallel.Pool'  ) as mock_multiprocessing_pool:
    UpperCamelCase_: Any   				=    {F'''{i}''': i for i in range(UpperCAmelCase__  )}
    UpperCamelCase_: int   				=    map_nested(lambda UpperCAmelCase__  : x + 1_0     ,   UpperCAmelCase__     ,   num_proc=UpperCAmelCase__     ,   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     _lowerCAmelCase(	UpperCAmelCase_ ):





       """simple docstring"""







       @require_tf
       def       _a   ( self							):
         import tensorflow as tf
         from tensorflow.keras import layers

         UpperCamelCase_: Dict   				=    layers.Dense(2							)

         def gen_random_output():
           UpperCamelCase_: Optional[Any]   				=    tf.random.uniform((1, 3)							)
           return model(_lowerCamelCase							).numpy()

         with temp_seed(4_2  ,      set_tensorflow=_lowerCamelCase							):
           UpperCamelCase_: int   				=    gen_random_output()
         with temp_seed(4_2  ,      set_tensorflow=_lowerCamelCase							):
           UpperCamelCase_: List[str]   				=    gen_random_output()
         UpperCamelCase_: str   				=    gen_random_output()

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





       @require_torch
       def       _a   ( self							):
         import torch

         def gen_random_output():
           UpperCamelCase_: Any   				=    torch.nn.Linear(3  ,      2							)
           UpperCamelCase_: Optional[Any]   				=    torch.rand(1  ,      3							)
           return model(_lowerCamelCase							).detach().numpy()

         with temp_seed(4_2  ,      set_pytorch=_lowerCamelCase							):
           UpperCamelCase_: Dict   				=    gen_random_output()
         with temp_seed(4_2  ,      set_pytorch=_lowerCamelCase							):
           UpperCamelCase_: str   				=    gen_random_output()
         UpperCamelCase_: str   				=    gen_random_output()

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





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

         with temp_seed(4_2							):
           UpperCamelCase_: Optional[Any]   				=    gen_random_output()
         with temp_seed(4_2							):
           UpperCamelCase_: Tuple   				=    gen_random_output()
         UpperCamelCase_: Optional[int]   				=    gen_random_output()

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







@pytest.mark.parametrize('input_data'     ,   [{}]  )
def   snake_case (UpperCAmelCase__  )	->    Dict:
  UpperCamelCase_: str   				=    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__  )	->    Dict:
  UpperCamelCase_: Optional[Any]   				=    NestedDataStructure(UpperCAmelCase__  ).flatten()
  assert output == expected_output

def   snake_case ()	->    Optional[int]:
  UpperCamelCase_: List[Any]   				=    A(x=1     ,   y='foobar'  )
  UpperCamelCase_: Optional[int]   				=    {'x': 1, 'y': 'foobar'}
  assert asdict(UpperCAmelCase__  ) == expected_output

  UpperCamelCase_: List[str]   				=    {'a': {'b': A(x=1_0     ,   y='foo'  )}, 'c': [A(x=2_0     ,   y='bar'  )]}
  UpperCamelCase_: Tuple   				=    {'a': {'b': {'x': 1_0, 'y': 'foo'}}, 'c': [{'x': 2_0, 'y': 'bar'}]}
  assert asdict(UpperCAmelCase__  ) == expected_output

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

def   snake_case (UpperCAmelCase__  )	->    Optional[Any]:
  return text.split()

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

def   snake_case ()	->    Union[str, Any]:
  with Pool(2  ) as pool:
    UpperCamelCase_: Optional[Any]   				=    list(iflatmap_unordered(UpperCAmelCase__     ,   _split_text     ,   kwargs_iterable=[{'text': 'hello there'}] * 1_0  )  )
    assert out.count('hello'  ) == 1_0
    assert out.count('there'  ) == 1_0
    assert len(UpperCAmelCase__  ) == 2_0

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

  # check that we get items as fast as possible
  with Pool(2  ) as pool:
    UpperCamelCase_: Any   				=    []
    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