





from typing import List

from packaging import version
from sklearn.metrics import fa_score

import datasets
from datasets.config import PY_VERSION


if PY_VERSION < version.parse('''3.8'''):
	import importlib_metadata
else:
	import importlib.metadata as importlib_metadata


# TODO(Patrick/Anton)
lowerCAmelCase_												=       '''\
'''

lowerCAmelCase_												=       '''\
    XTREME-S is a benchmark to evaluate universal cross-lingual speech representations in many languages.
    XTREME-S covers four task families: speech recognition, classification, speech-to-text translation and retrieval.
'''

lowerCAmelCase_												=       '''
Compute XTREME-S evaluation metric associated to each XTREME-S dataset.
Args:
    predictions: list of predictions to score.
        Each translation should be tokenized into a list of tokens.
    references: list of lists of references for each translation.
        Each reference should be tokenized into a list of tokens.
    bleu_kwargs: optional dict of keywords to be passed when computing \'bleu\'.
        Keywords include Dict can be one of \'smooth_method\', \'smooth_value\', \'force\', \'lowercase\',
        \'tokenize\', \'use_effective_order\'.
    wer_kwargs: optional dict of keywords to be passed when computing \'wer\' and \'cer\'.
        Keywords include \'concatenate_texts\'.
Returns: depending on the XTREME-S task, one or several of:
    "accuracy": Accuracy - for \'fleurs-lang_id\', \'minds14\'
    "f1": F1 score - for \'minds14\'
    "wer": Word error rate - for \'mls\', \'fleurs-asr\', \'voxpopuli\', \'babel\'
    "cer": Character error rate - for \'mls\', \'fleurs-asr\', \'voxpopuli\', \'babel\'
    "bleu": BLEU score according to the `sacrebleu` metric - for \'covost2\'
Examples:

    >>> xtreme_s_metric = datasets.load_metric(\'xtreme_s\', \'mls\')  # \'mls\', \'voxpopuli\', \'fleurs-asr\' or \'babel\'
    >>> references = ["it is sunny here", "paper and pen are essentials"]
    >>> predictions = ["it\'s sunny", "paper pen are essential"]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {\'wer\': 0.56, \'cer\': 0.27}

    >>> xtreme_s_metric = datasets.load_metric(\'xtreme_s\', \'covost2\')
    >>> references = ["bonjour paris", "il est necessaire de faire du sport de temps en temp"]
    >>> predictions = ["bonjour paris", "il est important de faire du sport souvent"]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {\'bleu\': 31.65}

    >>> xtreme_s_metric = datasets.load_metric(\'xtreme_s\', \'fleurs-lang_id\')
    >>> references = [0, 1, 0, 0, 1]
    >>> predictions = [0, 1, 1, 0, 0]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {\'accuracy\': 0.6}

    >>> xtreme_s_metric = datasets.load_metric(\'xtreme_s\', \'minds14\')
    >>> references = [0, 1, 0, 0, 1]
    >>> predictions = [0, 1, 1, 0, 0]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {\'f1\': 0.58, \'accuracy\': 0.6}
'''

lowerCAmelCase_												=       ['''fleurs-asr''', '''mls''', '''voxpopuli''', '''babel''', '''covost2''', '''fleurs-lang_id''', '''minds14''']
lowerCAmelCase_												=       ''''''

try:
	from jiwer import transforms as tr

	lowerCAmelCase_												=       True
except ImportError:
	lowerCAmelCase_												=       False

if _jiwer_available and version.parse(importlib_metadata.version('''jiwer''')) < version.parse('''2.3.0'''):



	class       __lowerCAmelCase (    tr.AbstractTransform  ):






				def __init__(self      ,							__magic_name__ = " "					)  ->     List[str]:


								'''simple docstring'''



								snake_case_		:					Dict			    =						sentence_delimiter






				def     lowerCamelCase						(self      ,							__magic_name__					)  ->     int:


								'''simple docstring'''



								return list(__magic_name__					)






				def     lowerCamelCase						(self      ,							__magic_name__					)  ->     str:


								'''simple docstring'''



								snake_case_		:					List[Any]			    =						[]
								for sent_idx, sentence in enumerate(__magic_name__					):
												chars.extend(self.process_string(__magic_name__					)					)
												if self.sentence_delimiter is not None and self.sentence_delimiter != "" and sent_idx < len(__magic_name__					) - 1:
																chars.append(self.sentence_delimiter					)
								return chars




	lowerCAmelCase_												=       tr.Compose(
	    [tr.RemoveMultipleSpaces(), tr.Strip(), SentencesToListOfCharacters(SENTENCE_DELIMITER)]
	)
