codingutf8 2022 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 tensorflow groupvit model from future import annotations import inspect import os import random import tempfile import unittest from importlib import importmodule import numpy as np import requests from transformers import groupvitconfig groupvittextconfig groupvitvisionconfig from transformers testingutils import ispttfcrosstest requiretensorflowprobability requiretf requirevision slow from transformers utils import istfavailable isvisionavailable from testconfigurationcommon import configtester from testmodelingtfcommon import tfmodeltestermixin floatstensor idstensor randomattentionmask from testpipelinemixin import pipelinetestermixin if istfavailable import tensorflow as tf from transformers import tfgroupvitmodel tfgroupvittextmodel tfgroupvitvisionmodel tfsharedembeddings from transformers models groupvit modelingtfgroupvit import tfgroupvitpretrainedmodelarchivelist if isvisionavailable from pil import image from transformers import clipprocessor class tfgroupvitvisionmodeltester def init self parent batchsize12 imagesize30 patchsize2 numchannels3 istrainingtrue hiddensize32 depths6 3 3 numgrouptokens64 8 0 numoutputgroups64 8 8 numattentionheads4 intermediatesize37 dropout0 1 attentiondropout0 1 initializerrange0 02 scopenone self parent parent self batchsize batchsize self imagesize imagesize self patchsize patchsize self numchannels numchannels self istraining istraining self hiddensize hiddensize self depths depths self numhiddenlayers sumdepths self expectednumhiddenlayers lendepths 1 self numgrouptokens numgrouptokens self numoutputgroups numoutputgroups self numattentionheads numattentionheads self intermediatesize intermediatesize self dropout dropout self attentiondropout attentiondropout self initializerrange initializerrange self scope scope numpatches imagesize patchsize 2 no cls token for groupvit self seqlength numpatches def prepareconfigandinputsself rng random random0 pixelvalues floatstensorself batchsize self numchannels self imagesize self imagesize rngrng config self getconfig return config pixelvalues def getconfigself return groupvitvisionconfig imagesizeself imagesize patchsizeself patchsize numchannelsself numchannels hiddensizeself hiddensize depthsself depths numgrouptokensself numgrouptokens numoutputgroupsself numoutputgroups numattentionheadsself numattentionheads intermediatesizeself intermediatesize dropoutself dropout attentiondropoutself attentiondropout initializerrangeself initializerrange def createandcheckmodelself config pixelvalues model tfgroupvitvisionmodelconfigconfig result modelpixelvalues trainingfalse self parent assertequal result lasthiddenstate shape self batchsize self numoutputgroups1 self hiddensize self parent assertequalresult pooleroutput shape self batchsize self hiddensize def prepareconfigandinputsforcommonself configandinputs self prepareconfigandinputs config pixelvalues configandinputs inputsdict pixelvalues pixelvalues return config inputsdict requiretf class tfgroupvitvisionmodeltesttfmodeltestermixin unittest testcase allmodelclasses tfgroupvitvisionmodel if istfavailable else testpruning false testresizeembeddings false testheadmasking false testonnx false def checkpttfoutputsself tfoutputs ptoutputs modelclass tol1e4 nameoutputs attributesnone we override with a slightly higher tol value as this model tends to diverge a bit more super checkpttfoutputstfoutputs ptoutputs modelclass tol name attributes def setupself self modeltester tfgroupvitvisionmodeltesterself self configtester configtester self configclassgroupvitvisionconfig hastextmodalityfalse hiddensize37 def testconfigself self configtester runcommontests unittest skipreasongroupvit does not use inputsembeds def testinputsembedsself pass requiretensorflowprobability slow def testsavedmodelcreationself super testsavedmodelcreation unittest skipreasongroupvit does not use inputsembeds def testgraphmodewithinputsembedsself pass def testmodelcommonattributesself config self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig self assertisinstancemodel getinputembeddings tf keras layers layer x model getoutputembeddings self asserttruex is none or isinstancex tf keras layers layer def testforwardsignatureself config self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig signature inspect signaturemodel call signature parameters is an ordereddict so argnames order is deterministic argnames signature parameters keys expectedargnames pixelvalues self assertlistequalargnames 1 expectedargnames def testmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfigandinputs def testattentionoutputsself config inputsdict self modeltester prepareconfigandinputsforcommon config returndict true seqlen getattrself modeltester seqlength none expectednumattentionoutputs sumg 0 for g in self modeltester numgrouptokens for modelclass in self allmodelclasses inputsdictoutputattentions true inputsdictoutputhiddenstates false config returndict true model modelclassconfig outputs modelself prepareforclassinputsdict modelclass trainingfalse attentions outputs attentions groupvit returns attention grouping of each stage self assertequallenattentions sumg 0 for g in self modeltester numgrouptokens check that outputattentions also work using config del inputsdictoutputattentions config outputattentions true model modelclassconfig outputs modelself prepareforclassinputsdict modelclass trainingfalse attentions outputs attentions groupvit returns attention grouping of each stage self assertequallenattentions expectednumattentionoutputs outlen lenoutputs check attention is always last and order is fine inputsdictoutputattentions true inputsdictoutputhiddenstates true model modelclassconfig outputs modelself prepareforclassinputsdict modelclass trainingfalse addedhiddenstates 1 self assertequaloutlen addedhiddenstates lenoutputs selfattentions outputs attentions groupvit returns attention grouping of each stage self assertequallenselfattentions expectednumattentionoutputs for i selfattn in enumerateselfattentions if selfattn is none continue self assertlistequal listselfattentionsi shape2 self modeltester numoutputgroupsi self modeltester numoutputgroupsi 1 if i 0 else seqlen def testhiddenstatesoutputself def checkhiddenstatesoutputinputsdict config modelclass model modelclassconfig outputs modelself prepareforclassinputsdict modelclass trainingfalse hiddenstates outputs encoderhiddenstates if config isencoderdecoder else outputs hiddenstates expectednumlayers getattr self modeltester expectednumhiddenlayers self modeltester numhiddenlayers 1 self assertequallenhiddenstates expectednumlayers seqlength getattrself modeltester seqlength none self assertlistequal listhiddenstates0 shape2 seqlength self modeltester hiddensize config inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses inputsdictoutputhiddenstates true checkhiddenstatesoutputinputsdict config modelclass check that outputhiddenstates also work using config del inputsdictoutputhiddenstates config outputhiddenstates true checkhiddenstatesoutputinputsdict config modelclass ispttfcrosstest def testpttfmodelequivalenceself groupvit computes some indices using argmax uses them as onehot encoding for further computation the problem is while pttf have very small difference in ysoft 1e9 the argmax could be totally different if there are at least 2 indices with almost identical values this leads to very large difference in the outputs we need specific seeds to avoid almost identical values happening in ysoft import torch seed 338 random seedseed np random seedseed torch manualseedseed torch cuda manualseedallseed tf random setseedseed return super testpttfmodelequivalence slow def testmodelfrompretrainedself for modelname in tfgroupvitpretrainedmodelarchivelist 1 model tfgroupvitvisionmodel frompretrainedmodelname self assertisnotnonemodel unittest skip tfgroupvitvisionmodel does not convert hiddenstates and attentions to tensors as they are all of different dimensions and we get got a nontensor value error when saving the model slow def testsavedmodelcreationextendedself config inputsdict self modeltester prepareconfigandinputsforcommon config outputhiddenstates true config outputattentions true if hasattrconfig usecache config usecache true seqlen getattrself modeltester seqlength none for modelclass in self allmodelclasses classinputsdict self prepareforclassinputsdict modelclass model modelclassconfig numout lenmodelclassinputsdict with tempfile temporarydirectory as tmpdirname model savepretrainedtmpdirname savedmodeltrue savedmodeldir os path jointmpdirname savedmodel 1 model tf keras models loadmodelsavedmodeldir outputs modelclassinputsdict outputhiddenstates outputshiddenstates outputattentions outputsattentions check num outputs self assertequallenoutputs numout check num layers expectednumlayers getattr self modeltester expectednumhiddenlayers self modeltester numhiddenlayers 1 self assertequallenoutputhiddenstates expectednumlayers self assertequallenoutputattentions self modeltester numhiddenlayers check attention outputs imagesize self modeltester imagesize self modeltester imagesize patchsize self modeltester patchsize self modeltester patchsize numpatches imagesize1 patchsize1 imagesize0 patchsize0 seqlen numpatches 1 self assertlistequal listoutputattentions0 shape3 self modeltester numattentionheads seqlen seqlen check hidden states self assertlistequal listoutputhiddenstates0 shape2 seqlen self modeltester hiddensize class tfgroupvittextmodeltester def init self parent batchsize12 seqlength7 istrainingtrue useinputmasktrue uselabelstrue vocabsize99 hiddensize32 numhiddenlayers2 numattentionheads4 intermediatesize37 dropout0 1 attentiondropout0 1 maxpositionembeddings512 initializerrange0 02 scopenone self parent parent self batchsize batchsize self seqlength seqlength self istraining istraining self useinputmask useinputmask self uselabels uselabels self vocabsize vocabsize self hiddensize hiddensize self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads self intermediatesize intermediatesize self dropout dropout self attentiondropout attentiondropout self maxpositionembeddings maxpositionembeddings self initializerrange initializerrange self scope scope def prepareconfigandinputsself rng random random0 inputids idstensorself batchsize self seqlength self vocabsize rngrng inputmask none if self useinputmask inputmask randomattentionmaskself batchsize self seqlength make sure the first token has attention mask 1 to ensure that after combining the causal mask there is still at least one token being attended to for each batch todo change randomattentionmask in pttfflax common test file after a discussion with the team inputmask tf concat tf oneslikeinputmask 1 dtypeinputmask dtype inputmask 1 axis1 config self getconfig return config inputids inputmask def getconfigself return groupvittextconfig vocabsizeself vocabsize hiddensizeself hiddensize numhiddenlayersself numhiddenlayers numattentionheadsself numattentionheads intermediatesizeself intermediatesize dropoutself dropout attentiondropoutself attentiondropout maxpositionembeddingsself maxpositionembeddings initializerrangeself initializerrange def createandcheckmodelself config inputids inputmask model tfgroupvittextmodelconfigconfig result modelinputids attentionmaskinputmask trainingfalse result modelinputids trainingfalse self parent assertequalresult lasthiddenstate shape self batchsize self seqlength self hiddensize self parent assertequalresult pooleroutput shape self batchsize self hiddensize def prepareconfigandinputsforcommonself configandinputs self prepareconfigandinputs config inputids inputmask configandinputs inputsdict inputids inputids attentionmask inputmask return config inputsdict requiretf class tfgroupvittextmodeltesttfmodeltestermixin unittest testcase allmodelclasses tfgroupvittextmodel if istfavailable else testpruning false testheadmasking false testonnx false def checkpttfoutputsself tfoutputs ptoutputs modelclass tol1e4 nameoutputs attributesnone we override with a slightly higher tol value as this model tends to diverge a bit more super checkpttfoutputstfoutputs ptoutputs modelclass tol name attributes def setupself self modeltester tfgroupvittextmodeltesterself self configtester configtesterself configclassgroupvittextconfig hiddensize37 def testconfigself self configtester runcommontests def testmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfigandinputs unittest skipreasongroupvittextmodel does not use inputsembeds def testinputsembedsself pass slow def testmodelfrompretrainedself for modelname in tfgroupvitpretrainedmodelarchivelist 1 model tfgroupvittextmodel frompretrainedmodelname self assertisnotnonemodel slow def testsavedmodelcreationextendedself config inputsdict self modeltester prepareconfigandinputsforcommon config outputhiddenstates true config outputattentions true if hasattrconfig usecache config usecache true for modelclass in self allmodelclasses classinputsdict self prepareforclassinputsdict modelclass model modelclassconfig numout lenmodelclassinputsdict with tempfile temporarydirectory as tmpdirname model savepretrainedtmpdirname savedmodeltrue savedmodeldir os path jointmpdirname savedmodel 1 model tf keras models loadmodelsavedmodeldir outputs modelclassinputsdict outputhiddenstates outputshiddenstates outputattentions outputsattentions check number of outputs self assertequallenoutputs numout check number of layers expectednumlayers getattr self modeltester expectednumhiddenlayers self modeltester numhiddenlayers 1 check hidden states self assertequallenoutputhiddenstates expectednumlayers self assertlistequal listoutputhiddenstates0 shape2 self modeltester seqlength self modeltester hiddensize check attention outputs self assertequallenoutputattentions self modeltester numhiddenlayers seqlength self modeltester seqlength keylength getattrself modeltester keylength seqlength self assertlistequal listoutputattentions0 shape3 self modeltester numattentionheads seqlength keylength class tfgroupvitmodeltester def initself parent istrainingtrue self parent parent self textmodeltester tfgroupvittextmodeltesterparent self visionmodeltester tfgroupvitvisionmodeltesterparent self istraining istraining def prepareconfigandinputsself textconfig inputids attentionmask self textmodeltester prepareconfigandinputs visionconfig pixelvalues self visionmodeltester prepareconfigandinputs config self getconfig return config inputids attentionmask pixelvalues def getconfigself return groupvitconfig fromtextvisionconfigs self textmodeltester getconfig self visionmodeltester getconfig projectiondim64 def createandcheckmodelself config inputids attentionmask pixelvalues model tfgroupvitmodelconfig result modelinputids pixelvalues attentionmask trainingfalse self parent assertequal result logitsperimage shape self visionmodeltester batchsize self textmodeltester batchsize self parent assertequal result logitspertext shape self textmodeltester batchsize self visionmodeltester batchsize def prepareconfigandinputsforcommonself configandinputs self prepareconfigandinputs config inputids attentionmask pixelvalues configandinputs inputsdict inputids inputids attentionmask attentionmask pixelvalues pixelvalues returnloss true return config inputsdict requiretf class tfgroupvitmodeltesttfmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses tfgroupvitmodel if istfavailable else pipelinemodelmapping featureextraction tfgroupvitmodel if istfavailable else testheadmasking false testpruning false testresizeembeddings false testattentionoutputs false testonnx false def checkpttfoutputsself tfoutputs ptoutputs modelclass tol1e4 nameoutputs attributesnone we override with a slightly higher tol value as this model tends to diverge a bit more super checkpttfoutputstfoutputs ptoutputs modelclass tol name attributes def setupself self modeltester tfgroupvitmodeltesterself def testmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfigandinputs unittest skipreasonhiddenstates are tested in individual model tests def testhiddenstatesoutputself pass unittest skipreasoninputembeds are tested in individual model tests def testinputsembedsself pass unittest skipreasonclipmodel does not have inputoutput embeddings def testmodelcommonattributesself pass requiretensorflowprobability slow def testkerasfitself super testkerasfit ispttfcrosstest def testpttfmodelequivalenceself groupvit computes some indices using argmax uses them as onehot encoding for further computation the problem is while pttf have very small difference in ysoft 1e9 the argmax could be totally different if there are at least 2 indices with almost identical values this leads to very large difference in the outputs we need specific seeds to avoid almost identical values happening in ysoft import torch seed 158 random seedseed np random seedseed torch manualseedseed torch cuda manualseedallseed tf random setseedseed return super testpttfmodelequivalence overwrite from common since tfgroupvitmodeltester set returnloss to true and causes the preparation of symbolicinputs failed def testkerassaveloadself config inputsdict self modeltester prepareconfigandinputsforcommon remove returnloss to make code work if self class name tfgroupvitmodeltest inputsdict popreturnloss none tfmainlayerclasses modulemember for modelclass in self allmodelclasses for module in importmodulemodelclass module for modulemembername in dirmodule if modulemembername endswithmainlayer this condition is required since modelingtfclip py has 3 classes whose names end with mainlayer and modulemembername lenmainlayer modelclass name lenmodel for modulemember in getattrmodule modulemembername if isinstancemodulemember type and tf keras layers layer in modulemember bases and getattrmodulemember kerasserializable false for mainlayerclass in tfmainlayerclasses t5mainlayer needs an embedtokens parameter when called without the inputsembeds parameter if t5 in mainlayerclass name take the same values than in tft5modeltester for this shared layer shared tfsharedembeddings99 32 nameshared config usecache inputsdict popusecache none mainlayer mainlayerclassconfig embedtokensshared else mainlayer mainlayerclassconfig symbolicinputs name tf keras inputtensor shape1 dtypetensor dtype for name tensor in inputsdict items model tf keras modelsymbolicinputs outputsmainlayersymbolicinputs outputs modelinputsdict with tempfile temporarydirectory as tmpdirname filepath os path jointmpdirname kerasmodel h5 model savefilepath if t5 in mainlayerclass name model tf keras models loadmodel filepath customobjects mainlayerclass name mainlayerclass tfsharedembeddings tfsharedembeddings else model tf keras models loadmodel filepath customobjectsmainlayerclass name mainlayerclass assert isinstancemodel tf keras model afteroutputs modelinputsdict self assertoutputssameafteroutputs outputs slow def testmodelfrompretrainedself for modelname in tfgroupvitpretrainedmodelarchivelist 1 model tfgroupvitmodel frompretrainedmodelname self assertisnotnonemodel unittest skipreasoncurrently savedmodel doesn t work with nested outputs slow def testsavedmodelcreationself pass unittest skipreasonsavedmodel doesn t work with nested outputs so no preparation happens slow def testprepareservingoutputself pass we will verify our results on an image of cute cats def prepareimg url http images cocodataset orgval2017000000039769 jpg im image openrequests geturl streamtrue raw return im requirevision requiretf class tfgroupvitmodelintegrationtestunittest testcase slow def testinferenceself modelname nvidiagroupvitgccyfcc model tfgroupvitmodel frompretrainedmodelname processor clipprocessor frompretrainedmodelname image prepareimg inputs processor texta photo of a cat a photo of a dog imagesimage paddingtrue returntensorstf outputs modelinputs trainingfalse verify the logits self assertequal outputs logitsperimage shape tf tensorshapeinputs pixelvalues shape0 inputs inputids shape0 self assertequal outputs logitspertext shape tf tensorshapeinputs inputids shape0 inputs pixelvalues shape0 expectedlogits tf constant13 3523 6 3629 tf debugging assertnearoutputs logitsperimage expectedlogits atol1e3 coding utf 8 2022 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 tensorflow groupvit model no cls token for groupvit here we also overwrite some of the tests of test_modeling_common py as groupvit does not use input_ids inputs_embeds attention_mask and seq_length we override with a slightly higher tol value as this model tends to diverge a bit more during saving tensorflow will also run with training true which trigger gumbel_softmax that requires tensorflow probability signature parameters is an ordereddict so arg_names order is deterministic groupvit returns attention grouping of each stage check that output_attentions also work using config groupvit returns attention grouping of each stage check attention is always last and order is fine groupvit returns attention grouping of each stage check that output_hidden_states also work using config groupvit computes some indices using argmax uses them as one hot encoding for further computation the problem is while pt tf have very small difference in y_soft 1e 9 the argmax could be totally different if there are at least 2 indices with almost identical values this leads to very large difference in the outputs we need specific seeds to avoid almost identical values happening in y_soft check num outputs check num layers check attention outputs check hidden states make sure the first token has attention mask 1 to ensure that after combining the causal mask there is still at least one token being attended to for each batch todo change random_attention_mask in pt tf flax common test file after a discussion with the team we override with a slightly higher tol value as this model tends to diverge a bit more check number of outputs check number of layers check hidden states check attention outputs we override with a slightly higher tol value as this model tends to diverge a bit more groupvit computes some indices using argmax uses them as one hot encoding for further computation the problem is while pt tf have very small difference in y_soft 1e 9 the argmax could be totally different if there are at least 2 indices with almost identical values this leads to very large difference in the outputs we need specific seeds to avoid almost identical values happening in y_soft overwrite from common since tfgroupvitmodeltester set return_loss to true and causes the preparation of symbolic_inputs failed remove return_loss to make code work this condition is required since modeling_tf_clip py has 3 classes whose names end with mainlayer t5mainlayer needs an embed_tokens parameter when called without the inputs_embeds parameter take the same values than in tft5modeltester for this shared layer we will verify our results on an image of cute cats verify the logits
from __future__ import annotations

