



'''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    lowerCamelCase__ (		_A       ):  # picklable for multiprocessing
  return x.sum()
def    lowerCamelCase__ (		_A       ):  # picklable for multiprocessing
  return i + 1

@dataclass
class 							a__:
    lowercase__				=					42
    lowercase__				=					42

class 							a__(   UpperCAmelCase__			):







    def     lowercase_ (   self				:						str			):
      a		:      List[str]       =  {}
      a		:      int       =  []
      a		:      Dict       =  1
      a		:      Optional[int]       =  [1, 2]
      a		:      Optional[int]       =  {'a': 1, 'b': 2}
      a		:      List[str]       =  {'a': [1, 2], 'b': [3, 4]}
      a		:      str       =  {'a': {'1': 1}, 'b': 2}
      a		:      int       =  {'a': 1, 'b': 2, 'c': 3, 'd': 4}
      a		:      List[Any]       =  {}
      a		:      Any       =  []
      a		:      int       =  2
      a		:      str       =  [2, 3]
      a		:      Optional[int]       =  {'a': 2, 'b': 3}
      a		:      Tuple       =  {'a': [2, 3], 'b': [4, 5]}
      a		:      Tuple       =  {'a': {'1': 2}, 'b': 3}
      a		:      List[Any]       =  {'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			)

      a		:      List[Any]       =  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			)

      a		:      Optional[int]       =  {'a': np.eye(2			), 'b': np.zeros(3			), 'c': np.ones(2			)}
      a		:      str       =  {'a': 2, 'b': 0, 'c': 2}
      a		:      Any       =  {
          '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 __snake_case			: x + 1		,    __lowerCAmelCase		,    num_proc=__lowerCAmelCase			)







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







    def     lowercase_ (   self				:						Dict			):

      class 							a__:
          lowercase__				=					"""bar"""
      a		:      Any       =  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    lowerCamelCase__ (		_A				,    _A				,    _A       ):

  with patch('datasets.utils.py_utils._single_map_nested'       ) as mock_single_map_nested, patch(
      'datasets.parallel.parallel.Pool'       ) as mock_multiprocessing_pool:
    a		:      str       =  {f"""{i}""": i for i in range(UpperCAmelCase__       )}
    a		:      Union[str, Any]       =  map_nested(lambda _A       : 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     lowercase_ (   self				:						Dict			):
      import tensorflow as tf
      from tensorflow.keras import layers

      a		:      Dict       =  layers.Dense(2			)

      def gen_random_output():
        a		:      Dict       =  tf.random.uniform((1, 3)			)
        return model(__lowerCAmelCase			).numpy()

      with temp_seed(42		,    set_tensorflow=__lowerCAmelCase			):
        a		:      Tuple       =  gen_random_output()
      with temp_seed(42		,    set_tensorflow=__lowerCAmelCase			):
        a		:      Optional[Any]       =  gen_random_output()
      a		:      List[str]       =  gen_random_output()

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







    @require_torch
    def     lowercase_ (   self				:						Dict			):
      import torch

      def gen_random_output():
        a		:      Union[str, Any]       =  torch.nn.Linear(3		,    2			)
        a		:      List[Any]       =  torch.rand(1		,    3			)
        return model(__lowerCAmelCase			).detach().numpy()

      with temp_seed(42		,    set_pytorch=__lowerCAmelCase			):
        a		:      Dict       =  gen_random_output()
      with temp_seed(42		,    set_pytorch=__lowerCAmelCase			):
        a		:      List[Any]       =  gen_random_output()
      a		:      Any       =  gen_random_output()

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







    def     lowercase_ (   self				:						Tuple			):






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

      with temp_seed(42			):
        a		:      int       =  gen_random_output()
      with temp_seed(42			):
        a		:      List[str]       =  gen_random_output()
      a		:      Any       =  gen_random_output()

      np.testing.assert_equal(__lowerCAmelCase		,    __lowerCAmelCase			)
      self.assertGreater(np.abs(outa - outa			).sum()		,    0			)
@pytest.mark.parametrize('input_data'				,    [{}]       )
def    lowerCamelCase__ (		_A       ):
  a		:      Dict       =  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    lowerCamelCase__ (		_A				,    _A       ):
  a		:      Optional[int]       =  NestedDataStructure(UpperCAmelCase__       ).flatten()
  assert output == expected_output
def    lowerCamelCase__ (		):
  a		:      int       =  A(x=1				,    y='foobar'       )
  a		:      Optional[Any]       =  {'x': 1, 'y': 'foobar'}
  assert asdict(UpperCAmelCase__       ) == expected_output

  a		:      Tuple       =  {'a': {'b': A(x=10				,    y='foo'       )}, 'c': [A(x=20				,    y='bar'       )]}
  a		:      Tuple       =  {'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    lowerCamelCase__ (		_A       ):
  return text.split()
def    lowerCamelCase__ (		_A       ):
  yield (time.time(), content)
  time.sleep(2       )
  yield (time.time(), content)

def    lowerCamelCase__ (		):

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




'''simple docstring'''






from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional

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


if TYPE_CHECKING:
     from ... import FeatureExtractionMixin, TensorType

lowerCAmelCase:      List[Any]		 =		logging.get_logger(__name__)

lowerCAmelCase:      List[str]		 =		{
    'openai/imagegpt-small': '',
    'openai/imagegpt-medium': '',
    'openai/imagegpt-large': '',
}

class 							a__(      lowerCamelCase__			):
    lowercase__				=					"""imagegpt"""
    lowercase__				=					["""past_key_values"""]
    lowercase__				=					{
        """hidden_size""": """n_embd""",
        """max_position_embeddings""": """n_positions""",
        """num_attention_heads""": """n_head""",
        """num_hidden_layers""": """n_layer""",
    }







    def __init__(   self				:						List[str]		,    __snake_case				:						int=5_12 + 1		,    __snake_case				:						Optional[Any]=32 * 32		,    __snake_case				:						Optional[Any]=5_12		,    __snake_case				:						List[str]=24		,    __snake_case				:						Tuple=8		,    __snake_case				:						Optional[int]=None		,    __snake_case				:						Dict="quick_gelu"		,    __snake_case				:						Dict=0.1		,    __snake_case				:						List[str]=0.1		,    __snake_case				:						Optional[Any]=0.1		,    __snake_case				:						Tuple=1e-5		,    __snake_case				:						Union[str, Any]=0.02		,    __snake_case				:						Optional[int]=True		,    __snake_case				:						Dict=True		,    __snake_case				:						str=False		,    __snake_case				:						Any=False		,    __snake_case				:						List[Any]=False		,    **__snake_case				:						List[str]		,    ):
      a		:      Optional[Any]       =  vocab_size
      a		:      Union[str, Any]       =  n_positions
      a		:      int       =  n_embd
      a		:      int       =  n_layer
      a		:      str       =  n_head
      a		:      int       =  n_inner
      a		:      int       =  activation_function
      a		:      Any       =  resid_pdrop
      a		:      Tuple       =  embd_pdrop
      a		:      List[Any]       =  attn_pdrop
      a		:      Tuple       =  layer_norm_epsilon
      a		:      Any       =  initializer_range
      a		:      str       =  scale_attn_weights
      a		:      List[Any]       =  use_cache
      a		:      Any       =  scale_attn_by_inverse_layer_idx
      a		:      Optional[Any]       =  reorder_and_upcast_attn
      a		:      str       =  tie_word_embeddings

      super().__init__(tie_word_embeddings=__snake_case		,    **__snake_case			)

class 							a__(      lowerCamelCase__			):







    @property
    def     lowercase_ (   self				:						Union[str, Any]			):
      return OrderedDict(
          [
              ('input_ids', {0: 'batch', 1: 'sequence'}),
          ]			)







    def     lowercase_ (   self				:						Union[str, Any]		,    __snake_case				:						"FeatureExtractionMixin"		,    __snake_case				:						int = 1		,    __snake_case				:						int = -1		,    __snake_case				:						bool = False		,    __snake_case				:						Optional["TensorType"] = None		,    __snake_case				:						int = 3		,    __snake_case				:						int = 32		,    __snake_case				:						int = 32		,    ):
      a		:      Dict       =  self._generate_dummy_images(__snake_case		,    __snake_case		,    __snake_case		,    __snake_case			)
      a		:      Union[str, Any]       =  dict(preprocessor(images=__snake_case		,    return_tensors=__snake_case			)			)

      return inputs