elif _jiwer_available:
	lowerCAmelCase_												=       tr.Compose(
	    [
	        tr.RemoveMultipleSpaces(),
	        tr.Strip(),
	        tr.ReduceToSingleSentence(SENTENCE_DELIMITER),
	        tr.ReduceToListOfListOfChars(),
	    ]
	)
else:
	lowerCAmelCase_												=       None






def   lowerCamelCase_		(       _UpperCamelCase				,       _UpperCamelCase				)	->		Any:






				"""simple docstring"""


				return float((preds == labels).mean()				)






def   lowerCamelCase_		(       _UpperCamelCase				,       _UpperCamelCase				)	->		Tuple:






				"""simple docstring"""


				return {
				    "f1": float(fa_score(y_true=_UpperCamelCase				,       y_pred=_UpperCamelCase				,       average='''macro'''				)				),
				    "accuracy": simple_accuracy(_UpperCamelCase				,       _UpperCamelCase				),
				}






def   lowerCamelCase_		(       _UpperCamelCase				,       _UpperCamelCase				,       _UpperCamelCase="exp"				,       _UpperCamelCase=None				,       _UpperCamelCase=False				,       _UpperCamelCase=False				,       _UpperCamelCase=None				,       _UpperCamelCase=False				,       )	->		Any:






				"""simple docstring"""


				snake_case_		:					Any			    =						[[label] for label in labels]
				snake_case_		:					int			    =						list(_UpperCamelCase				)
				try:
								import sacrebleu as scb
				except ImportError:
								raise ValueError(
								    '''sacrebleu has to be installed in order to apply the bleu metric for covost2.'''
								    '''You can install it via `pip install sacrebleu`.'''				)

				if version.parse(scb.__version__				) < version.parse('''1.4.12'''				):
								raise ImportWarning(
								    '''To use `sacrebleu`, the module `sacrebleu>=1.4.12` is required, and the current version of `sacrebleu` doesn\'t match this condition.\n'''
								    '''You can install it with `pip install "sacrebleu>=1.4.12"`.'''				)

				snake_case_		:					Optional[int]			    =						len(labels[0]				)
				if any(len(_UpperCamelCase				) != references_per_prediction for refs in labels				):
								raise ValueError('''Sacrebleu requires the same number of references for each prediction'''				)
				snake_case_		:					Optional[Any]			    =						[[refs[i] for refs in labels] for i in range(_UpperCamelCase				)]
				snake_case_		:					List[Any]			    =						scb.corpus_bleu(
				    _UpperCamelCase				,       _UpperCamelCase				,       smooth_method=_UpperCamelCase				,       smooth_value=_UpperCamelCase				,       force=_UpperCamelCase				,       lowercase=_UpperCamelCase				,       use_effective_order=_UpperCamelCase				,       **({'''tokenize''': tokenize} if tokenize else {})				,       )
				return {"bleu": output.score}






def   lowerCamelCase_		(       _UpperCamelCase				,       _UpperCamelCase				,       _UpperCamelCase				,       _UpperCamelCase				)	->		Tuple:






				"""simple docstring"""


				try:
								from jiwer import compute_measures
				except ImportError:
								raise ValueError(
								    f'''jiwer has to be installed in order to apply the wer metric for {config_name}.'''
								    '''You can install it via `pip install jiwer`.'''				)

				if concatenate_texts:
								snake_case_		:					Dict			    =						compute_measures(_UpperCamelCase				,       _UpperCamelCase				)['''wer''']

								snake_case_		:					Optional[int]			    =						compute_measures(_UpperCamelCase				,       _UpperCamelCase				,       truth_transform=_UpperCamelCase				,       hypothesis_transform=_UpperCamelCase				)['''wer''']
								return {"wer": wer, "cer": cer}
				else:

								def compute_score(_UpperCamelCase				,       _UpperCamelCase				,       _UpperCamelCase="wer"				):
												snake_case_		:					str			    =						0
												snake_case_		:					List[str]			    =						0
												for prediction, reference in zip(_UpperCamelCase				,       _UpperCamelCase				):
																if score_type == "wer":
																				snake_case_		:					Dict			    =						compute_measures(_UpperCamelCase				,       _UpperCamelCase				)
																elif score_type == "cer":
																				snake_case_		:					Optional[int]			    =						compute_measures(
																				    _UpperCamelCase				,       _UpperCamelCase				,       truth_transform=_UpperCamelCase				,       hypothesis_transform=_UpperCamelCase				)
																incorrect += measures["substitutions"] + measures["deletions"] + measures["insertions"]
																total += measures["substitutions"] + measures["deletions"] + measures["hits"]
												return incorrect / total

								return {"wer": compute_score(_UpperCamelCase				,       _UpperCamelCase				,       '''wer'''				), "cer": compute_score(_UpperCamelCase				,       _UpperCamelCase				,       '''cer'''				)}