import inspect
import os
import random
import tempfile
import unittest
from importlib import import_module

import numpy as np
import requests

from transformers import GroupViTConfig, GroupViTTextConfig, GroupViTVisionConfig
from transformers.testing_utils import (
    is_pt_tf_cross_test,
    require_tensorflow_probability,
    require_tf,
    require_vision,
    slow,
)
from transformers.utils import is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import TFGroupViTModel, TFGroupViTTextModel, TFGroupViTVisionModel, TFSharedEmbeddings
    from transformers.models.groupvit.modeling_tf_groupvit import TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import CLIPProcessor


class TFGroupViTVisionModelTester:
    def __init__(
        self,
        parent,
        batch_size=12,
        image_size=30,
        patch_size=2,
        num_channels=3,
        is_training=True,
        hidden_size=32,
        depths=[6, 3, 3],
        num_group_tokens=[64, 8, 0],
        num_output_groups=[64, 8, 8],
        num_attention_heads=4,
        intermediate_size=37,
        dropout=0.1,
        attention_dropout=0.1,
        initializer_range=0.02,
        scope=None,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_channels = num_channels
        self.is_training = is_training
        self.hidden_size = hidden_size
        self.depths = depths
        self.num_hidden_layers = sum(depths)
        self.expected_num_hidden_layers = len(depths) + 1
        self.num_group_tokens = num_group_tokens
        self.num_output_groups = num_output_groups
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.dropout = dropout
        self.attention_dropout = attention_dropout
        self.initializer_range = initializer_range
        self.scope = scope

        num_patches = (image_size // patch_size) ** 2
        
        self.seq_length = num_patches

    def prepare_config_and_inputs(self):
        rng = random.Random(0)
        pixel_values = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size], rng=rng)
        config = self.get_config()

        return config, pixel_values

    def get_config(self):
        return GroupViTVisionConfig(
            image_size=self.image_size,
            patch_size=self.patch_size,
            num_channels=self.num_channels,
            hidden_size=self.hidden_size,
            depths=self.depths,
            num_group_tokens=self.num_group_tokens,
            num_output_groups=self.num_output_groups,
            num_attention_heads=self.num_attention_heads,
            intermediate_size=self.intermediate_size,
            dropout=self.dropout,
            attention_dropout=self.attention_dropout,
            initializer_range=self.initializer_range,
        )

    def create_and_check_model(self, config, pixel_values):
        model = TFGroupViTVisionModel(config=config)
        result = model(pixel_values, training=False)
        self.parent.assertEqual(
            result.last_hidden_state.shape, (self.batch_size, self.num_output_groups[-1], self.hidden_size)
        )
        self.parent.assertEqual(result.pooler_output.shape, (self.batch_size, self.hidden_size))

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        config, pixel_values = config_and_inputs
        inputs_dict = {"pixel_values": pixel_values}
        return config, inputs_dict


