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 flax visiontextdualencoder model from typing import optional tuple import flax linen as nn import jax import jax numpy as jnp from flax core frozendict import frozendict freeze unfreeze from flax traverseutil import flattendict unflattendict from modelingflaxutils import flaxpretrainedmodel appendreplacereturndocstrings overwritecalldocstring from utils import addstartdocstrings logging from auto configurationauto import autoconfig from auto modelingflaxauto import flaxmodelmapping flaxautomodel from clip modelingflaxclip import flaxclipoutput flaxclipvisionmodel from configurationvisiontextdualencoder import visiontextdualencoderconfig logger logging getloggername configfordoc visiontextdualencoderconfig visiontextdualencoderstartdocstring r this class can be used to initialize a visiontext dual encoder model with any pretrained vision autoencoding model as the vision encoder and any pretrained text model as the text encoder the vision and text encoders are loaded via the flaxautomodel frompretrained method the projection layers are automatically added to the model and should be finetuned on a downstream task like contrastive imagetext modeling in lit zeroshot transfer with lockedimage text tuninghttps arxiv orgabs2111 07991 it is shown how leveraging pretrained lockedfrozen image and text model for contrastive learning yields significant improvment on new zeroshot vision tasks such as image classification or retrieval after such a visiontextdualencoder model has been trainedfinetuned it can be savedloaded just like any other models see the examples for more information this model inherits from pretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a flax linen modulehttps flax readthedocs ioenlatestapireferenceflax linenmodule html subclass use it as a regular flax linen module and refer to the flax documentation for all matter related to general usage and behavior finally this model supports inherent jax features such as justintime jit compilationhttps jax readthedocs ioenlatestjax htmljustintimecompilationjit automatic differentiationhttps jax readthedocs ioenlatestjax htmlautomaticdifferentiation vectorizationhttps jax readthedocs ioenlatestjax htmlvectorizationvmap parallelizationhttps jax readthedocs ioenlatestjax htmlparallelizationpmap parameters config visiontextdualencoderconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the flaxpretrainedmodel frompretrained method to load the model weights dtype jax numpy dtype optional defaults to jax numpy float32 the data type of the computation can be one of jax numpy float32 jax numpy float16 on gpus and jax numpy bfloat16 on tpus this can be used to enable mixedprecision training or halfprecision inference on gpus or tpus if specified all the computation will be performed with the given dtype note that this only specifies the dtype of the computation and does not influence the dtype of model parameters if you wish to change the dtype of the model parameters see flaxpretrainedmodel tofp16 and flaxpretrainedmodel tobf16 class flaxvisiontextdualencodermodulenn module config visiontextdualencoderconfig dtype jnp dtype jnp float32 def setupself visionconfig self config visionconfig textconfig self config textconfig self visionembeddim visionconfig hiddensize self textembeddim textconfig hiddensize self projectiondim self config projectiondim visionmodule flaxmodelmapping getself config visionconfig class flaxclipvisionmodel moduleclass textmodule flaxmodelmappingself config textconfig class moduleclass self visionmodel visionmodulevisionconfig dtypeself dtype self textmodel textmoduletextconfig dtypeself dtype self visualprojection nn dense self projectiondim dtypeself dtype kernelinitjax nn initializers normal0 02 usebiasfalse self textprojection nn dense self projectiondim dtypeself dtype kernelinitjax nn initializers normal0 02 usebiasfalse self logitscale self param logitscale lambda shape jnp onesshape self config logitscaleinitvalue def call self inputidsnone pixelvaluesnone attentionmasknone positionidsnone tokentypeidsnone deterministic bool true outputattentionsnone outputhiddenstatesnone returndictnone returndict returndict if returndict is not none else self config returndict visionoutputs self visionmodel pixelvaluespixelvalues deterministicdeterministic outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict textoutputs self textmodel inputidsinputids attentionmaskattentionmask tokentypeidstokentypeids positionidspositionids deterministicdeterministic outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict imageembeds visionoutputs1 imageembeds self visualprojectionimageembeds textembeds textoutputs1 textembeds self textprojectiontextembeds normalized features imageembeds imageembeds jnp linalg normimageembeds axis1 keepdimstrue textembeds textembeds jnp linalg normtextembeds axis1 keepdimstrue cosine similarity as logits logitscale jnp expself logitscale logitspertext jnp matmultextembeds imageembeds t logitscale logitsperimage logitspertext t if not returndict return logitsperimage logitspertext textembeds imageembeds textoutputs visionoutputs return flaxclipoutput logitsperimagelogitsperimage logitspertextlogitspertext textembedstextembeds imageembedsimageembeds textmodeloutputtextoutputs visionmodeloutputvisionoutputs addstartdocstringsvisiontextdualencoderstartdocstring class flaxvisiontextdualencodermodelflaxpretrainedmodel configclass visiontextdualencoderconfig moduleclass flaxvisiontextdualencodermodule def init self config visiontextdualencoderconfig inputshape optionaltuple none seed int 0 dtype jnp dtype jnp float32 doinit bool true kwargs if not doinit raise valueerror flaxvisiontextdualencodermodel cannot be created without initializing doinit must be true if inputshape is none inputshape 1 1 1 config visionconfig imagesize config visionconfig imagesize 3 module self moduleclassconfigconfig dtypedtype kwargs super initconfig module inputshapeinputshape seedseed dtypedtype def initweightsself rng jax random prngkey inputshape tuple params frozendict none frozendict init input tensor inputids jnp zerosinputshape0 dtypei4 positionids jnp broadcasttojnp arangejnp atleast2dinputids shape1 inputshape0 tokentypeids jnp oneslikeinputids attentionmask jnp oneslikeinputids pixelvalues jax random normalrng inputshape1 paramsrng dropoutrng jax random splitrng rngs params paramsrng dropout dropoutrng randomparams self module initrngs inputids pixelvalues attentionmask positionids tokentypeids params if params is not none randomparams flattendictunfreezerandomparams params flattendictunfreezeparams for missingkey in self missingkeys paramsmissingkey randomparamsmissingkey self missingkeys set return freezeunflattendictparams else return randomparams def call self inputids pixelvalues attentionmasknone positionidsnone tokentypeidsnone params dict none dropoutrng jax random prngkey none train bool false outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config returndict pixelvalues jnp transposepixelvalues 0 2 3 1 if positionids is none positionids jnp broadcasttojnp arangejnp atleast2dinputids shape1 inputids shape if tokentypeids is none tokentypeids jnp zeroslikeinputids if attentionmask is none attentionmask jnp oneslikeinputids handle any prng if needed rngs if dropoutrng is not none rngsdropout dropoutrng return self module apply params params or self params jnp arrayinputids dtypei4 jnp arraypixelvalues dtypejnp float32 jnp arrayattentionmask dtypei4 jnp arraypositionids dtypei4 jnp arraytokentypeids dtypei4 not train outputattentions outputhiddenstates returndict rngsrngs def gettextfeatures self inputids attentionmasknone positionidsnone tokentypeidsnone params dict none dropoutrng jax random prngkey none trainfalse r args inputids numpy ndarray of shape batchsize sequencelength indices of input sequence tokens in the vocabulary padding will be ignored by default should you provide it indices can be obtained using pretrainedtokenizer see pretrainedtokenizer encode and pretrainedtokenizer call for details what are input ids glossaryinputids returns textfeatures jnp ndarray of shape batchsize outputdim the text embeddings obtained by applying the projection layer to the pooled output of text model handle any prng if needed handle any prng if needed rngs if dropoutrng is not none rngsdropout dropoutrng def getfeaturesmodule pixelvalues deterministic visionoutputs module visionmodelpixelvaluespixelvalues deterministicdeterministic pooledoutput visionoutputs1 pooledoutput imagefeatures module visualprojectionpooledoutput return imagefeatures return self module apply params params or self params jnp arraypixelvalues dtypejnp float32 not train methodgetfeatures rngsrngs classmethod def fromvisiontextpretrained cls visionmodelnameorpath str none textmodelnameorpath str none modelargs kwargs flaxpretrainedmodel remove text vision kwargs from kwargs load and initialize the text and vision model instantiate config with corresponding kwargs init model the projection layers are always newly initialized when loading the model using pretrained vision and text model overwritecalldocstring flaxvisiontextdualencodermodel visiontextdualencoderinputsdocstring visiontextdualencodermodeldocstring appendreplacereturndocstrings flaxvisiontextdualencodermodel outputtypeflaxclipoutput configclassconfigfordoc 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 flax visiontextdualencoder model this class can be used to initialize a vision text dual encoder model with any pretrained vision autoencoding model as the vision encoder and any pretrained text model as the text encoder the vision and text encoders are loaded via the flaxautomodel from_pretrained method the projection layers are automatically added to the model and should be fine tuned on a downstream task like contrastive image text modeling in lit zero shot transfer with locked image text tuning https arxiv org abs 2111 07991 it is shown how leveraging pre trained locked frozen image and text model for contrastive learning yields significant improvment on new zero shot vision tasks such as image classification or retrieval after such a vision text dual encoder model has been trained fine tuned it can be saved loaded just like any other models see the examples for more information this model inherits from pretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a flax linen module https flax readthedocs io en latest api_reference flax linen module html subclass use it as a regular flax linen module and refer to the flax documentation for all matter related to general usage and behavior finally this model supports inherent jax features such as just in time jit compilation https jax readthedocs io en latest jax html just in time compilation jit automatic differentiation https jax readthedocs io en latest jax html automatic differentiation vectorization https jax readthedocs io en latest jax html vectorization vmap parallelization https jax readthedocs io en latest jax html parallelization pmap parameters config visiontextdualencoderconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the flaxpretrainedmodel from_pretrained method to load the model weights dtype jax numpy dtype optional defaults to jax numpy float32 the data type of the computation can be one of jax numpy float32 jax numpy float16 on gpus and jax numpy bfloat16 on tpus this can be used to enable mixed precision training or half precision inference on gpus or tpus if specified all the computation will be performed with the given dtype note that this only specifies the dtype of the computation and does not influence the dtype of model parameters if you wish to change the dtype of the model parameters see flaxpretrainedmodel to_fp16 and flaxpretrainedmodel to_bf16 args input_ids numpy ndarray of shape batch_size sequence_length indices of input sequence tokens in the vocabulary padding will be ignored by default should you provide it indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask position_ids numpy ndarray of shape batch_size sequence_length optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config max_position_embeddings 1 what are position ids glossary position ids pixel_values torch floattensor of shape batch_size num_channels height width pixel values padding will be ignored by default should you provide it pixel values can be obtained using an image processor e g if you use vit as the encoder you should use autoimageprocessor see vitimageprocessor __call__ for details output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple normalized features cosine similarity as logits init input tensor handle any prng if needed args input_ids numpy ndarray of shape batch_size sequence_length indices of input sequence tokens in the vocabulary padding will be ignored by default should you provide it indices can be obtained using pretrainedtokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids returns text_features jnp ndarray of shape batch_size output_dim the text embeddings obtained by applying the projection layer to the pooled output of text model handle any prng if needed args pixel_values numpy ndarray of shape batch_size num_channels height width pixel values padding will be ignored by default should you provide it pixel values can be obtained using imagefeatureextractionmixin see imagefeatureextractionmixin __call__ for details returns image_features jnp ndarray of shape batch_size output_dim the image embeddings obtained by applying the projection layer to the pooled output of vision model handle any prng if needed pooled_output params vision_model_name_or_path str optional defaults to none information necessary to initiate the vision model can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the root level like bert base uncased or namespaced under a user or organization name like dbmdz bert base german cased a path to a directory containing model weights saved using flaxpretrainedmodel save_pretrained e g my_model_directory a path or url to a pytorch checkpoint folder e g pt_model in this case from_pt should be set to true and a configuration object should be provided as config argument this loading path is slower than converting the pytorch checkpoint in a flax model using the provided conversion scripts and loading the flax model afterwards text_model_name_or_path str optional information necessary to initiate the text model can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the root level like bert base uncased or namespaced under a user or organization name like dbmdz bert base german cased a path to a directory containing model weights saved using flaxpretrainedmodel save_pretrained e g my_model_directory a path or url to a pytorch checkpoint folder e g pt_model in this case from_pt should be set to true and a configuration object should be provided as config argument this loading path is slower than converting the pytorch checkpoint in a flax model using the provided conversion scripts and loading the flax model afterwards model_args remaining positional arguments optional all remaning positional arguments will be passed to the underlying model s __init__ method kwargs remaining dictionary of keyword arguments optional can be used to update the configuration object after it being loaded and initiate the model e g output_attentions true to update the text configuration use the prefix text_ for each configuration parameter to update the vision configuration use the prefix vision_ for each configuration parameter to update the parent model configuration do not use a prefix for each configuration parameter behaves differently depending on whether a config is provided or automatically loaded example python from transformers import flaxvisiontextdualencodermodel initialize a model from pretrained vit and bert models note that the projection layers will be randomly initialized model flaxvisiontextdualencodermodel from_vision_text_pretrained google vit base patch16 224 bert base uncased saving model after fine tuning model save_pretrained vit bert load fine tuned model model flaxvisiontextdualencodermodel from_pretrained vit bert remove text vision kwargs from kwargs load and initialize the text and vision model instantiate config with corresponding kwargs init model the projection layers are always newly initialized when loading the model using pre trained vision and text model returns examples python from pil import image import requests import jax from transformers import flaxvisiontextdualencodermodel visiontextdualencoderprocessor autoimageprocessor autotokenizer tokenizer autotokenizer from_pretrained bert base uncased image_processor autoimageprocesor from_pretrained google vit base patch16 224 processor visiontextdualencoderprocessor image_processor tokenizer model flaxvisiontextdualencodermodel from_vision_text_pretrained google vit base patch16 224 bert base uncased contrastive training urls http images cocodataset org val2017 000000039769 jpg https farm3 staticflickr com 2674 5850229113_4fe05d5265_z jpg images image open requests get url stream true raw for url in urls inputs processor text a photo of a cat a photo of a dog images images return_tensors np padding true outputs model input_ids inputs input_ids attention_mask inputs attention_mask pixel_values inputs pixel_values logits_per_image outputs logits_per_image this is the image text similarity score save and load from pretrained model save_pretrained vit bert model flaxvisiontextdualencodermodel from_pretrained vit bert inference outputs model inputs logits_per_image outputs logits_per_image this is the image text similarity score probs jax nn softmax logits_per_image axis 1 we can take the softmax to get the label probabilities
from typing import Optional, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
from flax.traverse_util import flatten_dict, unflatten_dict

