"""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      __UpperCAmelCase       (      lowercase		):  # picklable for multiprocessing







  """simple docstring"""
  return x.sum()






def      __UpperCAmelCase       (      lowercase		):  # picklable for multiprocessing







  """simple docstring"""
  return i + 1


@dataclass
class 							a					:
    _snake_case  :	int
    _snake_case  :	str



class 							a					(		lowerCAmelCase_ ):


    def   lowerCAmelCase_	(       self						:	Optional[int]       ):
      _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(__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    			=   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    			=   {"""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(__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   lowerCAmelCase_	(       self						:	Any       ):
      _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(__lowerCAmelCase ,					__lowerCAmelCase ,					__lowerCAmelCase       )       ) ,					__lowerCAmelCase       )


    def   lowerCAmelCase_	(       self						:	Tuple       ):


      class 							a					:
          _snake_case  :	Union[str, Any]  				=				'bar'



      _UpperCAmelCase    			=   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      __UpperCAmelCase       (      lowercase ,lowercase ,lowercase		):







  """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(lowercase		)}
    _UpperCAmelCase    			=   map_nested(lambda lowercase		: x + 10 ,lowercase ,num_proc=lowercase ,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					(		lowerCAmelCase_ ):


    @require_tf
    def   lowerCAmelCase_	(       self						:	Dict       ):
      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(__lowerCAmelCase       ).numpy()

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

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


    @require_torch
    def   lowerCAmelCase_	(       self						:	Union[str, Any]       ):
      import torch

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

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

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


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






@pytest.mark.parametrize("""input_data""" ,[{}]		)
def      __UpperCAmelCase       (      lowercase		):







  """simple docstring"""
  _UpperCAmelCase    			=   NestedDataStructure(lowercase		).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      __UpperCAmelCase       (      lowercase ,lowercase		):







  """simple docstring"""
  _UpperCAmelCase    			=   NestedDataStructure(lowercase		).flatten()
  assert output == expected_output






def      __UpperCAmelCase       (      ):







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

  with pytest.raises(lowercase		):
    asdict([1, A(x=10 ,y="""foo"""		)]		)






def      __UpperCAmelCase       (      lowercase		):







  """simple docstring"""
  return text.split()






def      __UpperCAmelCase       (      lowercase		):







  """simple docstring"""
  yield (time.time(), content)
  time.sleep(2		)
  yield (time.time(), content)






def      __UpperCAmelCase       (      ):







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

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

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

"""simple docstring"""


from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase__			=    logging.get_logger(__name__)

UpperCAmelCase__			=    {
    """google/realm-cc-news-pretrained-embedder""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/config.json"""
    ),
    """google/realm-cc-news-pretrained-encoder""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/config.json"""
    ),
    """google/realm-cc-news-pretrained-scorer""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/config.json"""
    ),
    """google/realm-cc-news-pretrained-openqa""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/config.json"""
    ),
    """google/realm-orqa-nq-openqa""": """https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/config.json""",
    """google/realm-orqa-nq-reader""": """https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/config.json""",
    """google/realm-orqa-wq-openqa""": """https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/config.json""",
    """google/realm-orqa-wq-reader""": """https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/config.json""",
    # See all REALM models at https://huggingface.co/models?filter=realm
}


class 							a					(		lowerCAmelCase_ ):
    _snake_case  :	List[Any]  				=				'realm'


    def __init__(       self						:	List[Any] ,					__lowerCAmelCase						:	List[str]=3_0522 ,					__lowerCAmelCase						:	Tuple=768 ,					__lowerCAmelCase						:	List[Any]=128 ,					__lowerCAmelCase						:	Tuple=12 ,					__lowerCAmelCase						:	List[str]=12 ,					__lowerCAmelCase						:	Union[str, Any]=8 ,					__lowerCAmelCase						:	List[str]=3072 ,					__lowerCAmelCase						:	Union[str, Any]="gelu_new" ,					__lowerCAmelCase						:	str=0.1 ,					__lowerCAmelCase						:	Tuple=0.1 ,					__lowerCAmelCase						:	List[str]=512 ,					__lowerCAmelCase						:	Dict=2 ,					__lowerCAmelCase						:	List[Any]=0.02 ,					__lowerCAmelCase						:	Optional[int]=1e-1_2 ,					__lowerCAmelCase						:	Any=256 ,					__lowerCAmelCase						:	Union[str, Any]=10 ,					__lowerCAmelCase						:	Dict=1e-3 ,					__lowerCAmelCase						:	str=5 ,					__lowerCAmelCase						:	str=320 ,					__lowerCAmelCase						:	Tuple=1335_3718 ,					__lowerCAmelCase						:	Optional[Any]=5000 ,					__lowerCAmelCase						:	int=1 ,					__lowerCAmelCase						:	Dict=0 ,					__lowerCAmelCase						:	Any=2 ,					**__lowerCAmelCase						:	Optional[int] ,					):
      super().__init__(pad_token_id=__lowerCAmelCase ,					bos_token_id=__lowerCAmelCase ,					eos_token_id=__lowerCAmelCase ,					**__lowerCAmelCase       )

      # Common config
      _UpperCAmelCase    			=   vocab_size
      _UpperCAmelCase    			=   max_position_embeddings
      _UpperCAmelCase    			=   hidden_size
      _UpperCAmelCase    			=   retriever_proj_size
      _UpperCAmelCase    			=   num_hidden_layers
      _UpperCAmelCase    			=   num_attention_heads
      _UpperCAmelCase    			=   num_candidates
      _UpperCAmelCase    			=   intermediate_size
      _UpperCAmelCase    			=   hidden_act
      _UpperCAmelCase    			=   hidden_dropout_prob
      _UpperCAmelCase    			=   attention_probs_dropout_prob
      _UpperCAmelCase    			=   initializer_range
      _UpperCAmelCase    			=   type_vocab_size
      _UpperCAmelCase    			=   layer_norm_eps

      # Reader config
      _UpperCAmelCase    			=   span_hidden_size
      _UpperCAmelCase    			=   max_span_width
      _UpperCAmelCase    			=   reader_layer_norm_eps
      _UpperCAmelCase    			=   reader_beam_size
      _UpperCAmelCase    			=   reader_seq_len

      # Retrieval config
      _UpperCAmelCase    			=   num_block_records
      _UpperCAmelCase    			=   searcher_beam_size