@require_tf
class TFGroupViTVisionModelTest(TFModelTesterMixin, unittest.TestCase):
    

    all_model_classes = (TFGroupViTVisionModel,) if is_tf_available() else ()

    test_pruning = False
    test_resize_embeddings = False
    test_head_masking = False
    test_onnx = False

    def check_pt_tf_outputs(self, tf_outputs, pt_outputs, model_class, tol=1e-4, name="outputs", attributes=None):
        
        super().check_pt_tf_outputs(tf_outputs, pt_outputs, model_class, tol, name, attributes)

    def setUp(self):
        self.model_tester = TFGroupViTVisionModelTester(self)
        self.config_tester = ConfigTester(
            self, config_class=GroupViTVisionConfig, has_text_modality=False, hidden_size=37
        )

    def test_config(self):
        self.config_tester.run_common_tests()

    @unittest.skip(reason="GroupViT does not use inputs_embeds")
    def test_inputs_embeds(self):
        pass

    

    @require_tensorflow_probability
    @slow
    def test_saved_model_creation(self):
        super().test_saved_model_creation()

    @unittest.skip(reason="GroupViT does not use inputs_embeds")
    def test_graph_mode_with_inputs_embeds(self):
        pass

    def test_model_common_attributes(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            self.assertIsInstance(model.get_input_embeddings(), (tf.keras.layers.Layer))
            x = model.get_output_embeddings()
            self.assertTrue(x is None or isinstance(x, tf.keras.layers.Layer))

    def test_forward_signature(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            signature = inspect.signature(model.call)
            
            arg_names = [*signature.parameters.keys()]

            expected_arg_names = ["pixel_values"]
            self.assertListEqual(arg_names[:1], expected_arg_names)

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)

    def test_attention_outputs(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.return_dict = True

        seq_len = getattr(self.model_tester, "seq_length", None)

        expected_num_attention_outputs = sum(g > 0 for g in self.model_tester.num_group_tokens)

        for model_class in self.all_model_classes:
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = False
            config.return_dict = True
            model = model_class(config)
            outputs = model(**self._prepare_for_class(inputs_dict, model_class), training=False)
            attentions = outputs.attentions
            
            self.assertEqual(len(attentions), sum(g > 0 for g in self.model_tester.num_group_tokens))

            
            del inputs_dict["output_attentions"]
            config.output_attentions = True
            model = model_class(config)
            outputs = model(**self._prepare_for_class(inputs_dict, model_class), training=False)
            attentions = outputs.attentions
            
            self.assertEqual(len(attentions), expected_num_attention_outputs)

            out_len = len(outputs)

            
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = True
            model = model_class(config)
            outputs = model(**self._prepare_for_class(inputs_dict, model_class), training=False)

            added_hidden_states = 1
            self.assertEqual(out_len + added_hidden_states, len(outputs))

            self_attentions = outputs.attentions

            
            self.assertEqual(len(self_attentions), expected_num_attention_outputs)
            for i, self_attn in enumerate(self_attentions):
                if self_attn is None:
                    continue

                self.assertListEqual(
                    list(self_attentions[i].shape[-2:]),
                    [
                        self.model_tester.num_output_groups[i],
                        self.model_tester.num_output_groups[i - 1] if i > 0 else seq_len,
                    ],
                )

    def test_hidden_states_output(self):
        def check_hidden_states_output(inputs_dict, config, model_class):
            model = model_class(config)

            outputs = model(**self._prepare_for_class(inputs_dict, model_class), training=False)

            hidden_states = outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

            expected_num_layers = getattr(
                self.model_tester, "expected_num_hidden_layers", self.model_tester.num_hidden_layers + 1
            )
            self.assertEqual(len(hidden_states), expected_num_layers)

            seq_length = getattr(self.model_tester, "seq_length", None)

            self.assertListEqual(
                list(hidden_states[0].shape[-2:]),
                [seq_length, self.model_tester.hidden_size],
            )

        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            inputs_dict["output_hidden_states"] = True
            check_hidden_states_output(inputs_dict, config, model_class)

            
            del inputs_dict["output_hidden_states"]
            config.output_hidden_states = True

            check_hidden_states_output(inputs_dict, config, model_class)

    @is_pt_tf_cross_test
    def test_pt_tf_model_equivalence(self):
        
        
        
        
        
        
        
        import torch

        seed = 338
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        tf.random.set_seed(seed)
        return super().test_pt_tf_model_equivalence()

    @slow
    def test_model_from_pretrained(self):
        for model_name in TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            model = TFGroupViTVisionModel.from_pretrained(model_name)
            self.assertIsNotNone(model)

    @unittest.skip(
        "TFGroupViTVisionModel does not convert `hidden_states` and `attentions` to tensors as they are all of"
        " different dimensions, and we get `Got a non-Tensor value` error when saving the model."
    )
    @slow
    def test_saved_model_creation_extended(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.output_hidden_states = True
        config.output_attentions = True

        if hasattr(config, "use_cache"):
            config.use_cache = True

        seq_len = getattr(self.model_tester, "seq_length", None)

        for model_class in self.all_model_classes:
            class_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
            model = model_class(config)
            num_out = len(model(class_inputs_dict))

            with tempfile.TemporaryDirectory() as tmpdirname:
                model.save_pretrained(tmpdirname, saved_model=True)
                saved_model_dir = os.path.join(tmpdirname, "saved_model", "1")
                model = tf.keras.models.load_model(saved_model_dir)
                outputs = model(class_inputs_dict)
                output_hidden_states = outputs["hidden_states"]
                output_attentions = outputs["attentions"]

                
                self.assertEqual(len(outputs), num_out)

                
                expected_num_layers = getattr(
                    self.model_tester, "expected_num_hidden_layers", self.model_tester.num_hidden_layers + 1
                )

                self.assertEqual(len(output_hidden_states), expected_num_layers)
                self.assertEqual(len(output_attentions), self.model_tester.num_hidden_layers)

                
                image_size = (self.model_tester.image_size, self.model_tester.image_size)
                patch_size = (self.model_tester.patch_size, self.model_tester.patch_size)
                num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
                seq_len = num_patches + 1

                self.assertListEqual(
                    list(output_attentions[0].shape[-3:]),
                    [self.model_tester.num_attention_heads, seq_len, seq_len],
                )

                
                self.assertListEqual(
                    list(output_hidden_states[0].shape[-2:]),
                    [seq_len, self.model_tester.hidden_size],
                )


class TFGroupViTTextModelTester:
    def __init__(
        self,
        parent,
        batch_size=12,
        seq_length=7,
        is_training=True,
        use_input_mask=True,
        use_labels=True,
        vocab_size=99,
        hidden_size=32,
        num_hidden_layers=2,
        num_attention_heads=4,
        intermediate_size=37,
        dropout=0.1,
        attention_dropout=0.1,
        max_position_embeddings=512,
        initializer_range=0.02,
        scope=None,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        self.use_input_mask = use_input_mask
        self.use_labels = use_labels
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.dropout = dropout
        self.attention_dropout = attention_dropout
        self.max_position_embeddings = max_position_embeddings
        self.initializer_range = initializer_range
        self.scope = scope

    def prepare_config_and_inputs(self):
        rng = random.Random(0)
        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size, rng=rng)

        input_mask = None
        if self.use_input_mask:
            input_mask = random_attention_mask([self.batch_size, self.seq_length])
            
            
            
            input_mask = tf.concat(
                [tf.ones_like(input_mask[:, :1], dtype=input_mask.dtype), input_mask[:, 1:]], axis=-1
            )

        config = self.get_config()

        return config, input_ids, input_mask

    def get_config(self):
        return GroupViTTextConfig(
            vocab_size=self.vocab_size,
            hidden_size=self.hidden_size,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            intermediate_size=self.intermediate_size,
            dropout=self.dropout,
            attention_dropout=self.attention_dropout,
            max_position_embeddings=self.max_position_embeddings,
            initializer_range=self.initializer_range,
        )

    def create_and_check_model(self, config, input_ids, input_mask):
        model = TFGroupViTTextModel(config=config)
        result = model(input_ids, attention_mask=input_mask, training=False)
        result = model(input_ids, training=False)
        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))
        self.parent.assertEqual(result.pooler_output.shape, (self.batch_size, self.hidden_size))

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        config, input_ids, input_mask = config_and_inputs
        inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask}
        return config, inputs_dict


