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  snake_case						(__lowercase						)	->	Optional[Any]:

						'''simple docstring'''



						if isinstance(__lowercase ,      collections.abc.Iterable						):
												return x
						return (x, x)



@require_tf
class        lowercase_   :



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_    ):
													pass



							def 	UpperCamelCase	(						self    ):
													pass



							def 	UpperCamelCase	(						self    ):
													pass



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_				,    lowercase_				,    lowercase_				,    lowercase_=None				,    **lowercase_    ):
													_snake_case						:					Optional[Any]    =   VisionTextDualEncoderConfig.from_vision_text_configs(lowercase_				,    lowercase_    )

													_snake_case						:					int    =   TFVisionTextDualEncoderModel(lowercase_    )

													_snake_case						:					Union[str, Any]    =   model(input_ids=lowercase_				,    pixel_values=lowercase_				,    attention_mask=lowercase_    )

													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 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_				,    lowercase_				,    lowercase_				,    lowercase_=None				,    **lowercase_    ):
													_snake_case    ,_snake_case						:					List[Any]    =   self.get_vision_text_model(lowercase_				,    lowercase_    )
													_snake_case						:					List[str]    =   TFVisionTextDualEncoderModel(vision_model=lowercase_				,    text_model=lowercase_    )

													_snake_case						:					Union[str, Any]    =   model(input_ids=lowercase_				,    pixel_values=lowercase_				,    attention_mask=lowercase_    )

													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 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_				,    lowercase_				,    lowercase_				,    lowercase_=None				,    **lowercase_    ):
													_snake_case    ,_snake_case						:					Any    =   self.get_vision_text_model(lowercase_				,    lowercase_    )
													_snake_case						:					Optional[Any]    =   {"vision_model": vision_model, "text_model": text_model}
													_snake_case						:					Union[str, Any]    =   TFVisionTextDualEncoderModel.from_vision_text_pretrained(**lowercase_    )

													_snake_case						:					List[Any]    =   model(input_ids=lowercase_				,    pixel_values=lowercase_				,    attention_mask=lowercase_    )

													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 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_				,    lowercase_				,    lowercase_				,    lowercase_=None				,    **lowercase_    ):
													_snake_case    ,_snake_case						:					Tuple    =   self.get_vision_text_model(lowercase_				,    lowercase_    )
													_snake_case						:					Optional[Any]    =   TFVisionTextDualEncoderModel(vision_model=lowercase_				,    text_model=lowercase_    )

													_snake_case						:					List[str]    =   model(input_ids=lowercase_				,    pixel_values=lowercase_				,    attention_mask=lowercase_    )
													_snake_case						:					int    =   output[0].numpy()

													with tempfile.TemporaryDirectory() as tmpdirname:
																			model.save_pretrained(lowercase_    )
																			_snake_case						:					str    =   TFVisionTextDualEncoderModel.from_pretrained(lowercase_    )

																			_snake_case						:					Any    =   model(input_ids=lowercase_				,    pixel_values=lowercase_				,    attention_mask=lowercase_    )
																			_snake_case						:					Tuple    =   after_output[0].numpy()
																			_snake_case						:					Any    =   np.amax(np.abs(out_a - out_a    )    )
																			self.assertLessEqual(lowercase_				,    1e-5    )



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_				,    lowercase_				,    lowercase_				,    lowercase_=None				,    **lowercase_    ):
													_snake_case    ,_snake_case						:					Tuple    =   self.get_vision_text_model(lowercase_				,    lowercase_    )
													_snake_case						:					Dict    =   TFVisionTextDualEncoderModel(vision_model=lowercase_				,    text_model=lowercase_    )

													_snake_case						:					List[Any]    =   model(
													    input_ids=lowercase_				,    pixel_values=lowercase_				,    attention_mask=lowercase_				,    output_attentions=lowercase_    )

													_snake_case						:					List[str]    =   output.vision_model_output.attentions
													self.assertEqual(len(lowercase_    )				,    vision_config.num_hidden_layers    )

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

													_snake_case						:					Dict    =   output.text_model_output.attentions
													self.assertEqual(len(lowercase_    )				,    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 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_				,    lowercase_    ):
													_snake_case						:					Union[str, Any]    =   np.abs((a - b)    ).max()
													self.assertLessEqual(lowercase_				,    lowercase_				,    f"""Difference between torch and flax is {diff} (>= {tol})."""    )



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					List[Any]    =   self.prepare_config_and_inputs()
													self.check_vision_text_dual_encoder_model(**lowercase_    )



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					str    =   self.prepare_config_and_inputs()
													self.check_model_from_pretrained_configs(**lowercase_    )



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					Any    =   self.prepare_config_and_inputs()
													self.check_vision_text_dual_encoder_from_pretrained(**lowercase_    )



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					List[str]    =   self.prepare_config_and_inputs()
													self.check_save_load(**lowercase_    )



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					Dict    =   self.prepare_config_and_inputs()
													self.check_vision_text_output_attention(**lowercase_    )



							@slow
							def 	UpperCamelCase	(						self    ):
													_snake_case    ,_snake_case						:					Union[str, Any]    =   self.get_pretrained_model_and_inputs()

													_snake_case						:					Optional[int]    =   model_a(**lowercase_    )
													_snake_case						:					int    =   outputs[0].numpy()

													with tempfile.TemporaryDirectory() as tmp_dirname:
																			model_a.save_pretrained(lowercase_    )
																			_snake_case						:					Union[str, Any]    =   TFVisionTextDualEncoderModel.from_pretrained(lowercase_    )

																			_snake_case						:					Optional[Any]    =   model_a(**lowercase_    )
																			_snake_case						:					Optional[int]    =   after_outputs[0].numpy()
																			_snake_case						:					Tuple    =   np.amax(np.abs(out_a - out_a    )    )
																			self.assertLessEqual(lowercase_				,    1e-5    )