@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION,  _KWARGS_DESCRIPTION  )
class       __lowerCAmelCase (    datasets.Metric  ):






			def     lowerCamelCase						(self					)  ->     List[str]:


							'''simple docstring'''



							if self.config_name not in _CONFIG_NAMES:
											raise KeyError(F'''You should supply a configuration name selected in {_CONFIG_NAMES}'''					)

							snake_case_		:					List[str]			    =						'''int64''' if self.config_name in ['''fleurs-lang_id''', '''minds14'''] else '''string'''

							return datasets.MetricInfo(
							    description=_DESCRIPTION      ,							citation=_CITATION      ,							inputs_description=_KWARGS_DESCRIPTION      ,							features=datasets.Features(
							        {'''predictions''': datasets.Value(__magic_name__					), '''references''': datasets.Value(__magic_name__					)}					)      ,							codebase_urls=[]      ,							reference_urls=[]      ,							format='''numpy'''      ,							)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__=None      ,							__magic_name__=None					)  ->     int:


							'''simple docstring'''



							snake_case_		:					List[str]			    =						bleu_kwargs if bleu_kwargs is not None else {}
							snake_case_		:					List[str]			    =						wer_kwargs if wer_kwargs is not None else {}

							if self.config_name == "fleurs-lang_id":
											return {"accuracy": simple_accuracy(__magic_name__      ,							__magic_name__					)}
							elif self.config_name == "minds14":
											return fa_and_simple_accuracy(__magic_name__      ,							__magic_name__					)
							elif self.config_name == "covost2":
											snake_case_		:					Dict			    =						bleu_kwargs.pop('''smooth_method'''      ,							'''exp'''					)
											snake_case_		:					Any			    =						bleu_kwargs.pop('''smooth_value'''      ,							__magic_name__					)
											snake_case_		:					Dict			    =						bleu_kwargs.pop('''force'''      ,							__magic_name__					)
											snake_case_		:					Any			    =						bleu_kwargs.pop('''lowercase'''      ,							__magic_name__					)
											snake_case_		:					List[str]			    =						bleu_kwargs.pop('''tokenize'''      ,							__magic_name__					)
											snake_case_		:					str			    =						bleu_kwargs.pop('''use_effective_order'''      ,							__magic_name__					)
											return bleu(
											    preds=__magic_name__      ,							labels=__magic_name__      ,							smooth_method=__magic_name__      ,							smooth_value=__magic_name__      ,							force=__magic_name__      ,							lowercase=__magic_name__      ,							tokenize=__magic_name__      ,							use_effective_order=__magic_name__      ,							)
							elif self.config_name in ["fleurs-asr", "mls", "voxpopuli", "babel"]:
											snake_case_		:					List[str]			    =						wer_kwargs.pop('''concatenate_texts'''      ,							__magic_name__					)
											return wer_and_cer(__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							self.config_name					)
							else:
											raise KeyError(F'''You should supply a configuration name selected in {_CONFIG_NAMES}'''					)







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


