






'''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 		__SCREAMING_SNAKE_CASE       (	_UpperCamelCase				):


   """simple docstring"""




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


@require_flax
class 							_UpperCAmelCase    :




    def 					lowerCAmelCase__					(      self	:							List[str]	,  a	:							Tuple	,  a	:							int      ):






       '''simple docstring'''






       pass




    def 					lowerCAmelCase__					(      self	:							int      ):






       '''simple docstring'''






       pass




    def 					lowerCAmelCase__					(      self	:							Optional[Any]      ):






       '''simple docstring'''






       pass




    def 					lowerCAmelCase__					(      self	:							int	,  a	:							np.ndarray	,  a	:							np.ndarray	,  a	:							float      ):






       '''simple docstring'''






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




    def 					lowerCAmelCase__					(      self	:							str	,  a	:							int	,  a	:							List[str]	,  a	:							Tuple	,  a	:							Any	,  a	:							Optional[Any]=None	,  **a	:							Tuple      ):






       '''simple docstring'''






       lowercase_						:     List[Any]							     =  VisionTextDualEncoderConfig.from_vision_text_configs(a	,  a      )

       lowercase_						:     Any							     =  FlaxVisionTextDualEncoderModel(a      )

       lowercase_						:     List[str]							     =  model(input_ids=a	,  pixel_values=a	,  attention_mask=a      )

       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 					lowerCAmelCase__					(      self	:							Dict	,  a	:							List[str]	,  a	:							Optional[int]	,  a	:							Optional[int]	,  a	:							str	,  a	:							str=None	,  **a	:							int      ):






       '''simple docstring'''






       lowercase_     ,		lowercase_						:     str							     =  self.get_vision_text_model(a	,  a      )
       lowercase_						:     List[str]							     =  {"vision_model": vision_model, "text_model": text_model}
       lowercase_						:     Optional[Any]							     =  FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**a      )

       lowercase_						:     int							     =  model(input_ids=a	,  pixel_values=a	,  attention_mask=a      )

       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 					lowerCAmelCase__					(      self	:							str	,  a	:							str	,  a	:							Dict	,  a	:							str	,  a	:							Dict	,  a	:							List[str]=None	,  **a	:							Optional[Any]      ):






       '''simple docstring'''






       lowercase_     ,		lowercase_						:     Dict							     =  self.get_vision_text_model(a	,  a      )
       lowercase_						:     List[str]							     =  {"vision_model": vision_model, "text_model": text_model}
       lowercase_						:     Optional[Any]							     =  FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**a      )

       lowercase_						:     Dict							     =  model(input_ids=a	,  pixel_values=a	,  attention_mask=a      )
       lowercase_						:     Optional[Any]							     =  output[0]

       with tempfile.TemporaryDirectory() as tmpdirname:
          model.save_pretrained(a      )
          lowercase_						:     Union[str, Any]							     =  FlaxVisionTextDualEncoderModel.from_pretrained(a      )

          lowercase_						:     List[str]							     =  model(input_ids=a	,  pixel_values=a	,  attention_mask=a      )
          lowercase_						:     int							     =  after_output[0]
          lowercase_						:     List[str]							     =  np.amax(np.abs(out_a - out_a      )      )
          self.assertLessEqual(a	,  1e-3      )




    def 					lowerCAmelCase__					(      self	:							str	,  a	:							Any	,  a	:							Any	,  a	:							Dict	,  a	:							Optional[Any]	,  a	:							Optional[int]=None	,  **a	:							Optional[int]      ):






       '''simple docstring'''






       lowercase_     ,		lowercase_						:     List[str]							     =  self.get_vision_text_model(a	,  a      )
       lowercase_						:     str							     =  {"vision_model": vision_model, "text_model": text_model}
       lowercase_						:     int							     =  FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**a      )

       lowercase_						:     Tuple							     =  model(
           input_ids=a	,  pixel_values=a	,  attention_mask=a	,  output_attentions=a      )

       lowercase_						:     Any							     =  output.vision_model_output.attentions
       self.assertEqual(len(a      )	,  vision_config.num_hidden_layers      )

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

       lowercase_						:     str							     =  output.text_model_output.attentions
       self.assertEqual(len(a      )	,  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 					lowerCAmelCase__					(      self	:							int	,  a	:							List[str]	,  a	:							Union[str, Any]	,  a	:							Any      ):






       '''simple docstring'''






       pt_model.to(a      )
       pt_model.eval()

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

       with torch.no_grad():
          lowercase_						:     Union[str, Any]							     =  pt_model(**a      ).to_tuple()

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

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

       lowercase_						:     Union[str, Any]							     =  fx_model_loaded(**a      ).to_tuple()
       self.assertEqual(len(a      )	,  len(a      )	,  "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(a	,  pt_output.numpy()	,  4e-2      )

       # Flax -> PT
       with tempfile.TemporaryDirectory() as tmpdirname:
          fx_model.save_pretrained(a      )
          lowercase_						:     List[Any]							     =  VisionTextDualEncoderModel.from_pretrained(a	,  from_flax=a      )

       pt_model_loaded.to(a      )
       pt_model_loaded.eval()

       with torch.no_grad():
          lowercase_						:     int							     =  pt_model_loaded(**a      ).to_tuple()

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




    def 					lowerCAmelCase__					(      self	:							int	,  a	:							Tuple	,  a	:							Dict	,  a	:							str      ):






       '''simple docstring'''






       lowercase_						:     Tuple							     =  VisionTextDualEncoderConfig.from_vision_text_configs(a	,  a      )

       lowercase_						:     int							     =  VisionTextDualEncoderModel(a      )
       lowercase_						:     Any							     =  FlaxVisionTextDualEncoderModel(a      )

       lowercase_						:     Any							     =  convert_pytorch_state_dict_to_flax(pt_model.state_dict()	,  a      )
       lowercase_						:     Optional[Any]							     =  fx_state

       self.check_pt_flax_equivalence(a	,  a	,  a      )




    def 					lowerCAmelCase__					(      self	:							Union[str, Any]	,  a	:							Any	,  a	:							List[str]	,  a	:							Optional[int]      ):






       '''simple docstring'''






       lowercase_						:     Any							     =  VisionTextDualEncoderConfig.from_vision_text_configs(a	,  a      )

       lowercase_						:     int							     =  VisionTextDualEncoderModel(a      )
       lowercase_						:     Optional[int]							     =  FlaxVisionTextDualEncoderModel(a      )

       lowercase_						:     Optional[int]							     =  load_flax_weights_in_pytorch_model(a	,  fx_model.params      )

       self.check_pt_flax_equivalence(a	,  a	,  a      )




    def 					lowerCAmelCase__					(      self	:							List[Any]      ):






       '''simple docstring'''






       lowercase_						:     str							     =  self.prepare_config_and_inputs()
       self.check_model_from_pretrained_configs(**a      )




    def 					lowerCAmelCase__					(      self	:							Tuple      ):






       '''simple docstring'''






       lowercase_						:     Union[str, Any]							     =  self.prepare_config_and_inputs()
       self.check_vision_text_dual_encoder_from_pretrained(**a      )




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






       '''simple docstring'''






       lowercase_						:     Optional[int]							     =  self.prepare_config_and_inputs()
       self.check_save_load(**a      )




    def 					lowerCAmelCase__					(      self	:							Dict      ):






       '''simple docstring'''






       lowercase_						:     Dict							     =  self.prepare_config_and_inputs()
       self.check_vision_text_output_attention(**a      )




    @is_pt_flax_cross_test
    def 					lowerCAmelCase__					(      self	:							int      ):






       '''simple docstring'''






       lowercase_						:     Any							     =  self.prepare_config_and_inputs()
       lowercase_						:     int							     =  config_inputs_dict.pop("vision_config"      )
       lowercase_						:     Any							     =  config_inputs_dict.pop("text_config"      )

       lowercase_						:     Optional[int]							     =  config_inputs_dict

       self.check_equivalence_pt_to_flax(a	,  a	,  a      )
       self.check_equivalence_flax_to_pt(a	,  a	,  a      )





    @slow
    def 					lowerCAmelCase__					(      self	:							Any      ):






       '''simple docstring'''






       lowercase_     ,		lowercase_						:     Tuple							     =  self.get_pretrained_model_and_inputs()

       lowercase_						:     Any							     =  model_a(**a      )
       lowercase_						:     str							     =  outputs[0]

       with tempfile.TemporaryDirectory() as tmp_dirname:
          model_a.save_pretrained(a      )
          lowercase_						:     str							     =  FlaxVisionTextDualEncoderModel.from_pretrained(a      )

          lowercase_						:     Optional[Any]							     =  model_a(**a      )
          lowercase_						:     Any							     =  after_outputs[0]
          lowercase_						:     List[str]							     =  np.amax(np.abs(out_a - out_a      )      )
          self.assertLessEqual(a	,  1e-5      )


@require_flax
class 							_UpperCAmelCase    (		snake_case					,       unittest.TestCase      ):




    def 					lowerCAmelCase__					(      self	:							List[str]      ):






       '''simple docstring'''






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

       return model, inputs




    def 					lowerCAmelCase__					(      self	:							Tuple	,  a	:							Optional[Any]	,  a	:							List[Any]      ):






       '''simple docstring'''






       lowercase_						:     Any							     =  FlaxViTModel(a      )
       lowercase_						:     Tuple							     =  FlaxBertModel(a      )
       return vision_model, text_model





    def 					lowerCAmelCase__					(      self	:							int      ):






       '''simple docstring'''






       lowercase_						:     Dict							     =  FlaxViTModelTester(self      )
       lowercase_						:     Union[str, Any]							     =  FlaxBertModelTester(self      )
       lowercase_						:     int							     =  vit_model_tester.prepare_config_and_inputs()
       lowercase_						:     Optional[Any]							     =  bert_model_tester.prepare_config_and_inputs()

       lowercase_     ,		lowercase_						:     Any							     =  vision_config_and_inputs

       lowercase_     ,		lowercase_     ,		lowercase_     ,		lowercase_						:     List[str]							     =  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 							_UpperCAmelCase    (		snake_case					,       unittest.TestCase      ):




    def 					lowerCAmelCase__					(      self	:							Optional[Any]      ):






       '''simple docstring'''






       lowercase_						:     Tuple							     =  FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
           "hf-internal-testing/tiny-random-clip"	,  "hf-internal-testing/tiny-bert"	,  vision_from_pt=a	,  text_from_pt=a	,  )
       lowercase_						:     Dict							     =  1_3
       lowercase_						:     Optional[Any]							     =  floats_tensor(
           [
               batch_size,
               model.config.vision_config.num_channels,
               model.config.vision_config.image_size,
               model.config.vision_config.image_size,
           ]      )
       lowercase_						:     List[str]							     =  ids_tensor([batch_size, 4]	,  model.config.text_config.vocab_size      )
       lowercase_						:     str							     =  random_attention_mask([batch_size, 4]      )
       lowercase_						:     str							     =  {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

       return model, inputs




    def 					lowerCAmelCase__					(      self	:							List[str]	,  a	:							str	,  a	:							int      ):






       '''simple docstring'''






       lowercase_						:     Dict							     =  FlaxCLIPVisionModel(a      )
       lowercase_						:     List[str]							     =  FlaxBertModel(a      )
       return vision_model, text_model





    def 					lowerCAmelCase__					(      self	:							Dict      ):






       '''simple docstring'''






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

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

       lowercase_     ,		lowercase_     ,		lowercase_     ,		lowercase_						:     List[Any]							     =  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 							_UpperCAmelCase    (		unittest.TestCase      ):




    @slow
    def 					lowerCAmelCase__					(      self	:							Tuple      ):






       '''simple docstring'''






       lowercase_						:     Optional[Any]							     =  FlaxVisionTextDualEncoderModel.from_pretrained("clip-italian/clip-italian"	,  logit_scale_init_value=1.0      )
       lowercase_						:     Dict							     =  VisionTextDualEncoderProcessor.from_pretrained("clip-italian/clip-italian"      )

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

       lowercase_						:     Optional[Any]							     =  model(**a      )

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

       lowercase_						:     int							     =  np.array([[1.228_4727, 0.310_4122]]      )

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












'''simple docstring'''
import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto import CONFIG_MAPPING


UpperCamelCase__			=      logging.get_logger(__name__)

UpperCamelCase__			=      {
    'ut/deta': 'https://huggingface.co/ut/deta/resolve/main/config.json',
}


class 							_UpperCAmelCase    (		snake_case      ):
    __lowerCamelCase: Union[str, Any]									=     'deta'
    __lowerCamelCase: str									=     {
        'hidden_size': 'd_model',
        'num_attention_heads': 'encoder_attention_heads',
    }




    def __init__(      self	:							List[str]	,  a	:							Union[str, Any]=None	,  a	:							Dict=9_0_0	,  a	:							str=2_0_4_8	,  a	:							Union[str, Any]=6	,  a	:							int=2_0_4_8	,  a	:							Optional[int]=8	,  a	:							Optional[Any]=6	,  a	:							Optional[Any]=1_0_2_4	,  a	:							Any=8	,  a	:							Dict=0.0	,  a	:							Optional[int]=True	,  a	:							Union[str, Any]="relu"	,  a	:							str=2_5_6	,  a	:							str=0.1	,  a	:							Any=0.0	,  a	:							Optional[Any]=0.0	,  a	:							str=0.02	,  a	:							int=1.0	,  a	:							Tuple=True	,  a	:							Optional[int]=False	,  a	:							str="sine"	,  a	:							List[Any]=5	,  a	:							str=4	,  a	:							Any=4	,  a	:							Optional[Any]=True	,  a	:							Union[str, Any]=3_0_0	,  a	:							List[Any]=True	,  a	:							Tuple=True	,  a	:							Any=1	,  a	:							int=5	,  a	:							Optional[int]=2	,  a	:							Tuple=1	,  a	:							List[str]=1	,  a	:							Tuple=5	,  a	:							Tuple=2	,  a	:							List[str]=0.1	,  a	:							Tuple=0.25	,  **a	:							Optional[int]	,  ):






       '''simple docstring'''






       if backbone_config is None:
          logger.info("`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone."      )
          lowercase_						:     Optional[int]							     =  CONFIG_MAPPING["resnet"](out_features=["stage2", "stage3", "stage4"]      )
       else:
          if isinstance(a	,  a      ):
             lowercase_						:     Any							     =  backbone_config.pop("model_type"      )
             lowercase_						:     str							     =  CONFIG_MAPPING[backbone_model_type]
             lowercase_						:     Any							     =  config_class.from_dict(a      )

       lowercase_						:     Optional[int]							     =  backbone_config
       lowercase_						:     Dict							     =  num_queries
       lowercase_						:     Tuple							     =  max_position_embeddings
       lowercase_						:     Union[str, Any]							     =  d_model
       lowercase_						:     Optional[Any]							     =  encoder_ffn_dim
       lowercase_						:     Any							     =  encoder_layers
       lowercase_						:     int							     =  encoder_attention_heads
       lowercase_						:     Tuple							     =  decoder_ffn_dim
       lowercase_						:     Dict							     =  decoder_layers
       lowercase_						:     Dict							     =  decoder_attention_heads
       lowercase_						:     List[Any]							     =  dropout
       lowercase_						:     Union[str, Any]							     =  attention_dropout
       lowercase_						:     List[str]							     =  activation_dropout
       lowercase_						:     Optional[Any]							     =  activation_function
       lowercase_						:     str							     =  init_std
       lowercase_						:     Dict							     =  init_xavier_std
       lowercase_						:     List[Any]							     =  encoder_layerdrop
       lowercase_						:     Optional[int]							     =  auxiliary_loss
       lowercase_						:     str							     =  position_embedding_type
       # deformable attributes
       lowercase_						:     Optional[int]							     =  num_feature_levels
       lowercase_						:     List[Any]							     =  encoder_n_points
       lowercase_						:     Optional[Any]							     =  decoder_n_points
       lowercase_						:     List[Any]							     =  two_stage
       lowercase_						:     Any							     =  two_stage_num_proposals
       lowercase_						:     List[str]							     =  with_box_refine
       lowercase_						:     Optional[int]							     =  assign_first_stage
       if two_stage is True and with_box_refine is False:
          raise ValueError("If two_stage is True, with_box_refine must be True."      )
       # Hungarian matcher
       lowercase_						:     int							     =  class_cost
       lowercase_						:     Optional[int]							     =  bbox_cost
       lowercase_						:     str							     =  giou_cost
       # Loss coefficients
       lowercase_						:     Optional[int]							     =  mask_loss_coefficient
       lowercase_						:     Union[str, Any]							     =  dice_loss_coefficient
       lowercase_						:     Optional[Any]							     =  bbox_loss_coefficient
       lowercase_						:     Optional[Any]							     =  giou_loss_coefficient
       lowercase_						:     Any							     =  eos_coefficient
       lowercase_						:     str							     =  focal_alpha
       super().__init__(is_encoder_decoder=a	,  **a      )




    @property
    def 					lowerCAmelCase__					(      self	:							Union[str, Any]      ):






       '''simple docstring'''






       return self.encoder_attention_heads




    @property
    def 					lowerCAmelCase__					(      self	:							Union[str, Any]      ):






       '''simple docstring'''






       return self.d_model





    def 					lowerCAmelCase__					(      self	:							Optional[Any]      ):






       '''simple docstring'''






       lowercase_						:     List[Any]							     =  copy.deepcopy(self.__dict__      )
       lowercase_						:     Union[str, Any]							     =  self.backbone_config.to_dict()
       lowercase_						:     str							     =  self.__class__.model_type
       return output




