





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     ( lowerCAmelCase__						):
		if isinstance(lowerCAmelCase__							,collections.abc.Iterable						):
				return x
		return (x, x)






@require_flax
class 						UpperCamelCase_  :



					def   __a     (		self     :     Any       ,      lowerCamelCase     :     Tuple       ,      lowerCamelCase     :     Union[str, Any]						):
							pass



					def   __a     (		self     :     Union[str, Any]						):
							pass



					def   __a     (		self     :     Tuple						):
							pass



					def   __a     (		self     :     Union[str, Any]       ,      lowerCamelCase     :     np.ndarray       ,      lowerCamelCase     :     np.ndarray       ,      lowerCamelCase     :     float						):
							lowerCamelCase_	:     Tuple 						=      np.abs((a - b)						).max()
							self.assertLessEqual(lowerCamelCase       ,      lowerCamelCase       ,      F"Difference between torch and flax is {diff} (>= {tol})."						)



					def   __a     (		self     :     Optional[Any]       ,      lowerCamelCase     :     Dict       ,      lowerCamelCase     :     Any       ,      lowerCamelCase     :     str       ,      lowerCamelCase     :     Dict       ,      lowerCamelCase     :     Tuple=None       ,      **lowerCamelCase     :     Optional[int]						):
							lowerCamelCase_	:     Union[str, Any] 						=      VisionTextDualEncoderConfig.from_vision_text_configs(lowerCamelCase       ,      lowerCamelCase						)

							lowerCamelCase_	:     int 						=      FlaxVisionTextDualEncoderModel(lowerCamelCase						)

							lowerCamelCase_	:     str 						=      model(input_ids=lowerCamelCase       ,      pixel_values=lowerCamelCase       ,      attention_mask=lowerCamelCase						)

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

							lowerCamelCase_	:     Tuple 						=      model(input_ids=lowerCamelCase       ,      pixel_values=lowerCamelCase       ,      attention_mask=lowerCamelCase						)

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

							lowerCamelCase_	:     Tuple 						=      model(input_ids=lowerCamelCase       ,      pixel_values=lowerCamelCase       ,      attention_mask=lowerCamelCase						)
							lowerCamelCase_	:     Tuple 						=      output[0]

							with tempfile.TemporaryDirectory() as tmpdirname:
									model.save_pretrained(lowerCamelCase						)
									lowerCamelCase_	:     Optional[int] 						=      FlaxVisionTextDualEncoderModel.from_pretrained(lowerCamelCase						)

									lowerCamelCase_	:     Dict 						=      model(input_ids=lowerCamelCase       ,      pixel_values=lowerCamelCase       ,      attention_mask=lowerCamelCase						)
									lowerCamelCase_	:     str 						=      after_output[0]
									lowerCamelCase_	:     Tuple 						=      np.amax(np.abs(out_a - out_a						)						)
									self.assertLessEqual(lowerCamelCase       ,      1E-3						)



					def   __a     (		self     :     int       ,      lowerCamelCase     :     Tuple       ,      lowerCamelCase     :     Dict       ,      lowerCamelCase     :     List[str]       ,      lowerCamelCase     :     Union[str, Any]       ,      lowerCamelCase     :     Union[str, Any]=None       ,      **lowerCamelCase     :     List[str]						):
							lowerCamelCase_						,      lowerCamelCase_	:     List[Any] 						=      self.get_vision_text_model(lowerCamelCase       ,      lowerCamelCase						)
							lowerCamelCase_	:     str 						=      {'vision_model': vision_model, 'text_model': text_model}
							lowerCamelCase_	:     str 						=      FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCamelCase						)

							lowerCamelCase_	:     List[Any] 						=      model(
							    input_ids=lowerCamelCase       ,      pixel_values=lowerCamelCase       ,      attention_mask=lowerCamelCase       ,      output_attentions=lowerCamelCase						)

							lowerCamelCase_	:     Union[str, Any] 						=      output.vision_model_output.attentions
							self.assertEqual(len(lowerCamelCase						)       ,      vision_config.num_hidden_layers						)

							# in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
							lowerCamelCase_	:     str 						=      to_atuple(vision_model.config.image_size						)
							lowerCamelCase_	:     Dict 						=      to_atuple(vision_model.config.patch_size						)
							lowerCamelCase_	:     List[str] 						=      (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_	:     Optional[Any] 						=      output.text_model_output.attentions
							self.assertEqual(len(lowerCamelCase						)       ,      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   __a     (		self     :     Union[str, Any]       ,      lowerCamelCase     :     Union[str, Any]       ,      lowerCamelCase     :     Any       ,      lowerCamelCase     :     Optional[Any]						):
							pt_model.to(lowerCamelCase						)
							pt_model.eval()

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

							with torch.no_grad():
									lowerCamelCase_	:     Any 						=      pt_model(**lowerCamelCase						).to_tuple()

							lowerCamelCase_	:     Union[str, Any] 						=      fx_model(**lowerCamelCase						).to_tuple()
							self.assertEqual(len(lowerCamelCase						)       ,      len(lowerCamelCase						)       ,      'Output lengths differ between Flax and PyTorch'						)
							for fx_output, pt_output in zip(fx_outputs[:4]       ,      pt_outputs[:4]						):
									self.assert_almost_equals(lowerCamelCase       ,      pt_output.numpy()       ,      4E-2						)

							# PT -> Flax
							with tempfile.TemporaryDirectory() as tmpdirname:
									pt_model.save_pretrained(lowerCamelCase						)
									lowerCamelCase_	:     Optional[Any] 						=      FlaxVisionTextDualEncoderModel.from_pretrained(lowerCamelCase       ,      from_pt=lowerCamelCase						)

							lowerCamelCase_	:     str 						=      fx_model_loaded(**lowerCamelCase						).to_tuple()
							self.assertEqual(len(lowerCamelCase						)       ,      len(lowerCamelCase						)       ,      '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(lowerCamelCase       ,      pt_output.numpy()       ,      4E-2						)

							# Flax -> PT
							with tempfile.TemporaryDirectory() as tmpdirname:
									fx_model.save_pretrained(lowerCamelCase						)
									lowerCamelCase_	:     Optional[int] 						=      VisionTextDualEncoderModel.from_pretrained(lowerCamelCase       ,      from_flax=lowerCamelCase						)

							pt_model_loaded.to(lowerCamelCase						)
							pt_model_loaded.eval()

							with torch.no_grad():
									lowerCamelCase_	:     Optional[Any] 						=      pt_model_loaded(**lowerCamelCase						).to_tuple()

							self.assertEqual(len(lowerCamelCase						)       ,      len(lowerCamelCase						)       ,      '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(lowerCamelCase       ,      pt_output_loaded.numpy()       ,      4E-2						)



					def   __a     (		self     :     Dict       ,      lowerCamelCase     :     Dict       ,      lowerCamelCase     :     Optional[Any]       ,      lowerCamelCase     :     List[str]						):
							lowerCamelCase_	:     Tuple 						=      VisionTextDualEncoderConfig.from_vision_text_configs(lowerCamelCase       ,      lowerCamelCase						)

							lowerCamelCase_	:     List[str] 						=      VisionTextDualEncoderModel(lowerCamelCase						)
							lowerCamelCase_	:     int 						=      FlaxVisionTextDualEncoderModel(lowerCamelCase						)

							lowerCamelCase_	:     Any 						=      convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,      lowerCamelCase						)
							lowerCamelCase_	:     Any 						=      fx_state

							self.check_pt_flax_equivalence(lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase						)



					def   __a     (		self     :     int       ,      lowerCamelCase     :     Any       ,      lowerCamelCase     :     Any       ,      lowerCamelCase     :     Optional[int]						):
							lowerCamelCase_	:     Union[str, Any] 						=      VisionTextDualEncoderConfig.from_vision_text_configs(lowerCamelCase       ,      lowerCamelCase						)

							lowerCamelCase_	:     List[str] 						=      VisionTextDualEncoderModel(lowerCamelCase						)
							lowerCamelCase_	:     List[str] 						=      FlaxVisionTextDualEncoderModel(lowerCamelCase						)

							lowerCamelCase_	:     Optional[Any] 						=      load_flax_weights_in_pytorch_model(lowerCamelCase       ,      fx_model.params						)

							self.check_pt_flax_equivalence(lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase						)



					def   __a     (		self     :     Union[str, Any]						):
							lowerCamelCase_	:     Any 						=      self.prepare_config_and_inputs()
							self.check_model_from_pretrained_configs(**lowerCamelCase						)



					def   __a     (		self     :     List[str]						):
							lowerCamelCase_	:     Optional[int] 						=      self.prepare_config_and_inputs()
							self.check_vision_text_dual_encoder_from_pretrained(**lowerCamelCase						)



					def   __a     (		self     :     Union[str, Any]						):
							lowerCamelCase_	:     Optional[int] 						=      self.prepare_config_and_inputs()
							self.check_save_load(**lowerCamelCase						)



					def   __a     (		self     :     Tuple						):
							lowerCamelCase_	:     List[Any] 						=      self.prepare_config_and_inputs()
							self.check_vision_text_output_attention(**lowerCamelCase						)



					@is_pt_flax_cross_test
					def   __a     (		self     :     Optional[int]						):
							lowerCamelCase_	:     List[Any] 						=      self.prepare_config_and_inputs()
							lowerCamelCase_	:     Any 						=      config_inputs_dict.pop('vision_config'						)
							lowerCamelCase_	:     Any 						=      config_inputs_dict.pop('text_config'						)

							lowerCamelCase_	:     List[str] 						=      config_inputs_dict

							self.check_equivalence_pt_to_flax(lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase						)
							self.check_equivalence_flax_to_pt(lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase						)





					@slow
					def   __a     (		self     :     Any						):
							lowerCamelCase_						,      lowerCamelCase_	:     List[Any] 						=      self.get_pretrained_model_and_inputs()

							lowerCamelCase_	:     List[Any] 						=      model_a(**lowerCamelCase						)
							lowerCamelCase_	:     Union[str, Any] 						=      outputs[0]

							with tempfile.TemporaryDirectory() as tmp_dirname:
									model_a.save_pretrained(lowerCamelCase						)
									lowerCamelCase_	:     Dict 						=      FlaxVisionTextDualEncoderModel.from_pretrained(lowerCamelCase						)

									lowerCamelCase_	:     str 						=      model_a(**lowerCamelCase						)
									lowerCamelCase_	:     List[str] 						=      after_outputs[0]
									lowerCamelCase_	:     Tuple 						=      np.amax(np.abs(out_a - out_a						)						)
									self.assertLessEqual(lowerCamelCase       ,      1E-5						)






@require_flax
class 						UpperCamelCase_  (  snake_case__  ,     unittest.TestCase							):



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

							return model, inputs



					def   __a     (		self     :     Optional[int]       ,      lowerCamelCase     :     Tuple       ,      lowerCamelCase     :     Any						):
							lowerCamelCase_	:     Union[str, Any] 						=      FlaxViTModel(lowerCamelCase						)
							lowerCamelCase_	:     List[Any] 						=      FlaxBertModel(lowerCamelCase						)
							return vision_model, text_model





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

							lowerCamelCase_						,      lowerCamelCase_	:     Optional[int] 						=      vision_config_and_inputs

							lowerCamelCase_						,      lowerCamelCase_						,      lowerCamelCase_						,      lowerCamelCase_	:     Dict 						=      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 						UpperCamelCase_  (  snake_case__  ,     unittest.TestCase							):



					def   __a     (		self     :     Optional[int]						):
							lowerCamelCase_	:     Tuple 						=      FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
							    'hf-internal-testing/tiny-random-clip'       ,      'hf-internal-testing/tiny-bert'       ,      vision_from_pt=lowerCamelCase       ,      text_from_pt=lowerCamelCase       ,      )
							lowerCamelCase_	:     Optional[Any] 						=      13
							lowerCamelCase_	:     str 						=      floats_tensor(
							    [
							        batch_size,
							        model.config.vision_config.num_channels,
							        model.config.vision_config.image_size,
							        model.config.vision_config.image_size,
							    ]						)
							lowerCamelCase_	:     str 						=      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   __a     (		self     :     int       ,      lowerCamelCase     :     Optional[int]       ,      lowerCamelCase     :     str						):
							lowerCamelCase_	:     List[str] 						=      FlaxCLIPVisionModel(lowerCamelCase						)
							lowerCamelCase_	:     Tuple 						=      FlaxBertModel(lowerCamelCase						)
							return vision_model, text_model





					def   __a     (		self     :     Any						):
							lowerCamelCase_	:     int 						=      FlaxCLIPVisionModelTester(self						)
							lowerCamelCase_	:     Union[str, Any] 						=      FlaxBertModelTester(self						)
							lowerCamelCase_	:     Optional[Any] 						=      clip_model_tester.prepare_config_and_inputs()
							lowerCamelCase_	:     Any 						=      bert_model_tester.prepare_config_and_inputs()

							lowerCamelCase_						,      lowerCamelCase_	:     Union[str, Any] 						=      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 						UpperCamelCase_  (  unittest.TestCase							):



					@slow
					def   __a     (		self     :     Optional[int]						):
							lowerCamelCase_	:     str 						=      FlaxVisionTextDualEncoderModel.from_pretrained('clip-italian/clip-italian'       ,      logit_scale_init_value=1.0						)
							lowerCamelCase_	:     Dict 						=      VisionTextDualEncoderProcessor.from_pretrained('clip-italian/clip-italian'						)

							lowerCamelCase_	:     List[Any] 						=      Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'						)
							lowerCamelCase_	:     List[str] 						=      processor(
							    text=['una foto di un gatto', 'una foto di un cane']       ,      images=lowerCamelCase       ,      padding=lowerCamelCase       ,      return_tensors='np'						)

							lowerCamelCase_	:     Tuple 						=      model(**lowerCamelCase						)

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

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







def 		_SCREAMING_SNAKE_CASE     ( lowerCAmelCase__=2_81_23						):
		lowerCamelCase_	:     Any 						=      [1] * (limit + 1)

		for i in range(2							,int(limit**0.5						) + 1						):
				sum_divs[i * i] += i
				for k in range(i + 1							,limit // i + 1						):
						sum_divs[k * i] += k + i

		lowerCamelCase_	:     List[str] 						=      set()
		lowerCamelCase_	:     List[str] 						=      0

		for n in range(1							,limit + 1						):
				if sum_divs[n] > n:
						abundants.add(lowerCAmelCase__						)

				if not any((n - a in abundants) for a in abundants						):
						res += n

		return res


if __name__ == "__main__":
		print(solution())