lowerCAmelCase_												=       logging.get_logger(__name__)



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






			def     lowerCamelCase						(self					)  ->     int:


							'''simple docstring'''



							snake_case_		:					Any			    =						self.config.text_config
							snake_case_		:					Any			    =						self.config.vision_config

							snake_case_		:					Union[str, Any]			    =						self.config.projection_dim
							snake_case_		:					Dict			    =						text_config.hidden_size
							snake_case_		:					Union[str, Any]			    =						vision_config.hidden_size

							snake_case_		:					Union[str, Any]			    =						FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
							snake_case_		:					Dict			    =						FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__      ,							__magic_name__					).module_class

							snake_case_		:					Optional[Any]			    =						text_module(__magic_name__      ,							dtype=self.dtype					)
							snake_case_		:					Dict			    =						vision_module(__magic_name__      ,							dtype=self.dtype					)

							snake_case_		:					Dict			    =						nn.Dense(
							    self.projection_dim      ,							dtype=self.dtype      ,							kernel_init=jax.nn.initializers.normal(0.02					)      ,							use_bias=__magic_name__      ,							)
							snake_case_		:					Optional[Any]			    =						nn.Dense(
							    self.projection_dim      ,							dtype=self.dtype      ,							kernel_init=jax.nn.initializers.normal(0.02					)      ,							use_bias=__magic_name__      ,							)
							snake_case_		:					Optional[Any]			    =						self.param('''logit_scale'''      ,							jax.nn.initializers.ones      ,							[]					)






			def __call__(self      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__ = True      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							)  ->     Any:


							'''simple docstring'''



							snake_case_		:					Tuple			    =						return_dict if return_dict is not None else self.config.return_dict

							snake_case_		:					Tuple			    =						self.vision_model(
							    pixel_values=__magic_name__      ,							deterministic=__magic_name__      ,							output_attentions=__magic_name__      ,							output_hidden_states=__magic_name__      ,							return_dict=__magic_name__      ,							)

							snake_case_		:					Tuple			    =						self.text_model(
							    input_ids=__magic_name__      ,							attention_mask=__magic_name__      ,							token_type_ids=__magic_name__      ,							position_ids=__magic_name__      ,							deterministic=__magic_name__      ,							output_attentions=__magic_name__      ,							output_hidden_states=__magic_name__      ,							return_dict=__magic_name__      ,							)

							snake_case_		:					Dict			    =						vision_outputs[1]
							snake_case_		:					Optional[Any]			    =						self.visual_projection(__magic_name__					)

							snake_case_		:					List[str]			    =						text_outputs[1]
							snake_case_		:					Union[str, Any]			    =						self.text_projection(__magic_name__					)

							# normalized features
							snake_case_		:					List[str]			    =						image_embeds / jnp.linalg.norm(__magic_name__      ,							axis=-1      ,							keepdims=__magic_name__					)
							snake_case_		:					Any			    =						text_embeds / jnp.linalg.norm(__magic_name__      ,							axis=-1      ,							keepdims=__magic_name__					)

							# cosine similarity as logits
							snake_case_		:					Optional[Any]			    =						jnp.exp(self.logit_scale					)
							snake_case_		:					str			    =						jnp.matmul(__magic_name__      ,							image_embeds.T					) * logit_scale
							snake_case_		:					List[str]			    =						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=__magic_name__      ,							logits_per_text=__magic_name__      ,							text_embeds=__magic_name__      ,							image_embeds=__magic_name__      ,							text_model_output=__magic_name__      ,							vision_model_output=__magic_name__      ,							)




