




"""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       UpperCAmelCase__ (snake_case__							:      List[Any]					):
  """simple docstring"""

  if isinstance(_SCREAMING_SNAKE_CASE	,     collections.abc.Iterable					):
    return x
  return (x, x)






@require_flax
class       lowercase:







 '''simple docstring'''





 def 				UpperCamelCase_     (       self:	List[str],   a_:	List[Any],   a_:	Optional[Any]	):




   '''simple docstring'''



   pass


 def 				UpperCamelCase_     (       self:	Union[str, Any]	):




   '''simple docstring'''



   pass


 def 				UpperCamelCase_     (       self:	List[Any]	):




   '''simple docstring'''



   pass


 def 				UpperCamelCase_     (       self:	Dict,   a_:	np.ndarray,   a_:	np.ndarray,   a_:	float	):




   '''simple docstring'''



   _snake_case							:	Union[str, Any]			=						np.abs((a - b)	).max()
   self.assertLessEqual(__UpperCamelCase,   __UpperCamelCase,   f"Difference between torch and flax is {diff} (>= {tol})."	)


 def 				UpperCamelCase_     (       self:	Optional[int],   a_:	Dict,   a_:	List[str],   a_:	List[Any],   a_:	Tuple,   a_:	Union[str, Any]=None,   **a_:	Optional[int]	):




   '''simple docstring'''



   _snake_case							:	str			=						VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCamelCase,   __UpperCamelCase	)

   _snake_case							:	List[Any]			=						FlaxVisionTextDualEncoderModel(__UpperCamelCase	)

   _snake_case							:	List[str]			=						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 				UpperCamelCase_     (       self:	str,   a_:	List[str],   a_:	str,   a_:	Dict,   a_:	Dict,   a_:	List[str]=None,   **a_:	str	):




   '''simple docstring'''



   _snake_case	,      _snake_case							:	str			=						self.get_vision_text_model(__UpperCamelCase,   __UpperCamelCase	)
   _snake_case							:	str			=						{"""vision_model""": vision_model, """text_model""": text_model}
   _snake_case							:	Any			=						FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCamelCase	)

   _snake_case							:	int			=						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 				UpperCamelCase_     (       self:	str,   a_:	str,   a_:	Optional[int],   a_:	Union[str, Any],   a_:	Optional[int],   a_:	List[str]=None,   **a_:	List[Any]	):




   '''simple docstring'''



   _snake_case	,      _snake_case							:	Tuple			=						self.get_vision_text_model(__UpperCamelCase,   __UpperCamelCase	)
   _snake_case							:	Optional[int]			=						{"""vision_model""": vision_model, """text_model""": text_model}
   _snake_case							:	List[Any]			=						FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCamelCase	)

   _snake_case							:	Dict			=						model(input_ids=__UpperCamelCase,   pixel_values=__UpperCamelCase,   attention_mask=__UpperCamelCase	)
   _snake_case							:	List[Any]			=						output[0]

   with tempfile.TemporaryDirectory() as tmpdirname:
     model.save_pretrained(__UpperCamelCase	)
     _snake_case							:	List[Any]			=						FlaxVisionTextDualEncoderModel.from_pretrained(__UpperCamelCase	)

     _snake_case							:	Union[str, Any]			=						model(input_ids=__UpperCamelCase,   pixel_values=__UpperCamelCase,   attention_mask=__UpperCamelCase	)
     _snake_case							:	Tuple			=						after_output[0]
     _snake_case							:	Any			=						np.amax(np.abs(out_a - out_a	)	)
     self.assertLessEqual(__UpperCamelCase,   1E-3	)


 def 				UpperCamelCase_     (       self:	Any,   a_:	List[Any],   a_:	Tuple,   a_:	Dict,   a_:	str,   a_:	List[str]=None,   **a_:	int	):




   '''simple docstring'''



   _snake_case	,      _snake_case							:	List[Any]			=						self.get_vision_text_model(__UpperCamelCase,   __UpperCamelCase	)
   _snake_case							:	Any			=						{"""vision_model""": vision_model, """text_model""": text_model}
   _snake_case							:	List[Any]			=						FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCamelCase	)

   _snake_case							:	List[Any]			=						model(
       input_ids=__UpperCamelCase,   pixel_values=__UpperCamelCase,   attention_mask=__UpperCamelCase,   output_attentions=__UpperCamelCase	)

   _snake_case							:	Tuple			=						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)
   _snake_case							:	Any			=						to_atuple(vision_model.config.image_size	)
   _snake_case							:	Any			=						to_atuple(vision_model.config.patch_size	)
   _snake_case							:	str			=						(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
   _snake_case							:	List[Any]			=						num_patches + 1
   self.assertEqual(vision_attentions[0].shape[-3:],   (vision_config.num_attention_heads, seq_len, seq_len)	)

   _snake_case							:	Dict			=						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 				UpperCamelCase_     (       self:	int,   a_:	Dict,   a_:	Dict,   a_:	List[str]	):




   '''simple docstring'''



   pt_model.to(__UpperCamelCase	)
   pt_model.eval()

   # prepare inputs
   _snake_case							:	List[str]			=						inputs_dict
   _snake_case							:	Optional[int]			=						{k: torch.tensor(v.tolist()	) for k, v in flax_inputs.items()}

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

   _snake_case							:	str			=						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	)
     _snake_case							:	int			=						FlaxVisionTextDualEncoderModel.from_pretrained(__UpperCamelCase,   from_pt=__UpperCamelCase	)

   _snake_case							:	List[Any]			=						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	)
     _snake_case							:	Optional[Any]			=						VisionTextDualEncoderModel.from_pretrained(__UpperCamelCase,   from_flax=__UpperCamelCase	)

   pt_model_loaded.to(__UpperCamelCase	)
   pt_model_loaded.eval()

   with torch.no_grad():
     _snake_case							:	str			=						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 				UpperCamelCase_     (       self:	List[Any],   a_:	Tuple,   a_:	Optional[int],   a_:	Any	):




   '''simple docstring'''



   _snake_case							:	Optional[int]			=						VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCamelCase,   __UpperCamelCase	)

   _snake_case							:	List[str]			=						VisionTextDualEncoderModel(__UpperCamelCase	)
   _snake_case							:	Any			=						FlaxVisionTextDualEncoderModel(__UpperCamelCase	)

   _snake_case							:	Union[str, Any]			=						convert_pytorch_state_dict_to_flax(pt_model.state_dict(),   __UpperCamelCase	)
   _snake_case							:	int			=						fx_state

   self.check_pt_flax_equivalence(__UpperCamelCase,   __UpperCamelCase,   __UpperCamelCase	)


 def 				UpperCamelCase_     (       self:	str,   a_:	List[str],   a_:	List[str],   a_:	Optional[int]	):




   '''simple docstring'''



   _snake_case							:	List[str]			=						VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCamelCase,   __UpperCamelCase	)

   _snake_case							:	List[Any]			=						VisionTextDualEncoderModel(__UpperCamelCase	)
   _snake_case							:	int			=						FlaxVisionTextDualEncoderModel(__UpperCamelCase	)

   _snake_case							:	List[str]			=						load_flax_weights_in_pytorch_model(__UpperCamelCase,   fx_model.params	)

   self.check_pt_flax_equivalence(__UpperCamelCase,   __UpperCamelCase,   __UpperCamelCase	)


 def 				UpperCamelCase_     (       self:	List[Any]	):




   '''simple docstring'''



   _snake_case							:	int			=						self.prepare_config_and_inputs()
   self.check_model_from_pretrained_configs(**__UpperCamelCase	)


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   _snake_case							:	List[str]			=						self.prepare_config_and_inputs()
   self.check_vision_text_dual_encoder_from_pretrained(**__UpperCamelCase	)


 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   _snake_case							:	List[str]			=						self.prepare_config_and_inputs()
   self.check_save_load(**__UpperCamelCase	)


 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   _snake_case							:	List[str]			=						self.prepare_config_and_inputs()
   self.check_vision_text_output_attention(**__UpperCamelCase	)


 @is_pt_flax_cross_test
 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   _snake_case							:	str			=						self.prepare_config_and_inputs()
   _snake_case							:	Dict			=						config_inputs_dict.pop("""vision_config"""	)
   _snake_case							:	Union[str, Any]			=						config_inputs_dict.pop("""text_config"""	)

   _snake_case							:	Optional[Any]			=						config_inputs_dict

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


 @slow
 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   _snake_case	,      _snake_case							:	Tuple			=						self.get_pretrained_model_and_inputs()

   _snake_case							:	Tuple			=						model_a(**__UpperCamelCase	)
   _snake_case							:	List[Any]			=						outputs[0]

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

     _snake_case							:	str			=						model_a(**__UpperCamelCase	)
     _snake_case							:	int			=						after_outputs[0]
     _snake_case							:	List[str]			=						np.amax(np.abs(out_a - out_a	)	)
     self.assertLessEqual(__UpperCamelCase,   1E-5	)







@require_flax
class       lowercase(	__a ,    unittest.TestCase				):







 '''simple docstring'''





 def 				UpperCamelCase_     (       self:	Tuple	):




   '''simple docstring'''



   _snake_case							:	str			=						FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
       """hf-internal-testing/tiny-random-vit""",   """hf-internal-testing/tiny-bert""",   vision_from_pt=__UpperCamelCase,   text_from_pt=__UpperCamelCase,   )
   _snake_case							:	Union[str, Any]			=						13
   _snake_case							:	Any			=						floats_tensor(
       [
           batch_size,
           model.config.vision_config.num_channels,
           model.config.vision_config.image_size,
           model.config.vision_config.image_size,
       ]	)
   _snake_case							:	Tuple			=						ids_tensor([batch_size, 4],   model.config.text_config.vocab_size	)
   _snake_case							:	Union[str, Any]			=						random_attention_mask([batch_size, 4]	)
   _snake_case							:	Any			=						{"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

   return model, inputs


 def 				UpperCamelCase_     (       self:	List[Any],   a_:	str,   a_:	Any	):




   '''simple docstring'''



   _snake_case							:	int			=						FlaxViTModel(__UpperCamelCase	)
   _snake_case							:	Tuple			=						FlaxBertModel(__UpperCamelCase	)
   return vision_model, text_model


 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						FlaxViTModelTester(self	)
   _snake_case							:	Tuple			=						FlaxBertModelTester(self	)
   _snake_case							:	Tuple			=						vit_model_tester.prepare_config_and_inputs()
   _snake_case							:	List[Any]			=						bert_model_tester.prepare_config_and_inputs()

   _snake_case	,      _snake_case							:	Optional[Any]			=						vision_config_and_inputs

   _snake_case	,      _snake_case	,      _snake_case	,      _snake_case							:	int			=						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       lowercase(	__a ,    unittest.TestCase				):







 '''simple docstring'''





 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   _snake_case							:	Tuple			=						FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
       """hf-internal-testing/tiny-random-clip""",   """hf-internal-testing/tiny-bert""",   vision_from_pt=__UpperCamelCase,   text_from_pt=__UpperCamelCase,   )
   _snake_case							:	str			=						13
   _snake_case							:	Union[str, Any]			=						floats_tensor(
       [
           batch_size,
           model.config.vision_config.num_channels,
           model.config.vision_config.image_size,
           model.config.vision_config.image_size,
       ]	)
   _snake_case							:	Any			=						ids_tensor([batch_size, 4],   model.config.text_config.vocab_size	)
   _snake_case							:	Any			=						random_attention_mask([batch_size, 4]	)
   _snake_case							:	str			=						{"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

   return model, inputs


 def 				UpperCamelCase_     (       self:	List[Any],   a_:	int,   a_:	List[Any]	):




   '''simple docstring'''



   _snake_case							:	Tuple			=						FlaxCLIPVisionModel(__UpperCamelCase	)
   _snake_case							:	str			=						FlaxBertModel(__UpperCamelCase	)
   return vision_model, text_model


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						FlaxCLIPVisionModelTester(self	)
   _snake_case							:	Any			=						FlaxBertModelTester(self	)
   _snake_case							:	str			=						clip_model_tester.prepare_config_and_inputs()
   _snake_case							:	Union[str, Any]			=						bert_model_tester.prepare_config_and_inputs()

   _snake_case	,      _snake_case							:	Dict			=						vision_config_and_inputs

   _snake_case	,      _snake_case	,      _snake_case	,      _snake_case							:	Tuple			=						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       lowercase(	unittest.TestCase				):







 '''simple docstring'''





 @slow
 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   _snake_case							:	str			=						FlaxVisionTextDualEncoderModel.from_pretrained("""clip-italian/clip-italian""",   logit_scale_init_value=1.0	)
   _snake_case							:	Union[str, Any]			=						VisionTextDualEncoderProcessor.from_pretrained("""clip-italian/clip-italian"""	)

   _snake_case							:	Optional[Any]			=						Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""	)
   _snake_case							:	int			=						processor(
       text=["""una foto di un gatto""", """una foto di un cane"""],   images=__UpperCamelCase,   padding=__UpperCamelCase,   return_tensors="""np"""	)

   _snake_case							:	List[str]			=						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]),   )

   _snake_case							:	Any			=						np.array([[1.2_284_727, 0.3_104_122]]	)

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








"""simple docstring"""


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


A_      			=       {
    '''configuration_owlvit''': [
        '''OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''OwlViTConfig''',
        '''OwlViTOnnxConfig''',
        '''OwlViTTextConfig''',
        '''OwlViTVisionConfig''',
    ],
    '''processing_owlvit''': ['''OwlViTProcessor'''],
}


try:
     if not is_vision_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     A_      			=       ['''OwlViTFeatureExtractor''']
     A_      			=       ['''OwlViTImageProcessor''']

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     A_      			=       [
         '''OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''OwlViTModel''',
         '''OwlViTPreTrainedModel''',
         '''OwlViTTextModel''',
         '''OwlViTVisionModel''',
         '''OwlViTForObjectDetection''',
     ]

if TYPE_CHECKING:
     from .configuration_owlvit import (
         OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP,
         OwlViTConfig,
         OwlViTOnnxConfig,
         OwlViTTextConfig,
         OwlViTVisionConfig,
     )
     from .processing_owlvit import OwlViTProcessor

     try:
          if not is_vision_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .feature_extraction_owlvit import OwlViTFeatureExtractor
          from .image_processing_owlvit import OwlViTImageProcessor

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_owlvit import (
              OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
              OwlViTForObjectDetection,
              OwlViTModel,
              OwlViTPreTrainedModel,
              OwlViTTextModel,
              OwlViTVisionModel,
          )

else:
     import sys

     A_      			=       _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)


