






import copy
import os
from typing import Union

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


A_		:		Optional[int]		     =			logging.get_logger(__name__)

A_		:		str		     =			{
    'BridgeTower/bridgetower-base': 'https://huggingface.co/BridgeTower/bridgetower-base/blob/main/config.json',
    'BridgeTower/bridgetower-base-itm-mlm': (
        'https://huggingface.co/BridgeTower/bridgetower-base-itm-mlm/blob/main/config.json'
    ),
}




class     _lowerCAmelCase(	UpperCAmelCase_ ):





       """simple docstring"""







       a    :		List[Any]									='''bridgetower_vision_model'''





       def __init__( self  ,      _lowerCamelCase=7_6_8  ,      _lowerCamelCase=1_2  ,      _lowerCamelCase=3  ,      _lowerCamelCase=1_6  ,      _lowerCamelCase=2_8_8  ,      _lowerCamelCase=1  ,      _lowerCamelCase=1e-05  ,      _lowerCamelCase=False  ,      _lowerCamelCase=True  ,      _lowerCamelCase=False  ,      **_lowerCamelCase  ,      ):
         super().__init__(**_lowerCamelCase							)
         UpperCamelCase_: Any   				=    hidden_size
         UpperCamelCase_: Union[str, Any]   				=    num_hidden_layers
         UpperCamelCase_: List[str]   				=    num_channels
         UpperCamelCase_: int   				=    patch_size
         UpperCamelCase_: List[Any]   				=    image_size
         UpperCamelCase_: str   				=    initializer_factor
         UpperCamelCase_: str   				=    layer_norm_eps
         UpperCamelCase_: int   				=    stop_gradient
         UpperCamelCase_: int   				=    share_layernorm
         UpperCamelCase_: Union[str, Any]   				=    remove_last_layer





       @classmethod
       def       _a   ( cls  ,      _lowerCamelCase  ,      **_lowerCamelCase							):
         UpperCamelCase_				,UpperCamelCase_: Optional[Any]   				=    cls.get_config_dict(_lowerCamelCase  ,      **_lowerCamelCase							)

         if config_dict.get('model_type'							) == "bridgetower":
           UpperCamelCase_: int   				=    config_dict['text_config']

         if "model_type" in config_dict and hasattr(cls  ,      'model_type'							) and config_dict["model_type"] != cls.model_type:
           logger.warning(
               f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
               f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

         return cls.from_dict(_lowerCamelCase  ,      **_lowerCamelCase							)







class     _lowerCAmelCase(	UpperCAmelCase_ ):





       """simple docstring"""







       a    :		List[str]									='''bridgetower_text_model'''





       def __init__( self  ,      _lowerCamelCase=5_0_2_6_5  ,      _lowerCamelCase=7_6_8  ,      _lowerCamelCase=1_2  ,      _lowerCamelCase=1_2  ,      _lowerCamelCase=1  ,      _lowerCamelCase=3_0_7_2  ,      _lowerCamelCase="gelu"  ,      _lowerCamelCase=0.1  ,      _lowerCamelCase=0.1  ,      _lowerCamelCase=5_1_4  ,      _lowerCamelCase=1  ,      _lowerCamelCase=1e-05  ,      _lowerCamelCase=1  ,      _lowerCamelCase=0  ,      _lowerCamelCase=2  ,      _lowerCamelCase="absolute"  ,      _lowerCamelCase=True  ,      **_lowerCamelCase  ,      ):
         super().__init__(**_lowerCamelCase							)

         UpperCamelCase_: Tuple   				=    vocab_size
         UpperCamelCase_: str   				=    hidden_size
         UpperCamelCase_: List[str]   				=    num_hidden_layers
         UpperCamelCase_: Tuple   				=    num_attention_heads
         UpperCamelCase_: Optional[Any]   				=    hidden_act
         UpperCamelCase_: int   				=    initializer_factor
         UpperCamelCase_: Optional[int]   				=    intermediate_size
         UpperCamelCase_: Optional[Any]   				=    hidden_dropout_prob
         UpperCamelCase_: Dict   				=    attention_probs_dropout_prob
         UpperCamelCase_: int   				=    max_position_embeddings
         UpperCamelCase_: int   				=    type_vocab_size
         UpperCamelCase_: Optional[Any]   				=    layer_norm_eps
         UpperCamelCase_: Any   				=    position_embedding_type
         UpperCamelCase_: int   				=    use_cache
         UpperCamelCase_: Optional[Any]   				=    pad_token_id
         UpperCamelCase_: Optional[Any]   				=    bos_token_id
         UpperCamelCase_: List[str]   				=    eos_token_id





       @classmethod
       def       _a   ( cls  ,      _lowerCamelCase  ,      **_lowerCamelCase							):
         UpperCamelCase_				,UpperCamelCase_: int   				=    cls.get_config_dict(_lowerCamelCase  ,      **_lowerCamelCase							)

         if config_dict.get('model_type'							) == "bridgetower":
           UpperCamelCase_: Dict   				=    config_dict['text_config']

         if "model_type" in config_dict and hasattr(cls  ,      'model_type'							) and config_dict["model_type"] != cls.model_type:
           logger.warning(
               f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
               f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

         return cls.from_dict(_lowerCamelCase  ,      **_lowerCamelCase							)







class     _lowerCAmelCase(	UpperCAmelCase_ ):





       """simple docstring"""







       a    :		Optional[Any]									='''bridgetower'''





       def __init__( self  ,      _lowerCamelCase=True  ,      _lowerCamelCase="gelu"  ,      _lowerCamelCase=7_6_8  ,      _lowerCamelCase=1  ,      _lowerCamelCase=1e-05  ,      _lowerCamelCase=False  ,      _lowerCamelCase="add"  ,      _lowerCamelCase=1_2  ,      _lowerCamelCase=6  ,      _lowerCamelCase=False  ,      _lowerCamelCase=False  ,      _lowerCamelCase=None  ,      _lowerCamelCase=None  ,      **_lowerCamelCase  ,      ):
         # TODO: remove this once the Hub files are updated.
         UpperCamelCase_: Optional[int]   				=    kwargs.pop('text_config_dict'  ,      _lowerCamelCase							)
         UpperCamelCase_: Union[str, Any]   				=    kwargs.pop('vision_config_dict'  ,      _lowerCamelCase							)

         super().__init__(**_lowerCamelCase							)
         UpperCamelCase_: List[Any]   				=    share_cross_modal_transformer_layers
         UpperCamelCase_: List[str]   				=    hidden_act
         UpperCamelCase_: Optional[Any]   				=    hidden_size
         UpperCamelCase_: Optional[Any]   				=    initializer_factor
         UpperCamelCase_: Any   				=    layer_norm_eps
         UpperCamelCase_: Dict   				=    share_link_tower_layers
         UpperCamelCase_: int   				=    link_tower_type
         UpperCamelCase_: Union[str, Any]   				=    num_attention_heads
         UpperCamelCase_: Optional[int]   				=    num_hidden_layers
         UpperCamelCase_: Tuple   				=    tie_word_embeddings
         UpperCamelCase_: Optional[int]   				=    init_layernorm_from_vision_encoder

         if text_config is None:
           UpperCamelCase_: Any   				=    {}
           logger.info('`text_config` is `None`. Initializing the `BridgeTowerTextConfig` with default values.'							)

         if vision_config is None:
           UpperCamelCase_: Any   				=    {}
           logger.info('`vision_config` is `None`. Initializing the `BridgeTowerVisionConfig` with default values.'							)

         UpperCamelCase_: Optional[Any]   				=    BridgeTowerTextConfig(**_lowerCamelCase							)
         UpperCamelCase_: Optional[Any]   				=    BridgeTowerVisionConfig(**_lowerCamelCase							)





       @classmethod
       def       _a   ( cls  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      **_lowerCamelCase							):
         return cls(text_config=text_config.to_dict()  ,      vision_config=vision_config.to_dict()  ,      **_lowerCamelCase							)





       def       _a   ( self							):
         UpperCamelCase_: Tuple   				=    copy.deepcopy(self.__dict__							)
         UpperCamelCase_: Optional[Any]   				=    self.text_config.to_dict()
         UpperCamelCase_: Union[str, Any]   				=    self.vision_config.to_dict()
         UpperCamelCase_: Tuple   				=    self.__class__.model_type
         return output







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