


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  (				SCREAMING_SNAKE_CASE	:				int     ):







 """simple docstring"""



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

@require_flax
class       __magic_name__	:


     def      UpperCAmelCase__	(      self							:     Optional[Any]							,	lowerCamelCase__							:     Optional[int]							,	lowerCamelCase__							:     Optional[int]			)  ->      Union[str, Any]:





      '''simple docstring'''






      pass


     def      UpperCAmelCase__	(      self							:     List[Any]			)  ->      Optional[int]:





      '''simple docstring'''






      pass


     def      UpperCAmelCase__	(      self							:     Any			)  ->      Optional[int]:





      '''simple docstring'''






      pass


     def      UpperCAmelCase__	(      self							:     List[Any]							,	lowerCamelCase__							:     np.ndarray							,	lowerCamelCase__							:     np.ndarray							,	lowerCamelCase__							:     float			)  ->      int:





      '''simple docstring'''






      UpperCamelCase__					:    Optional[Any]          =   np.abs((a - b)			).max()
      self.assertLessEqual(lowerCamelCase__							,	lowerCamelCase__							,	F"Difference between torch and flax is {diff} (>= {tol})."			)


     def      UpperCAmelCase__	(      self							:     Dict							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     Any							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     Optional[Any]=None							,	**lowerCamelCase__							:     Dict			)  ->      List[str]:





      '''simple docstring'''






      UpperCamelCase__					:    List[str]          =   VisionTextDualEncoderConfig.from_vision_text_configs(lowerCamelCase__							,	lowerCamelCase__			)

      UpperCamelCase__					:    List[str]          =   FlaxVisionTextDualEncoderModel(lowerCamelCase__			)

      UpperCamelCase__					:    Union[str, Any]          =   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							:     int							,	lowerCamelCase__							:     Any							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     Tuple=None							,	**lowerCamelCase__							:     Optional[Any]			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__       ,				UpperCamelCase__					:    Dict          =   self.get_vision_text_model(lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    Optional[int]          =   {'''vision_model''': vision_model, '''text_model''': text_model}
      UpperCamelCase__					:    Optional[int]          =   FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCamelCase__			)

      UpperCamelCase__					:    List[str]          =   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							:     Optional[int]							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     str							,	lowerCamelCase__							:     str							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     int=None							,	**lowerCamelCase__							:     List[str]			)  ->      Optional[int]:





      '''simple docstring'''






      UpperCamelCase__       ,				UpperCamelCase__					:    int          =   self.get_vision_text_model(lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    Union[str, Any]          =   {'''vision_model''': vision_model, '''text_model''': text_model}
      UpperCamelCase__					:    str          =   FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCamelCase__			)

      UpperCamelCase__					:    Dict          =   model(input_ids=lowerCamelCase__							,	pixel_values=lowerCamelCase__							,	attention_mask=lowerCamelCase__			)
      UpperCamelCase__					:    str          =   output[0]

      with tempfile.TemporaryDirectory() as tmpdirname:
       model.save_pretrained(lowerCamelCase__			)
       UpperCamelCase__					:    Dict          =   FlaxVisionTextDualEncoderModel.from_pretrained(lowerCamelCase__			)

       UpperCamelCase__					:    List[Any]          =   model(input_ids=lowerCamelCase__							,	pixel_values=lowerCamelCase__							,	attention_mask=lowerCamelCase__			)
       UpperCamelCase__					:    Tuple          =   after_output[0]
       UpperCamelCase__					:    Any          =   np.amax(np.abs(out_a - out_a			)			)
       self.assertLessEqual(lowerCamelCase__							,	1E-3			)


     def      UpperCAmelCase__	(      self							:     Dict							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     str							,	lowerCamelCase__							:     Optional[int]							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     Any=None							,	**lowerCamelCase__							:     List[Any]			)  ->      Any:





      '''simple docstring'''






      UpperCamelCase__       ,				UpperCamelCase__					:    List[Any]          =   self.get_vision_text_model(lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    int          =   {'''vision_model''': vision_model, '''text_model''': text_model}
      UpperCamelCase__					:    Optional[Any]          =   FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCamelCase__			)

      UpperCamelCase__					:    int          =   model(
          input_ids=lowerCamelCase__							,	pixel_values=lowerCamelCase__							,	attention_mask=lowerCamelCase__							,	output_attentions=lowerCamelCase__			)

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

      UpperCamelCase__					:    Union[str, Any]          =   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							:     Optional[Any]							,	lowerCamelCase__							:     Tuple							,	lowerCamelCase__							:     Tuple							,	lowerCamelCase__							:     Optional[int]			)  ->      Tuple:





      '''simple docstring'''






      pt_model.to(lowerCamelCase__			)
      pt_model.eval()

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

      with torch.no_grad():
       UpperCamelCase__					:    Optional[int]          =   pt_model(**lowerCamelCase__			).to_tuple()

      UpperCamelCase__					:    Optional[Any]          =   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__			)
       UpperCamelCase__					:    Optional[Any]          =   FlaxVisionTextDualEncoderModel.from_pretrained(lowerCamelCase__							,	from_pt=lowerCamelCase__			)

      UpperCamelCase__					:    Tuple          =   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__			)
       UpperCamelCase__					:    List[str]          =   VisionTextDualEncoderModel.from_pretrained(lowerCamelCase__							,	from_flax=lowerCamelCase__			)

      pt_model_loaded.to(lowerCamelCase__			)
      pt_model_loaded.eval()

      with torch.no_grad():
       UpperCamelCase__					:    int          =   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							:     Union[str, Any]							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     List[str]			)  ->      List[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    Optional[Any]          =   VisionTextDualEncoderConfig.from_vision_text_configs(lowerCamelCase__							,	lowerCamelCase__			)

      UpperCamelCase__					:    Tuple          =   VisionTextDualEncoderModel(lowerCamelCase__			)
      UpperCamelCase__					:    Optional[Any]          =   FlaxVisionTextDualEncoderModel(lowerCamelCase__			)

      UpperCamelCase__					:    List[Any]          =   convert_pytorch_state_dict_to_flax(pt_model.state_dict()							,	lowerCamelCase__			)
      UpperCamelCase__					:    Dict          =   fx_state

      self.check_pt_flax_equivalence(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__			)


     def      UpperCAmelCase__	(      self							:     Tuple							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     Optional[int]							,	lowerCamelCase__							:     List[str]			)  ->      List[str]:





      '''simple docstring'''






      UpperCamelCase__					:    Union[str, Any]          =   VisionTextDualEncoderConfig.from_vision_text_configs(lowerCamelCase__							,	lowerCamelCase__			)

      UpperCamelCase__					:    Tuple          =   VisionTextDualEncoderModel(lowerCamelCase__			)
      UpperCamelCase__					:    Any          =   FlaxVisionTextDualEncoderModel(lowerCamelCase__			)

      UpperCamelCase__					:    str          =   load_flax_weights_in_pytorch_model(lowerCamelCase__							,	fx_model.params			)

      self.check_pt_flax_equivalence(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__			)


     def      UpperCAmelCase__	(      self							:     Any			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   self.prepare_config_and_inputs()
      self.check_model_from_pretrained_configs(**lowerCamelCase__			)


     def      UpperCAmelCase__	(      self							:     int			)  ->      Any:





      '''simple docstring'''






      UpperCamelCase__					:    Tuple          =   self.prepare_config_and_inputs()
      self.check_vision_text_dual_encoder_from_pretrained(**lowerCamelCase__			)


     def      UpperCAmelCase__	(      self							:     Optional[int]			)  ->      Optional[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    str          =   self.prepare_config_and_inputs()
      self.check_save_load(**lowerCamelCase__			)


     def      UpperCAmelCase__	(      self							:     str			)  ->      Optional[int]:





      '''simple docstring'''






      UpperCamelCase__					:    Union[str, Any]          =   self.prepare_config_and_inputs()
      self.check_vision_text_output_attention(**lowerCamelCase__			)


     @is_pt_flax_cross_test
     def      UpperCAmelCase__	(      self							:     Optional[Any]			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   self.prepare_config_and_inputs()
      UpperCamelCase__					:    str          =   config_inputs_dict.pop('''vision_config'''			)
      UpperCamelCase__					:    Any          =   config_inputs_dict.pop('''text_config'''			)

      UpperCamelCase__					:    Optional[Any]          =   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[Any]			)  ->      str:





      '''simple docstring'''






      UpperCamelCase__       ,				UpperCamelCase__					:    Optional[int]          =   self.get_pretrained_model_and_inputs()

      UpperCamelCase__					:    Optional[int]          =   model_a(**lowerCamelCase__			)
      UpperCamelCase__					:    Optional[int]          =   outputs[0]

      with tempfile.TemporaryDirectory() as tmp_dirname:
       model_a.save_pretrained(lowerCamelCase__			)
       UpperCamelCase__					:    Tuple          =   FlaxVisionTextDualEncoderModel.from_pretrained(lowerCamelCase__			)

       UpperCamelCase__					:    List[str]          =   model_a(**lowerCamelCase__			)
       UpperCamelCase__					:    List[str]          =   after_outputs[0]
       UpperCamelCase__					:    Union[str, Any]          =   np.amax(np.abs(out_a - out_a			)			)
       self.assertLessEqual(lowerCamelCase__							,	1E-5			)



@require_flax
class       __magic_name__	(     __lowerCAmelCase					,				unittest.TestCase):


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





      '''simple docstring'''






      UpperCamelCase__					:    Optional[Any]          =   FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
          '''hf-internal-testing/tiny-random-vit'''							,	'''hf-internal-testing/tiny-bert'''							,	vision_from_pt=lowerCamelCase__							,	text_from_pt=lowerCamelCase__							,	)
      UpperCamelCase__					:    Dict          =   13
      UpperCamelCase__					:    Any          =   floats_tensor(
          [
              batch_size,
              model.config.vision_config.num_channels,
              model.config.vision_config.image_size,
              model.config.vision_config.image_size,
          ]			)
      UpperCamelCase__					:    int          =   ids_tensor([batch_size, 4]							,	model.config.text_config.vocab_size			)
      UpperCamelCase__					:    Dict          =   random_attention_mask([batch_size, 4]			)
      UpperCamelCase__					:    str          =   {'''pixel_values''': pixel_values, '''input_ids''': input_ids, '''attention_mask''': attention_mask}

      return model, inputs


     def      UpperCAmelCase__	(      self							:     List[str]							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     List[Any]			)  ->      Optional[int]:





      '''simple docstring'''






      UpperCamelCase__					:    List[str]          =   FlaxViTModel(lowerCamelCase__			)
      UpperCamelCase__					:    Optional[int]          =   FlaxBertModel(lowerCamelCase__			)
      return vision_model, text_model


     def      UpperCAmelCase__	(      self							:     int			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    List[str]          =   FlaxViTModelTester(self			)
      UpperCamelCase__					:    Optional[int]          =   FlaxBertModelTester(self			)
      UpperCamelCase__					:    str          =   vit_model_tester.prepare_config_and_inputs()
      UpperCamelCase__					:    Tuple          =   bert_model_tester.prepare_config_and_inputs()

      UpperCamelCase__       ,				UpperCamelCase__					:    Union[str, Any]          =   vision_config_and_inputs

      UpperCamelCase__       ,				UpperCamelCase__       ,				UpperCamelCase__       ,				UpperCamelCase__					:    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_torch
class       __magic_name__	(     __lowerCAmelCase					,				unittest.TestCase):


     def      UpperCAmelCase__	(      self							:     Optional[int]			)  ->      Optional[int]:





      '''simple docstring'''






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

      return model, inputs


     def      UpperCAmelCase__	(      self							:     List[str]							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     Any			)  ->      Optional[int]:





      '''simple docstring'''






      UpperCamelCase__					:    Any          =   FlaxCLIPVisionModel(lowerCamelCase__			)
      UpperCamelCase__					:    Dict          =   FlaxBertModel(lowerCamelCase__			)
      return vision_model, text_model


     def      UpperCAmelCase__	(      self							:     int			)  ->      str:





      '''simple docstring'''






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

      UpperCamelCase__       ,				UpperCamelCase__					:    List[str]          =   vision_config_and_inputs

      UpperCamelCase__       ,				UpperCamelCase__       ,				UpperCamelCase__       ,				UpperCamelCase__					:    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       __magic_name__	(     unittest.TestCase):


     @slow
     def      UpperCAmelCase__	(      self							:     str			)  ->      Optional[int]:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   FlaxVisionTextDualEncoderModel.from_pretrained('''clip-italian/clip-italian'''							,	logit_scale_init_value=1.0			)
      UpperCamelCase__					:    Optional[Any]          =   VisionTextDualEncoderProcessor.from_pretrained('''clip-italian/clip-italian'''			)

      UpperCamelCase__					:    Dict          =   Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''			)
      UpperCamelCase__					:    Optional[Any]          =   processor(
          text=['''una foto di un gatto''', '''una foto di un cane''']							,	images=lowerCamelCase__							,	padding=lowerCamelCase__							,	return_tensors='''np'''			)

      UpperCamelCase__					:    Optional[Any]          =   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])							,	)

      UpperCamelCase__					:    Union[str, Any]          =   np.array([[1.228_4727, 0.310_4122]]			)

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






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_torch_available,
)


__UpperCamelCase       :       Dict						       =				{
    "configuration_trocr": ["TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP", "TrOCRConfig"],
    "processing_trocr": ["TrOCRProcessor"],
}


try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase       :       Optional[int]						       =				[
           "TROCR_PRETRAINED_MODEL_ARCHIVE_LIST",
           "TrOCRForCausalLM",
           "TrOCRPreTrainedModel",
       ]


if TYPE_CHECKING:
       from .configuration_trocr import TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP, TrOCRConfig
       from .processing_trocr import TrOCRProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_trocr import TROCR_PRETRAINED_MODEL_ARCHIVE_LIST, TrOCRForCausalLM, TrOCRPreTrainedModel

else:
       import sys

       __UpperCamelCase       :       Tuple						       =				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


