codingutf8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch visiontextdualencoder model import collections import tempfile import unittest import numpy as np from transformers testingutils import isptflaxcrosstest requireflax requiretorch requirevision slow torchdevice from transformers utils import isflaxavailable istorchavailable isvisionavailable from testmodelingflaxcommon import floatstensor idstensor randomattentionmask from bert testmodelingflaxbert import flaxbertmodeltester from clip testmodelingflaxclip import flaxclipvisionmodeltester from vit testmodelingflaxvit import flaxvitmodeltester if isflaxavailable from transformers import flaxbertmodel flaxclipvisionmodel flaxvisiontextdualencodermodel flaxvitmodel visiontextdualencoderconfig visiontextdualencoderprocessor from transformers modelingflaxpytorchutils import convertpytorchstatedicttoflax loadflaxweightsinpytorchmodel if istorchavailable import torch from transformers import visiontextdualencodermodel if isvisionavailable from pil import image inspired by https github comrwightmanpytorchimagemodelsblobb9bd960a032c75ca6b808ddeed76bee5f3ed4972timmmodelslayershelpers py from pytorch internals def to2tuplex if isinstancex collections abc iterable return x return x x requireflax class visiontextdualencodermixin def getvisiontextmodelself config textconfig pass def prepareconfigandinputsself pass def getpretrainedmodelandinputsself pass def assertalmostequalsself a np ndarray b np ndarray tol float diff np absa b max self assertlessequaldiff tol fdifference between torch and flax is diff tol def checkmodelfrompretrainedconfigs self textconfig inputids attentionmask visionconfig pixelvaluesnone kwargs config visiontextdualencoderconfig fromvisiontextconfigsvisionconfig textconfig model flaxvisiontextdualencodermodelconfig output modelinputidsinputids pixelvaluespixelvalues attentionmaskattentionmask self assertequaloutputtextembeds shape inputids shape0 config projectiondim self assertequaloutputimageembeds shape pixelvalues shape0 config projectiondim def checkvisiontextdualencoderfrompretrained self textconfig inputids attentionmask visionconfig pixelvaluesnone kwargs visionmodel textmodel self getvisiontextmodelvisionconfig textconfig kwargs visionmodel visionmodel textmodel textmodel model flaxvisiontextdualencodermodel fromvisiontextpretrainedkwargs output modelinputidsinputids pixelvaluespixelvalues attentionmaskattentionmask self assertequaloutputtextembeds shape inputids shape0 model config projectiondim self assertequaloutputimageembeds shape pixelvalues shape0 model config projectiondim def checksaveloadself textconfig inputids attentionmask visionconfig pixelvaluesnone kwargs visionmodel textmodel self getvisiontextmodelvisionconfig textconfig kwargs visionmodel visionmodel textmodel textmodel model flaxvisiontextdualencodermodel fromvisiontextpretrainedkwargs output modelinputidsinputids pixelvaluespixelvalues attentionmaskattentionmask out1 output0 with tempfile temporarydirectory as tmpdirname model savepretrainedtmpdirname model flaxvisiontextdualencodermodel frompretrainedtmpdirname afteroutput modelinputidsinputids pixelvaluespixelvalues attentionmaskattentionmask out2 afteroutput0 maxdiff np amaxnp absout2 out1 self assertlessequalmaxdiff 1e3 def checkvisiontextoutputattention self textconfig inputids attentionmask visionconfig pixelvaluesnone kwargs visionmodel textmodel self getvisiontextmodelvisionconfig textconfig kwargs visionmodel visionmodel textmodel textmodel model flaxvisiontextdualencodermodel fromvisiontextpretrainedkwargs output model inputidsinputids pixelvaluespixelvalues attentionmaskattentionmask outputattentionstrue visionattentions output visionmodeloutput attentions self assertequallenvisionattentions visionconfig numhiddenlayers in vit the seqlen equals the number of patches 1 we add 1 for the cls token imagesize to2tuplevisionmodel config imagesize patchsize to2tuplevisionmodel config patchsize numpatches imagesize1 patchsize1 imagesize0 patchsize0 seqlen numpatches 1 self assertequalvisionattentions0 shape3 visionconfig numattentionheads seqlen seqlen textattentions output textmodeloutput attentions self assertequallentextattentions textconfig numhiddenlayers self assertequal textattentions0 shape3 textconfig numattentionheads inputids shape1 inputids shape1 def checkptflaxequivalenceself ptmodel fxmodel inputsdict ptmodel totorchdevice ptmodel eval prepare inputs flaxinputs inputsdict ptinputs k torch tensorv tolist for k v in flaxinputs items with torch nograd ptoutputs ptmodelptinputs totuple fxoutputs fxmodelinputsdict totuple self assertequallenfxoutputs lenptoutputs output lengths differ between flax and pytorch for fxoutput ptoutput in zipfxoutputs 4 ptoutputs 4 self assertalmostequalsfxoutput ptoutput numpy 4e2 pt flax with tempfile temporarydirectory as tmpdirname ptmodel savepretrainedtmpdirname fxmodelloaded flaxvisiontextdualencodermodel frompretrainedtmpdirname frompttrue fxoutputsloaded fxmodelloadedinputsdict totuple self assertequallenfxoutputsloaded lenptoutputs output lengths differ between flax and pytorch for fxoutputloaded ptoutput in zipfxoutputsloaded 4 ptoutputs 4 self assertalmostequalsfxoutputloaded ptoutput numpy 4e2 flax pt with tempfile temporarydirectory as tmpdirname fxmodel savepretrainedtmpdirname ptmodelloaded visiontextdualencodermodel frompretrainedtmpdirname fromflaxtrue ptmodelloaded totorchdevice ptmodelloaded eval with torch nograd ptoutputsloaded ptmodelloadedptinputs totuple self assertequallenfxoutputs lenptoutputsloaded output lengths differ between flax and pytorch for fxoutput ptoutputloaded in zipfxoutputs 4 ptoutputsloaded 4 self assertalmostequalsfxoutput ptoutputloaded numpy 4e2 def checkequivalencepttoflaxself visionconfig textconfig inputsdict config visiontextdualencoderconfig fromvisiontextconfigsvisionconfig textconfig ptmodel visiontextdualencodermodelconfig fxmodel flaxvisiontextdualencodermodelconfig fxstate convertpytorchstatedicttoflaxptmodel statedict fxmodel fxmodel params fxstate self checkptflaxequivalenceptmodel fxmodel inputsdict def checkequivalenceflaxtoptself visionconfig textconfig inputsdict config visiontextdualencoderconfig fromvisiontextconfigsvisionconfig textconfig ptmodel visiontextdualencodermodelconfig fxmodel flaxvisiontextdualencodermodelconfig ptmodel loadflaxweightsinpytorchmodelptmodel fxmodel params self checkptflaxequivalenceptmodel fxmodel inputsdict def testmodelfrompretrainedconfigsself inputsdict self prepareconfigandinputs self checkmodelfrompretrainedconfigsinputsdict def testvisiontextdualencoderfrompretrainedself inputsdict self prepareconfigandinputs self checkvisiontextdualencoderfrompretrainedinputsdict def testsaveloadself inputsdict self prepareconfigandinputs self checksaveloadinputsdict def testvisiontextoutputattentionself inputsdict self prepareconfigandinputs self checkvisiontextoutputattentioninputsdict isptflaxcrosstest def testptflaxequivalenceself configinputsdict self prepareconfigandinputs visionconfig configinputsdict popvisionconfig textconfig configinputsdict poptextconfig inputsdict configinputsdict self checkequivalencepttoflaxvisionconfig textconfig inputsdict self checkequivalenceflaxtoptvisionconfig textconfig inputsdict slow def testrealmodelsaveloadfrompretrainedself model2 inputs self getpretrainedmodelandinputs outputs model2inputs out2 outputs0 with tempfile temporarydirectory as tmpdirname model2 savepretrainedtmpdirname model1 flaxvisiontextdualencodermodel frompretrainedtmpdirname afteroutputs model1inputs out1 afteroutputs0 maxdiff np amaxnp absout1 out2 self assertlessequalmaxdiff 1e5 requireflax class flaxvitbertmodeltestvisiontextdualencodermixin unittest testcase def getpretrainedmodelandinputsself model flaxvisiontextdualencodermodel fromvisiontextpretrained hfinternaltestingtinyrandomvit hfinternaltestingtinybert visionfrompttrue textfrompttrue batchsize 13 pixelvalues floatstensor batchsize model config visionconfig numchannels model config visionconfig imagesize model config visionconfig imagesize inputids idstensorbatchsize 4 model config textconfig vocabsize attentionmask randomattentionmaskbatchsize 4 inputs pixelvalues pixelvalues inputids inputids attentionmask attentionmask return model inputs def getvisiontextmodelself visionconfig textconfig visionmodel flaxvitmodelvisionconfig textmodel flaxbertmodeltextconfig return visionmodel textmodel def prepareconfigandinputsself vitmodeltester flaxvitmodeltesterself bertmodeltester flaxbertmodeltesterself visionconfigandinputs vitmodeltester prepareconfigandinputs textconfigandinputs bertmodeltester prepareconfigandinputs visionconfig pixelvalues visionconfigandinputs textconfig inputids tokentypeids attentionmask textconfigandinputs make sure that cross attention layers are added return textconfig textconfig visionconfig visionconfig pixelvalues pixelvalues attentionmask attentionmask inputids inputids tokentypeids tokentypeids requiretorch class flaxclipvisionbertmodeltestvisiontextdualencodermixin unittest testcase def getpretrainedmodelandinputsself model flaxvisiontextdualencodermodel fromvisiontextpretrained hfinternaltestingtinyrandomclip hfinternaltestingtinybert visionfrompttrue textfrompttrue batchsize 13 pixelvalues floatstensor batchsize model config visionconfig numchannels model config visionconfig imagesize model config visionconfig imagesize inputids idstensorbatchsize 4 model config textconfig vocabsize attentionmask randomattentionmaskbatchsize 4 inputs pixelvalues pixelvalues inputids inputids attentionmask attentionmask return model inputs def getvisiontextmodelself visionconfig textconfig visionmodel flaxclipvisionmodelvisionconfig textmodel flaxbertmodeltextconfig return visionmodel textmodel def prepareconfigandinputsself clipmodeltester flaxclipvisionmodeltesterself bertmodeltester flaxbertmodeltesterself visionconfigandinputs clipmodeltester prepareconfigandinputs textconfigandinputs bertmodeltester prepareconfigandinputs visionconfig pixelvalues visionconfigandinputs textconfig inputids tokentypeids attentionmask textconfigandinputs make sure that cross attention layers are added return textconfig textconfig visionconfig visionconfig pixelvalues pixelvalues attentionmask attentionmask inputids inputids tokentypeids tokentypeids requireflax requirevision class flaxvisiontextdualencoderintegrationtestunittest testcase slow def testinferenceself model flaxvisiontextdualencodermodel frompretrainedclipitalianclipitalian logitscaleinitvalue1 0 processor visiontextdualencoderprocessor frompretrainedclipitalianclipitalian image image open testsfixturestestssamplescoco000000039769 png inputs processor textuna foto di un gatto una foto di un cane imagesimage paddingtrue returntensorsnp outputs modelinputs verify the logits self assertequaloutputs logitsperimage shape inputs pixelvalues shape0 inputs inputids shape0 self assertequal outputs logitspertext shape inputs inputids shape0 inputs pixelvalues shape0 expectedlogits np array1 2284727 0 3104122 self asserttruenp allcloseoutputs logitsperimage expectedlogits atol1e3 coding utf 8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch visiontextdualencoder model inspired by https github com rwightman pytorch image models blob b9bd960a032c75ca6b808ddeed76bee5f3ed4972 timm models layers helpers py from pytorch internals in vit the seq_len equals the number of patches 1 we add 1 for the cls token prepare inputs pt flax flax pt make sure that cross attention layers are added make sure that cross attention layers are added verify the logits
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 to_2tuple(x):
    if isinstance(x, collections.abc.Iterable):
        return x
    return (x, x)


