



"""simple docstring"""


import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_flax_cross_test,
    require_flax,
    require_torch,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import is_flax_available, is_torch_available, is_vision_available

from ...test_modeling_flax_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_flax_bert import FlaxBertModelTester
from ..clip.test_modeling_flax_clip import FlaxCLIPVisionModelTester
from ..vit.test_modeling_flax_vit import FlaxViTModelTester


if is_flax_available():
					from transformers import (
					    FlaxBertModel,
					    FlaxCLIPVisionModel,
					    FlaxVisionTextDualEncoderModel,
					    FlaxViTModel,
					    VisionTextDualEncoderConfig,
					    VisionTextDualEncoderProcessor,
					)
					from transformers.modeling_flax_pytorch_utils import (
					    convert_pytorch_state_dict_to_flax,
					    load_flax_weights_in_pytorch_model,
					)


if is_torch_available():
					import torch

					from transformers import VisionTextDualEncoderModel

if is_vision_available():
					from PIL import Image







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						List[Any]				):






				'''simple docstring'''



				if isinstance(_SCREAMING_SNAKE_CASE					,     collections.abc.Iterable				):
								return x
				return (x, x)
@require_flax
class 			_a :
			"""simple docstring"""



			def  lowercase__						(		self     :    List[str]				,   __UpperCamelCase     :    List[Any]				,   __UpperCamelCase     :    Optional[Any]   )->Dict:
							pass



			def  lowercase__						(		self     :    Union[str, Any]   )->Dict:
							pass



			def  lowercase__						(		self     :    List[Any]   )->Tuple:
							pass



			def  lowercase__						(		self     :    Dict				,   __UpperCamelCase     :    np.ndarray				,   __UpperCamelCase     :    np.ndarray				,   __UpperCamelCase     :    float   )->int:
							_UpperCAmelCase					=    np.abs((a - b)   ).max()
							self.assertLessEqual(__UpperCamelCase				,   __UpperCamelCase				,   F'Difference between torch and flax is {diff} (>= {tol}).'   )



			def  lowercase__						(		self     :    Optional[int]				,   __UpperCamelCase     :    Dict				,   __UpperCamelCase     :    List[str]				,   __UpperCamelCase     :    List[Any]				,   __UpperCamelCase     :    Tuple				,   __UpperCamelCase     :    Union[str, Any]=None				,   **__UpperCamelCase     :    Optional[int]   )->Union[str, Any]:
							_UpperCAmelCase					=    VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCamelCase				,   __UpperCamelCase   )

							_UpperCAmelCase					=    FlaxVisionTextDualEncoderModel(__UpperCamelCase   )

							_UpperCAmelCase					=    model(input_ids=__UpperCamelCase				,   pixel_values=__UpperCamelCase				,   attention_mask=__UpperCamelCase   )

							self.assertEqual(output['''text_embeds'''].shape				,   (input_ids.shape[0], config.projection_dim)   )
							self.assertEqual(output['''image_embeds'''].shape				,   (pixel_values.shape[0], config.projection_dim)   )



			def  lowercase__						(		self     :    str				,   __UpperCamelCase     :    List[str]				,   __UpperCamelCase     :    str				,   __UpperCamelCase     :    Dict				,   __UpperCamelCase     :    Dict				,   __UpperCamelCase     :    List[str]=None				,   **__UpperCamelCase     :    str   )->Tuple:
							_UpperCAmelCase		,						_UpperCAmelCase					=    self.get_vision_text_model(__UpperCamelCase				,   __UpperCamelCase   )
							_UpperCAmelCase					=    {'''vision_model''': vision_model, '''text_model''': text_model}
							_UpperCAmelCase					=    FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCamelCase   )

							_UpperCAmelCase					=    model(input_ids=__UpperCamelCase				,   pixel_values=__UpperCamelCase				,   attention_mask=__UpperCamelCase   )

							self.assertEqual(output['''text_embeds'''].shape				,   (input_ids.shape[0], model.config.projection_dim)   )
							self.assertEqual(output['''image_embeds'''].shape				,   (pixel_values.shape[0], model.config.projection_dim)   )



			def  lowercase__						(		self     :    str				,   __UpperCamelCase     :    str				,   __UpperCamelCase     :    Optional[int]				,   __UpperCamelCase     :    Union[str, Any]				,   __UpperCamelCase     :    Optional[int]				,   __UpperCamelCase     :    List[str]=None				,   **__UpperCamelCase     :    List[Any]   )->Tuple:
							_UpperCAmelCase		,						_UpperCAmelCase					=    self.get_vision_text_model(__UpperCamelCase				,   __UpperCamelCase   )
							_UpperCAmelCase					=    {'''vision_model''': vision_model, '''text_model''': text_model}
							_UpperCAmelCase					=    FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCamelCase   )

							_UpperCAmelCase					=    model(input_ids=__UpperCamelCase				,   pixel_values=__UpperCamelCase				,   attention_mask=__UpperCamelCase   )
							_UpperCAmelCase					=    output[0]

							with tempfile.TemporaryDirectory() as tmpdirname:
											model.save_pretrained(__UpperCamelCase   )
											_UpperCAmelCase					=    FlaxVisionTextDualEncoderModel.from_pretrained(__UpperCamelCase   )

											_UpperCAmelCase					=    model(input_ids=__UpperCamelCase				,   pixel_values=__UpperCamelCase				,   attention_mask=__UpperCamelCase   )
											_UpperCAmelCase					=    after_output[0]
											_UpperCAmelCase					=    np.amax(np.abs(out_a - out_a   )   )
											self.assertLessEqual(__UpperCamelCase				,   1e-3   )



			def  lowercase__						(		self     :    Any				,   __UpperCamelCase     :    List[Any]				,   __UpperCamelCase     :    Tuple				,   __UpperCamelCase     :    Dict				,   __UpperCamelCase     :    str				,   __UpperCamelCase     :    List[str]=None				,   **__UpperCamelCase     :    int   )->Optional[Any]:
							_UpperCAmelCase		,						_UpperCAmelCase					=    self.get_vision_text_model(__UpperCamelCase				,   __UpperCamelCase   )
							_UpperCAmelCase					=    {'''vision_model''': vision_model, '''text_model''': text_model}
							_UpperCAmelCase					=    FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCamelCase   )

							_UpperCAmelCase					=    model(
							    input_ids=__UpperCamelCase				,   pixel_values=__UpperCamelCase				,   attention_mask=__UpperCamelCase				,   output_attentions=__UpperCamelCase   )

							_UpperCAmelCase					=    output.vision_model_output.attentions
							self.assertEqual(len(__UpperCamelCase   )				,   vision_config.num_hidden_layers   )

							# in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
							_UpperCAmelCase					=    to_atuple(vision_model.config.image_size   )
							_UpperCAmelCase					=    to_atuple(vision_model.config.patch_size   )
							_UpperCAmelCase					=    (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
							_UpperCAmelCase					=    num_patches + 1
							self.assertEqual(vision_attentions[0].shape[-3:]				,   (vision_config.num_attention_heads, seq_len, seq_len)   )

							_UpperCAmelCase					=    output.text_model_output.attentions
							self.assertEqual(len(__UpperCamelCase   )				,   text_config.num_hidden_layers   )

							self.assertEqual(
							    text_attentions[0].shape[-3:]				,   (text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])				,   )



			def  lowercase__						(		self     :    int				,   __UpperCamelCase     :    Dict				,   __UpperCamelCase     :    Dict				,   __UpperCamelCase     :    List[str]   )->Any:
							pt_model.to(__UpperCamelCase   )
							pt_model.eval()

							# prepare inputs
							_UpperCAmelCase					=    inputs_dict
							_UpperCAmelCase					=    {k: torch.tensor(v.tolist()   ) for k, v in flax_inputs.items()}

							with torch.no_grad():
											_UpperCAmelCase					=    pt_model(**__UpperCamelCase   ).to_tuple()

							_UpperCAmelCase					=    fx_model(**__UpperCamelCase   ).to_tuple()
							self.assertEqual(len(__UpperCamelCase   )				,   len(__UpperCamelCase   )				,   '''Output lengths differ between Flax and PyTorch'''   )
							for fx_output, pt_output in zip(fx_outputs[:4]				,   pt_outputs[:4]   ):
											self.assert_almost_equals(__UpperCamelCase				,   pt_output.numpy()				,   4e-2   )

							# PT -> Flax
							with tempfile.TemporaryDirectory() as tmpdirname:
											pt_model.save_pretrained(__UpperCamelCase   )
											_UpperCAmelCase					=    FlaxVisionTextDualEncoderModel.from_pretrained(__UpperCamelCase				,   from_pt=__UpperCamelCase   )

							_UpperCAmelCase					=    fx_model_loaded(**__UpperCamelCase   ).to_tuple()
							self.assertEqual(len(__UpperCamelCase   )				,   len(__UpperCamelCase   )				,   '''Output lengths differ between Flax and PyTorch'''   )
							for fx_output_loaded, pt_output in zip(fx_outputs_loaded[:4]				,   pt_outputs[:4]   ):
											self.assert_almost_equals(__UpperCamelCase				,   pt_output.numpy()				,   4e-2   )

							# Flax -> PT
							with tempfile.TemporaryDirectory() as tmpdirname:
											fx_model.save_pretrained(__UpperCamelCase   )
											_UpperCAmelCase					=    VisionTextDualEncoderModel.from_pretrained(__UpperCamelCase				,   from_flax=__UpperCamelCase   )

							pt_model_loaded.to(__UpperCamelCase   )
							pt_model_loaded.eval()

							with torch.no_grad():
											_UpperCAmelCase					=    pt_model_loaded(**__UpperCamelCase   ).to_tuple()

							self.assertEqual(len(__UpperCamelCase   )				,   len(__UpperCamelCase   )				,   '''Output lengths differ between Flax and PyTorch'''   )
							for fx_output, pt_output_loaded in zip(fx_outputs[:4]				,   pt_outputs_loaded[:4]   ):
											self.assert_almost_equals(__UpperCamelCase				,   pt_output_loaded.numpy()				,   4e-2   )



			def  lowercase__						(		self     :    List[Any]				,   __UpperCamelCase     :    Tuple				,   __UpperCamelCase     :    Optional[int]				,   __UpperCamelCase     :    Any   )->List[Any]:
							_UpperCAmelCase					=    VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCamelCase				,   __UpperCamelCase   )

							_UpperCAmelCase					=    VisionTextDualEncoderModel(__UpperCamelCase   )
							_UpperCAmelCase					=    FlaxVisionTextDualEncoderModel(__UpperCamelCase   )

							_UpperCAmelCase					=    convert_pytorch_state_dict_to_flax(pt_model.state_dict()				,   __UpperCamelCase   )
							_UpperCAmelCase					=    fx_state

							self.check_pt_flax_equivalence(__UpperCamelCase				,   __UpperCamelCase				,   __UpperCamelCase   )



			def  lowercase__						(		self     :    str				,   __UpperCamelCase     :    List[str]				,   __UpperCamelCase     :    List[str]				,   __UpperCamelCase     :    Optional[int]   )->Union[str, Any]:
							_UpperCAmelCase					=    VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCamelCase				,   __UpperCamelCase   )

							_UpperCAmelCase					=    VisionTextDualEncoderModel(__UpperCamelCase   )
							_UpperCAmelCase					=    FlaxVisionTextDualEncoderModel(__UpperCamelCase   )

							_UpperCAmelCase					=    load_flax_weights_in_pytorch_model(__UpperCamelCase				,   fx_model.params   )

							self.check_pt_flax_equivalence(__UpperCamelCase				,   __UpperCamelCase				,   __UpperCamelCase   )



			def  lowercase__						(		self     :    List[Any]   )->Optional[Any]:
							_UpperCAmelCase					=    self.prepare_config_and_inputs()
							self.check_model_from_pretrained_configs(**__UpperCamelCase   )



			def  lowercase__						(		self     :    Optional[Any]   )->List[Any]:
							_UpperCAmelCase					=    self.prepare_config_and_inputs()
							self.check_vision_text_dual_encoder_from_pretrained(**__UpperCamelCase   )



			def  lowercase__						(		self     :    int   )->List[str]:
							_UpperCAmelCase					=    self.prepare_config_and_inputs()
							self.check_save_load(**__UpperCamelCase   )



			def  lowercase__						(		self     :    Dict   )->Optional[int]:
							_UpperCAmelCase					=    self.prepare_config_and_inputs()
							self.check_vision_text_output_attention(**__UpperCamelCase   )



			@is_pt_flax_cross_test
			def  lowercase__						(		self     :    List[str]   )->Optional[int]:
							_UpperCAmelCase					=    self.prepare_config_and_inputs()
							_UpperCAmelCase					=    config_inputs_dict.pop('''vision_config'''   )
							_UpperCAmelCase					=    config_inputs_dict.pop('''text_config'''   )

							_UpperCAmelCase					=    config_inputs_dict

							self.check_equivalence_pt_to_flax(__UpperCamelCase				,   __UpperCamelCase				,   __UpperCamelCase   )
							self.check_equivalence_flax_to_pt(__UpperCamelCase				,   __UpperCamelCase				,   __UpperCamelCase   )



			@slow
			def  lowercase__						(		self     :    List[str]   )->List[Any]:
							_UpperCAmelCase		,						_UpperCAmelCase					=    self.get_pretrained_model_and_inputs()

							_UpperCAmelCase					=    model_a(**__UpperCamelCase   )
							_UpperCAmelCase					=    outputs[0]

							with tempfile.TemporaryDirectory() as tmp_dirname:
											model_a.save_pretrained(__UpperCamelCase   )
											_UpperCAmelCase					=    FlaxVisionTextDualEncoderModel.from_pretrained(__UpperCamelCase   )

											_UpperCAmelCase					=    model_a(**__UpperCamelCase   )
											_UpperCAmelCase					=    after_outputs[0]
											_UpperCAmelCase					=    np.amax(np.abs(out_a - out_a   )   )
											self.assertLessEqual(__UpperCamelCase				,   1e-5   )