@require_tf
class        lowercase_   (					__snake_case						,			unittest.TestCase		):



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					Any    =   TFVisionTextDualEncoderModel.from_vision_text_pretrained(
													    "hf-internal-testing/tiny-random-vit"				,    "hf-internal-testing/tiny-random-bert"    )
													_snake_case						:					Optional[int]    =   13
													_snake_case						:					List[str]    =   floats_tensor(
													    [
													        batch_size,
													        model.vision_model.config.num_channels,
													        model.vision_model.config.image_size,
													        model.vision_model.config.image_size,
													    ]    )
													_snake_case						:					str    =   ids_tensor([batch_size, 4]				,    model.text_model.config.vocab_size    )
													_snake_case						:					Tuple    =   random_attention_mask([batch_size, 4]    )
													_snake_case						:					List[str]    =   {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

													return model, inputs



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_    ):
													_snake_case						:					int    =   TFViTModel(lowercase_				,    name="vision_model"    )
													_snake_case						:					List[str]    =   TFBertModel(lowercase_				,    name="text_model"    )
													return vision_model, text_model



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					List[Any]    =   TFViTModelTester(self    )
													_snake_case						:					List[str]    =   TFBertModelTester(self    )
													_snake_case						:					Any    =   vit_model_tester.prepare_config_and_inputs()
													_snake_case						:					Union[str, Any]    =   bert_model_tester.prepare_config_and_inputs()

													_snake_case    ,_snake_case    ,_snake_case						:					int    =   vision_config_and_inputs

													(
													    (
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,
													)						:					Any    =   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        lowercase_   (					__snake_case						,			unittest.TestCase		):



							def 	UpperCamelCase	(						self    ):
													# DeiT repo doesn't have TF weights, but we don't actually use the weights at all so let's
													# just reinitialize it.
													_snake_case						:					Tuple    =   TFVisionTextDualEncoderModel.from_vision_text_pretrained(
													    "Rocketknight1/tiny-random-deit-tf"				,    "hf-internal-testing/tiny-random-roberta"    )
													_snake_case						:					Optional[int]    =   13
													_snake_case						:					Optional[Any]    =   floats_tensor(
													    [
													        batch_size,
													        model.vision_model.config.num_channels,
													        model.vision_model.config.image_size,
													        model.vision_model.config.image_size,
													    ]    )
													_snake_case						:					Union[str, Any]    =   ids_tensor([batch_size, 4]				,    model.text_model.config.vocab_size    )
													_snake_case						:					List[Any]    =   random_attention_mask([batch_size, 4]    )
													_snake_case						:					Optional[Any]    =   {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

													return model, inputs



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_				,    lowercase_				,    lowercase_				,    lowercase_=None				,    **lowercase_    ):
													_snake_case    ,_snake_case						:					Dict    =   self.get_vision_text_model(lowercase_				,    lowercase_    )
													_snake_case						:					Optional[Any]    =   TFVisionTextDualEncoderModel(vision_model=lowercase_				,    text_model=lowercase_    )

													_snake_case						:					int    =   model(
													    input_ids=lowercase_				,    pixel_values=lowercase_				,    attention_mask=lowercase_				,    output_attentions=lowercase_    )

													_snake_case						:					Dict    =   output.vision_model_output.attentions
													self.assertEqual(len(lowercase_    )				,    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)
													_snake_case						:					int    =   to_atuple(vision_model.config.image_size    )
													_snake_case						:					List[Any]    =   to_atuple(vision_model.config.patch_size    )
													_snake_case						:					Optional[Any]    =   (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
													_snake_case						:					Union[str, Any]    =   num_patches + 2
													self.assertEqual(vision_attentions[0].shape[-3:]				,    (vision_config.num_attention_heads, seq_len, seq_len)    )

													_snake_case						:					List[str]    =   output.text_model_output.attentions
													self.assertEqual(len(lowercase_    )				,    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 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_    ):
													_snake_case						:					List[Any]    =   TFDeiTModel(lowercase_				,    name="vision_model"    )
													_snake_case						:					int    =   TFRobertaModel(lowercase_				,    name="text_model"    )
													return vision_model, text_model



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					Union[str, Any]    =   TFDeiTModelTester(self    )
													_snake_case						:					List[Any]    =   TFRobertaModelTester(self    )
													_snake_case						:					List[str]    =   vit_model_tester.prepare_config_and_inputs()
													_snake_case						:					List[Any]    =   bert_model_tester.prepare_config_and_inputs()

													_snake_case    ,_snake_case    ,_snake_case						:					str    =   vision_config_and_inputs

													(
													    (
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,
													)						:					Any    =   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        lowercase_   (					__snake_case						,			unittest.TestCase		):



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					List[Any]    =   TFVisionTextDualEncoderModel.from_vision_text_pretrained(
													    "Rocketknight1/tiny-random-clip-tf"				,    "hf-internal-testing/tiny-random-bert"    )
													_snake_case						:					Tuple    =   13
													_snake_case						:					Any    =   floats_tensor(
													    [
													        batch_size,
													        model.vision_model.config.num_channels,
													        model.vision_model.config.image_size,
													        model.vision_model.config.image_size,
													    ]    )
													_snake_case						:					Union[str, Any]    =   ids_tensor([batch_size, 4]				,    model.text_model.config.vocab_size    )
													_snake_case						:					Tuple    =   random_attention_mask([batch_size, 4]    )
													_snake_case						:					List[str]    =   {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

													return model, inputs



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_    ):
													_snake_case						:					List[Any]    =   TFCLIPVisionModel(lowercase_				,    name="vision_model"    )
													_snake_case						:					Dict    =   TFBertModel(lowercase_				,    name="text_model"    )
													return vision_model, text_model



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					Dict    =   TFCLIPVisionModelTester(self    )
													_snake_case						:					Union[str, Any]    =   TFBertModelTester(self    )
													_snake_case						:					List[Any]    =   clip_model_tester.prepare_config_and_inputs()
													_snake_case						:					int    =   bert_model_tester.prepare_config_and_inputs()

													_snake_case    ,_snake_case						:					List[Any]    =   vision_config_and_inputs

													(
													    (
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,
													)						:					int    =   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        lowercase_   (					unittest.TestCase		):



							@slow
							def 	UpperCamelCase	(						self    ):
													_snake_case						:					Dict    =   TFVisionTextDualEncoderModel.from_pretrained(
													    "clip-italian/clip-italian"				,    logit_scale_init_value=1.0				,    from_pt=lowercase_    )
													_snake_case						:					List[str]    =   VisionTextDualEncoderProcessor.from_pretrained("clip-italian/clip-italian"    )

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

													_snake_case						:					List[Any]    =   model(**lowercase_    )

													# 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])				,    )

													_snake_case						:					Union[str, Any]    =   np.array([[1.2_284_727, 0.3_104_122]]    )

													self.assertTrue(np.allclose(outputs.logits_per_image.numpy()				,    lowercase_				,    atol=1e-3    )    )
from __future__ import annotations

from collections import deque
from collections.abc import Iterator
from dataclasses import dataclass



@dataclass
class        lowercase_   :
							_lowerCamelCase						=		42
							_lowerCamelCase						=		42




class        lowercase_   :



							def __init__(						self				,    lowercase_    ):
													_snake_case						:					list[list[Edge]]    =   [[] for _ in range(lowercase_    )]
													_snake_case						:					Union[str, Any]    =   size



							def __getitem__(						self				,    lowercase_    ):
													return iter(self._graph[vertex]    )



							@property
							def 	UpperCamelCase	(						self    ):
													return self._size



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_				,    lowercase_    ):
													if weight not in (0, 1):
																			raise ValueError("Edge weight must be either 0 or 1."    )

													if to_vertex < 0 or to_vertex >= self.size:
																			raise ValueError("Vertex indexes must be in [0; size)."    )

													self._graph[from_vertex].append(Edge(lowercase_				,    lowercase_    )    )



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_    ):
													_snake_case						:					Optional[int]    =   deque([start_vertex]    )
													_snake_case						:					list[int | None]    =   [None] * self.size
													_snake_case						:					Tuple    =   0

													while queue:
																			_snake_case						:					List[Any]    =   queue.popleft()
																			_snake_case						:					Tuple    =   distances[current_vertex]
																			if current_distance is None:
																									continue

																			for edge in self[current_vertex]:
																									_snake_case						:					Dict    =   current_distance + edge.weight
																									_snake_case						:					str    =   distances[edge.destination_vertex]
																									if (
																									    isinstance(lowercase_				,    lowercase_    )
																									    and new_distance >= dest_vertex_distance
																									):
																															continue
																									_snake_case						:					List[Any]    =   new_distance
																									if edge.weight == 0:
																															queue.appendleft(edge.destination_vertex    )
																									else:
																															queue.append(edge.destination_vertex    )

													if distances[finish_vertex] is None:
																			raise ValueError("No path from start_vertex to finish_vertex."    )

													return distances[finish_vertex]


if __name__ == "__main__":
				import doctest

				doctest.testmod()