



"""simple docstring"""






import os
from shutil import copyfile
from typing import List, Optional, Tuple

from tokenizers import processors

from ...tokenization_utils import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
	from .tokenization_nllb import NllbTokenizer
else:
	__snake_case							:   Optional[Any]      =    None


__snake_case							:   Tuple      =    logging.get_logger(__name__)


__snake_case							:   str      =    {'vocab_file': 'sentencepiece.bpe.model', 'tokenizer_file': 'tokenizer.json'}

__snake_case							:   int      =    {
    'vocab_file': {
        'facebook/nllb-200-distilled-600M': (
            'https://huggingface.co/facebook/nllb-200-distilled-600M/resolve/main/sentencepiece.bpe.model'
        ),
    },
    'tokenizer_file': {
        'facebook/nllb-200-distilled-600M': (
            'https://huggingface.co/facebook/nllb-200-distilled-600M/resolve/main/tokenizer.json'
        ),
    },
}

__snake_case							:   Any      =    {
    'facebook/nllb-large-en-ro': 1_024,
    'facebook/nllb-200-distilled-600M': 1_024,
}

# fmt: off
__snake_case							:   List[Any]      =    ['ace_Arab', 'ace_Latn', 'acm_Arab', 'acq_Arab', 'aeb_Arab', 'afr_Latn', 'ajp_Arab', 'aka_Latn', 'amh_Ethi', 'apc_Arab', 'arb_Arab', 'ars_Arab', 'ary_Arab', 'arz_Arab', 'asm_Beng', 'ast_Latn', 'awa_Deva', 'ayr_Latn', 'azb_Arab', 'azj_Latn', 'bak_Cyrl', 'bam_Latn', 'ban_Latn', 'bel_Cyrl', 'bem_Latn', 'ben_Beng', 'bho_Deva', 'bjn_Arab', 'bjn_Latn', 'bod_Tibt', 'bos_Latn', 'bug_Latn', 'bul_Cyrl', 'cat_Latn', 'ceb_Latn', 'ces_Latn', 'cjk_Latn', 'ckb_Arab', 'crh_Latn', 'cym_Latn', 'dan_Latn', 'deu_Latn', 'dik_Latn', 'dyu_Latn', 'dzo_Tibt', 'ell_Grek', 'eng_Latn', 'epo_Latn', 'est_Latn', 'eus_Latn', 'ewe_Latn', 'fao_Latn', 'pes_Arab', 'fij_Latn', 'fin_Latn', 'fon_Latn', 'fra_Latn', 'fur_Latn', 'fuv_Latn', 'gla_Latn', 'gle_Latn', 'glg_Latn', 'grn_Latn', 'guj_Gujr', 'hat_Latn', 'hau_Latn', 'heb_Hebr', 'hin_Deva', 'hne_Deva', 'hrv_Latn', 'hun_Latn', 'hye_Armn', 'ibo_Latn', 'ilo_Latn', 'ind_Latn', 'isl_Latn', 'ita_Latn', 'jav_Latn', 'jpn_Jpan', 'kab_Latn', 'kac_Latn', 'kam_Latn', 'kan_Knda', 'kas_Arab', 'kas_Deva', 'kat_Geor', 'knc_Arab', 'knc_Latn', 'kaz_Cyrl', 'kbp_Latn', 'kea_Latn', 'khm_Khmr', 'kik_Latn', 'kin_Latn', 'kir_Cyrl', 'kmb_Latn', 'kon_Latn', 'kor_Hang', 'kmr_Latn', 'lao_Laoo', 'lvs_Latn', 'lij_Latn', 'lim_Latn', 'lin_Latn', 'lit_Latn', 'lmo_Latn', 'ltg_Latn', 'ltz_Latn', 'lua_Latn', 'lug_Latn', 'luo_Latn', 'lus_Latn', 'mag_Deva', 'mai_Deva', 'mal_Mlym', 'mar_Deva', 'min_Latn', 'mkd_Cyrl', 'plt_Latn', 'mlt_Latn', 'mni_Beng', 'khk_Cyrl', 'mos_Latn', 'mri_Latn', 'zsm_Latn', 'mya_Mymr', 'nld_Latn', 'nno_Latn', 'nob_Latn', 'npi_Deva', 'nso_Latn', 'nus_Latn', 'nya_Latn', 'oci_Latn', 'gaz_Latn', 'ory_Orya', 'pag_Latn', 'pan_Guru', 'pap_Latn', 'pol_Latn', 'por_Latn', 'prs_Arab', 'pbt_Arab', 'quy_Latn', 'ron_Latn', 'run_Latn', 'rus_Cyrl', 'sag_Latn', 'san_Deva', 'sat_Beng', 'scn_Latn', 'shn_Mymr', 'sin_Sinh', 'slk_Latn', 'slv_Latn', 'smo_Latn', 'sna_Latn', 'snd_Arab', 'som_Latn', 'sot_Latn', 'spa_Latn', 'als_Latn', 'srd_Latn', 'srp_Cyrl', 'ssw_Latn', 'sun_Latn', 'swe_Latn', 'swh_Latn', 'szl_Latn', 'tam_Taml', 'tat_Cyrl', 'tel_Telu', 'tgk_Cyrl', 'tgl_Latn', 'tha_Thai', 'tir_Ethi', 'taq_Latn', 'taq_Tfng', 'tpi_Latn', 'tsn_Latn', 'tso_Latn', 'tuk_Latn', 'tum_Latn', 'tur_Latn', 'twi_Latn', 'tzm_Tfng', 'uig_Arab', 'ukr_Cyrl', 'umb_Latn', 'urd_Arab', 'uzn_Latn', 'vec_Latn', 'vie_Latn', 'war_Latn', 'wol_Latn', 'xho_Latn', 'ydd_Hebr', 'yor_Latn', 'yue_Hant', 'zho_Hans', 'zho_Hant', 'zul_Latn']