class       __lowerCAmelCase (    _a  ):
			lowerCamelCase_  :  Dict     =   HybridCLIPConfig
			lowerCamelCase_  :  int     =   FlaxHybridCLIPModule






			def __init__(self      ,							__magic_name__      ,							__magic_name__ = None      ,							__magic_name__ = 0      ,							__magic_name__ = jnp.floataa      ,							**__magic_name__      ,							)  ->     str:


							'''simple docstring'''



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

							snake_case_		:					Union[str, Any]			    =						self.module_class(config=__magic_name__      ,							dtype=__magic_name__      ,							**__magic_name__					)
							super().__init__(__magic_name__      ,							__magic_name__      ,							input_shape=__magic_name__      ,							seed=__magic_name__      ,							dtype=__magic_name__					)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__ = None					)  ->     FrozenDict:


							'''simple docstring'''



							snake_case_		:					List[Any]			    =						jnp.zeros(input_shape[0]      ,							dtype='''i4'''					)
							snake_case_		:					Optional[Any]			    =						jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__magic_name__					).shape[-1]					)      ,							input_shape[0]					)
							snake_case_		:					Union[str, Any]			    =						jnp.ones_like(__magic_name__					)
							snake_case_		:					List[Any]			    =						jnp.ones_like(__magic_name__					)

							snake_case_		:					Union[str, Any]			    =						jax.random.normal(__magic_name__      ,							input_shape[1]					)

							snake_case_    ,							snake_case_		:					Optional[int]			    =						jax.random.split(__magic_name__					)
							snake_case_		:					int			    =						{'''params''': params_rng, '''dropout''': dropout_rng}

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






			def __call__(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							__magic_name__ = False      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							)  ->     Any:


							'''simple docstring'''



							snake_case_		:					Optional[int]			    =						output_attentions if output_attentions is not None else self.config.output_attentions
							snake_case_		:					Tuple			    =						(
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							snake_case_		:					List[Any]			    =						return_dict if return_dict is not None else self.config.return_dict

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

							if token_type_ids is None:
											snake_case_		:					Tuple			    =						jnp.zeros_like(__magic_name__					)

							if attention_mask is None:
											snake_case_		:					Optional[Any]			    =						jnp.ones_like(__magic_name__					)

							# Handle any PRNG if needed
							snake_case_		:					List[str]			    =						{}
							if dropout_rng is not None:
											snake_case_		:					Tuple			    =						dropout_rng

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






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							__magic_name__=False      ,							)  ->     Union[str, Any]:


							'''simple docstring'''



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

							if token_type_ids is None:
											snake_case_		:					List[str]			    =						jnp.zeros_like(__magic_name__					)

							if attention_mask is None:
											snake_case_		:					Union[str, Any]			    =						jnp.ones_like(__magic_name__					)

							# Handle any PRNG if needed
							snake_case_		:					Optional[int]			    =						{}
							if dropout_rng is not None:
											snake_case_		:					int			    =						dropout_rng

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

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






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							__magic_name__=False					)  ->     Dict:


							'''simple docstring'''



							snake_case_		:					Any			    =						{}
							if dropout_rng is not None:
											snake_case_		:					Optional[int]			    =						dropout_rng

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

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






			@classmethod
			def     lowerCamelCase						(cls      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							*__magic_name__      ,							**__magic_name__      ,							)  ->     FlaxPreTrainedModel:


							'''simple docstring'''



							snake_case_		:					Optional[Any]			    =						{
							    argument[len('''text_'''					) :]: value for argument, value in kwargs.items() if argument.startswith('''text_'''					)
							}

							snake_case_		:					Dict			    =						{
							    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
							snake_case_		:					str			    =						kwargs_text.pop('''model'''      ,							__magic_name__					)
							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

															snake_case_		:					Dict			    =						AutoConfig.from_pretrained(__magic_name__					)
															snake_case_		:					int			    =						text_config

											snake_case_		:					Optional[int]			    =						FlaxAutoModel.from_pretrained(__magic_name__      ,							*__magic_name__      ,							**__magic_name__					)

							snake_case_		:					Any			    =						kwargs_vision.pop('''model'''      ,							__magic_name__					)
							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

															snake_case_		:					Optional[int]			    =						AutoConfig.from_pretrained(__magic_name__					)
															snake_case_		:					List[str]			    =						vision_config

											snake_case_		:					List[Any]			    =						FlaxAutoModel.from_pretrained(__magic_name__      ,							*__magic_name__      ,							**__magic_name__					)

							# instantiate config with corresponding kwargs
							snake_case_		:					Dict			    =						kwargs.pop('''dtype'''      ,							jnp.floataa					)
							snake_case_		:					str			    =						HybridCLIPConfig.from_text_vision_configs(text_model.config      ,							vision_model.config      ,							**__magic_name__					)

							# init model
							snake_case_		:					int			    =						cls(__magic_name__      ,							*__magic_name__      ,							dtype=__magic_name__      ,							**__magic_name__					)

							if vision_config.model_type == "clip":
											snake_case_		:					int			    =						vision_model.params['''vision_model''']
											snake_case_		:					Any			    =						vision_model.params['''visual_projection''']['''kernel''']
							else:
											snake_case_		:					List[Any]			    =						vision_model.params

							snake_case_		:					Union[str, Any]			    =						text_model.params

							return model
