

"""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   _A					(	_a	:						Tuple				):



      """simple docstring"""






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





@require_flax
class lowerCamelCase__   :



    '''simple docstring'''




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Tuple:
          pass




    def UpperCamelCase__		( self    )    ->      Tuple:
          pass




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




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Tuple:
          A				  =		np.abs((a - b)    ).max()
          self.assertLessEqual(lowerCamelCase_ ,lowerCamelCase_ ,f'Difference between torch and flax is {diff} (>= {tol}).'    )




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_=None ,**lowerCamelCase_    )    ->      Any:
          A				  =		VisionTextDualEncoderConfig.from_vision_text_configs(lowerCamelCase_ ,lowerCamelCase_    )

          A				  =		FlaxVisionTextDualEncoderModel(lowerCamelCase_    )

          A				  =		model(input_ids=lowerCamelCase_ ,pixel_values=lowerCamelCase_ ,attention_mask=lowerCamelCase_    )

          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 ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_=None ,**lowerCamelCase_    )    ->      int:
          A  ,			A				  =		self.get_vision_text_model(lowerCamelCase_ ,lowerCamelCase_    )
          A				  =		{"""vision_model""": vision_model, """text_model""": text_model}
          A				  =		FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCamelCase_    )

          A				  =		model(input_ids=lowerCamelCase_ ,pixel_values=lowerCamelCase_ ,attention_mask=lowerCamelCase_    )

          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 ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_=None ,**lowerCamelCase_    )    ->      Union[str, Any]:
          A  ,			A				  =		self.get_vision_text_model(lowerCamelCase_ ,lowerCamelCase_    )
          A				  =		{"""vision_model""": vision_model, """text_model""": text_model}
          A				  =		FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCamelCase_    )

          A				  =		model(input_ids=lowerCamelCase_ ,pixel_values=lowerCamelCase_ ,attention_mask=lowerCamelCase_    )
          A				  =		output[0]

          with tempfile.TemporaryDirectory() as tmpdirname:
                model.save_pretrained(lowerCamelCase_    )
                A				  =		FlaxVisionTextDualEncoderModel.from_pretrained(lowerCamelCase_    )

                A				  =		model(input_ids=lowerCamelCase_ ,pixel_values=lowerCamelCase_ ,attention_mask=lowerCamelCase_    )
                A				  =		after_output[0]
                A				  =		np.amax(np.abs(out_a - out_a    )    )
                self.assertLessEqual(lowerCamelCase_ ,1E-3    )




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_=None ,**lowerCamelCase_    )    ->      Any:
          A  ,			A				  =		self.get_vision_text_model(lowerCamelCase_ ,lowerCamelCase_    )
          A				  =		{"""vision_model""": vision_model, """text_model""": text_model}
          A				  =		FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCamelCase_    )

          A				  =		model(
              input_ids=lowerCamelCase_ ,pixel_values=lowerCamelCase_ ,attention_mask=lowerCamelCase_ ,output_attentions=lowerCamelCase_    )

          A				  =		output.vision_model_output.attentions
          self.assertEqual(len(lowerCamelCase_    ) ,vision_config.num_hidden_layers    )

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

          A				  =		output.text_model_output.attentions
          self.assertEqual(len(lowerCamelCase_    ) ,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 ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      str:
          pt_model.to(lowerCamelCase_    )
          pt_model.eval()

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

          with torch.no_grad():
                A				  =		pt_model(**lowerCamelCase_    ).to_tuple()

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

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

          A				  =		fx_model_loaded(**lowerCamelCase_    ).to_tuple()
          self.assertEqual(len(lowerCamelCase_    ) ,len(lowerCamelCase_    ) ,"""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(lowerCamelCase_ ,pt_output.numpy() ,4E-2    )

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

          pt_model_loaded.to(lowerCamelCase_    )
          pt_model_loaded.eval()

          with torch.no_grad():
                A				  =		pt_model_loaded(**lowerCamelCase_    ).to_tuple()

          self.assertEqual(len(lowerCamelCase_    ) ,len(lowerCamelCase_    ) ,"""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(lowerCamelCase_ ,pt_output_loaded.numpy() ,4E-2    )




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Optional[Any]:
          A				  =		VisionTextDualEncoderConfig.from_vision_text_configs(lowerCamelCase_ ,lowerCamelCase_    )

          A				  =		VisionTextDualEncoderModel(lowerCamelCase_    )
          A				  =		FlaxVisionTextDualEncoderModel(lowerCamelCase_    )

          A				  =		convert_pytorch_state_dict_to_flax(pt_model.state_dict() ,lowerCamelCase_    )
          A				  =		fx_state

          self.check_pt_flax_equivalence(lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      List[str]:
          A				  =		VisionTextDualEncoderConfig.from_vision_text_configs(lowerCamelCase_ ,lowerCamelCase_    )

          A				  =		VisionTextDualEncoderModel(lowerCamelCase_    )
          A				  =		FlaxVisionTextDualEncoderModel(lowerCamelCase_    )

          A				  =		load_flax_weights_in_pytorch_model(lowerCamelCase_ ,fx_model.params    )

          self.check_pt_flax_equivalence(lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      List[Any]:
          A				  =		self.prepare_config_and_inputs()
          self.check_model_from_pretrained_configs(**lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      List[Any]:
          A				  =		self.prepare_config_and_inputs()
          self.check_vision_text_dual_encoder_from_pretrained(**lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      Tuple:
          A				  =		self.prepare_config_and_inputs()
          self.check_save_load(**lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      List[Any]:
          A				  =		self.prepare_config_and_inputs()
          self.check_vision_text_output_attention(**lowerCamelCase_    )




    @is_pt_flax_cross_test
    def UpperCamelCase__		( self    )    ->      Optional[Any]:
          A				  =		self.prepare_config_and_inputs()
          A				  =		config_inputs_dict.pop("""vision_config"""    )
          A				  =		config_inputs_dict.pop("""text_config"""    )

          A				  =		config_inputs_dict

          self.check_equivalence_pt_to_flax(lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )
          self.check_equivalence_flax_to_pt(lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )




    @slow
    def UpperCamelCase__		( self    )    ->      Optional[int]:
          A  ,			A				  =		self.get_pretrained_model_and_inputs()

          A				  =		model_a(**lowerCamelCase_    )
          A				  =		outputs[0]

          with tempfile.TemporaryDirectory() as tmp_dirname:
                model_a.save_pretrained(lowerCamelCase_    )
                A				  =		FlaxVisionTextDualEncoderModel.from_pretrained(lowerCamelCase_    )

                A				  =		model_a(**lowerCamelCase_    )
                A				  =		after_outputs[0]
                A				  =		np.amax(np.abs(out_a - out_a    )    )
                self.assertLessEqual(lowerCamelCase_ ,1E-5    )





@require_flax
class lowerCamelCase__   (				SCREAMING_SNAKE_CASE     ,   unittest.TestCase    ):



    '''simple docstring'''




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

          return model, inputs




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Any:
          A				  =		FlaxViTModel(lowerCamelCase_    )
          A				  =		FlaxBertModel(lowerCamelCase_    )
          return vision_model, text_model




    def UpperCamelCase__		( self    )    ->      int:
          A				  =		FlaxViTModelTester(self    )
          A				  =		FlaxBertModelTester(self    )
          A				  =		vit_model_tester.prepare_config_and_inputs()
          A				  =		bert_model_tester.prepare_config_and_inputs()

          A  ,			A				  =		vision_config_and_inputs

          A  ,			A  ,			A  ,			A				  =		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 lowerCamelCase__   (				SCREAMING_SNAKE_CASE     ,   unittest.TestCase    ):



    '''simple docstring'''




    def UpperCamelCase__		( self    )    ->      List[str]:
          A				  =		FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
              """hf-internal-testing/tiny-random-clip""" ,"""hf-internal-testing/tiny-bert""" ,vision_from_pt=lowerCamelCase_ ,text_from_pt=lowerCamelCase_ ,)
          A				  =		1_3
          A				  =		floats_tensor(
              [
                  batch_size,
                  model.config.vision_config.num_channels,
                  model.config.vision_config.image_size,
                  model.config.vision_config.image_size,
              ]    )
          A				  =		ids_tensor([batch_size, 4] ,model.config.text_config.vocab_size    )
          A				  =		random_attention_mask([batch_size, 4]    )
          A				  =		{"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

          return model, inputs




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_    )    ->      List[str]:
          A				  =		FlaxCLIPVisionModel(lowerCamelCase_    )
          A				  =		FlaxBertModel(lowerCamelCase_    )
          return vision_model, text_model




    def UpperCamelCase__		( self    )    ->      int:
          A				  =		FlaxCLIPVisionModelTester(self    )
          A				  =		FlaxBertModelTester(self    )
          A				  =		clip_model_tester.prepare_config_and_inputs()
          A				  =		bert_model_tester.prepare_config_and_inputs()

          A  ,			A				  =		vision_config_and_inputs

          A  ,			A  ,			A  ,			A				  =		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 lowerCamelCase__   (				unittest.TestCase    ):



    '''simple docstring'''




    @slow
    def UpperCamelCase__		( self    )    ->      str:
          A				  =		FlaxVisionTextDualEncoderModel.from_pretrained("""clip-italian/clip-italian""" ,logit_scale_init_value=1.0    )
          A				  =		VisionTextDualEncoderProcessor.from_pretrained("""clip-italian/clip-italian"""    )

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

          A				  =		model(**lowerCamelCase_    )

          # 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]) ,)

          A				  =		np.array([[1.2_28_47_27, 0.3_10_41_22]]    )

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







"""simple docstring"""


from typing import Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
  from PIL import Image

  from ..image_utils import load_image

if is_torch_available():
  from ..models.auto.modeling_auto import MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING

UpperCAmelCase													=logging.get_logger(__name__)





@add_end_docstrings(SCREAMING_SNAKE_CASE    )
class lowerCamelCase__   (				SCREAMING_SNAKE_CASE    ):



    '''simple docstring'''




    def __init__( self ,*lowerCamelCase_ ,**lowerCamelCase_    )    ->      str:
          super().__init__(*lowerCamelCase_ ,**lowerCamelCase_    )
          self.check_model_type(lowerCamelCase_    )




    def UpperCamelCase__		( self ,lowerCamelCase_=None ,lowerCamelCase_=None ,lowerCamelCase_=None ,**lowerCamelCase_    )    ->      List[Any]:
          A  ,			A				  =		{}, {}
          if padding is not None:
                A				  =		padding
          if truncation is not None:
                A				  =		truncation
          if top_k is not None:
                A				  =		top_k
          return preprocess_params, {}, postprocess_params




    def __call__( self ,lowerCamelCase_ ,lowerCamelCase_ = None ,**lowerCamelCase_    )    ->      Dict:
          if isinstance(lowerCamelCase_ ,(Image.Image, str)    ) and isinstance(lowerCamelCase_ ,lowerCamelCase_    ):
                A				  =		{"""image""": image, """question""": question}
          else:
                A				  =		image
          A				  =		super().__call__(lowerCamelCase_ ,**lowerCamelCase_    )
          return results




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_=False ,lowerCamelCase_=False    )    ->      List[Any]:
          A				  =		load_image(inputs["""image"""]    )
          A				  =		self.tokenizer(
              inputs["""question"""] ,return_tensors=self.framework ,padding=lowerCamelCase_ ,truncation=lowerCamelCase_    )
          A				  =		self.image_processor(images=lowerCamelCase_ ,return_tensors=self.framework    )
          model_inputs.update(lowerCamelCase_    )
          return model_inputs




    def UpperCamelCase__		( self ,lowerCamelCase_    )    ->      Union[str, Any]:
          A				  =		self.model(**lowerCamelCase_    )
          return model_outputs




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_=5    )    ->      int:
          if top_k > self.model.config.num_labels:
                A				  =		self.model.config.num_labels

          if self.framework == "pt":
                A				  =		model_outputs.logits.sigmoid()[0]
                A  ,			A				  =		probs.topk(lowerCamelCase_    )
          else:
                raise ValueError(f'Unsupported framework: {self.framework}'    )

          A				  =		scores.tolist()
          A				  =		ids.tolist()
          return [{"score": score, "answer": self.model.config.idalabel[_id]} for score, _id in zip(lowerCamelCase_ ,lowerCamelCase_    )]




