



import argparse
import os

import align
import numpy as np
import requests
import tensorflow as tf
import torch
from PIL import Image
from tokenizer import Tokenizer

from transformers import (
    AlignConfig,
    AlignModel,
    AlignProcessor,
    BertConfig,
    BertTokenizer,
    EfficientNetConfig,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
__snake_case		:Optional[int] 						=	logging.get_logger(__name__)





def 	__snake_case      (     _UpperCAmelCase							):
    __a					=							tf.image.resize(snake_case_  ,					(346, 346)							)
    __a					=							tf.image.crop_to_bounding_box(snake_case_  ,					(346 - 289) // 2  ,					(346 - 289) // 2  ,					289  ,					289							)
    return image





def 	__snake_case      (     ):
    __a					=							EfficientNetConfig.from_pretrained('''google/efficientnet-b7'''							)
    __a					=							289
    __a					=							640
    __a					=							{'''0''': '''LABEL_0''', '''1''': '''LABEL_1'''}
    __a					=							{'''LABEL_0''': 0, '''LABEL_1''': 1}
    __a					=							[]

    __a					=							BertConfig()
    __a					=							AlignConfig.from_text_vision_configs(
        text_config=snake_case_  ,					vision_config=snake_case_  ,					projection_dim=640							)
    return config





def 	__snake_case      (     ):
    __a					=							'''http://images.cocodataset.org/val2017/000000039769.jpg'''
    __a					=							Image.open(requests.get(snake_case_  ,					stream=snake_case_							).raw							)
    return im





def 	__snake_case      (     ):
    __a					=							EfficientNetImageProcessor(
        do_center_crop=snake_case_  ,					rescale_factor=1 / 1_27.5  ,					rescale_offset=snake_case_  ,					do_normalize=snake_case_  ,					include_top=snake_case_  ,					resample=Image.BILINEAR  ,					)
    __a					=							BertTokenizer.from_pretrained('''bert-base-uncased'''							)
    __a					=							64
    __a					=							AlignProcessor(image_processor=snake_case_  ,					tokenizer=snake_case_							)
    return processor





def 	__snake_case      (     _UpperCAmelCase							):
    __a					=							[v.split('''_'''							)[0].split('''block'''							)[1] for v in original_param_names if v.startswith('''block'''							)]
    __a					=							list(set(snake_case_							)							)
    __a					=							sorted(snake_case_							)
    __a					=							len(snake_case_							)
    __a					=							{b: str(snake_case_							) for b, i in zip(snake_case_  ,					range(snake_case_							)							)}

    __a					=							[]
    rename_keys.append(('''stem_conv/kernel:0''', '''embeddings.convolution.weight''')							)
    rename_keys.append(('''stem_bn/gamma:0''', '''embeddings.batchnorm.weight''')							)
    rename_keys.append(('''stem_bn/beta:0''', '''embeddings.batchnorm.bias''')							)
    rename_keys.append(('''stem_bn/moving_mean:0''', '''embeddings.batchnorm.running_mean''')							)
    rename_keys.append(('''stem_bn/moving_variance:0''', '''embeddings.batchnorm.running_var''')							)

    for b in block_names:
        __a					=							block_name_mapping[b]
        rename_keys.append((f'block{b}_expand_conv/kernel:0', f'encoder.blocks.{hf_b}.expansion.expand_conv.weight')							)
        rename_keys.append((f'block{b}_expand_bn/gamma:0', f'encoder.blocks.{hf_b}.expansion.expand_bn.weight')							)
        rename_keys.append((f'block{b}_expand_bn/beta:0', f'encoder.blocks.{hf_b}.expansion.expand_bn.bias')							)
        rename_keys.append(
            (f'block{b}_expand_bn/moving_mean:0', f'encoder.blocks.{hf_b}.expansion.expand_bn.running_mean')							)
        rename_keys.append(
            (f'block{b}_expand_bn/moving_variance:0', f'encoder.blocks.{hf_b}.expansion.expand_bn.running_var')							)
        rename_keys.append(
            (f'block{b}_dwconv/depthwise_kernel:0', f'encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight')							)
        rename_keys.append((f'block{b}_bn/gamma:0', f'encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight')							)
        rename_keys.append((f'block{b}_bn/beta:0', f'encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias')							)
        rename_keys.append(
            (f'block{b}_bn/moving_mean:0', f'encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean')							)
        rename_keys.append(
            (f'block{b}_bn/moving_variance:0', f'encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var')							)

        rename_keys.append((f'block{b}_se_reduce/kernel:0', f'encoder.blocks.{hf_b}.squeeze_excite.reduce.weight')							)
        rename_keys.append((f'block{b}_se_reduce/bias:0', f'encoder.blocks.{hf_b}.squeeze_excite.reduce.bias')							)
        rename_keys.append((f'block{b}_se_expand/kernel:0', f'encoder.blocks.{hf_b}.squeeze_excite.expand.weight')							)
        rename_keys.append((f'block{b}_se_expand/bias:0', f'encoder.blocks.{hf_b}.squeeze_excite.expand.bias')							)
        rename_keys.append(
            (f'block{b}_project_conv/kernel:0', f'encoder.blocks.{hf_b}.projection.project_conv.weight')							)
        rename_keys.append((f'block{b}_project_bn/gamma:0', f'encoder.blocks.{hf_b}.projection.project_bn.weight')							)
        rename_keys.append((f'block{b}_project_bn/beta:0', f'encoder.blocks.{hf_b}.projection.project_bn.bias')							)
        rename_keys.append(
            (f'block{b}_project_bn/moving_mean:0', f'encoder.blocks.{hf_b}.projection.project_bn.running_mean')							)
        rename_keys.append(
            (f'block{b}_project_bn/moving_variance:0', f'encoder.blocks.{hf_b}.projection.project_bn.running_var')							)

    __a					=							{}
    for item in rename_keys:
        if item[0] in original_param_names:
            __a					=							'''vision_model.''' + item[1]

    # BERT text encoder
    __a					=							[]
    __a					=							'''tf_bert_model/bert'''
    __a					=							'''text_model'''
    for i in range(12							):
        rename_keys.append(
            (
                f'{old}/encoder/layer_._{i}/attention/self/query/kernel:0',
                f'{new}.encoder.layer.{i}.attention.self.query.weight',
            )							)
        rename_keys.append(
            (
                f'{old}/encoder/layer_._{i}/attention/self/query/bias:0',
                f'{new}.encoder.layer.{i}.attention.self.query.bias',
            )							)
        rename_keys.append(
            (
                f'{old}/encoder/layer_._{i}/attention/self/key/kernel:0',
                f'{new}.encoder.layer.{i}.attention.self.key.weight',
            )							)
        rename_keys.append(
            (
                f'{old}/encoder/layer_._{i}/attention/self/key/bias:0',
                f'{new}.encoder.layer.{i}.attention.self.key.bias',
            )							)
        rename_keys.append(
            (
                f'{old}/encoder/layer_._{i}/attention/self/value/kernel:0',
                f'{new}.encoder.layer.{i}.attention.self.value.weight',
            )							)
        rename_keys.append(
            (
                f'{old}/encoder/layer_._{i}/attention/self/value/bias:0',
                f'{new}.encoder.layer.{i}.attention.self.value.bias',
            )							)
        rename_keys.append(
            (
                f'{old}/encoder/layer_._{i}/attention/output/dense/kernel:0',
                f'{new}.encoder.layer.{i}.attention.output.dense.weight',
            )							)
        rename_keys.append(
            (
                f'{old}/encoder/layer_._{i}/attention/output/dense/bias:0',
                f'{new}.encoder.layer.{i}.attention.output.dense.bias',
            )							)
        rename_keys.append(
            (
                f'{old}/encoder/layer_._{i}/attention/output/LayerNorm/gamma:0',
                f'{new}.encoder.layer.{i}.attention.output.LayerNorm.weight',
            )							)
        rename_keys.append(
            (
                f'{old}/encoder/layer_._{i}/attention/output/LayerNorm/beta:0',
                f'{new}.encoder.layer.{i}.attention.output.LayerNorm.bias',
            )							)
        rename_keys.append(
            (
                f'{old}/encoder/layer_._{i}/intermediate/dense/kernel:0',
                f'{new}.encoder.layer.{i}.intermediate.dense.weight',
            )							)
        rename_keys.append(
            (
                f'{old}/encoder/layer_._{i}/intermediate/dense/bias:0',
                f'{new}.encoder.layer.{i}.intermediate.dense.bias',
            )							)
        rename_keys.append(
            (f'{old}/encoder/layer_._{i}/output/dense/kernel:0', f'{new}.encoder.layer.{i}.output.dense.weight')							)
        rename_keys.append(
            (f'{old}/encoder/layer_._{i}/output/dense/bias:0', f'{new}.encoder.layer.{i}.output.dense.bias')							)
        rename_keys.append(
            (f'{old}/encoder/layer_._{i}/output/LayerNorm/gamma:0', f'{new}.encoder.layer.{i}.output.LayerNorm.weight')							)
        rename_keys.append(
            (f'{old}/encoder/layer_._{i}/output/LayerNorm/beta:0', f'{new}.encoder.layer.{i}.output.LayerNorm.bias')							)

    rename_keys.append((f'{old}/embeddings/word_embeddings/weight:0', f'{new}.embeddings.word_embeddings.weight')							)
    rename_keys.append(
        (f'{old}/embeddings/position_embeddings/embeddings:0', f'{new}.embeddings.position_embeddings.weight')							)
    rename_keys.append(
        (f'{old}/embeddings/token_type_embeddings/embeddings:0', f'{new}.embeddings.token_type_embeddings.weight')							)
    rename_keys.append((f'{old}/embeddings/LayerNorm/gamma:0', f'{new}.embeddings.LayerNorm.weight')							)
    rename_keys.append((f'{old}/embeddings/LayerNorm/beta:0', f'{new}.embeddings.LayerNorm.bias')							)

    rename_keys.append((f'{old}/pooler/dense/kernel:0', f'{new}.pooler.dense.weight')							)
    rename_keys.append((f'{old}/pooler/dense/bias:0', f'{new}.pooler.dense.bias')							)
    rename_keys.append(('''dense/kernel:0''', '''text_projection.weight''')							)
    rename_keys.append(('''dense/bias:0''', '''text_projection.bias''')							)
    rename_keys.append(('''dense/bias:0''', '''text_projection.bias''')							)
    rename_keys.append(('''temperature:0''', '''temperature''')							)

    for item in rename_keys:
        if item[0] in original_param_names:
            __a					=							item[1]
    return key_mapping





def 	__snake_case      (     _UpperCAmelCase  ,					_UpperCAmelCase  ,					_UpperCAmelCase							):
    list(hf_params.keys()							)

    for key, value in tf_params.items():
        if key not in key_mapping:
            continue

        __a					=							key_mapping[key]
        if "_conv" in key and "kernel" in key:
            __a					=							torch.from_numpy(snake_case_							).permute(3  ,					2  ,					0  ,					1							)
        elif "embeddings" in key:
            __a					=							torch.from_numpy(snake_case_							)
        elif "depthwise_kernel" in key:
            __a					=							torch.from_numpy(snake_case_							).permute(2  ,					3  ,					0  ,					1							)
        elif "kernel" in key:
            __a					=							torch.from_numpy(np.transpose(snake_case_							)							)
        elif "temperature" in key:
            __a					=							value
        elif "bn/gamma" or "bn/beta" in key:
            __a					=							torch.from_numpy(np.transpose(snake_case_							)							).squeeze()
        else:
            __a					=							torch.from_numpy(snake_case_							)

        # Replace HF parameters with original TF model parameters
        hf_params[hf_key].copy_(snake_case_							)





@torch.no_grad()
def 	__snake_case      (     _UpperCAmelCase  ,					_UpperCAmelCase  ,					_UpperCAmelCase  ,					_UpperCAmelCase							):
    __a					=							64
    __a					=							Tokenizer(snake_case_							)
    __a					=							align.Align('''efficientnet-b7'''  ,					'''bert-base'''  ,					640  ,					snake_case_  ,					tok.get_vocab_size()							)
    original_model.compile()
    original_model.load_weights(snake_case_							)

    __a					=							original_model.trainable_variables
    __a					=							original_model.non_trainable_variables
    __a					=							{param.name: param.numpy() for param in tf_params}
    for param in tf_non_train_params:
        __a					=							param.numpy()
    __a					=							list(tf_params.keys()							)

    # Load HuggingFace model
    __a					=							get_align_config()
    __a					=							AlignModel(snake_case_							).eval()
    __a					=							hf_model.state_dict()

    # Create src-to-dst parameter name mapping dictionary
    print('''Converting parameters...'''							)
    __a					=							rename_keys(snake_case_							)
    replace_params(snake_case_  ,					snake_case_  ,					snake_case_							)

    # Initialize processor
    __a					=							get_processor()
    __a					=							processor(
        images=prepare_img()  ,					text='''A picture of a cat'''  ,					padding='''max_length'''  ,					max_length=64  ,					return_tensors='''pt'''							)

    # HF model inference
    hf_model.eval()
    with torch.no_grad():
        __a					=							hf_model(**snake_case_							)

    __a					=							outputs.image_embeds.detach().numpy()
    __a					=							outputs.text_embeds.detach().numpy()

    # Original model inference
    __a					=							False
    __a					=							EfficientNetImageProcessor(
        do_center_crop=snake_case_  ,					do_rescale=snake_case_  ,					do_normalize=snake_case_  ,					include_top=snake_case_  ,					resample=Image.BILINEAR  ,					)
    __a					=							tf_image_processor(images=prepare_img()  ,					return_tensors='''tf'''  ,					data_format='''channels_last'''							)['''pixel_values''']
    __a					=							tok(tf.constant(['''A picture of a cat''']							)							)

    __a					=							original_model.image_encoder(snake_case_  ,					training=snake_case_							)
    __a					=							original_model.text_encoder(snake_case_  ,					training=snake_case_							)

    __a					=							tf.nn.la_normalize(snake_case_  ,					axis=-1							)
    __a					=							tf.nn.la_normalize(snake_case_  ,					axis=-1							)

    # Check whether original and HF model outputs match  -> np.allclose
    if not np.allclose(snake_case_  ,					snake_case_  ,					atol=1E-3							):
        raise ValueError('''The predicted image features are not the same.'''							)
    if not np.allclose(snake_case_  ,					snake_case_  ,					atol=1E-3							):
        raise ValueError('''The predicted text features are not the same.'''							)
    print('''Model outputs match!'''							)

    if save_model:
        # Create folder to save model
        if not os.path.isdir(snake_case_							):
            os.mkdir(snake_case_							)
        # Save converted model and image processor
        hf_model.save_pretrained(snake_case_							)
        processor.save_pretrained(snake_case_							)

    if push_to_hub:
        # Push model and image processor to hub
        print('''Pushing converted ALIGN to the hub...'''							)
        processor.push_to_hub('''align-base'''							)
        hf_model.push_to_hub('''align-base'''							)


if __name__ == "__main__":
 __snake_case		:str 						=	argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--checkpoint_path''',
     default='''./weights/model-weights''',
     type=str,
     help='''Path to the pretrained TF ALIGN checkpoint.''',
 )
 parser.add_argument(
     '''--pytorch_dump_folder_path''',
     default='''hf_model''',
     type=str,
     help='''Path to the output PyTorch model directory.''',
 )
 parser.add_argument('''--save_model''', action='''store_true''', help='''Save model to local''')
 parser.add_argument('''--push_to_hub''', action='''store_true''', help='''Push model and image processor to the hub''')

 __snake_case		:int 						=	parser.parse_args()
 convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)











from typing import Optional, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
from configuration_hybrid_clip import HybridCLIPConfig
from flax.core.frozen_dict import FrozenDict

from transformers import FLAX_MODEL_MAPPING, FlaxCLIPVisionModel
from transformers.modeling_flax_utils import FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPOutput
from transformers.utils import logging


lowercase_      :							List[str]    							=					logging.get_logger(__name__)






class        __lowerCAmelCase					(			nn.Module					):
			snake_case_ :					HybridCLIPConfig
			snake_case_ :					jnp.dtype						=	jnp.floataa







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




									"""simple docstring"""


									_UpperCAmelCase						=	self.config.text_config
									_UpperCAmelCase						=	self.config.vision_config

									_UpperCAmelCase						=	self.config.projection_dim
									_UpperCAmelCase						=	text_config.hidden_size
									_UpperCAmelCase						=	vision_config.hidden_size

									_UpperCAmelCase						=	FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
									_UpperCAmelCase						=	FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__       ,       snake_case__ ).module_class

									_UpperCAmelCase						=	text_module(snake_case__       ,       dtype=self.dtype )
									_UpperCAmelCase						=	vision_module(snake_case__       ,       dtype=self.dtype )

									_UpperCAmelCase						=	nn.Dense(
									    self.projection_dim       ,       dtype=self.dtype       ,       kernel_init=jax.nn.initializers.normal(0.02 )       ,       use_bias=snake_case__       ,       )
									_UpperCAmelCase						=	nn.Dense(
									    self.projection_dim       ,       dtype=self.dtype       ,       kernel_init=jax.nn.initializers.normal(0.02 )       ,       use_bias=snake_case__       ,       )
									_UpperCAmelCase						=	self.param("logit_scale"       ,       jax.nn.initializers.ones       ,       [] )







			def __call__(			self       :     List[Any]       ,       snake_case__       :     List[Any]=None       ,       snake_case__       :     List[str]=None       ,       snake_case__       :     Optional[Any]=None       ,       snake_case__       :     Dict=None       ,       snake_case__       :     Tuple=None       ,       snake_case__       :     bool = True       ,       snake_case__       :     Dict=None       ,       snake_case__       :     Dict=None       ,       snake_case__       :     Dict=None       ,       ):




									"""simple docstring"""


									_UpperCAmelCase						=	return_dict if return_dict is not None else self.config.return_dict

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

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

									_UpperCAmelCase						=	vision_outputs[1]
									_UpperCAmelCase						=	self.visual_projection(snake_case__ )

									_UpperCAmelCase						=	text_outputs[1]
									_UpperCAmelCase						=	self.text_projection(snake_case__ )

									# normalized features
									_UpperCAmelCase						=	image_embeds / jnp.linalg.norm(snake_case__       ,       axis=-1       ,       keepdims=snake_case__ )
									_UpperCAmelCase						=	text_embeds / jnp.linalg.norm(snake_case__       ,       axis=-1       ,       keepdims=snake_case__ )

									# cosine similarity as logits
									_UpperCAmelCase						=	jnp.exp(self.logit_scale )
									_UpperCAmelCase						=	jnp.matmul(snake_case__       ,       image_embeds.T ) * logit_scale
									_UpperCAmelCase						=	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=snake_case__       ,       logits_per_text=snake_case__       ,       text_embeds=snake_case__       ,       image_embeds=snake_case__       ,       text_model_output=snake_case__       ,       vision_model_output=snake_case__       ,       )






class        __lowerCAmelCase					(			UpperCAmelCase__					):
			snake_case_ :					Union[str, Any]						=	HybridCLIPConfig
			snake_case_ :					Optional[Any]						=	FlaxHybridCLIPModule







			def __init__(			self       :     List[Any]       ,       snake_case__       :     HybridCLIPConfig       ,       snake_case__       :     Optional[Tuple] = None       ,       snake_case__       :     int = 0       ,       snake_case__       :     jnp.dtype = jnp.floataa       ,       **snake_case__       :     Any       ,       ):




									"""simple docstring"""


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

									_UpperCAmelCase						=	self.module_class(config=snake_case__       ,       dtype=snake_case__       ,       **snake_case__ )
									super().__init__(snake_case__       ,       snake_case__       ,       input_shape=snake_case__       ,       seed=snake_case__       ,       dtype=snake_case__ )







			def UpperCamelCase     (			self       :     Union[str, Any]       ,       snake_case__       :     jax.random.PRNGKey       ,       snake_case__       :     Tuple       ,       snake_case__       :     FrozenDict = None ):




									"""simple docstring"""


									_UpperCAmelCase						=	jnp.zeros(input_shape[0]       ,       dtype="i4" )
									_UpperCAmelCase						=	jnp.broadcast_to(jnp.arange(jnp.atleast_ad(snake_case__ ).shape[-1] )       ,       input_shape[0] )
									_UpperCAmelCase						=	jnp.ones_like(snake_case__ )
									_UpperCAmelCase						=	jnp.ones_like(snake_case__ )

									_UpperCAmelCase						=	jax.random.normal(snake_case__       ,       input_shape[1] )

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

									return self.module.init(snake_case__       ,       snake_case__       ,       snake_case__       ,       snake_case__       ,       snake_case__       ,       snake_case__ )["params"]







			def __call__(			self       :     Union[str, Any]       ,       snake_case__       :     Optional[int]       ,       snake_case__       :     Tuple       ,       snake_case__       :     Optional[int]=None       ,       snake_case__       :     Tuple=None       ,       snake_case__       :     Optional[Any]=None       ,       snake_case__       :     dict = None       ,       snake_case__       :     jax.random.PRNGKey = None       ,       snake_case__       :     bool = False       ,       snake_case__       :     Optional[bool] = None       ,       snake_case__       :     Optional[bool] = None       ,       snake_case__       :     Optional[bool] = None       ,       ):




									"""simple docstring"""


									_UpperCAmelCase						=	output_attentions if output_attentions is not None else self.config.output_attentions
									_UpperCAmelCase						=	(
									    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
									)
									_UpperCAmelCase						=	return_dict if return_dict is not None else self.config.return_dict

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

									if token_type_ids is None:
															_UpperCAmelCase						=	jnp.zeros_like(snake_case__ )

									if attention_mask is None:
															_UpperCAmelCase						=	jnp.ones_like(snake_case__ )

									# Handle any PRNG if needed
									_UpperCAmelCase						=	{}
									if dropout_rng is not None:
															_UpperCAmelCase						=	dropout_rng

									return self.module.apply(
									    {"params": params or self.params}       ,       jnp.array(snake_case__       ,       dtype="i4" )       ,       jnp.array(snake_case__       ,       dtype=jnp.floataa )       ,       jnp.array(snake_case__       ,       dtype="i4" )       ,       jnp.array(snake_case__       ,       dtype="i4" )       ,       jnp.array(snake_case__       ,       dtype="i4" )       ,       not train       ,       snake_case__       ,       snake_case__       ,       snake_case__       ,       rngs=snake_case__       ,       )







			def UpperCamelCase     (			self       :     int       ,       snake_case__       :     int       ,       snake_case__       :     List[str]=None       ,       snake_case__       :     Union[str, Any]=None       ,       snake_case__       :     str=None       ,       snake_case__       :     dict = None       ,       snake_case__       :     jax.random.PRNGKey = None       ,       snake_case__       :     int=False       ,       ):




									"""simple docstring"""


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

									if token_type_ids is None:
															_UpperCAmelCase						=	jnp.zeros_like(snake_case__ )

									if attention_mask is None:
															_UpperCAmelCase						=	jnp.ones_like(snake_case__ )

									# Handle any PRNG if needed
									_UpperCAmelCase						=	{}
									if dropout_rng is not None:
															_UpperCAmelCase						=	dropout_rng

									def _get_features(snake_case__       :     Union[str, Any]       ,       snake_case__       :     Optional[Any]       ,       snake_case__       :     List[Any]       ,       snake_case__       :     int       ,       snake_case__       :     Dict       ,       snake_case__       :     List[str] ):
															_UpperCAmelCase						=	module.text_model(
															    input_ids=snake_case__       ,       attention_mask=snake_case__       ,       position_ids=snake_case__       ,       token_type_ids=snake_case__       ,       deterministic=snake_case__       ,       )
															_UpperCAmelCase						=	text_outputs[1]
															_UpperCAmelCase						=	module.text_projection(snake_case__ )
															return text_features

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







			def UpperCamelCase     (			self       :     List[Any]       ,       snake_case__       :     Optional[Any]       ,       snake_case__       :     dict = None       ,       snake_case__       :     jax.random.PRNGKey = None       ,       snake_case__       :     str=False ):




									"""simple docstring"""


									_UpperCAmelCase						=	{}
									if dropout_rng is not None:
															_UpperCAmelCase						=	dropout_rng

									def _get_features(snake_case__       :     Optional[Any]       ,       snake_case__       :     Tuple       ,       snake_case__       :     int ):
															_UpperCAmelCase						=	module.vision_model(pixel_values=snake_case__       ,       deterministic=snake_case__ )
															_UpperCAmelCase						=	vision_outputs[1]  # pooled_output
															_UpperCAmelCase						=	module.visual_projection(snake_case__ )
															return image_features

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







			@classmethod
			def UpperCamelCase     (			cls       :     Tuple       ,       snake_case__       :     str = None       ,       snake_case__       :     str = None       ,       *snake_case__       :     Any       ,       **snake_case__       :     Tuple       ,       ):




									"""simple docstring"""


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

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

									# remove text, vision kwargs from kwargs
									for key in kwargs_text.keys():
															del kwargs["text_" + key]
									for key in kwargs_vision.keys():
															del kwargs["vision_" + key]

									# Load and initialize the text and vision model
									_UpperCAmelCase						=	kwargs_text.pop("model"       ,       snake_case__ )
									if text_model is None:
															assert (
															    text_model_name_or_path is not None
															), "If `model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
															from transformers import FlaxAutoModel

															if "config" not in kwargs_text:
																					from transformers import AutoConfig

																					_UpperCAmelCase						=	AutoConfig.from_pretrained(snake_case__ )
																					_UpperCAmelCase						=	text_config

															_UpperCAmelCase						=	FlaxAutoModel.from_pretrained(snake_case__       ,       *snake_case__       ,       **snake_case__ )

									_UpperCAmelCase						=	kwargs_vision.pop("model"       ,       snake_case__ )
									if vision_model is None:
															assert (
															    vision_model_name_or_path is not None
															), "If `model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
															from transformers import FlaxAutoModel

															if "config" not in kwargs_vision:
																					from transformers import AutoConfig

																					_UpperCAmelCase						=	AutoConfig.from_pretrained(snake_case__ )
																					_UpperCAmelCase						=	vision_config

															_UpperCAmelCase						=	FlaxAutoModel.from_pretrained(snake_case__       ,       *snake_case__       ,       **snake_case__ )

									# instantiate config with corresponding kwargs
									_UpperCAmelCase						=	kwargs.pop("dtype"       ,       jnp.floataa )
									_UpperCAmelCase						=	HybridCLIPConfig.from_text_vision_configs(text_model.config       ,       vision_model.config       ,       **snake_case__ )

									# init model
									_UpperCAmelCase						=	cls(snake_case__       ,       *snake_case__       ,       dtype=snake_case__       ,       **snake_case__ )

									if vision_config.model_type == "clip":
															_UpperCAmelCase						=	vision_model.params["vision_model"]
															_UpperCAmelCase						=	vision_model.params["visual_projection"]["kernel"]
									else:
															_UpperCAmelCase						=	vision_model.params

									_UpperCAmelCase						=	text_model.params

									return model




