






'''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  SCREAMING_SNAKE_CASE_       (UpperCamelCase      )      ->       Union[str, Any]:
							if isinstance(UpperCamelCase				,			collections.abc.Iterable      ):
														return x
							return (x, x)

@require_flax
class _lowercase	:




							def 		lowerCamelCase_				(							self:				Dict			,							UpperCamelCase__:				str			,							UpperCamelCase__:				Optional[Any]     ):
														pass




							def 		lowerCamelCase_				(							self:				Dict     ):
														pass




							def 		lowerCamelCase_				(							self:				Tuple     ):
														pass




							def 		lowerCamelCase_				(							self:				Tuple			,							UpperCamelCase__:				np.ndarray			,							UpperCamelCase__:				np.ndarray			,							UpperCamelCase__:				float     ):
														lowerCamelCase__					:	List[str]         =       np.abs((a - b)     ).max()
														self.assertLessEqual(UpperCamelCase__			,							UpperCamelCase__			,							F'''Difference between torch and flax is {diff} (>= {tol}).'''     )




							def 		lowerCamelCase_				(							self:				Dict			,							UpperCamelCase__:				Tuple			,							UpperCamelCase__:				Optional[Any]			,							UpperCamelCase__:				List[str]			,							UpperCamelCase__:				Tuple			,							UpperCamelCase__:				int=None			,							**UpperCamelCase__:				Any     ):
														lowerCamelCase__					:	Dict         =       VisionTextDualEncoderConfig.from_vision_text_configs(UpperCamelCase__			,							UpperCamelCase__     )

														lowerCamelCase__					:	Tuple         =       FlaxVisionTextDualEncoderModel(UpperCamelCase__     )

														lowerCamelCase__					:	Union[str, Any]         =       model(input_ids=UpperCamelCase__			,							pixel_values=UpperCamelCase__			,							attention_mask=UpperCamelCase__     )

														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 		lowerCamelCase_				(							self:				Any			,							UpperCamelCase__:				Dict			,							UpperCamelCase__:				Dict			,							UpperCamelCase__:				Union[str, Any]			,							UpperCamelCase__:				Dict			,							UpperCamelCase__:				List[str]=None			,							**UpperCamelCase__:				Optional[int]     ):
														lowerCamelCase__  ,		lowerCamelCase__					:	Union[str, Any]         =       self.get_vision_text_model(UpperCamelCase__			,							UpperCamelCase__     )
														lowerCamelCase__					:	Dict         =       {"""vision_model""": vision_model, """text_model""": text_model}
														lowerCamelCase__					:	int         =       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**UpperCamelCase__     )

														lowerCamelCase__					:	Optional[Any]         =       model(input_ids=UpperCamelCase__			,							pixel_values=UpperCamelCase__			,							attention_mask=UpperCamelCase__     )

														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 		lowerCamelCase_				(							self:				Any			,							UpperCamelCase__:				Optional[Any]			,							UpperCamelCase__:				Any			,							UpperCamelCase__:				List[Any]			,							UpperCamelCase__:				Optional[Any]			,							UpperCamelCase__:				Optional[Any]=None			,							**UpperCamelCase__:				List[Any]     ):
														lowerCamelCase__  ,		lowerCamelCase__					:	Tuple         =       self.get_vision_text_model(UpperCamelCase__			,							UpperCamelCase__     )
														lowerCamelCase__					:	Optional[Any]         =       {"""vision_model""": vision_model, """text_model""": text_model}
														lowerCamelCase__					:	Any         =       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**UpperCamelCase__     )

														lowerCamelCase__					:	Dict         =       model(input_ids=UpperCamelCase__			,							pixel_values=UpperCamelCase__			,							attention_mask=UpperCamelCase__     )
														lowerCamelCase__					:	Tuple         =       output[0]

														with tempfile.TemporaryDirectory() as tmpdirname:
																					model.save_pretrained(UpperCamelCase__     )
																					lowerCamelCase__					:	List[Any]         =       FlaxVisionTextDualEncoderModel.from_pretrained(UpperCamelCase__     )

																					lowerCamelCase__					:	List[Any]         =       model(input_ids=UpperCamelCase__			,							pixel_values=UpperCamelCase__			,							attention_mask=UpperCamelCase__     )
																					lowerCamelCase__					:	List[Any]         =       after_output[0]
																					lowerCamelCase__					:	int         =       np.amax(np.abs(out_a - out_a     )     )
																					self.assertLessEqual(UpperCamelCase__			,							1e-3     )




							def 		lowerCamelCase_				(							self:				List[str]			,							UpperCamelCase__:				Union[str, Any]			,							UpperCamelCase__:				List[Any]			,							UpperCamelCase__:				Optional[Any]			,							UpperCamelCase__:				Any			,							UpperCamelCase__:				Optional[Any]=None			,							**UpperCamelCase__:				Any     ):
														lowerCamelCase__  ,		lowerCamelCase__					:	str         =       self.get_vision_text_model(UpperCamelCase__			,							UpperCamelCase__     )
														lowerCamelCase__					:	Optional[Any]         =       {"""vision_model""": vision_model, """text_model""": text_model}
														lowerCamelCase__					:	Dict         =       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**UpperCamelCase__     )

														lowerCamelCase__					:	str         =       model(
														    input_ids=UpperCamelCase__			,							pixel_values=UpperCamelCase__			,							attention_mask=UpperCamelCase__			,							output_attentions=UpperCamelCase__     )

														lowerCamelCase__					:	int         =       output.vision_model_output.attentions
														self.assertEqual(len(UpperCamelCase__     )			,							vision_config.num_hidden_layers     )

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

														lowerCamelCase__					:	List[str]         =       output.text_model_output.attentions
														self.assertEqual(len(UpperCamelCase__     )			,							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 		lowerCamelCase_				(							self:				Optional[Any]			,							UpperCamelCase__:				Union[str, Any]			,							UpperCamelCase__:				Dict			,							UpperCamelCase__:				List[Any]     ):
														pt_model.to(UpperCamelCase__     )
														pt_model.eval()

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

														with torch.no_grad():
																					lowerCamelCase__					:	List[str]         =       pt_model(**UpperCamelCase__     ).to_tuple()

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

														# PT -> Flax
														with tempfile.TemporaryDirectory() as tmpdirname:
																					pt_model.save_pretrained(UpperCamelCase__     )
																					lowerCamelCase__					:	Dict         =       FlaxVisionTextDualEncoderModel.from_pretrained(UpperCamelCase__			,							from_pt=UpperCamelCase__     )

														lowerCamelCase__					:	Union[str, Any]         =       fx_model_loaded(**UpperCamelCase__     ).to_tuple()
														self.assertEqual(len(UpperCamelCase__     )			,							len(UpperCamelCase__     )			,							"""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(UpperCamelCase__			,							pt_output.numpy()			,							4e-2     )

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

														pt_model_loaded.to(UpperCamelCase__     )
														pt_model_loaded.eval()

														with torch.no_grad():
																					lowerCamelCase__					:	Optional[Any]         =       pt_model_loaded(**UpperCamelCase__     ).to_tuple()

														self.assertEqual(len(UpperCamelCase__     )			,							len(UpperCamelCase__     )			,							"""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(UpperCamelCase__			,							pt_output_loaded.numpy()			,							4e-2     )




							def 		lowerCamelCase_				(							self:				List[Any]			,							UpperCamelCase__:				Any			,							UpperCamelCase__:				int			,							UpperCamelCase__:				Any     ):
														lowerCamelCase__					:	int         =       VisionTextDualEncoderConfig.from_vision_text_configs(UpperCamelCase__			,							UpperCamelCase__     )

														lowerCamelCase__					:	Union[str, Any]         =       VisionTextDualEncoderModel(UpperCamelCase__     )
														lowerCamelCase__					:	Union[str, Any]         =       FlaxVisionTextDualEncoderModel(UpperCamelCase__     )

														lowerCamelCase__					:	str         =       convert_pytorch_state_dict_to_flax(pt_model.state_dict()			,							UpperCamelCase__     )
														lowerCamelCase__					:	Tuple         =       fx_state

														self.check_pt_flax_equivalence(UpperCamelCase__			,							UpperCamelCase__			,							UpperCamelCase__     )




							def 		lowerCamelCase_				(							self:				Any			,							UpperCamelCase__:				Dict			,							UpperCamelCase__:				int			,							UpperCamelCase__:				str     ):
														lowerCamelCase__					:	Union[str, Any]         =       VisionTextDualEncoderConfig.from_vision_text_configs(UpperCamelCase__			,							UpperCamelCase__     )

														lowerCamelCase__					:	int         =       VisionTextDualEncoderModel(UpperCamelCase__     )
														lowerCamelCase__					:	Tuple         =       FlaxVisionTextDualEncoderModel(UpperCamelCase__     )

														lowerCamelCase__					:	List[Any]         =       load_flax_weights_in_pytorch_model(UpperCamelCase__			,							fx_model.params     )

														self.check_pt_flax_equivalence(UpperCamelCase__			,							UpperCamelCase__			,							UpperCamelCase__     )




							def 		lowerCamelCase_				(							self:				Dict     ):
														lowerCamelCase__					:	Tuple         =       self.prepare_config_and_inputs()
														self.check_model_from_pretrained_configs(**UpperCamelCase__     )




							def 		lowerCamelCase_				(							self:				Tuple     ):
														lowerCamelCase__					:	Dict         =       self.prepare_config_and_inputs()
														self.check_vision_text_dual_encoder_from_pretrained(**UpperCamelCase__     )




							def 		lowerCamelCase_				(							self:				int     ):
														lowerCamelCase__					:	Union[str, Any]         =       self.prepare_config_and_inputs()
														self.check_save_load(**UpperCamelCase__     )




							def 		lowerCamelCase_				(							self:				Optional[int]     ):
														lowerCamelCase__					:	List[str]         =       self.prepare_config_and_inputs()
														self.check_vision_text_output_attention(**UpperCamelCase__     )




							@is_pt_flax_cross_test
							def 		lowerCamelCase_				(							self:				Dict     ):
														lowerCamelCase__					:	Any         =       self.prepare_config_and_inputs()
														lowerCamelCase__					:	int         =       config_inputs_dict.pop("""vision_config"""     )
														lowerCamelCase__					:	str         =       config_inputs_dict.pop("""text_config"""     )

														lowerCamelCase__					:	Dict         =       config_inputs_dict

														self.check_equivalence_pt_to_flax(UpperCamelCase__			,							UpperCamelCase__			,							UpperCamelCase__     )
														self.check_equivalence_flax_to_pt(UpperCamelCase__			,							UpperCamelCase__			,							UpperCamelCase__     )




							@slow
							def 		lowerCamelCase_				(							self:				Dict     ):
														lowerCamelCase__  ,		lowerCamelCase__					:	Any         =       self.get_pretrained_model_and_inputs()

														lowerCamelCase__					:	Dict         =       model_a(**UpperCamelCase__     )
														lowerCamelCase__					:	str         =       outputs[0]

														with tempfile.TemporaryDirectory() as tmp_dirname:
																					model_a.save_pretrained(UpperCamelCase__     )
																					lowerCamelCase__					:	Union[str, Any]         =       FlaxVisionTextDualEncoderModel.from_pretrained(UpperCamelCase__     )

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



@require_flax
class _lowercase	(				_lowercase				,      unittest.TestCase		):




							def 		lowerCamelCase_				(							self:				List[Any]     ):
														lowerCamelCase__					:	List[Any]         =       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
														    """hf-internal-testing/tiny-random-vit"""			,							"""hf-internal-testing/tiny-bert"""			,							vision_from_pt=UpperCamelCase__			,							text_from_pt=UpperCamelCase__			,							)
														lowerCamelCase__					:	Optional[int]         =       13
														lowerCamelCase__					:	Optional[Any]         =       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__					:	Any         =       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 		lowerCamelCase_				(							self:				Union[str, Any]			,							UpperCamelCase__:				Tuple			,							UpperCamelCase__:				Tuple     ):
														lowerCamelCase__					:	Dict         =       FlaxViTModel(UpperCamelCase__     )
														lowerCamelCase__					:	Optional[Any]         =       FlaxBertModel(UpperCamelCase__     )
														return vision_model, text_model




							def 		lowerCamelCase_				(							self:				Tuple     ):
														lowerCamelCase__					:	List[Any]         =       FlaxViTModelTester(self     )
														lowerCamelCase__					:	List[Any]         =       FlaxBertModelTester(self     )
														lowerCamelCase__					:	List[Any]         =       vit_model_tester.prepare_config_and_inputs()
														lowerCamelCase__					:	str         =       bert_model_tester.prepare_config_and_inputs()

														lowerCamelCase__  ,		lowerCamelCase__					:	Union[str, Any]         =       vision_config_and_inputs

														lowerCamelCase__  ,		lowerCamelCase__  ,		lowerCamelCase__  ,		lowerCamelCase__					:	Tuple         =       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 _lowercase	(				_lowercase				,      unittest.TestCase		):




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

														return model, inputs




							def 		lowerCamelCase_				(							self:				Tuple			,							UpperCamelCase__:				int			,							UpperCamelCase__:				List[Any]     ):
														lowerCamelCase__					:	List[Any]         =       FlaxCLIPVisionModel(UpperCamelCase__     )
														lowerCamelCase__					:	Optional[int]         =       FlaxBertModel(UpperCamelCase__     )
														return vision_model, text_model




							def 		lowerCamelCase_				(							self:				int     ):
														lowerCamelCase__					:	List[str]         =       FlaxCLIPVisionModelTester(self     )
														lowerCamelCase__					:	Tuple         =       FlaxBertModelTester(self     )
														lowerCamelCase__					:	Any         =       clip_model_tester.prepare_config_and_inputs()
														lowerCamelCase__					:	List[str]         =       bert_model_tester.prepare_config_and_inputs()

														lowerCamelCase__  ,		lowerCamelCase__					:	str         =       vision_config_and_inputs

														lowerCamelCase__  ,		lowerCamelCase__  ,		lowerCamelCase__  ,		lowerCamelCase__					:	str         =       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 _lowercase	(				unittest.TestCase		):




							@slow
							def 		lowerCamelCase_				(							self:				str     ):
														lowerCamelCase__					:	Optional[Any]         =       FlaxVisionTextDualEncoderModel.from_pretrained("""clip-italian/clip-italian"""			,							logit_scale_init_value=1.0     )
														lowerCamelCase__					:	str         =       VisionTextDualEncoderProcessor.from_pretrained("""clip-italian/clip-italian"""     )

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

														lowerCamelCase__					:	Union[str, Any]         =       model(**UpperCamelCase__     )

														# 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__					:	str         =       np.array([[1.2_284_727, 0.3_104_122]]     )

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









'''simple docstring'''







from __future__ import annotations

_A	:			Any         ={
    '''A''': ['''B''', '''C''', '''E'''],
    '''B''': ['''A''', '''D''', '''E'''],
    '''C''': ['''A''', '''F''', '''G'''],
    '''D''': ['''B'''],
    '''E''': ['''A''', '''B''', '''D'''],
    '''F''': ['''C'''],
    '''G''': ['''C'''],
}

class _lowercase	:




							def __init__(							self:				Tuple			,							UpperCamelCase__:				dict[str, list[str]]			,							UpperCamelCase__:				str     ):
														lowerCamelCase__					:	str         =       graph
														# mapping node to its parent in resulting breadth first tree
														lowerCamelCase__					:	dict[str, str | None]         =       {}
														lowerCamelCase__					:	Any         =       source_vertex




							def 		lowerCamelCase_				(							self:				List[str]     ):
														lowerCamelCase__					:	List[str]         =       {self.source_vertex}
														lowerCamelCase__					:	List[str]         =       None
														lowerCamelCase__					:	Tuple         =       [self.source_vertex]  # first in first out queue

														while queue:
																					lowerCamelCase__					:	Tuple         =       queue.pop(0     )
																					for adjacent_vertex in self.graph[vertex]:
																												if adjacent_vertex not in visited:
																																			visited.add(UpperCamelCase__     )
																																			lowerCamelCase__					:	List[str]         =       vertex
																																			queue.append(UpperCamelCase__     )




							def 		lowerCamelCase_				(							self:				str			,							UpperCamelCase__:				str     ):
														if target_vertex == self.source_vertex:
																					return self.source_vertex

														lowerCamelCase__					:	Tuple         =       self.parent.get(UpperCamelCase__     )
														if target_vertex_parent is None:
																					lowerCamelCase__					:	int         =       (
																					    F'''No path from vertex: {self.source_vertex} to vertex: {target_vertex}'''
																					)
																					raise ValueError(UpperCamelCase__     )

														return self.shortest_path(UpperCamelCase__     ) + F'''->{target_vertex}'''


if __name__ == "__main__":
			_A	:			int         =Graph(graph, '''G''')
			g.breath_first_search()
			print(g.shortest_path('''D'''))
			print(g.shortest_path('''G'''))
			print(g.shortest_path('''Foo'''))