@require_tf
class TFGroupViTTextModelTest(TFModelTesterMixin, unittest.TestCase):
    all_model_classes = (TFGroupViTTextModel,) if is_tf_available() else ()
    test_pruning = False
    test_head_masking = False
    test_onnx = False

    def check_pt_tf_outputs(self, tf_outputs, pt_outputs, model_class, tol=1e-4, name="outputs", attributes=None):
        
        super().check_pt_tf_outputs(tf_outputs, pt_outputs, model_class, tol, name, attributes)

    def setUp(self):
        self.model_tester = TFGroupViTTextModelTester(self)
        self.config_tester = ConfigTester(self, config_class=GroupViTTextConfig, hidden_size=37)

    def test_config(self):
        self.config_tester.run_common_tests()

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)

    @unittest.skip(reason="GroupViTTextModel does not use inputs_embeds")
    def test_inputs_embeds(self):
        pass

    @slow
    def test_model_from_pretrained(self):
        for model_name in TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            model = TFGroupViTTextModel.from_pretrained(model_name)
            self.assertIsNotNone(model)

    @slow
    def test_saved_model_creation_extended(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.output_hidden_states = True
        config.output_attentions = True

        if hasattr(config, "use_cache"):
            config.use_cache = True

        for model_class in self.all_model_classes:
            class_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
            model = model_class(config)
            num_out = len(model(class_inputs_dict))

            with tempfile.TemporaryDirectory() as tmpdirname:
                model.save_pretrained(tmpdirname, saved_model=True)
                saved_model_dir = os.path.join(tmpdirname, "saved_model", "1")
                model = tf.keras.models.load_model(saved_model_dir)
                outputs = model(class_inputs_dict)
                output_hidden_states = outputs["hidden_states"]
                output_attentions = outputs["attentions"]

                
                self.assertEqual(len(outputs), num_out)

                
                expected_num_layers = getattr(
                    self.model_tester, "expected_num_hidden_layers", self.model_tester.num_hidden_layers + 1
                )

                
                self.assertEqual(len(output_hidden_states), expected_num_layers)
                self.assertListEqual(
                    list(output_hidden_states[0].shape[-2:]),
                    [self.model_tester.seq_length, self.model_tester.hidden_size],
                )

                
                self.assertEqual(len(output_attentions), self.model_tester.num_hidden_layers)

                seq_length = self.model_tester.seq_length
                key_length = getattr(self.model_tester, "key_length", seq_length)

                self.assertListEqual(
                    list(output_attentions[0].shape[-3:]),
                    [self.model_tester.num_attention_heads, seq_length, key_length],
                )


class TFGroupViTModelTester:
    def __init__(self, parent, is_training=True):
        self.parent = parent
        self.text_model_tester = TFGroupViTTextModelTester(parent)
        self.vision_model_tester = TFGroupViTVisionModelTester(parent)
        self.is_training = is_training

    def prepare_config_and_inputs(self):
        text_config, input_ids, attention_mask = self.text_model_tester.prepare_config_and_inputs()
        vision_config, pixel_values = self.vision_model_tester.prepare_config_and_inputs()

        config = self.get_config()

        return config, input_ids, attention_mask, pixel_values

    def get_config(self):
        return GroupViTConfig.from_text_vision_configs(
            self.text_model_tester.get_config(), self.vision_model_tester.get_config(), projection_dim=64
        )

    def create_and_check_model(self, config, input_ids, attention_mask, pixel_values):
        model = TFGroupViTModel(config)
        result = model(input_ids, pixel_values, attention_mask, training=False)
        self.parent.assertEqual(
            result.logits_per_image.shape, (self.vision_model_tester.batch_size, self.text_model_tester.batch_size)
        )
        self.parent.assertEqual(
            result.logits_per_text.shape, (self.text_model_tester.batch_size, self.vision_model_tester.batch_size)
        )

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        config, input_ids, attention_mask, pixel_values = config_and_inputs
        inputs_dict = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "pixel_values": pixel_values,
            "return_loss": True,
        }
        return config, inputs_dict


