






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

def      lowerCAmelCase__	(    _a       :			Optional[Any]    ):  # picklable for multiprocessing
  return i + 1




@dataclass
class 		UpperCAmelCase_						:


      '''simple docstring'''




      A     :   int
      A     :   str




class 		UpperCAmelCase_						(	_UpperCAmelCase					):


      '''simple docstring'''




      def       _lowerCAmelCase       (	self				) ->				Dict:
        snake_case_				:							Tuple						 =       {}
        snake_case_				:							Any						 =       []
        snake_case_				:							Any						 =       1
        snake_case_				:							List[str]						 =       [1, 2]
        snake_case_				:							Dict						 =       {"""a""": 1, """b""": 2}
        snake_case_				:							Tuple						 =       {"""a""": [1, 2], """b""": [3, 4]}
        snake_case_				:							Optional[Any]						 =       {"""a""": {"""1""": 1}, """b""": 2}
        snake_case_				:							List[str]						 =       {"""a""": 1, """b""": 2, """c""": 3, """d""": 4}
        snake_case_				:							List[str]						 =       {}
        snake_case_				:							List[Any]						 =       []
        snake_case_				:							Optional[int]						 =       2
        snake_case_				:							List[Any]						 =       [2, 3]
        snake_case_				:							List[str]						 =       {"""a""": 2, """b""": 3}
        snake_case_				:							Tuple						 =       {"""a""": [2, 3], """b""": [4, 5]}
        snake_case_				:							Optional[Any]						 =       {"""a""": {"""1""": 2}, """b""": 3}
        snake_case_				:							Optional[int]						 =       {"""a""": 2, """b""": 3, """c""": 4, """d""": 5}
        self.assertEqual(map_nested(lowercase__		,				lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__				)		,				lowercase__				)

        snake_case_				:							Optional[Any]						 =       2
        self.assertEqual(map_nested(lowercase__		,				lowercase__		,				num_proc=lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__		,				num_proc=lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__		,				num_proc=lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__		,				num_proc=lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__		,				num_proc=lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__		,				num_proc=lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__		,				num_proc=lowercase__				)		,				lowercase__				)
        self.assertEqual(map_nested(lowercase__		,				lowercase__		,				num_proc=lowercase__				)		,				lowercase__				)

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






      def       _lowerCAmelCase       (	self				) ->				str:
        snake_case_				:							List[str]						 =       {"""a""": 1, """b""": 2}
        snake_case_				:							Dict						 =       {"""a""": 3, """b""": 4}
        snake_case_				:							Dict						 =       {"""a""": 5, """b""": 6}
        snake_case_				:							List[Any]						 =       sorted([("a", (1, 3, 5)), ("b", (2, 4, 6))]				)
        self.assertEqual(sorted(zip_dict(lowercase__		,				lowercase__		,				lowercase__				)				)		,				lowercase__				)






      def       _lowerCAmelCase       (	self				) ->				Optional[int]:




        class 		UpperCAmelCase_						:


              '''simple docstring'''




              A     :   str						=					"""bar"""




        snake_case_				:							Dict						 =       Foo()
        self.assertEqual(foo.my_attr		,				"bar"				)
        with temporary_assignment(lowercase__		,				"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      lowerCAmelCase__	(    _a       :			Tuple  ,      _a       :			int  ,      _a       :			Optional[Any]    ):





  with patch("datasets.utils.py_utils._single_map_nested"    ) as mock_single_map_nested, patch(
      "datasets.parallel.parallel.Pool"    ) as mock_multiprocessing_pool:
    snake_case_				:							Optional[Any]						 =       {F'''{i}''': i for i in range(SCREAMING_SNAKE_CASE__    )}
    snake_case_				:							Union[str, Any]						 =       map_nested(lambda _a    : 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_						(	_UpperCAmelCase					):


      '''simple docstring'''




      @require_tf
      def       _lowerCAmelCase       (	self				) ->				Optional[int]:
        import tensorflow as tf
        from tensorflow.keras import layers

        snake_case_				:							Tuple						 =       layers.Dense(2				)

        def gen_random_output():
          snake_case_				:							List[Any]						 =       tf.random.uniform((1, 3)				)
          return model(lowercase__				).numpy()

        with temp_seed(42		,				set_tensorflow=lowercase__				):
          snake_case_				:							Tuple						 =       gen_random_output()
        with temp_seed(42		,				set_tensorflow=lowercase__				):
          snake_case_				:							str						 =       gen_random_output()
        snake_case_				:							Optional[int]						 =       gen_random_output()

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






      @require_torch
      def       _lowerCAmelCase       (	self				) ->				Optional[int]:
        import torch

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

        with temp_seed(42		,				set_pytorch=lowercase__				):
          snake_case_				:							Optional[Any]						 =       gen_random_output()
        with temp_seed(42		,				set_pytorch=lowercase__				):
          snake_case_				:							Optional[Any]						 =       gen_random_output()
        snake_case_				:							Tuple						 =       gen_random_output()

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






      def       _lowerCAmelCase       (	self				) ->				Union[str, Any]:
        def gen_random_output():
          return np.random.rand(1		,				3				)

        with temp_seed(42				):
          snake_case_				:							List[Any]						 =       gen_random_output()
        with temp_seed(42				):
          snake_case_				:							Union[str, Any]						 =       gen_random_output()
        snake_case_				:							Optional[Any]						 =       gen_random_output()

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




@pytest.mark.parametrize("input_data"  ,      [{}]    )
def      lowerCAmelCase__	(    _a       :			str    ):
  snake_case_				:							Union[str, Any]						 =       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      lowerCAmelCase__	(    _a       :			Tuple  ,      _a       :			List[Any]    ):
  snake_case_				:							Dict						 =       NestedDataStructure(SCREAMING_SNAKE_CASE__    ).flatten()
  assert output == expected_output

def      lowerCAmelCase__	(    ):
  snake_case_				:							Union[str, Any]						 =       A(x=1  ,      y="foobar"    )
  snake_case_				:							List[Any]						 =       {"""x""": 1, """y""": """foobar"""}
  assert asdict(SCREAMING_SNAKE_CASE__    ) == expected_output

  snake_case_				:							Tuple						 =       {"""a""": {"""b""": A(x=10  ,      y="foo"    )}, """c""": [A(x=20  ,      y="bar"    )]}
  snake_case_				:							Union[str, Any]						 =       {"""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      lowerCAmelCase__	(    _a       :			str    ):
  return text.split()

def      lowerCAmelCase__	(    _a       :			Dict    ):
  yield (time.time(), content)
  time.sleep(2    )
  yield (time.time(), content)

def      lowerCAmelCase__	(    ):





  with Pool(2    ) as pool:
    snake_case_				:							Optional[int]						 =       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:
    snake_case_				:							Union[str, Any]						 =       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:
    snake_case_				:							Optional[int]						 =       []
    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 warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType




class 		UpperCAmelCase_						(	SCREAMING_SNAKE_CASE__					):


						'''simple docstring'''




						A     :   Tuple						=					['image_processor', 'tokenizer']
						A     :   List[Any]						=					'ViltImageProcessor'
						A     :   Optional[Any]						=					('BertTokenizer', 'BertTokenizerFast')






						def __init__(	self		,				_SCREAMING_SNAKE_CASE=None		,				_SCREAMING_SNAKE_CASE=None		,				**_SCREAMING_SNAKE_CASE				) ->				int:
								snake_case_				:							str						 =       None
								if "feature_extractor" in kwargs:
										warnings.warn(
										    "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
										    " instead."		,				_SCREAMING_SNAKE_CASE		,				)
										snake_case_				:							str						 =       kwargs.pop("feature_extractor"				)

								snake_case_				:							Tuple						 =       image_processor if image_processor is not None else feature_extractor
								if image_processor is None:
										raise ValueError("You need to specify an `image_processor`."				)
								if tokenizer is None:
										raise ValueError("You need to specify a `tokenizer`."				)

								super().__init__(_SCREAMING_SNAKE_CASE		,				_SCREAMING_SNAKE_CASE				)
								snake_case_				:							str						 =       self.image_processor






						def __call__(	self		,				_SCREAMING_SNAKE_CASE		,				_SCREAMING_SNAKE_CASE = None		,				_SCREAMING_SNAKE_CASE = True		,				_SCREAMING_SNAKE_CASE = False		,				_SCREAMING_SNAKE_CASE = None		,				_SCREAMING_SNAKE_CASE = None		,				_SCREAMING_SNAKE_CASE = 0		,				_SCREAMING_SNAKE_CASE = None		,				_SCREAMING_SNAKE_CASE = None		,				_SCREAMING_SNAKE_CASE = None		,				_SCREAMING_SNAKE_CASE = False		,				_SCREAMING_SNAKE_CASE = False		,				_SCREAMING_SNAKE_CASE = False		,				_SCREAMING_SNAKE_CASE = False		,				_SCREAMING_SNAKE_CASE = True		,				_SCREAMING_SNAKE_CASE = None		,				**_SCREAMING_SNAKE_CASE		,				) ->				BatchEncoding:
								snake_case_				:							List[Any]						 =       self.tokenizer(
								    text=_SCREAMING_SNAKE_CASE		,				add_special_tokens=_SCREAMING_SNAKE_CASE		,				padding=_SCREAMING_SNAKE_CASE		,				truncation=_SCREAMING_SNAKE_CASE		,				max_length=_SCREAMING_SNAKE_CASE		,				stride=_SCREAMING_SNAKE_CASE		,				pad_to_multiple_of=_SCREAMING_SNAKE_CASE		,				return_token_type_ids=_SCREAMING_SNAKE_CASE		,				return_attention_mask=_SCREAMING_SNAKE_CASE		,				return_overflowing_tokens=_SCREAMING_SNAKE_CASE		,				return_special_tokens_mask=_SCREAMING_SNAKE_CASE		,				return_offsets_mapping=_SCREAMING_SNAKE_CASE		,				return_length=_SCREAMING_SNAKE_CASE		,				verbose=_SCREAMING_SNAKE_CASE		,				return_tensors=_SCREAMING_SNAKE_CASE		,				**_SCREAMING_SNAKE_CASE		,				)
								# add pixel_values + pixel_mask
								snake_case_				:							str						 =       self.image_processor(_SCREAMING_SNAKE_CASE		,				return_tensors=_SCREAMING_SNAKE_CASE				)
								encoding.update(_SCREAMING_SNAKE_CASE				)

								return encoding






						def       _lowerCAmelCase       (	self		,				*_SCREAMING_SNAKE_CASE		,				**_SCREAMING_SNAKE_CASE				) ->				List[str]:
								return self.tokenizer.batch_decode(*_SCREAMING_SNAKE_CASE		,				**_SCREAMING_SNAKE_CASE				)






						def       _lowerCAmelCase       (	self		,				*_SCREAMING_SNAKE_CASE		,				**_SCREAMING_SNAKE_CASE				) ->				Dict:
								return self.tokenizer.decode(*_SCREAMING_SNAKE_CASE		,				**_SCREAMING_SNAKE_CASE				)






						@property
						def       _lowerCAmelCase       (	self				) ->				Tuple:
								snake_case_				:							List[Any]						 =       self.tokenizer.model_input_names
								snake_case_				:							Tuple						 =       self.image_processor.model_input_names
								return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names				)				)






						@property
						def       _lowerCAmelCase       (	self				) ->				Optional[Any]:
								warnings.warn(
								    "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."		,				_SCREAMING_SNAKE_CASE		,				)
								return self.image_processor_class






						@property
						def       _lowerCAmelCase       (	self				) ->				Optional[int]:
								warnings.warn(
								    "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."		,				_SCREAMING_SNAKE_CASE		,				)
								return self.image_processor