from ...modeling_flax_utils import FlaxPreTrainedModel, append_replace_return_docstrings, overwrite_call_docstring
from ...utils import add_start_docstrings, logging
from ..auto.configuration_auto import AutoConfig
from ..auto.modeling_flax_auto import FLAX_MODEL_MAPPING, FlaxAutoModel
from ..clip.modeling_flax_clip import FlaxCLIPOutput, FlaxCLIPVisionModel
from .configuration_vision_text_dual_encoder import VisionTextDualEncoderConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "VisionTextDualEncoderConfig"

VISION_TEXT_DUAL_ENCODER_START_DOCSTRING = r


VISION_TEXT_DUAL_ENCODER_INPUTS_DOCSTRING = r


class FlaxVisionTextDualEncoderModule(nn.Module):
    config: VisionTextDualEncoderConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self):
        vision_config = self.config.vision_config
        text_config = self.config.text_config

        self.vision_embed_dim = vision_config.hidden_size
        self.text_embed_dim = text_config.hidden_size
        self.projection_dim = self.config.projection_dim

        vision_module = FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__, FlaxCLIPVisionModel).module_class
        text_module = FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class

        self.vision_model = vision_module(vision_config, dtype=self.dtype)
        self.text_model = text_module(text_config, dtype=self.dtype)

        self.visual_projection = nn.Dense(
            self.projection_dim,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.normal(0.02),
            use_bias=False,
        )
        self.text_projection = nn.Dense(
            self.projection_dim,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.normal(0.02),
            use_bias=False,
        )

        self.logit_scale = self.param(
            "logit_scale", lambda _, shape: jnp.ones(shape) * self.config.logit_scale_init_value, []
        )

    def __call__(
        self,
        input_ids=None,
        pixel_values=None,
        attention_mask=None,
        position_ids=None,
        token_type_ids=None,
        deterministic: bool = True,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        vision_outputs = self.vision_model(
            pixel_values=pixel_values,
            deterministic=deterministic,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        text_outputs = self.text_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            deterministic=deterministic,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        image_embeds = vision_outputs[1]
        image_embeds = self.visual_projection(image_embeds)

        text_embeds = text_outputs[1]
        text_embeds = self.text_projection(text_embeds)

        
        image_embeds = image_embeds / jnp.linalg.norm(image_embeds, axis=-1, keepdims=True)
        text_embeds = text_embeds / jnp.linalg.norm(text_embeds, axis=-1, keepdims=True)

        
        logit_scale = jnp.exp(self.logit_scale)
        logits_per_text = jnp.matmul(text_embeds, image_embeds.T) * logit_scale
        logits_per_image = logits_per_text.T

        if not return_dict:
            return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)

        return FlaxCLIPOutput(
            logits_per_image=logits_per_image,
            logits_per_text=logits_per_text,
            text_embeds=text_embeds,
            image_embeds=image_embeds,
            text_model_output=text_outputs,
            vision_model_output=vision_outputs,
        )


@add_start_docstrings(VISION_TEXT_DUAL_ENCODER_START_DOCSTRING)
class FlaxVisionTextDualEncoderModel(FlaxPreTrainedModel):
    config_class = VisionTextDualEncoderConfig
    module_class = FlaxVisionTextDualEncoderModule

    def __init__(
        self,
        config: VisionTextDualEncoderConfig,
        input_shape: Optional[Tuple] = None,
        seed: int = 0,
        dtype: jnp.dtype = jnp.float32,
        _do_init: bool = True,
        **kwargs,
    ):
        if not _do_init:
            raise ValueError(
                "`FlaxVisionTextDualEncoderModel` cannot be created without initializing, `_do_init` must be `True`."
            )

        if input_shape is None:
            input_shape = ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

        module = self.module_class(config=config, dtype=dtype, **kwargs)
        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)

    def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
        
        input_ids = jnp.zeros(input_shape[0], dtype="i4")
        position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape[0])
        token_type_ids = jnp.ones_like(input_ids)
        attention_mask = jnp.ones_like(input_ids)

        pixel_values = jax.random.normal(rng, input_shape[1])

        params_rng, dropout_rng = jax.random.split(rng)
        rngs = {"params": params_rng, "dropout": dropout_rng}

        random_params = self.module.init(rngs, input_ids, pixel_values, attention_mask, position_ids, token_type_ids)[
            "params"
        ]

        if params is not None:
            random_params = flatten_dict(unfreeze(random_params))
            params = flatten_dict(unfreeze(params))
            for missing_key in self._missing_keys:
                params[missing_key] = random_params[missing_key]
            self._missing_keys = set()
            return freeze(unflatten_dict(params))
        else:
            return random_params

    def __call__(
        self,
        input_ids,
        pixel_values,
        attention_mask=None,
        position_ids=None,
        token_type_ids=None,
        params: dict = None,
        dropout_rng: jax.random.PRNGKey = None,
        train: bool = False,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        pixel_values = jnp.transpose(pixel_values, (0, 2, 3, 1))

        if position_ids is None:
            position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)

        if token_type_ids is None:
            token_type_ids = jnp.zeros_like(input_ids)

        if attention_mask is None:
            attention_mask = jnp.ones_like(input_ids)

        
        rngs = {}
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

        return self.module.apply(
            {"params": params or self.params},
            jnp.array(input_ids, dtype="i4"),
            jnp.array(pixel_values, dtype=jnp.float32),
            jnp.array(attention_mask, dtype="i4"),
            jnp.array(position_ids, dtype="i4"),
            jnp.array(token_type_ids, dtype="i4"),
            not train,
            output_attentions,
            output_hidden_states,
            return_dict,
            rngs=rngs,
        )

    def get_text_features(
        self,
        input_ids,
        attention_mask=None,
        position_ids=None,
        token_type_ids=None,
        params: dict = None,
        dropout_rng: jax.random.PRNGKey = None,
        train=False,
    ):
        r
        if position_ids is None:
            position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)

        if token_type_ids is None:
            token_type_ids = jnp.zeros_like(input_ids)

        if attention_mask is None:
            attention_mask = jnp.ones_like(input_ids)

        
        rngs = {}
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

        def _get_features(module, input_ids, attention_mask, position_ids, token_type_ids, deterministic):
            text_outputs = module.text_model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                token_type_ids=token_type_ids,
                deterministic=deterministic,
            )
            pooled_output = text_outputs[1]
            text_features = module.text_projection(pooled_output)
            return text_features

        return self.module.apply(
            {"params": params or self.params},
            jnp.array(input_ids, dtype="i4"),
            jnp.array(attention_mask, dtype="i4"),
            jnp.array(position_ids, dtype="i4"),
            jnp.array(token_type_ids, dtype="i4"),
            not train,
            method=_get_features,
            rngs=rngs,
        )

    def get_image_features(
        self, pixel_values, params: dict = None, dropout_rng: jax.random.PRNGKey = None, train=False
    ):
        r

        
        rngs = {}
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

        def _get_features(module, pixel_values, deterministic):
            vision_outputs = module.vision_model(pixel_values=pixel_values, deterministic=deterministic)
            pooled_output = vision_outputs[1]  
            image_features = module.visual_projection(pooled_output)
            return image_features

        return self.module.apply(
            {"params": params or self.params},
            jnp.array(pixel_values, dtype=jnp.float32),
            not train,
            method=_get_features,
            rngs=rngs,
        )

    @classmethod
    def from_vision_text_pretrained(
        cls,
        vision_model_name_or_path: str = None,
        text_model_name_or_path: str = None,
        *model_args,
        **kwargs,
    ) -> FlaxPreTrainedModel:
        

        kwargs_vision = {
            argument[len("vision_") :]: value for argument, value in kwargs.items() if argument.startswith("vision_")
        }

        kwargs_text = {
            argument[len("text_") :]: value for argument, value in kwargs.items() if argument.startswith("text_")
        }

        
        for key in kwargs_vision.keys():
            del kwargs["vision_" + key]
        for key in kwargs_text.keys():
            del kwargs["text_" + key]

        
        vision_model = kwargs_vision.pop("model", None)
        if vision_model is None:
            if vision_model_name_or_path is None:
                raise ValueError(
                    "If `vision_model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
                )

            if "config" not in kwargs_vision:
                vision_config = AutoConfig.from_pretrained(vision_model_name_or_path)

            if vision_config.model_type == "clip":
                kwargs_vision["config"] = vision_config.vision_config
                vision_model = FlaxCLIPVisionModel.from_pretrained(
                    vision_model_name_or_path, *model_args, **kwargs_vision
                )
            else:
                kwargs_vision["config"] = vision_config
                vision_model = FlaxAutoModel.from_pretrained(vision_model_name_or_path, *model_args, **kwargs_vision)

        text_model = kwargs_text.pop("model", None)
        if text_model is None:
            if text_model_name_or_path is None:
                raise ValueError(
                    "If `text_model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
                )

            if "config" not in kwargs_text:
                text_config = AutoConfig.from_pretrained(text_model_name_or_path)
                kwargs_text["config"] = text_config

            text_model = FlaxAutoModel.from_pretrained(text_model_name_or_path, *model_args, **kwargs_text)

        
        dtype = kwargs.pop("dtype", jnp.float32)
        config = VisionTextDualEncoderConfig.from_vision_text_configs(vision_model.config, text_model.config, **kwargs)

        
        model = cls(config, *model_args, dtype=dtype, **kwargs)

        model.params["vision_model"] = vision_model.params
        model.params["text_model"] = text_model.params

        
        
        logger.warning(
            "The projection layer and logit scale weights `[('visual_projection', 'kernel'), ('text_projection',"
            " 'kernel'), ('logit_scale',)]` are newly initialized. You should probably TRAIN this model on a"
            " down-stream task to be able to use it for predictions and inference."
        )

        return model


VISION_TEXT_DUAL_ENCODER_MODEL_DOCSTRING = r

overwrite_call_docstring(
    FlaxVisionTextDualEncoderModel,
    VISION_TEXT_DUAL_ENCODER_INPUTS_DOCSTRING + VISION_TEXT_DUAL_ENCODER_MODEL_DOCSTRING,
)
append_replace_return_docstrings(
    FlaxVisionTextDualEncoderModel, output_type=FlaxCLIPOutput, config_class=_CONFIG_FOR_DOC
)