@require_flax
class 			_a (					lowerCAmelCase			,							unittest.TestCase):
			"""simple docstring"""



			def  lowercase__						(		self     :    Tuple   )->List[Any]:
							_UpperCAmelCase					=    FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
							    '''hf-internal-testing/tiny-random-vit'''				,   '''hf-internal-testing/tiny-bert'''				,   vision_from_pt=__UpperCamelCase				,   text_from_pt=__UpperCamelCase				,   )
							_UpperCAmelCase					=    1_3
							_UpperCAmelCase					=    floats_tensor(
							    [
							        batch_size,
							        model.config.vision_config.num_channels,
							        model.config.vision_config.image_size,
							        model.config.vision_config.image_size,
							    ]   )
							_UpperCAmelCase					=    ids_tensor([batch_size, 4]				,   model.config.text_config.vocab_size   )
							_UpperCAmelCase					=    random_attention_mask([batch_size, 4]   )
							_UpperCAmelCase					=    {'''pixel_values''': pixel_values, '''input_ids''': input_ids, '''attention_mask''': attention_mask}

							return model, inputs



			def  lowercase__						(		self     :    List[Any]				,   __UpperCamelCase     :    str				,   __UpperCamelCase     :    Any   )->List[Any]:
							_UpperCAmelCase					=    FlaxViTModel(__UpperCamelCase   )
							_UpperCAmelCase					=    FlaxBertModel(__UpperCamelCase   )
							return vision_model, text_model



			def  lowercase__						(		self     :    List[str]   )->List[Any]:
							_UpperCAmelCase					=    FlaxViTModelTester(self   )
							_UpperCAmelCase					=    FlaxBertModelTester(self   )
							_UpperCAmelCase					=    vit_model_tester.prepare_config_and_inputs()
							_UpperCAmelCase					=    bert_model_tester.prepare_config_and_inputs()

							_UpperCAmelCase		,						_UpperCAmelCase					=    vision_config_and_inputs

							_UpperCAmelCase		,						_UpperCAmelCase		,						_UpperCAmelCase		,						_UpperCAmelCase					=    text_config_and_inputs

							# make sure that cross attention layers are added
							return {
							    "text_config": text_config,
							    "vision_config": vision_config,
							    "pixel_values": pixel_values,
							    "attention_mask": attention_mask,
							    "input_ids": input_ids,
							    "token_type_ids": token_type_ids,
							}



