




from __future__ import annotations

import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import is_tf_available, is_vision_available

from ...test_modeling_tf_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_tf_bert import TFBertModelTester
from ..clip.test_modeling_tf_clip import TFCLIPVisionModelTester
from ..deit.test_modeling_tf_deit import TFDeiTModelTester
from ..roberta.test_modeling_tf_roberta import TFRobertaModelTester
from ..vit.test_modeling_tf_vit import TFViTModelTester


if is_tf_available():
					from transformers import (
					    TFBertModel,
					    TFCLIPVisionModel,
					    TFDeiTModel,
					    TFRobertaModel,
					    TFVisionTextDualEncoderModel,
					    TFViTModel,
					    VisionTextDualEncoderConfig,
					)

if is_vision_available():
					from PIL import Image

					from transformers import VisionTextDualEncoderProcessor







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Optional[Any]  ):
					if isinstance(SCREAMING_SNAKE_CASE_   ,   collections.abc.Iterable  ):
										return x
					return (x, x)
@require_tf
class     a__					:






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A    ):




											"""simple docstring"""



											pass






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											pass






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											pass






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A	,    _A	,    _A=None	,    **_A    ):




											"""simple docstring"""



											__lowerCAmelCase            = VisionTextDualEncoderConfig.from_vision_text_configs(_A	,    _A    )

											__lowerCAmelCase            = TFVisionTextDualEncoderModel(_A    )

											__lowerCAmelCase            = model(input_ids=_A	,    pixel_values=_A	,    attention_mask=_A    )

											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	,    _A	,    _A	,    _A	,    _A	,    _A=None	,    **_A    ):




											"""simple docstring"""



											__lowerCAmelCase     ,							__lowerCAmelCase            = self.get_vision_text_model(_A	,    _A    )
											__lowerCAmelCase            = TFVisionTextDualEncoderModel(vision_model=_A	,    text_model=_A    )

											__lowerCAmelCase            = model(input_ids=_A	,    pixel_values=_A	,    attention_mask=_A    )

											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	,    _A	,    _A	,    _A	,    _A	,    _A=None	,    **_A    ):




											"""simple docstring"""



											__lowerCAmelCase     ,							__lowerCAmelCase            = self.get_vision_text_model(_A	,    _A    )
											__lowerCAmelCase            = {"vision_model": vision_model, "text_model": text_model}
											__lowerCAmelCase            = TFVisionTextDualEncoderModel.from_vision_text_pretrained(**_A    )

											__lowerCAmelCase            = model(input_ids=_A	,    pixel_values=_A	,    attention_mask=_A    )

											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	,    _A	,    _A	,    _A	,    _A	,    _A=None	,    **_A    ):




											"""simple docstring"""



											__lowerCAmelCase     ,							__lowerCAmelCase            = self.get_vision_text_model(_A	,    _A    )
											__lowerCAmelCase            = TFVisionTextDualEncoderModel(vision_model=_A	,    text_model=_A    )

											__lowerCAmelCase            = model(input_ids=_A	,    pixel_values=_A	,    attention_mask=_A    )
											__lowerCAmelCase            = output[0].numpy()

											with tempfile.TemporaryDirectory() as tmpdirname:
																model.save_pretrained(_A    )
																__lowerCAmelCase            = TFVisionTextDualEncoderModel.from_pretrained(_A    )

																__lowerCAmelCase            = model(input_ids=_A	,    pixel_values=_A	,    attention_mask=_A    )
																__lowerCAmelCase            = after_output[0].numpy()
																__lowerCAmelCase            = np.amax(np.abs(out_a - out_a    )    )
																self.assertLessEqual(_A	,    1E-5    )






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A	,    _A	,    _A=None	,    **_A    ):




											"""simple docstring"""



											__lowerCAmelCase     ,							__lowerCAmelCase            = self.get_vision_text_model(_A	,    _A    )
											__lowerCAmelCase            = TFVisionTextDualEncoderModel(vision_model=_A	,    text_model=_A    )

											__lowerCAmelCase            = model(
											    input_ids=_A	,    pixel_values=_A	,    attention_mask=_A	,    output_attentions=_A    )

											__lowerCAmelCase            = output.vision_model_output.attentions
											self.assertEqual(len(_A    )	,    vision_config.num_hidden_layers    )

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

											__lowerCAmelCase            = output.text_model_output.attentions
											self.assertEqual(len(_A    )	,    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	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



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






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.prepare_config_and_inputs()
											self.check_vision_text_dual_encoder_model(**_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.prepare_config_and_inputs()
											self.check_model_from_pretrained_configs(**_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.prepare_config_and_inputs()
											self.check_vision_text_dual_encoder_from_pretrained(**_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.prepare_config_and_inputs()
											self.check_save_load(**_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.prepare_config_and_inputs()
											self.check_vision_text_output_attention(**_A    )






						@slow
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase     ,							__lowerCAmelCase            = self.get_pretrained_model_and_inputs()

											__lowerCAmelCase            = model_a(**_A    )
											__lowerCAmelCase            = outputs[0].numpy()

											with tempfile.TemporaryDirectory() as tmp_dirname:
																model_a.save_pretrained(_A    )
																__lowerCAmelCase            = TFVisionTextDualEncoderModel.from_pretrained(_A    )

																__lowerCAmelCase            = model_a(**_A    )
																__lowerCAmelCase            = after_outputs[0].numpy()
																__lowerCAmelCase            = np.amax(np.abs(out_a - out_a    )    )
																self.assertLessEqual(_A	,    1E-5    )






@require_tf
class     a__					(      snake_case__				,       unittest.TestCase   ):






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFVisionTextDualEncoderModel.from_vision_text_pretrained(
											    "hf-internal-testing/tiny-random-vit"	,    "hf-internal-testing/tiny-random-bert"    )
											__lowerCAmelCase            = 1_3
											__lowerCAmelCase            = floats_tensor(
											    [
											        batch_size,
											        model.vision_model.config.num_channels,
											        model.vision_model.config.image_size,
											        model.vision_model.config.image_size,
											    ]    )
											__lowerCAmelCase            = ids_tensor([batch_size, 4]	,    model.text_model.config.vocab_size    )
											__lowerCAmelCase            = random_attention_mask([batch_size, 4]    )
											__lowerCAmelCase            = {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

											return model, inputs






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFViTModel(_A	,    name="vision_model"    )
											__lowerCAmelCase            = TFBertModel(_A	,    name="text_model"    )
											return vision_model, text_model






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFViTModelTester(self    )
											__lowerCAmelCase            = TFBertModelTester(self    )
											__lowerCAmelCase            = vit_model_tester.prepare_config_and_inputs()
											__lowerCAmelCase            = bert_model_tester.prepare_config_and_inputs()

											__lowerCAmelCase     ,							__lowerCAmelCase     ,							__lowerCAmelCase            = vision_config_and_inputs

											(
											    (
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							
											)            = text_config_and_inputs

											return {
											    "text_config": text_config,
											    "vision_config": vision_config,
											    "pixel_values": pixel_values,
											    "attention_mask": input_mask,
											    "input_ids": input_ids,
											    "text_token_type_ids": token_type_ids,
											    "text_sequence_labels": sequence_labels,
											    "text_token_labels": token_labels,
											    "text_choice_labels": choice_labels,
											}






@require_tf
class     a__					(      snake_case__				,       unittest.TestCase   ):






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFVisionTextDualEncoderModel.from_vision_text_pretrained(
											    "Rocketknight1/tiny-random-deit-tf"	,    "hf-internal-testing/tiny-random-roberta"    )
											__lowerCAmelCase            = 1_3
											__lowerCAmelCase            = floats_tensor(
											    [
											        batch_size,
											        model.vision_model.config.num_channels,
											        model.vision_model.config.image_size,
											        model.vision_model.config.image_size,
											    ]    )
											__lowerCAmelCase            = ids_tensor([batch_size, 4]	,    model.text_model.config.vocab_size    )
											__lowerCAmelCase            = random_attention_mask([batch_size, 4]    )
											__lowerCAmelCase            = {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

											return model, inputs






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A	,    _A	,    _A=None	,    **_A    ):




											"""simple docstring"""



											__lowerCAmelCase     ,							__lowerCAmelCase            = self.get_vision_text_model(_A	,    _A    )
											__lowerCAmelCase            = TFVisionTextDualEncoderModel(vision_model=_A	,    text_model=_A    )

											__lowerCAmelCase            = model(
											    input_ids=_A	,    pixel_values=_A	,    attention_mask=_A	,    output_attentions=_A    )

											__lowerCAmelCase            = output.vision_model_output.attentions
											self.assertEqual(len(_A    )	,    vision_config.num_hidden_layers    )

											# in DEiT, the seq_len equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
											__lowerCAmelCase            = to_atuple(vision_model.config.image_size    )
											__lowerCAmelCase            = to_atuple(vision_model.config.patch_size    )
											__lowerCAmelCase            = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
											__lowerCAmelCase            = num_patches + 2
											self.assertEqual(vision_attentions[0].shape[-3:]	,    (vision_config.num_attention_heads, seq_len, seq_len)    )

											__lowerCAmelCase            = output.text_model_output.attentions
											self.assertEqual(len(_A    )	,    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	,    _A	,    _A    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFDeiTModel(_A	,    name="vision_model"    )
											__lowerCAmelCase            = TFRobertaModel(_A	,    name="text_model"    )
											return vision_model, text_model






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFDeiTModelTester(self    )
											__lowerCAmelCase            = TFRobertaModelTester(self    )
											__lowerCAmelCase            = vit_model_tester.prepare_config_and_inputs()
											__lowerCAmelCase            = bert_model_tester.prepare_config_and_inputs()

											__lowerCAmelCase     ,							__lowerCAmelCase     ,							__lowerCAmelCase            = vision_config_and_inputs

											(
											    (
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							
											)            = text_config_and_inputs

											return {
											    "text_config": text_config,
											    "vision_config": vision_config,
											    "pixel_values": pixel_values,
											    "attention_mask": input_mask,
											    "input_ids": input_ids,
											    "text_token_type_ids": token_type_ids,
											    "text_sequence_labels": sequence_labels,
											    "text_token_labels": token_labels,
											    "text_choice_labels": choice_labels,
											}






@require_tf
class     a__					(      snake_case__				,       unittest.TestCase   ):






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFVisionTextDualEncoderModel.from_vision_text_pretrained(
											    "Rocketknight1/tiny-random-clip-tf"	,    "hf-internal-testing/tiny-random-bert"    )
											__lowerCAmelCase            = 1_3
											__lowerCAmelCase            = floats_tensor(
											    [
											        batch_size,
											        model.vision_model.config.num_channels,
											        model.vision_model.config.image_size,
											        model.vision_model.config.image_size,
											    ]    )
											__lowerCAmelCase            = ids_tensor([batch_size, 4]	,    model.text_model.config.vocab_size    )
											__lowerCAmelCase            = random_attention_mask([batch_size, 4]    )
											__lowerCAmelCase            = {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

											return model, inputs






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFCLIPVisionModel(_A	,    name="vision_model"    )
											__lowerCAmelCase            = TFBertModel(_A	,    name="text_model"    )
											return vision_model, text_model






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFCLIPVisionModelTester(self    )
											__lowerCAmelCase            = TFBertModelTester(self    )
											__lowerCAmelCase            = clip_model_tester.prepare_config_and_inputs()
											__lowerCAmelCase            = bert_model_tester.prepare_config_and_inputs()

											__lowerCAmelCase     ,							__lowerCAmelCase            = vision_config_and_inputs

											(
											    (
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							
											)            = text_config_and_inputs

											return {
											    "text_config": text_config,
											    "vision_config": vision_config,
											    "pixel_values": pixel_values,
											    "attention_mask": input_mask,
											    "input_ids": input_ids,
											    "text_token_type_ids": token_type_ids,
											    "text_sequence_labels": sequence_labels,
											    "text_token_labels": token_labels,
											    "text_choice_labels": choice_labels,
											}






@require_vision
@require_tf
class     a__					(      unittest.TestCase   ):






						@slow
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFVisionTextDualEncoderModel.from_pretrained(
											    "clip-italian/clip-italian"	,    logit_scale_init_value=1.0	,    from_pt=_A    )
											__lowerCAmelCase            = VisionTextDualEncoderProcessor.from_pretrained("clip-italian/clip-italian"    )

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

											__lowerCAmelCase            = model(**_A    )

											# 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            = np.array([[1.2_28_47_27, 0.3_10_41_22]]    )

											self.assertTrue(np.allclose(outputs.logits_per_image.numpy()	,    _A	,    atol=1E-3    )    )










from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__					    =    logging.get_logger(__name__)

UpperCamelCase__					    =    {
    """facebook/nllb-moe-54B""": """https://huggingface.co/facebook/nllb-moe-54b/resolve/main/config.json""",
}
class     a__					(      snake_case__   ):
						_a							:					Tuple      =					"""nllb-moe"""
						_a							:					Dict      =					["""past_key_values"""]
						_a							:					Optional[int]      =					{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}






						def __init__(     self	,    _A=1_2_8_1_1_2	,    _A=1_0_2_4	,    _A=1_2	,    _A=4_0_9_6	,    _A=1_6	,    _A=1_2	,    _A=4_0_9_6	,    _A=1_6	,    _A=0.05	,    _A=0.05	,    _A=True	,    _A=True	,    _A="relu"	,    _A=1_0_2_4	,    _A=0.1	,    _A=0.1	,    _A=0.0	,    _A=0.02	,    _A=2	,    _A=True	,    _A=False	,    _A="float32"	,    _A=False	,    _A=1_2_8	,    _A=6_4	,    _A=4	,    _A=4	,    _A=0.0_01	,    _A=0.0_01	,    _A="all"	,    _A=False	,    _A=False	,    _A=1.0	,    _A=0.2	,    _A=1	,    _A=0	,    _A=2	,    _A=False	,    **_A	,    ):




											"""simple docstring"""



											__lowerCAmelCase            = vocab_size
											__lowerCAmelCase            = max_position_embeddings
											__lowerCAmelCase            = d_model
											__lowerCAmelCase            = encoder_ffn_dim
											__lowerCAmelCase            = encoder_layers
											__lowerCAmelCase            = encoder_attention_heads
											__lowerCAmelCase            = decoder_ffn_dim
											__lowerCAmelCase            = decoder_layers
											__lowerCAmelCase            = decoder_attention_heads
											__lowerCAmelCase            = dropout
											__lowerCAmelCase            = attention_dropout
											__lowerCAmelCase            = activation_dropout
											__lowerCAmelCase            = activation_function
											__lowerCAmelCase            = init_std
											__lowerCAmelCase            = encoder_layerdrop
											__lowerCAmelCase            = decoder_layerdrop
											__lowerCAmelCase            = use_cache
											__lowerCAmelCase            = encoder_layers
											__lowerCAmelCase            = scale_embedding  # scale factor will be sqrt(d_model) if True
											__lowerCAmelCase            = router_z_loss_coef
											__lowerCAmelCase            = router_aux_loss_coef
											__lowerCAmelCase            = decoder_sparse_step
											__lowerCAmelCase            = encoder_sparse_step
											__lowerCAmelCase            = num_experts
											__lowerCAmelCase            = expert_capacity
											__lowerCAmelCase            = router_bias
											if router_dtype not in ["float32", "float16", "bfloat16"]:
																raise ValueError(f"""`router_dtype` must be one of 'float32', 'float16' or 'bfloat16', got {router_dtype}"""    )
											__lowerCAmelCase            = router_dtype

											__lowerCAmelCase            = router_ignore_padding_tokens
											__lowerCAmelCase            = batch_prioritized_routing
											__lowerCAmelCase            = second_expert_policy
											__lowerCAmelCase            = normalize_router_prob_before_dropping
											__lowerCAmelCase            = moe_eval_capacity_token_fraction
											__lowerCAmelCase            = moe_token_dropout
											__lowerCAmelCase            = output_router_logits
											super().__init__(
											    pad_token_id=_A	,    bos_token_id=_A	,    eos_token_id=_A	,    is_encoder_decoder=_A	,    decoder_start_token_id=_A	,    **_A	,    )