@require_flax
class VisionTextDualEncoderMixin:
    def get_vision_text_model(self, config, text_config):
        pass

    def prepare_config_and_inputs(self):
        pass

    def get_pretrained_model_and_inputs(self):
        pass

    def assert_almost_equals(self, a: np.ndarray, b: np.ndarray, tol: float):
        diff = np.abs((a - b)).max()
        self.assertLessEqual(diff, tol, f"Difference between torch and flax is {diff} (>= {tol}).")

    def check_model_from_pretrained_configs(
        self, text_config, input_ids, attention_mask, vision_config, pixel_values=None, **kwargs
    ):
        config = VisionTextDualEncoderConfig.from_vision_text_configs(vision_config, text_config)

        model = FlaxVisionTextDualEncoderModel(config)

        output = model(input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask)

        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 check_vision_text_dual_encoder_from_pretrained(
        self, text_config, input_ids, attention_mask, vision_config, pixel_values=None, **kwargs
    ):
        vision_model, text_model = self.get_vision_text_model(vision_config, text_config)
        kwargs = {"vision_model": vision_model, "text_model": text_model}
        model = FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**kwargs)

        output = model(input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask)

        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 check_save_load(self, text_config, input_ids, attention_mask, vision_config, pixel_values=None, **kwargs):
        vision_model, text_model = self.get_vision_text_model(vision_config, text_config)
        kwargs = {"vision_model": vision_model, "text_model": text_model}
        model = FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**kwargs)

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

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

            after_output = model(input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask)
            out_2 = after_output[0]
            max_diff = np.amax(np.abs(out_2 - out_1))
            self.assertLessEqual(max_diff, 1e-3)

    def check_vision_text_output_attention(
        self, text_config, input_ids, attention_mask, vision_config, pixel_values=None, **kwargs
    ):
        vision_model, text_model = self.get_vision_text_model(vision_config, text_config)
        kwargs = {"vision_model": vision_model, "text_model": text_model}
        model = FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**kwargs)

        output = model(
            input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask, output_attentions=True
        )

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

        
        image_size = to_2tuple(vision_model.config.image_size)
        patch_size = to_2tuple(vision_model.config.patch_size)
        num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
        seq_len = num_patches + 1
        self.assertEqual(vision_attentions[0].shape[-3:], (vision_config.num_attention_heads, seq_len, seq_len))

        text_attentions = output.text_model_output.attentions
        self.assertEqual(len(text_attentions), 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 check_pt_flax_equivalence(self, pt_model, fx_model, inputs_dict):
        pt_model.to(torch_device)
        pt_model.eval()

        
        flax_inputs = inputs_dict
        pt_inputs = {k: torch.tensor(v.tolist()) for k, v in flax_inputs.items()}

        with torch.no_grad():
            pt_outputs = pt_model(**pt_inputs).to_tuple()

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

        
        with tempfile.TemporaryDirectory() as tmpdirname:
            pt_model.save_pretrained(tmpdirname)
            fx_model_loaded = FlaxVisionTextDualEncoderModel.from_pretrained(tmpdirname, from_pt=True)

        fx_outputs_loaded = fx_model_loaded(**inputs_dict).to_tuple()
        self.assertEqual(len(fx_outputs_loaded), len(pt_outputs), "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(fx_output_loaded, pt_output.numpy(), 4e-2)

        
        with tempfile.TemporaryDirectory() as tmpdirname:
            fx_model.save_pretrained(tmpdirname)
            pt_model_loaded = VisionTextDualEncoderModel.from_pretrained(tmpdirname, from_flax=True)

        pt_model_loaded.to(torch_device)
        pt_model_loaded.eval()

        with torch.no_grad():
            pt_outputs_loaded = pt_model_loaded(**pt_inputs).to_tuple()

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

    def check_equivalence_pt_to_flax(self, vision_config, text_config, inputs_dict):
        config = VisionTextDualEncoderConfig.from_vision_text_configs(vision_config, text_config)

        pt_model = VisionTextDualEncoderModel(config)
        fx_model = FlaxVisionTextDualEncoderModel(config)

        fx_state = convert_pytorch_state_dict_to_flax(pt_model.state_dict(), fx_model)
        fx_model.params = fx_state

        self.check_pt_flax_equivalence(pt_model, fx_model, inputs_dict)

    def check_equivalence_flax_to_pt(self, vision_config, text_config, inputs_dict):
        config = VisionTextDualEncoderConfig.from_vision_text_configs(vision_config, text_config)

        pt_model = VisionTextDualEncoderModel(config)
        fx_model = FlaxVisionTextDualEncoderModel(config)

        pt_model = load_flax_weights_in_pytorch_model(pt_model, fx_model.params)

        self.check_pt_flax_equivalence(pt_model, fx_model, inputs_dict)

    def test_model_from_pretrained_configs(self):
        inputs_dict = self.prepare_config_and_inputs()
        self.check_model_from_pretrained_configs(**inputs_dict)

    def test_vision_text_dual_encoder_from_pretrained(self):
        inputs_dict = self.prepare_config_and_inputs()
        self.check_vision_text_dual_encoder_from_pretrained(**inputs_dict)

    def test_save_load(self):
        inputs_dict = self.prepare_config_and_inputs()
        self.check_save_load(**inputs_dict)

    def test_vision_text_output_attention(self):
        inputs_dict = self.prepare_config_and_inputs()
        self.check_vision_text_output_attention(**inputs_dict)

    @is_pt_flax_cross_test
    def test_pt_flax_equivalence(self):
        config_inputs_dict = self.prepare_config_and_inputs()
        vision_config = config_inputs_dict.pop("vision_config")
        text_config = config_inputs_dict.pop("text_config")

        inputs_dict = config_inputs_dict

        self.check_equivalence_pt_to_flax(vision_config, text_config, inputs_dict)
        self.check_equivalence_flax_to_pt(vision_config, text_config, inputs_dict)

    @slow
    def test_real_model_save_load_from_pretrained(self):
        model_2, inputs = self.get_pretrained_model_and_inputs()

        outputs = model_2(**inputs)
        out_2 = outputs[0]

        with tempfile.TemporaryDirectory() as tmp_dirname:
            model_2.save_pretrained(tmp_dirname)
            model_1 = FlaxVisionTextDualEncoderModel.from_pretrained(tmp_dirname)

            after_outputs = model_1(**inputs)
            out_1 = after_outputs[0]
            max_diff = np.amax(np.abs(out_1 - out_2))
            self.assertLessEqual(max_diff, 1e-5)


@require_flax
class FlaxViTBertModelTest(VisionTextDualEncoderMixin, unittest.TestCase):
    def get_pretrained_model_and_inputs(self):
        model = FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
            "hf-internal-testing/tiny-random-vit",
            "hf-internal-testing/tiny-bert",
            vision_from_pt=True,
            text_from_pt=True,
        )
        batch_size = 13
        pixel_values = floats_tensor(
            [
                batch_size,
                model.config.vision_config.num_channels,
                model.config.vision_config.image_size,
                model.config.vision_config.image_size,
            ]
        )
        input_ids = ids_tensor([batch_size, 4], model.config.text_config.vocab_size)
        attention_mask = random_attention_mask([batch_size, 4])
        inputs = {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

        return model, inputs

    def get_vision_text_model(self, vision_config, text_config):
        vision_model = FlaxViTModel(vision_config)
        text_model = FlaxBertModel(text_config)
        return vision_model, text_model

    def prepare_config_and_inputs(self):
        vit_model_tester = FlaxViTModelTester(self)
        bert_model_tester = FlaxBertModelTester(self)
        vision_config_and_inputs = vit_model_tester.prepare_config_and_inputs()
        text_config_and_inputs = bert_model_tester.prepare_config_and_inputs()

        vision_config, pixel_values = vision_config_and_inputs

        text_config, input_ids, token_type_ids, attention_mask = text_config_and_inputs

        
        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 FlaxCLIPVisionBertModelTest(VisionTextDualEncoderMixin, unittest.TestCase):
    def get_pretrained_model_and_inputs(self):
        model = FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
            "hf-internal-testing/tiny-random-clip",
            "hf-internal-testing/tiny-bert",
            vision_from_pt=True,
            text_from_pt=True,
        )
        batch_size = 13
        pixel_values = floats_tensor(
            [
                batch_size,
                model.config.vision_config.num_channels,
                model.config.vision_config.image_size,
                model.config.vision_config.image_size,
            ]
        )
        input_ids = ids_tensor([batch_size, 4], model.config.text_config.vocab_size)
        attention_mask = random_attention_mask([batch_size, 4])
        inputs = {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

        return model, inputs

    def get_vision_text_model(self, vision_config, text_config):
        vision_model = FlaxCLIPVisionModel(vision_config)
        text_model = FlaxBertModel(text_config)
        return vision_model, text_model

    def prepare_config_and_inputs(self):
        clip_model_tester = FlaxCLIPVisionModelTester(self)
        bert_model_tester = FlaxBertModelTester(self)
        vision_config_and_inputs = clip_model_tester.prepare_config_and_inputs()
        text_config_and_inputs = bert_model_tester.prepare_config_and_inputs()

        vision_config, pixel_values = vision_config_and_inputs

        text_config, input_ids, token_type_ids, attention_mask = text_config_and_inputs

        
        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 FlaxVisionTextDualEncoderIntegrationTest(unittest.TestCase):
    @slow
    def test_inference(self):
        model = FlaxVisionTextDualEncoderModel.from_pretrained("clip-italian/clip-italian", logit_scale_init_value=1.0)
        processor = VisionTextDualEncoderProcessor.from_pretrained("clip-italian/clip-italian")

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

        outputs = model(**inputs)

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

        expected_logits = np.array([[1.2284727, 0.3104122]])

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