@require_torch
class 			_a (					lowerCAmelCase			,							unittest.TestCase):
			"""simple docstring"""



			def  lowercase__						(		self     :    Any   )->Optional[Any]:
							_UpperCAmelCase					=    FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
							    '''hf-internal-testing/tiny-random-clip'''				,   '''hf-internal-testing/tiny-bert'''				,   vision_from_pt=__UpperCamelCase				,   text_from_pt=__UpperCamelCase				,   )
							_UpperCAmelCase					=    1_3
							_UpperCAmelCase					=    floats_tensor(
							    [
							        batch_size,
							        model.config.vision_config.num_channels,
							        model.config.vision_config.image_size,
							        model.config.vision_config.image_size,
							    ]   )
							_UpperCAmelCase					=    ids_tensor([batch_size, 4]				,   model.config.text_config.vocab_size   )
							_UpperCAmelCase					=    random_attention_mask([batch_size, 4]   )
							_UpperCAmelCase					=    {'''pixel_values''': pixel_values, '''input_ids''': input_ids, '''attention_mask''': attention_mask}

							return model, inputs



			def  lowercase__						(		self     :    List[Any]				,   __UpperCamelCase     :    int				,   __UpperCamelCase     :    List[Any]   )->Optional[int]:
							_UpperCAmelCase					=    FlaxCLIPVisionModel(__UpperCamelCase   )
							_UpperCAmelCase					=    FlaxBertModel(__UpperCamelCase   )
							return vision_model, text_model



			def  lowercase__						(		self     :    Optional[Any]   )->int:
							_UpperCAmelCase					=    FlaxCLIPVisionModelTester(self   )
							_UpperCAmelCase					=    FlaxBertModelTester(self   )
							_UpperCAmelCase					=    clip_model_tester.prepare_config_and_inputs()
							_UpperCAmelCase					=    bert_model_tester.prepare_config_and_inputs()

							_UpperCAmelCase		,						_UpperCAmelCase					=    vision_config_and_inputs

							_UpperCAmelCase		,						_UpperCAmelCase		,						_UpperCAmelCase		,						_UpperCAmelCase					=    text_config_and_inputs

							# make sure that cross attention layers are added
							return {
							    "text_config": text_config,
							    "vision_config": vision_config,
							    "pixel_values": pixel_values,
							    "attention_mask": attention_mask,
							    "input_ids": input_ids,
							    "token_type_ids": token_type_ids,
							}



