




'''simple docstring'''

import argparse
import json
import os

import tensorstore as ts
import torch
from flax import serialization
from flax.traverse_util import flatten_dict, unflatten_dict
from tensorflow.io import gfile

from transformers.modeling_utils import dtype_byte_size
from transformers.models.switch_transformers.convert_switch_transformers_original_flax_checkpoint_to_pytorch import (
    rename_keys,
)
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME
from transformers.utils.hub import convert_file_size_to_int


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase	)		->					int:

					'''simple docstring'''

					if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 3:
										# expert layer
										snake_case_									=				flax_key_tuple[:-1] + ('''weight''',)
										snake_case_									=				torch.permute(SCREAMING_SNAKE_CASE_,      (0, 2, 1)	)
					elif flax_key_tuple[-1] == "kernel" and ".".join(SCREAMING_SNAKE_CASE_	):
										# linear layer
										snake_case_									=				flax_key_tuple[:-1] + ('''weight''',)
										snake_case_									=				flax_tensor.T
					elif flax_key_tuple[-1] in ["scale", "embedding"]:
										snake_case_									=				flax_key_tuple[:-1] + ('''weight''',)

					return flax_key_tuple, flax_tensor


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase,      __UpperCAmelCase	)		->					List[str]:

					'''simple docstring'''

					if "metadata" in layer:
										snake_case_									=				layer.split('''metadata'''	)
										snake_case_									=				''''''.join(split_layer[0]	)[:-1]
										snake_case_									=				[tuple(('''metadata''' + split_layer[1]).split('''/'''	)	)]
					elif "kvstore" in layer:
										snake_case_									=				layer.split('''kvstore'''	)
										snake_case_									=				''''''.join(split_layer[0]	)[:-1]
										snake_case_									=				[tuple(('''kvstore''' + split_layer[1]).split('''/'''	)	)]

					else:
										snake_case_									=				layer.split('''/'''	)
										snake_case_									=				'''/'''.join(split_layer[:-1]	)
										snake_case_									=				(split_layer[-1],)

					if "kvstore/path" in layer:
										snake_case_									=				F"{switch_checkpoint_path}/{checkpoint_info[layer]}"
					elif "kvstore/driver" in layer:
										snake_case_									=				'''file'''
					else:
										snake_case_									=				checkpoint_info[layer]

					return curr_real_layer_name, split_layer, content


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase	)		->					Dict:

					'''simple docstring'''

					snake_case_									=				rename_keys(SCREAMING_SNAKE_CASE_	)
					snake_case_									=				{}
					for k, v in current_block.items():
										snake_case_									=				v
					snake_case_									=				new_current_block
					torch.save(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase,      __UpperCAmelCase,      __UpperCAmelCase,      __UpperCAmelCase = WEIGHTS_NAME	)		->					str:

					'''simple docstring'''

					snake_case_									=				convert_file_size_to_int(SCREAMING_SNAKE_CASE_	)
					snake_case_									=				[]
					snake_case_									=				{}
					snake_case_									=				0
					snake_case_									=				0

					os.makedirs(SCREAMING_SNAKE_CASE_,      exist_ok=SCREAMING_SNAKE_CASE_	)
					with gfile.GFile(switch_checkpoint_path + '''/checkpoint''',      '''rb'''	) as fp:
										snake_case_									=				serialization.msgpack_restore(fp.read()	)['''optimizer''']['''target''']
										snake_case_									=				flatten_dict(SCREAMING_SNAKE_CASE_,      sep='''/'''	)

					snake_case_									=				{}
					for layer in checkpoint_info.keys():
										snake_case_			,snake_case_			,snake_case_									=				get_key_and_tensorstore_dict(
										    SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)
										if curr_real_layer_name in all_layers:
															snake_case_									=				content
										else:
															snake_case_									=				{split_layer[-1]: content}

					for key in all_layers.keys():
										# open tensorstore file
										snake_case_									=				ts.open(unflatten_dict(all_layers[key]	)	).result().read().result()
										snake_case_									=				torch.tensor(SCREAMING_SNAKE_CASE_	)
										snake_case_									=				raw_weights.numel() * dtype_byte_size(raw_weights.dtype	)

										# use the renaming pattern from the small conversion scripts
										snake_case_			,snake_case_									=				rename_base_flax_keys(tuple(key.split('''/'''	)	),      SCREAMING_SNAKE_CASE_	)
										snake_case_									=				'''/'''.join(SCREAMING_SNAKE_CASE_	)

										# If this weight is going to tip up over the maximal size, we split.
										if current_block_size + weight_size > max_shard_size:
															snake_case_									=				os.path.join(
															    SCREAMING_SNAKE_CASE_,      weights_name.replace('''.bin''',      F"-{len(SCREAMING_SNAKE_CASE_	)+1:05d}-of-???.bin"	)	)
															rename_and_save_block(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)
															sharded_state_dicts.append(current_block.keys()	)
															del current_block
															snake_case_									=				{}
															snake_case_									=				0

										snake_case_									=				raw_weights.to(getattr(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)	)
										current_block_size += weight_size
										total_size += weight_size

					# Add the last block
					snake_case_									=				os.path.join(SCREAMING_SNAKE_CASE_,      weights_name.replace('''.bin''',      F"-{len(SCREAMING_SNAKE_CASE_	)+1:05d}-of-???.bin"	)	)
					rename_and_save_block(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)
					sharded_state_dicts.append(current_block.keys()	)

					# If we only have one shard, we return it
					if len(SCREAMING_SNAKE_CASE_	) == 1:
										return {weights_name: sharded_state_dicts[0]}, None

					# Otherwise, let's build the index
					snake_case_									=				{}
					snake_case_									=				{}
					for idx, shard in enumerate(SCREAMING_SNAKE_CASE_	):
										snake_case_									=				weights_name.replace(
										    '''.bin''',      F"-{idx+1:05d}-of-{len(SCREAMING_SNAKE_CASE_	):05d}.bin"	)  # len(sharded_state_dicts):05d}
										snake_case_									=				os.path.join(SCREAMING_SNAKE_CASE_,      weights_name.replace('''.bin''',      F"-{idx+1:05d}-of-???.bin"	)	)
										os.rename(SCREAMING_SNAKE_CASE_,      os.path.join(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)	)
										snake_case_									=				shard
										for key in shard:
															snake_case_									=				shard_file

    # Add the metadata
					snake_case_									=				{'''total_size''': total_size}
					snake_case_									=				{'''metadata''': metadata, '''weight_map''': weight_map}

					with open(os.path.join(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	),      '''w''',      encoding='''utf-8'''	) as f:
										snake_case_									=				json.dumps(SCREAMING_SNAKE_CASE_,      indent=2,      sort_keys=SCREAMING_SNAKE_CASE_	) + '''\n'''
										f.write(SCREAMING_SNAKE_CASE_	)

					return metadata, index