@require_tf
class TFGroupViTModelTest(TFModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (TFGroupViTModel,) if is_tf_available() else ()
    pipeline_model_mapping = {"feature-extraction": TFGroupViTModel} if is_tf_available() else {}
    test_head_masking = False
    test_pruning = False
    test_resize_embeddings = False
    test_attention_outputs = False
    test_onnx = False

    def check_pt_tf_outputs(self, tf_outputs, pt_outputs, model_class, tol=1e-4, name="outputs", attributes=None):
        
        super().check_pt_tf_outputs(tf_outputs, pt_outputs, model_class, tol, name, attributes)

    def setUp(self):
        self.model_tester = TFGroupViTModelTester(self)

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)

    @unittest.skip(reason="hidden_states are tested in individual model tests")
    def test_hidden_states_output(self):
        pass

    @unittest.skip(reason="input_embeds are tested in individual model tests")
    def test_inputs_embeds(self):
        pass

    @unittest.skip(reason="CLIPModel does not have input/output embeddings")
    def test_model_common_attributes(self):
        pass

    @require_tensorflow_probability
    @slow
    def test_keras_fit(self):
        super().test_keras_fit()

    @is_pt_tf_cross_test
    def test_pt_tf_model_equivalence(self):
        
        
        
        
        
        
        
        import torch

        seed = 158
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        tf.random.set_seed(seed)
        return super().test_pt_tf_model_equivalence()

    
    
    def test_keras_save_load(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        
        if self.__class__.__name__ == "TFGroupViTModelTest":
            inputs_dict.pop("return_loss", None)

        tf_main_layer_classes = {
            module_member
            for model_class in self.all_model_classes
            for module in (import_module(model_class.__module__),)
            for module_member_name in dir(module)
            if module_member_name.endswith("MainLayer")
            
            and module_member_name[: -len("MainLayer")] == model_class.__name__[: -len("Model")]
            for module_member in (getattr(module, module_member_name),)
            if isinstance(module_member, type)
            and tf.keras.layers.Layer in module_member.__bases__
            and getattr(module_member, "_keras_serializable", False)
        }
        for main_layer_class in tf_main_layer_classes:
            
            if "T5" in main_layer_class.__name__:
                
                shared = TFSharedEmbeddings(99, 32, name="shared")
                config.use_cache = inputs_dict.pop("use_cache", None)
                main_layer = main_layer_class(config, embed_tokens=shared)
            else:
                main_layer = main_layer_class(config)

            symbolic_inputs = {
                name: tf.keras.Input(tensor.shape[1:], dtype=tensor.dtype) for name, tensor in inputs_dict.items()
            }

            model = tf.keras.Model(symbolic_inputs, outputs=main_layer(symbolic_inputs))
            outputs = model(inputs_dict)

            with tempfile.TemporaryDirectory() as tmpdirname:
                filepath = os.path.join(tmpdirname, "keras_model.h5")
                model.save(filepath)
                if "T5" in main_layer_class.__name__:
                    model = tf.keras.models.load_model(
                        filepath,
                        custom_objects={
                            main_layer_class.__name__: main_layer_class,
                            "TFSharedEmbeddings": TFSharedEmbeddings,
                        },
                    )
                else:
                    model = tf.keras.models.load_model(
                        filepath, custom_objects={main_layer_class.__name__: main_layer_class}
                    )
                assert isinstance(model, tf.keras.Model)
                after_outputs = model(inputs_dict)
                self.assert_outputs_same(after_outputs, outputs)

    @slow
    def test_model_from_pretrained(self):
        for model_name in TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            model = TFGroupViTModel.from_pretrained(model_name)
            self.assertIsNotNone(model)

    @unittest.skip(reason="Currently `saved_model` doesn't work with nested outputs.")
    @slow
    def test_saved_model_creation(self):
        pass

    @unittest.skip(reason="`saved_model` doesn't work with nested outputs so no preparation happens.")
    @slow
    def test_prepare_serving_output(self):
        pass



def prepare_img():
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    im = Image.open(requests.get(url, stream=True).raw)
    return im


@require_vision
@require_tf
class TFGroupViTModelIntegrationTest(unittest.TestCase):
    @slow
    def test_inference(self):
        model_name = "nvidia/groupvit-gcc-yfcc"
        model = TFGroupViTModel.from_pretrained(model_name)
        processor = CLIPProcessor.from_pretrained(model_name)

        image = prepare_img()
        inputs = processor(
            text=["a photo of a cat", "a photo of a dog"], images=image, padding=True, return_tensors="tf"
        )

        outputs = model(**inputs, training=False)

        
        self.assertEqual(
            outputs.logits_per_image.shape,
            tf.TensorShape((inputs.pixel_values.shape[0], inputs.input_ids.shape[0])),
        )
        self.assertEqual(
            outputs.logits_per_text.shape,
            tf.TensorShape((inputs.input_ids.shape[0], inputs.pixel_values.shape[0])),
        )

        expected_logits = tf.constant([[13.3523, 6.3629]])

        tf.debugging.assert_near(outputs.logits_per_image, expected_logits, atol=1e-3)