@require_flax
@require_vision
class 			_a (					unittest.TestCase):
			"""simple docstring"""



			@slow
			def  lowercase__						(		self     :    Any   )->List[Any]:
							_UpperCAmelCase					=    FlaxVisionTextDualEncoderModel.from_pretrained('''clip-italian/clip-italian'''				,   logit_scale_init_value=1.0   )
							_UpperCAmelCase					=    VisionTextDualEncoderProcessor.from_pretrained('''clip-italian/clip-italian'''   )

							_UpperCAmelCase					=    Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''   )
							_UpperCAmelCase					=    processor(
							    text=['''una foto di un gatto''', '''una foto di un cane''']				,   images=__UpperCamelCase				,   padding=__UpperCamelCase				,   return_tensors='''np'''   )

							_UpperCAmelCase					=    model(**__UpperCamelCase   )

							# verify the logits
							self.assertEqual(outputs.logits_per_image.shape				,   (inputs.pixel_values.shape[0], inputs.input_ids.shape[0])   )
							self.assertEqual(
							    outputs.logits_per_text.shape				,   (inputs.input_ids.shape[0], inputs.pixel_values.shape[0])				,   )

							_UpperCAmelCase					=    np.array([[1.2_2_8_4_7_2_7, 0.3_1_0_4_1_2_2]]   )

							self.assertTrue(np.allclose(outputs.logits_per_image				,   __UpperCamelCase				,   atol=1e-3   )   )









