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






def 							A						(    snake_case				:Optional[Any]		)				->     Optional[int]:  # picklable for multiprocessing
 return i + 1

@dataclass
class __lowerCAmelCase		:
     lowercase						       =						42
     lowercase						       =						42


class __lowerCAmelCase		(				__SCREAMING_SNAKE_CASE      ):





     def 						UpperCAmelCase (     self				):

      '''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(__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				):

      '''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(__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				)				)     ,					__UpperCAmelCase				)







     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      class __lowerCAmelCase		:
           lowercase						       =						"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),
        (1_6, 1_6, 1_6),
        (1_6, 1_7, 1_6),
        (1_7, 1_6, 1_6),
    ]    , )
def 							A						(    snake_case				:Any    , snake_case				:Optional[Any]    , snake_case				:List[str]		)				->     List[str]:
 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(snake_case		)}
  __UpperCamelCase      =	map_nested(lambda snake_case		: x + 1_0    , snake_case    , num_proc=snake_case    , 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		(				__SCREAMING_SNAKE_CASE      ):





     @require_tf
     def 						UpperCAmelCase (     self				):

      '''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(__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				):

      '''simple docstring'''




      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				):

      '''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(__UpperCAmelCase     ,					__UpperCAmelCase				)
      self.assertGreater(np.abs(outa - outa				).sum()     ,					0				)






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






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

 __UpperCamelCase      =	{'a': {'b': A(x=1_0    , y='foo'		)}, 'c': [A(x=2_0    , y='bar'		)]}
 __UpperCamelCase      =	{'a': {'b': {'x': 1_0, 'y': 'foo'}}, 'c': [{'x': 2_0, 'y': 'bar'}]}
 assert asdict(snake_case		) == expected_output

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






def 							A						(    snake_case				:str		)				->     Optional[int]:
 return text.split()






def 							A						(    snake_case				:Tuple		)				->     Any:
 yield (time.time(), content)
 time.sleep(2		)
 yield (time.time(), content)






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

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

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





"""simple docstring"""

from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCamelCase     :    Union[str, Any]       =			logging.get_logger(__name__)

UpperCamelCase     :    int       =			{
    "distilbert-base-uncased": "https://huggingface.co/distilbert-base-uncased/resolve/main/config.json",
    "distilbert-base-uncased-distilled-squad": (
        "https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/config.json"
    ),
    "distilbert-base-cased": "https://huggingface.co/distilbert-base-cased/resolve/main/config.json",
    "distilbert-base-cased-distilled-squad": (
        "https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/config.json"
    ),
    "distilbert-base-german-cased": "https://huggingface.co/distilbert-base-german-cased/resolve/main/config.json",
    "distilbert-base-multilingual-cased": (
        "https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/config.json"
    ),
    "distilbert-base-uncased-finetuned-sst-2-english": (
        "https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/config.json"
    ),
}

class __lowerCAmelCase		(				__SCREAMING_SNAKE_CASE      ):
     lowercase						       =						"distilbert"
     lowercase						       =						{
         "hidden_size": "dim",
         "num_attention_heads": "n_heads",
         "num_hidden_layers": "n_layers",
     }





     def __init__(     self     ,					__UpperCAmelCase=3_0522     ,					__UpperCAmelCase=512     ,					__UpperCAmelCase=False     ,					__UpperCAmelCase=6     ,					__UpperCAmelCase=12     ,					__UpperCAmelCase=768     ,					__UpperCAmelCase=4 * 768     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase="gelu"     ,					__UpperCAmelCase=0.0_2     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=0.2     ,					__UpperCAmelCase=0     ,					**__UpperCAmelCase     ,					):

      '''simple docstring'''




      __UpperCamelCase      =	vocab_size
      __UpperCamelCase      =	max_position_embeddings
      __UpperCamelCase      =	sinusoidal_pos_embds
      __UpperCamelCase      =	n_layers
      __UpperCamelCase      =	n_heads
      __UpperCamelCase      =	dim
      __UpperCamelCase      =	hidden_dim
      __UpperCamelCase      =	dropout
      __UpperCamelCase      =	attention_dropout
      __UpperCamelCase      =	activation
      __UpperCamelCase      =	initializer_range
      __UpperCamelCase      =	qa_dropout
      __UpperCamelCase      =	seq_classif_dropout
      super().__init__(**__UpperCAmelCase     ,					pad_token_id=__UpperCAmelCase				)


class __lowerCAmelCase		(				__SCREAMING_SNAKE_CASE      ):





     @property
     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      if self.task == "multiple-choice":
       __UpperCamelCase      =	{0: 'batch', 1: 'choice', 2: 'sequence'}
      else:
       __UpperCamelCase      =	{0: 'batch', 1: 'sequence'}
      return OrderedDict(
          [
              ('input_ids', dynamic_axis),
              ('attention_mask', dynamic_axis),
          ]				)