if __name__ == "__main__":
							a  :				Any									=  argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--switch_t5x_checkpoint_path',
							    default='/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128/checkpoint_634600',
							    type=str,
							    required=False,
							    help='Path to a directory containing a folder per layer. Follows the original Google format.',
							)
							parser.add_argument('--max_shard_size', default='10GB', required=False, help='Max shard size')
							parser.add_argument('--dtype', default='bfloat16', type=str, required=False, help='dtype of the saved model')
							parser.add_argument(
							    '--pytorch_dump_folder_path',
							    default='/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128-converted',
							    type=str,
							    required=False,
							    help='Path to the output pytorch model.',
							)
							a  :				str									=  parser.parse_args()
							shard_on_the_fly(
							    args.switch_tax_checkpoint_path,
							    args.pytorch_dump_folder_path,
							    args.max_shard_size,
							    args.dtype,
							)


def 		__magic_name__    (	)		->					str:

					'''simple docstring'''

					from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration, TaTokenizer

					snake_case_									=				SwitchTransformersConfig.from_pretrained('''google/switch-base-8'''	)
					config.save_pretrained('''/home/arthur_huggingface_co/transformers/switch_converted'''	)
					snake_case_									=				SwitchTransformersForConditionalGeneration.from_pretrained(
					    '''/home/arthur_huggingface_co/transformers/switch_converted''',      device_map='''auto'''	)

					snake_case_									=				TaTokenizer.from_pretrained('''t5-small'''	)
					snake_case_									=				'''A <extra_id_0> walks into a bar a orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.'''

					snake_case_									=				tokenizer(SCREAMING_SNAKE_CASE_,      return_tensors='''pt'''	).input_ids
					snake_case_									=				model.generate(SCREAMING_SNAKE_CASE_,      decoder_start_token_id=0	)
					print(tokenizer.decode(out[0]	)	)







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_							:				Union[str, 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			)			)