"""simple docstring"""


from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


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

__A   :							Optional[Any]										=   {
    "google/bit-50": "https://huggingface.co/google/bit-50/resolve/main/config.json",
}
class 			_a (					lowerCAmelCase			,							lowerCAmelCase):
			"""simple docstring"""
			UpperCamelCase__			      =	"""bit"""
			UpperCamelCase__			      =	["""preactivation""", """bottleneck"""]
			UpperCamelCase__			      =	["""SAME""", """VALID"""]



			def __init__(		self     :    List[Any]				,   __UpperCamelCase     :    Optional[int]=3				,   __UpperCamelCase     :    Optional[int]=6_4				,   __UpperCamelCase     :    List[Any]=[2_5_6, 5_1_2, 1_0_2_4, 2_0_4_8]				,   __UpperCamelCase     :    Any=[3, 4, 6, 3]				,   __UpperCamelCase     :    str="preactivation"				,   __UpperCamelCase     :    Union[str, Any]="relu"				,   __UpperCamelCase     :    str=None				,   __UpperCamelCase     :    Dict=3_2				,   __UpperCamelCase     :    List[str]=0.0				,   __UpperCamelCase     :    Any=False				,   __UpperCamelCase     :    Union[str, Any]=3_2				,   __UpperCamelCase     :    List[str]=1				,   __UpperCamelCase     :    List[str]=None				,   __UpperCamelCase     :    Optional[Any]=None				,   **__UpperCamelCase     :    List[str]				,   )->List[Any]:
							super().__init__(**__UpperCamelCase   )
							if layer_type not in self.layer_types:
											raise ValueError(F'layer_type={layer_type} is not one of {",".join(self.layer_types   )}'   )
							if global_padding is not None:
											if global_padding.upper() in self.supported_padding:
															_UpperCAmelCase					=    global_padding.upper()
											else:
															raise ValueError(F'Padding strategy {global_padding} not supported'   )
							_UpperCAmelCase					=    num_channels
							_UpperCAmelCase					=    embedding_size
							_UpperCAmelCase					=    hidden_sizes
							_UpperCAmelCase					=    depths
							_UpperCAmelCase					=    layer_type
							_UpperCAmelCase					=    hidden_act
							_UpperCAmelCase					=    global_padding
							_UpperCAmelCase					=    num_groups
							_UpperCAmelCase					=    drop_path_rate
							_UpperCAmelCase					=    embedding_dynamic_padding
							_UpperCAmelCase					=    output_stride
							_UpperCAmelCase					=    width_factor

							_UpperCAmelCase					=    ['''stem'''] + [F'stage{idx}' for idx in range(1				,   len(__UpperCamelCase   ) + 1   )]
							_UpperCAmelCase		,						_UpperCAmelCase					=    get_aligned_output_features_output_indices(
							    out_features=__UpperCamelCase				,   out_indices=__UpperCamelCase				,   stage_names=self.stage_names   )