class 			A__   (  __SCREAMING_SNAKE_CASE							):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE          =	VOCAB_FILES_NAMES
	SCREAMING_SNAKE_CASE          =	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	SCREAMING_SNAKE_CASE          =	PRETRAINED_VOCAB_FILES_MAP
	SCREAMING_SNAKE_CASE          =	['input_ids', 'attention_mask']
	SCREAMING_SNAKE_CASE          =	NllbTokenizer

	SCREAMING_SNAKE_CASE          =	[]
	SCREAMING_SNAKE_CASE          =	[]

	def __init__(     self:		List[str]       ,			_SCREAMING_SNAKE_CASE:		str=None       ,			_SCREAMING_SNAKE_CASE:		Tuple=None       ,			_SCREAMING_SNAKE_CASE:		str="<s>"       ,			_SCREAMING_SNAKE_CASE:		Union[str, Any]="</s>"       ,			_SCREAMING_SNAKE_CASE:		List[str]="</s>"       ,			_SCREAMING_SNAKE_CASE:		Any="<s>"       ,			_SCREAMING_SNAKE_CASE:		List[str]="<unk>"       ,			_SCREAMING_SNAKE_CASE:		int="<pad>"       ,			_SCREAMING_SNAKE_CASE:		int="<mask>"       ,			_SCREAMING_SNAKE_CASE:		Dict=None       ,			_SCREAMING_SNAKE_CASE:		Any=None       ,			_SCREAMING_SNAKE_CASE:		Any=None       ,			_SCREAMING_SNAKE_CASE:		Optional[int]=False       ,			**_SCREAMING_SNAKE_CASE:		Union[str, Any]       ,			)				->       Tuple:

				"""simple docstring"""

				__lowerCAmelCase      :  Optional[int]												=       AddedToken(_SCREAMING_SNAKE_CASE       ,			lstrip=_SCREAMING_SNAKE_CASE       ,			rstrip=_SCREAMING_SNAKE_CASE) if isinstance(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE) else mask_token
				__lowerCAmelCase      :  Union[str, Any]												=       legacy_behaviour
				super().__init__(
				    vocab_file=_SCREAMING_SNAKE_CASE       ,			tokenizer_file=_SCREAMING_SNAKE_CASE       ,			bos_token=_SCREAMING_SNAKE_CASE       ,			eos_token=_SCREAMING_SNAKE_CASE       ,			sep_token=_SCREAMING_SNAKE_CASE       ,			cls_token=_SCREAMING_SNAKE_CASE       ,			unk_token=_SCREAMING_SNAKE_CASE       ,			pad_token=_SCREAMING_SNAKE_CASE       ,			mask_token=_SCREAMING_SNAKE_CASE       ,			src_lang=_SCREAMING_SNAKE_CASE       ,			tgt_lang=_SCREAMING_SNAKE_CASE       ,			additional_special_tokens=_SCREAMING_SNAKE_CASE       ,			legacy_behaviour=_SCREAMING_SNAKE_CASE       ,			**_SCREAMING_SNAKE_CASE       ,			)

				__lowerCAmelCase      :  Tuple												=       vocab_file
				__lowerCAmelCase      :  int												=       False if not self.vocab_file else True

				__lowerCAmelCase      :  Optional[Any]												=       FAIRSEQ_LANGUAGE_CODES.copy()

				if additional_special_tokens is not None:
							# Only add those special tokens if they are not already there.
							_additional_special_tokens.extend(
							    [t for t in additional_special_tokens if t not in _additional_special_tokens])

				self.add_special_tokens({"additional_special_tokens": _additional_special_tokens})
				__lowerCAmelCase      :  Dict												=       {
				    lang_code: self.convert_tokens_to_ids(_SCREAMING_SNAKE_CASE) for lang_code in FAIRSEQ_LANGUAGE_CODES
				}

				__lowerCAmelCase      :  Dict												=       src_lang if src_lang is not None else "eng_Latn"
				__lowerCAmelCase      :  Dict												=       self.convert_tokens_to_ids(self._src_lang)
				__lowerCAmelCase      :  Any												=       tgt_lang
				self.set_src_lang_special_tokens(self._src_lang)

	@property
	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any])				->       str:

				"""simple docstring"""

				return self._src_lang

	@src_lang.setter
	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		str)				->       None:

				"""simple docstring"""

				__lowerCAmelCase      :  Tuple												=       new_src_lang
				self.set_src_lang_special_tokens(self._src_lang)

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str]       ,			_SCREAMING_SNAKE_CASE:		List[int]       ,			_SCREAMING_SNAKE_CASE:		Optional[List[int]] = None)				->       List[int]:

				"""simple docstring"""

				if token_ids_a is None:
							return self.prefix_tokens + token_ids_a + self.suffix_tokens
				# We don't expect to process pairs, but leave the pair logic for API consistency
				return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		List[int]       ,			_SCREAMING_SNAKE_CASE:		Optional[List[int]] = None)				->       List[int]:

				"""simple docstring"""

				__lowerCAmelCase      :  List[Any]												=       [self.sep_token_id]
				__lowerCAmelCase      :  List[str]												=       [self.cls_token_id]

				if token_ids_a is None:
							return len(cls + token_ids_a + sep) * [0]
				return len(cls + token_ids_a + sep + sep + token_ids_a + sep) * [0]

	def 							_SCREAMING_SNAKE_CASE   (     self:		Any       ,			_SCREAMING_SNAKE_CASE:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		str       ,			_SCREAMING_SNAKE_CASE:		Optional[str]       ,			_SCREAMING_SNAKE_CASE:		Optional[str]       ,			**_SCREAMING_SNAKE_CASE:		List[Any])				->       int:

				"""simple docstring"""

				if src_lang is None or tgt_lang is None:
							raise ValueError("Translation requires a `src_lang` and a `tgt_lang` for this model")
				__lowerCAmelCase      :  Dict												=       src_lang
				__lowerCAmelCase      :  List[str]												=       self(_SCREAMING_SNAKE_CASE       ,			add_special_tokens=_SCREAMING_SNAKE_CASE       ,			return_tensors=_SCREAMING_SNAKE_CASE       ,			**_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Union[str, Any]												=       self.convert_tokens_to_ids(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Optional[Any]												=       tgt_lang_id
				return inputs

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str]       ,			_SCREAMING_SNAKE_CASE:		List[str]       ,			_SCREAMING_SNAKE_CASE:		str = "eng_Latn"       ,			_SCREAMING_SNAKE_CASE:		Optional[List[str]] = None       ,			_SCREAMING_SNAKE_CASE:		str = "fra_Latn"       ,			**_SCREAMING_SNAKE_CASE:		List[Any]       ,			)				->       BatchEncoding:

				"""simple docstring"""

				__lowerCAmelCase      :  Any												=       src_lang
				__lowerCAmelCase      :  Union[str, Any]												=       tgt_lang
				return super().prepare_seqaseq_batch(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			**_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[int])				->       Dict:

				"""simple docstring"""

				return self.set_src_lang_special_tokens(self.src_lang)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[Any])				->       List[Any]:

				"""simple docstring"""

				return self.set_tgt_lang_special_tokens(self.tgt_lang)

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[int])				->       None:

				"""simple docstring"""

				__lowerCAmelCase      :  List[str]												=       self.convert_tokens_to_ids(_SCREAMING_SNAKE_CASE)

				if self.legacy_behaviour:
							__lowerCAmelCase      :  str												=       []
							__lowerCAmelCase      :  Union[str, Any]												=       [self.eos_token_id, self.cur_lang_code]
				else:
							__lowerCAmelCase      :  str												=       [self.cur_lang_code]
							__lowerCAmelCase      :  List[str]												=       [self.eos_token_id]

				__lowerCAmelCase      :  Union[str, Any]												=       self.convert_ids_to_tokens(self.prefix_tokens)
				__lowerCAmelCase      :  List[str]												=       self.convert_ids_to_tokens(self.suffix_tokens)

				__lowerCAmelCase      :  Optional[int]												=       processors.TemplateProcessing(
				    single=prefix_tokens_str + ["$A"] + suffix_tokens_str       ,			pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str       ,			special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str       ,			self.prefix_tokens + self.suffix_tokens))       ,			)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Union[str, Any]       ,			_SCREAMING_SNAKE_CASE:		str)				->       None:

				"""simple docstring"""

				__lowerCAmelCase      :  Any												=       self.convert_tokens_to_ids(_SCREAMING_SNAKE_CASE)
				if self.legacy_behaviour:
							__lowerCAmelCase      :  str												=       []
							__lowerCAmelCase      :  str												=       [self.eos_token_id, self.cur_lang_code]
				else:
							__lowerCAmelCase      :  Tuple												=       [self.cur_lang_code]
							__lowerCAmelCase      :  List[Any]												=       [self.eos_token_id]

				__lowerCAmelCase      :  str												=       self.convert_ids_to_tokens(self.prefix_tokens)
				__lowerCAmelCase      :  Union[str, Any]												=       self.convert_ids_to_tokens(self.suffix_tokens)

				__lowerCAmelCase      :  Optional[int]												=       processors.TemplateProcessing(
				    single=prefix_tokens_str + ["$A"] + suffix_tokens_str       ,			pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str       ,			special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str       ,			self.prefix_tokens + self.suffix_tokens))       ,			)







	def 							_SCREAMING_SNAKE_CASE   (     self:		Union[str, Any]       ,			_SCREAMING_SNAKE_CASE:		str       ,			_SCREAMING_SNAKE_CASE:		Optional[str] = None)				->       Tuple[str]:

				"""simple docstring"""

				if not self.can_save_slow_tokenizer:
							raise ValueError(
							    "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
							    "tokenizer.")

				if not os.path.isdir(_SCREAMING_SNAKE_CASE):
							logger.error(F"""Vocabulary path ({save_directory}) should be a directory.""")
							return
				__lowerCAmelCase      :  str												=       os.path.join(
				    _SCREAMING_SNAKE_CASE       ,			(filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"])

				if os.path.abspath(self.vocab_file) != os.path.abspath(_SCREAMING_SNAKE_CASE):
							copyfile(self.vocab_file       ,			_SCREAMING_SNAKE_CASE)

				return (out_vocab_file,)




"""simple docstring"""






import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_flax_cross_test,
    require_flax,
    require_torch,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import is_flax_available, is_torch_available, is_vision_available

from ...test_modeling_flax_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_flax_bert import FlaxBertModelTester
from ..clip.test_modeling_flax_clip import FlaxCLIPVisionModelTester
from ..vit.test_modeling_flax_vit import FlaxViTModelTester


if is_flax_available():
	from transformers import (
	    FlaxBertModel,
	    FlaxCLIPVisionModel,
	    FlaxVisionTextDualEncoderModel,
	    FlaxViTModel,
	    VisionTextDualEncoderConfig,
	    VisionTextDualEncoderProcessor,
	)
	from transformers.modeling_flax_pytorch_utils import (
	    convert_pytorch_state_dict_to_flax,
	    load_flax_weights_in_pytorch_model,
	)


if is_torch_available():
	import torch

	from transformers import VisionTextDualEncoderModel

if is_vision_available():
	from PIL import Image


def 			_lowercase				(  __snake_case	)						->   List[str]:
			if isinstance(__snake_case					,collections.abc.Iterable	):
						return x
			return (x, x)



@require_flax
class 			A__   :




	'''simple docstring'''





	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		Dict       ,			_SCREAMING_SNAKE_CASE:		Optional[Any])				->       str:

				"""simple docstring"""

				pass

	def 							_SCREAMING_SNAKE_CASE   (     self:		str)				->       int:

				"""simple docstring"""

				pass

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str])				->       Tuple:

				"""simple docstring"""

				pass

	def 							_SCREAMING_SNAKE_CASE   (     self:		Any       ,			_SCREAMING_SNAKE_CASE:		np.ndarray       ,			_SCREAMING_SNAKE_CASE:		np.ndarray       ,			_SCREAMING_SNAKE_CASE:		float)				->       List[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  str												=       np.abs((a - b)).max()
				self.assertLessEqual(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			F"""Difference between torch and flax is {diff} (>= {tol}).""")

	def 							_SCREAMING_SNAKE_CASE   (     self:		str       ,			_SCREAMING_SNAKE_CASE:		List[str]       ,			_SCREAMING_SNAKE_CASE:		Dict       ,			_SCREAMING_SNAKE_CASE:		Any       ,			_SCREAMING_SNAKE_CASE:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		Any=None       ,			**_SCREAMING_SNAKE_CASE:		Tuple)				->       List[str]:

				"""simple docstring"""

				__lowerCAmelCase      :  Tuple												=       VisionTextDualEncoderConfig.from_vision_text_configs(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  List[str]												=       FlaxVisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Optional[int]												=       model(input_ids=_SCREAMING_SNAKE_CASE       ,			pixel_values=_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE)

				self.assertEqual(output["text_embeds"].shape       ,			(input_ids.shape[0], config.projection_dim))
				self.assertEqual(output["image_embeds"].shape       ,			(pixel_values.shape[0], config.projection_dim))

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		str       ,			_SCREAMING_SNAKE_CASE:		Dict       ,			_SCREAMING_SNAKE_CASE:		Dict       ,			_SCREAMING_SNAKE_CASE:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[Any]=None       ,			**_SCREAMING_SNAKE_CASE:		List[str])				->       Dict:

				"""simple docstring"""

				__lowerCAmelCase						,	__lowerCAmelCase      :  List[Any]												=       self.get_vision_text_model(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Union[str, Any]												=       {"vision_model": vision_model, "text_model": text_model}
				__lowerCAmelCase      :  Optional[Any]												=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Union[str, Any]												=       model(input_ids=_SCREAMING_SNAKE_CASE       ,			pixel_values=_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE)

				self.assertEqual(output["text_embeds"].shape       ,			(input_ids.shape[0], model.config.projection_dim))
				self.assertEqual(output["image_embeds"].shape       ,			(pixel_values.shape[0], model.config.projection_dim))

	def 							_SCREAMING_SNAKE_CASE   (     self:		Dict       ,			_SCREAMING_SNAKE_CASE:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		List[str]       ,			_SCREAMING_SNAKE_CASE:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		List[Any]=None       ,			**_SCREAMING_SNAKE_CASE:		Tuple)				->       str:

				"""simple docstring"""

				__lowerCAmelCase						,	__lowerCAmelCase      :  Optional[Any]												=       self.get_vision_text_model(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Optional[Any]												=       {"vision_model": vision_model, "text_model": text_model}
				__lowerCAmelCase      :  Any												=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Optional[Any]												=       model(input_ids=_SCREAMING_SNAKE_CASE       ,			pixel_values=_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Any												=       output[0]

				with tempfile.TemporaryDirectory() as tmpdirname:
							model.save_pretrained(_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase      :  Union[str, Any]												=       FlaxVisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase      :  int												=       model(input_ids=_SCREAMING_SNAKE_CASE       ,			pixel_values=_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase      :  Union[str, Any]												=       after_output[0]
							__lowerCAmelCase      :  Any												=       np.amax(np.abs(out_a - out_a))
							self.assertLessEqual(_SCREAMING_SNAKE_CASE       ,			1e-3)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Tuple       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		int       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		Any=None       ,			**_SCREAMING_SNAKE_CASE:		List[str])				->       Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase						,	__lowerCAmelCase      :  Any												=       self.get_vision_text_model(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Union[str, Any]												=       {"vision_model": vision_model, "text_model": text_model}
				__lowerCAmelCase      :  Optional[Any]												=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Optional[Any]												=       model(
				    input_ids=_SCREAMING_SNAKE_CASE       ,			pixel_values=_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE       ,			output_attentions=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Optional[Any]												=       output.vision_model_output.attentions
				self.assertEqual(len(_SCREAMING_SNAKE_CASE)       ,			vision_config.num_hidden_layers)

				# in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
				__lowerCAmelCase      :  List[str]												=       to_atuple(vision_model.config.image_size)
				__lowerCAmelCase      :  Any												=       to_atuple(vision_model.config.patch_size)
				__lowerCAmelCase      :  Union[str, Any]												=       (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
				__lowerCAmelCase      :  Tuple												=       num_patches + 1
				self.assertEqual(vision_attentions[0].shape[-3:]       ,			(vision_config.num_attention_heads, seq_len, seq_len))

				__lowerCAmelCase      :  Union[str, Any]												=       output.text_model_output.attentions
				self.assertEqual(len(_SCREAMING_SNAKE_CASE)       ,			text_config.num_hidden_layers)

				self.assertEqual(
				    text_attentions[0].shape[-3:]       ,			(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])       ,			)

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str]       ,			_SCREAMING_SNAKE_CASE:		Union[str, Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		int)				->       str:

				"""simple docstring"""

				pt_model.to(_SCREAMING_SNAKE_CASE)
				pt_model.eval()

				# prepare inputs
				__lowerCAmelCase      :  Union[str, Any]												=       inputs_dict
				__lowerCAmelCase      :  Union[str, Any]												=       {k: torch.tensor(v.tolist()) for k, v in flax_inputs.items()}

				with torch.no_grad():
							__lowerCAmelCase      :  Any												=       pt_model(**_SCREAMING_SNAKE_CASE).to_tuple()

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

				# PT -> Flax
				with tempfile.TemporaryDirectory() as tmpdirname:
							pt_model.save_pretrained(_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase      :  str												=       FlaxVisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE       ,			from_pt=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  int												=       fx_model_loaded(**_SCREAMING_SNAKE_CASE).to_tuple()
				self.assertEqual(len(_SCREAMING_SNAKE_CASE)       ,			len(_SCREAMING_SNAKE_CASE)       ,			"Output lengths differ between Flax and PyTorch")
				for fx_output_loaded, pt_output in zip(fx_outputs_loaded[:4]       ,			pt_outputs[:4]):
							self.assert_almost_equals(_SCREAMING_SNAKE_CASE       ,			pt_output.numpy()       ,			4e-2)

				# Flax -> PT
				with tempfile.TemporaryDirectory() as tmpdirname:
							fx_model.save_pretrained(_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase      :  Union[str, Any]												=       VisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE       ,			from_flax=_SCREAMING_SNAKE_CASE)

				pt_model_loaded.to(_SCREAMING_SNAKE_CASE)
				pt_model_loaded.eval()

				with torch.no_grad():
							__lowerCAmelCase      :  Optional[Any]												=       pt_model_loaded(**_SCREAMING_SNAKE_CASE).to_tuple()

				self.assertEqual(len(_SCREAMING_SNAKE_CASE)       ,			len(_SCREAMING_SNAKE_CASE)       ,			"Output lengths differ between Flax and PyTorch")
				for fx_output, pt_output_loaded in zip(fx_outputs[:4]       ,			pt_outputs_loaded[:4]):
							self.assert_almost_equals(_SCREAMING_SNAKE_CASE       ,			pt_output_loaded.numpy()       ,			4e-2)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Union[str, Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		Dict)				->       Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  Union[str, Any]												=       VisionTextDualEncoderConfig.from_vision_text_configs(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Dict												=       VisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  str												=       FlaxVisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  str												=       convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,			_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Optional[Any]												=       fx_state

				self.check_pt_flax_equivalence(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		Tuple       ,			_SCREAMING_SNAKE_CASE:		Dict)				->       str:

				"""simple docstring"""

				__lowerCAmelCase      :  Optional[int]												=       VisionTextDualEncoderConfig.from_vision_text_configs(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Dict												=       VisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Tuple												=       FlaxVisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  int												=       load_flax_weights_in_pytorch_model(_SCREAMING_SNAKE_CASE       ,			fx_model.params)

				self.check_pt_flax_equivalence(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any])				->       str:

				"""simple docstring"""

				__lowerCAmelCase      :  List[str]												=       self.prepare_config_and_inputs()
				self.check_model_from_pretrained_configs(**_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Dict)				->       int:

				"""simple docstring"""

				__lowerCAmelCase      :  List[Any]												=       self.prepare_config_and_inputs()
				self.check_vision_text_dual_encoder_from_pretrained(**_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any])				->       List[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  Dict												=       self.prepare_config_and_inputs()
				self.check_save_load(**_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		int)				->       Dict:

				"""simple docstring"""

				__lowerCAmelCase      :  List[str]												=       self.prepare_config_and_inputs()
				self.check_vision_text_output_attention(**_SCREAMING_SNAKE_CASE)

	@is_pt_flax_cross_test
	def 							_SCREAMING_SNAKE_CASE   (     self:		Union[str, Any])				->       Any:

				"""simple docstring"""

				__lowerCAmelCase      :  Dict												=       self.prepare_config_and_inputs()
				__lowerCAmelCase      :  List[Any]												=       config_inputs_dict.pop("vision_config")
				__lowerCAmelCase      :  str												=       config_inputs_dict.pop("text_config")

				__lowerCAmelCase      :  Union[str, Any]												=       config_inputs_dict

				self.check_equivalence_pt_to_flax(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)
				self.check_equivalence_flax_to_pt(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)







	@slow
	def 							_SCREAMING_SNAKE_CASE   (     self:		str)				->       Dict:

				"""simple docstring"""

				__lowerCAmelCase						,	__lowerCAmelCase      :  Dict												=       self.get_pretrained_model_and_inputs()

				__lowerCAmelCase      :  Union[str, Any]												=       model_a(**_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  List[str]												=       outputs[0]

				with tempfile.TemporaryDirectory() as tmp_dirname:
							model_a.save_pretrained(_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase      :  List[Any]												=       FlaxVisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase      :  Union[str, Any]												=       model_a(**_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase      :  List[str]												=       after_outputs[0]
							__lowerCAmelCase      :  List[Any]												=       np.amax(np.abs(out_a - out_a))
							self.assertLessEqual(_SCREAMING_SNAKE_CASE       ,			1e-5)






@require_flax
class 			A__   (  __SCREAMING_SNAKE_CASE  ,		unittest.TestCase							):




	'''simple docstring'''





	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str])				->       List[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  Any												=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
				    "hf-internal-testing/tiny-random-vit"       ,			"hf-internal-testing/tiny-bert"       ,			vision_from_pt=_SCREAMING_SNAKE_CASE       ,			text_from_pt=_SCREAMING_SNAKE_CASE       ,			)
				__lowerCAmelCase      :  Union[str, Any]												=       13
				__lowerCAmelCase      :  Optional[int]												=       floats_tensor(
				    [
				        batch_size,
				        model.config.vision_config.num_channels,
				        model.config.vision_config.image_size,
				        model.config.vision_config.image_size,
				    ])
				__lowerCAmelCase      :  List[Any]												=       ids_tensor([batch_size, 4]       ,			model.config.text_config.vocab_size)
				__lowerCAmelCase      :  List[Any]												=       random_attention_mask([batch_size, 4])
				__lowerCAmelCase      :  str												=       {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

				return model, inputs

	def 							_SCREAMING_SNAKE_CASE   (     self:		int       ,			_SCREAMING_SNAKE_CASE:		Union[str, Any]       ,			_SCREAMING_SNAKE_CASE:		Union[str, Any])				->       Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase      :  List[str]												=       FlaxViTModel(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Union[str, Any]												=       FlaxBertModel(_SCREAMING_SNAKE_CASE)
				return vision_model, text_model







	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[int])				->       int:

				"""simple docstring"""

				__lowerCAmelCase      :  List[Any]												=       FlaxViTModelTester(self)
				__lowerCAmelCase      :  Optional[Any]												=       FlaxBertModelTester(self)
				__lowerCAmelCase      :  int												=       vit_model_tester.prepare_config_and_inputs()
				__lowerCAmelCase      :  List[str]												=       bert_model_tester.prepare_config_and_inputs()

				__lowerCAmelCase						,	__lowerCAmelCase      :  Tuple												=       vision_config_and_inputs

				__lowerCAmelCase						,	__lowerCAmelCase						,	__lowerCAmelCase						,	__lowerCAmelCase      :  Union[str, Any]												=       text_config_and_inputs

				# make sure that cross attention layers are added
				return {
				    "text_config": text_config,
				    "vision_config": vision_config,
				    "pixel_values": pixel_values,
				    "attention_mask": attention_mask,
				    "input_ids": input_ids,
				    "token_type_ids": token_type_ids,
				}






@require_torch
class 			A__   (  __SCREAMING_SNAKE_CASE  ,		unittest.TestCase							):




	'''simple docstring'''





	def 							_SCREAMING_SNAKE_CASE   (     self:		Dict)				->       List[str]:

				"""simple docstring"""

				__lowerCAmelCase      :  Optional[int]												=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
				    "hf-internal-testing/tiny-random-clip"       ,			"hf-internal-testing/tiny-bert"       ,			vision_from_pt=_SCREAMING_SNAKE_CASE       ,			text_from_pt=_SCREAMING_SNAKE_CASE       ,			)
				__lowerCAmelCase      :  Optional[int]												=       13
				__lowerCAmelCase      :  List[str]												=       floats_tensor(
				    [
				        batch_size,
				        model.config.vision_config.num_channels,
				        model.config.vision_config.image_size,
				        model.config.vision_config.image_size,
				    ])
				__lowerCAmelCase      :  Any												=       ids_tensor([batch_size, 4]       ,			model.config.text_config.vocab_size)
				__lowerCAmelCase      :  str												=       random_attention_mask([batch_size, 4])
				__lowerCAmelCase      :  Optional[int]												=       {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

				return model, inputs

	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		int       ,			_SCREAMING_SNAKE_CASE:		Union[str, Any])				->       int:

				"""simple docstring"""

				__lowerCAmelCase      :  int												=       FlaxCLIPVisionModel(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  List[str]												=       FlaxBertModel(_SCREAMING_SNAKE_CASE)
				return vision_model, text_model







	def 							_SCREAMING_SNAKE_CASE   (     self:		Union[str, Any])				->       Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  List[Any]												=       FlaxCLIPVisionModelTester(self)
				__lowerCAmelCase      :  str												=       FlaxBertModelTester(self)
				__lowerCAmelCase      :  Optional[Any]												=       clip_model_tester.prepare_config_and_inputs()
				__lowerCAmelCase      :  Dict												=       bert_model_tester.prepare_config_and_inputs()

				__lowerCAmelCase						,	__lowerCAmelCase      :  Any												=       vision_config_and_inputs

				__lowerCAmelCase						,	__lowerCAmelCase						,	__lowerCAmelCase						,	__lowerCAmelCase      :  List[Any]												=       text_config_and_inputs

				# make sure that cross attention layers are added
				return {
				    "text_config": text_config,
				    "vision_config": vision_config,
				    "pixel_values": pixel_values,
				    "attention_mask": attention_mask,
				    "input_ids": input_ids,
				    "token_type_ids": token_type_ids,
				}






@require_flax
@require_vision
class 			A__   (  unittest.TestCase							):




	'''simple docstring'''





	@slow
	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str])				->       List[str]:

				"""simple docstring"""

				__lowerCAmelCase      :  Dict												=       FlaxVisionTextDualEncoderModel.from_pretrained("clip-italian/clip-italian"       ,			logit_scale_init_value=1.0)
				__lowerCAmelCase      :  str												=       VisionTextDualEncoderProcessor.from_pretrained("clip-italian/clip-italian")

				__lowerCAmelCase      :  Tuple												=       Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
				__lowerCAmelCase      :  Optional[int]												=       processor(
				    text=["una foto di un gatto", "una foto di un cane"]       ,			images=_SCREAMING_SNAKE_CASE       ,			padding=_SCREAMING_SNAKE_CASE       ,			return_tensors="np")

				__lowerCAmelCase      :  List[str]												=       model(**_SCREAMING_SNAKE_CASE)

				# verify the logits
				self.assertEqual(outputs.logits_per_image.shape       ,			(inputs.pixel_values.shape[0], inputs.input_ids.shape[0]))
				self.assertEqual(
				    outputs.logits_per_text.shape       ,			(inputs.input_ids.shape[0], inputs.pixel_values.shape[0])       ,			)

				__lowerCAmelCase      :  List[str]												=       np.array([[1.228_4727, 0.310_4122]])

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