



"""simple docstring"""





import argparse
import gc
import json
import os
import re

import torch
from huggingface_hub import hf_hub_download

from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerFast, RwkvConfig
from transformers.modeling_utils import WEIGHTS_INDEX_NAME, shard_checkpoint


_UpperCamelCase							=      {
    "169M": 12,
    "430M": 24,
    "1B5": 24,
    "3B": 32,
    "7B": 32,
    "14B": 40,
}

_UpperCamelCase							=      {
    "169M": 768,
    "430M": 1024,
    "1B5": 2048,
    "3B": 2560,
    "7B": 4096,
    "14B": 5120,
}



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Any   ):


							'''simple docstring'''







							__lowerCamelCase      :      int         =list(state_dict.keys()   )
							for name in state_dict_keys:
														__lowerCamelCase      :      List[str]         =state_dict.pop(UpperCAmelCase__   )
														# emb -> embedding
														if name.startswith('''emb.'''   ):
																					__lowerCamelCase      :      str         =name.replace('''emb.'''    ,						'''embeddings.'''   )
														# ln_0 -> pre_ln (only present at block 0)
														if name.startswith('''blocks.0.ln0'''   ):
																					__lowerCamelCase      :      int         =name.replace('''blocks.0.ln0'''    ,						'''blocks.0.pre_ln'''   )
														# att -> attention
														__lowerCamelCase      :      Union[str, Any]         =re.sub(r'''blocks\.(\d+)\.att'''    ,						r'''blocks.\1.attention'''    ,						UpperCAmelCase__   )
														# ffn -> feed_forward
														__lowerCamelCase      :      Optional[Any]         =re.sub(r'''blocks\.(\d+)\.ffn'''    ,						r'''blocks.\1.feed_forward'''    ,						UpperCAmelCase__   )
														# time_mix_k -> time_mix_key and reshape
														if name.endswith('''.time_mix_k'''   ):
																					__lowerCamelCase      :      Any         =name.replace('''.time_mix_k'''    ,						'''.time_mix_key'''   )
														# time_mix_v -> time_mix_value and reshape
														if name.endswith('''.time_mix_v'''   ):
																					__lowerCamelCase      :      int         =name.replace('''.time_mix_v'''    ,						'''.time_mix_value'''   )
														# time_mix_r -> time_mix_key and reshape
														if name.endswith('''.time_mix_r'''   ):
																					__lowerCamelCase      :      Any         =name.replace('''.time_mix_r'''    ,						'''.time_mix_receptance'''   )

														if name != "head.weight":
																					__lowerCamelCase      :      int         ="""rwkv.""" + name

														__lowerCamelCase      :      Any         =weight
							return state_dict



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							List[str]    ,						SCREAMING_SNAKE_CASE       :							Optional[int]    ,						SCREAMING_SNAKE_CASE       :							int    ,						SCREAMING_SNAKE_CASE       :							Optional[Any]=None    ,						SCREAMING_SNAKE_CASE       :							List[str]=None    ,						SCREAMING_SNAKE_CASE       :							str=False    ,						SCREAMING_SNAKE_CASE       :							Optional[Any]=None   ):


							'''simple docstring'''







							if tokenizer_file is None:
														print('''No `--tokenizer_file` provided, we will use the default tokenizer.'''   )
														__lowerCamelCase      :      Dict         =50277
														__lowerCamelCase      :      Optional[int]         =AutoTokenizer.from_pretrained('''EleutherAI/gpt-neox-20b'''   )
							else:
														__lowerCamelCase      :      Union[str, Any]         =PreTrainedTokenizerFast(tokenizer_file=UpperCAmelCase__   )
														__lowerCamelCase      :      str         =len(UpperCAmelCase__   )
							tokenizer.save_pretrained(UpperCAmelCase__   )

							# 2. Build the config
							__lowerCamelCase      :      Optional[Any]         =list(NUM_HIDDEN_LAYERS_MAPPING.keys()   )
							if size is None:
														# Try to infer size from the checkpoint name
														for candidate in possible_sizes:
																					if candidate in checkpoint_file:
																												__lowerCamelCase      :      List[str]         =candidate
																												break
														if size is None:
																					raise ValueError('''Could not infer the size, please provide it with the `--size` argument.'''   )
							if size not in possible_sizes:
														raise ValueError(F'`size` should be one of {possible_sizes}, got {size}.'   )

							__lowerCamelCase      :      Dict         =RwkvConfig(
							    vocab_size=UpperCAmelCase__    ,						num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size]    ,						hidden_size=HIDEN_SIZE_MAPPING[size]    ,						)
							config.save_pretrained(UpperCAmelCase__   )

							# 3. Download model file then convert state_dict
							__lowerCamelCase      :      Optional[Any]         =hf_hub_download(UpperCAmelCase__    ,						UpperCAmelCase__   )
							__lowerCamelCase      :      List[Any]         =torch.load(UpperCAmelCase__    ,						map_location='''cpu'''   )
							__lowerCamelCase      :      Dict         =convert_state_dict(UpperCAmelCase__   )

							# 4. Split in shards and save
							__lowerCamelCase      :      Dict         =shard_checkpoint(UpperCAmelCase__   )
							for shard_file, shard in shards.items():
														torch.save(UpperCAmelCase__    ,						os.path.join(UpperCAmelCase__    ,						UpperCAmelCase__   )   )

							if index is not None:
														__lowerCamelCase      :      Dict         =os.path.join(UpperCAmelCase__    ,						UpperCAmelCase__   )
														# Save the index as well
														with open(UpperCAmelCase__    ,						'''w'''    ,						encoding='''utf-8'''   ) as f:
																					__lowerCamelCase      :      Optional[Any]         =json.dumps(UpperCAmelCase__    ,						indent=2    ,						sort_keys=UpperCAmelCase__   ) + """\n"""
																					f.write(UpperCAmelCase__   )

														# 5. Clean up shards (for some reason the file PyTorch saves take the same space as the whole state_dict
														print(
														    '''Cleaning up shards. This may error with an OOM error, it this is the case don\'t worry you still have converted the model.'''   )
														__lowerCamelCase      :      Optional[Any]         =list(shards.keys()   )

														del state_dict
														del shards
														gc.collect()

														for shard_file in shard_files:
																					__lowerCamelCase      :      Optional[Any]         =torch.load(os.path.join(UpperCAmelCase__    ,						UpperCAmelCase__   )   )
																					torch.save({k: v.cpu().clone() for k, v in state_dict.items()}    ,						os.path.join(UpperCAmelCase__    ,						UpperCAmelCase__   )   )

							del state_dict
							gc.collect()

							if push_to_hub:
														if model_name is None:
																					raise ValueError('''Please provide a `model_name` to push the model to the Hub.'''   )
														__lowerCamelCase      :      Dict         =AutoModelForCausalLM.from_pretrained(UpperCAmelCase__   )
														model.push_to_hub(UpperCAmelCase__    ,						max_shard_size='''2GB'''   )
														tokenizer.push_to_hub(UpperCAmelCase__   )


if __name__ == "__main__":
							_UpperCamelCase							=      argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--repo_id', default=None, type=str, required=True, help='Repo ID from which to pull the checkpoint.'
							)
							parser.add_argument(
							    '--checkpoint_file', default=None, type=str, required=True, help='Name of the checkpoint file in the repo.'
							)
							parser.add_argument(
							    '--output_dir', default=None, type=str, required=True, help='Where to save the converted model.'
							)
							parser.add_argument(
							    '--tokenizer_file',
							    default=None,
							    type=str,
							    help='Path to the tokenizer file to use (if not provided, only the model is converted).',
							)
							parser.add_argument(
							    '--size',
							    default=None,
							    type=str,
							    help='Size of the model. Will be inferred from the `checkpoint_file` if not passed.',
							)
							parser.add_argument(
							    '--push_to_hub',
							    action='store_true',
							    help='Push to the Hub the converted model.',
							)
							parser.add_argument(
							    '--model_name',
							    default=None,
							    type=str,
							    help='Name of the pushed model on the Hub, including the username / organization.',
							)

							_UpperCamelCase							=      parser.parse_args()
							convert_rmkv_checkpoint_to_hf_format(
							    args.repo_id,
							    args.checkpoint_file,
							    args.output_dir,
							    size=args.size,
							    tokenizer_file=args.tokenizer_file,
							    push_to_hub=args.push_to_hub,
							    model_name=args.model_name,
							)











"""simple docstring"""





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 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Optional[int]   ):


							'''simple docstring'''







							if isinstance(SCREAMING_SNAKE_CASE    ,						collections.abc.Iterable   ):
														return x
							return (x, x)
@require_tf
class 		SCREAMING_SNAKE_CASE_							:
				"""simple docstring"""




				def 					__lowercase  (     self  :Optional[int]						,    __lowercase  :Union[str, Any]						,    __lowercase  :List[Any]			):
											pass




				def 					__lowercase  (     self  :Any			):
											pass




				def 					__lowercase  (     self  :str			):
											pass




				def 					__lowercase  (     self  :str						,    __lowercase  :Dict						,    __lowercase  :List[Any]						,    __lowercase  :Union[str, Any]						,    __lowercase  :Union[str, Any]						,    __lowercase  :Any=None						,    **__lowercase  :Tuple			):
											__lowerCamelCase      :      Optional[Any]         =VisionTextDualEncoderConfig.from_vision_text_configs(__lowercase						,    __lowercase			)

											__lowerCamelCase      :      Dict         =TFVisionTextDualEncoderModel(__lowercase			)

											__lowerCamelCase      :      List[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 					__lowercase  (     self  :str						,    __lowercase  :Any						,    __lowercase  :Dict						,    __lowercase  :Any						,    __lowercase  :int						,    __lowercase  :str=None						,    **__lowercase  :List[str]			):
											__lowerCamelCase    ,						__lowerCamelCase      :      Any         =self.get_vision_text_model(__lowercase						,    __lowercase			)
											__lowerCamelCase      :      List[str]         =TFVisionTextDualEncoderModel(vision_model=__lowercase						,    text_model=__lowercase			)

											__lowerCamelCase      :      Optional[int]         =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 					__lowercase  (     self  :int						,    __lowercase  :Dict						,    __lowercase  :List[Any]						,    __lowercase  :Any						,    __lowercase  :Optional[Any]						,    __lowercase  :Tuple=None						,    **__lowercase  :List[Any]			):
											__lowerCamelCase    ,						__lowerCamelCase      :      int         =self.get_vision_text_model(__lowercase						,    __lowercase			)
											__lowerCamelCase      :      Dict         ={'''vision_model''': vision_model, '''text_model''': text_model}
											__lowerCamelCase      :      List[Any]         =TFVisionTextDualEncoderModel.from_vision_text_pretrained(**__lowercase			)

											__lowerCamelCase      :      Optional[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 					__lowercase  (     self  :Optional[Any]						,    __lowercase  :Dict						,    __lowercase  :int						,    __lowercase  :Dict						,    __lowercase  :List[Any]						,    __lowercase  :Tuple=None						,    **__lowercase  :Optional[Any]			):
											__lowerCamelCase    ,						__lowerCamelCase      :      Tuple         =self.get_vision_text_model(__lowercase						,    __lowercase			)
											__lowerCamelCase      :      List[Any]         =TFVisionTextDualEncoderModel(vision_model=__lowercase						,    text_model=__lowercase			)

											__lowerCamelCase      :      Any         =model(input_ids=__lowercase						,    pixel_values=__lowercase						,    attention_mask=__lowercase			)
											__lowerCamelCase      :      Optional[Any]         =output[0].numpy()

											with tempfile.TemporaryDirectory() as tmpdirname:
																		model.save_pretrained(__lowercase			)
																		__lowerCamelCase      :      Dict         =TFVisionTextDualEncoderModel.from_pretrained(__lowercase			)

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




				def 					__lowercase  (     self  :int						,    __lowercase  :Any						,    __lowercase  :List[Any]						,    __lowercase  :Any						,    __lowercase  :str						,    __lowercase  :int=None						,    **__lowercase  :List[Any]			):
											__lowerCamelCase    ,						__lowerCamelCase      :      Any         =self.get_vision_text_model(__lowercase						,    __lowercase			)
											__lowerCamelCase      :      Union[str, Any]         =TFVisionTextDualEncoderModel(vision_model=__lowercase						,    text_model=__lowercase			)

											__lowerCamelCase      :      Any         =model(
											    input_ids=__lowercase						,    pixel_values=__lowercase						,    attention_mask=__lowercase						,    output_attentions=__lowercase			)

											__lowerCamelCase      :      Optional[Any]         =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)
											__lowerCamelCase      :      Dict         =to_atuple(vision_model.config.image_size			)
											__lowerCamelCase      :      Tuple         =to_atuple(vision_model.config.patch_size			)
											__lowerCamelCase      :      Tuple         =(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
											__lowerCamelCase      :      Any         =num_patches + 1
											self.assertEqual(vision_attentions[0].shape[-3:]						,    (vision_config.num_attention_heads, seq_len, seq_len)			)

											__lowerCamelCase      :      List[Any]         =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 					__lowercase  (     self  :List[Any]						,    __lowercase  :np.ndarray						,    __lowercase  :np.ndarray						,    __lowercase  :float			):
											__lowerCamelCase      :      Dict         =np.abs((a - b)			).max()
											self.assertLessEqual(__lowercase						,    __lowercase						,    f'Difference between torch and flax is {diff} (>= {tol}).'			)




				def 					__lowercase  (     self  :List[str]			):
											__lowerCamelCase      :      Any         =self.prepare_config_and_inputs()
											self.check_vision_text_dual_encoder_model(**__lowercase			)




				def 					__lowercase  (     self  :Tuple			):
											__lowerCamelCase      :      List[Any]         =self.prepare_config_and_inputs()
											self.check_model_from_pretrained_configs(**__lowercase			)




				def 					__lowercase  (     self  :Union[str, Any]			):
											__lowerCamelCase      :      Dict         =self.prepare_config_and_inputs()
											self.check_vision_text_dual_encoder_from_pretrained(**__lowercase			)




				def 					__lowercase  (     self  :Any			):
											__lowerCamelCase      :      List[Any]         =self.prepare_config_and_inputs()
											self.check_save_load(**__lowercase			)




				def 					__lowercase  (     self  :List[str]			):
											__lowerCamelCase      :      Any         =self.prepare_config_and_inputs()
											self.check_vision_text_output_attention(**__lowercase			)




				@slow
				def 					__lowercase  (     self  :Tuple			):
											__lowerCamelCase    ,						__lowerCamelCase      :      Optional[Any]         =self.get_pretrained_model_and_inputs()

											__lowerCamelCase      :      Optional[Any]         =model_a(**__lowercase			)
											__lowerCamelCase      :      int         =outputs[0].numpy()

											with tempfile.TemporaryDirectory() as tmp_dirname:
																		model_a.save_pretrained(__lowercase			)
																		__lowerCamelCase      :      List[Any]         =TFVisionTextDualEncoderModel.from_pretrained(__lowercase			)

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

@require_tf
class 		SCREAMING_SNAKE_CASE_							(    snake_case__       ,  unittest.TestCase   ):
				"""simple docstring"""




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

											return model, inputs




				def 					__lowercase  (     self  :Optional[Any]						,    __lowercase  :List[Any]						,    __lowercase  :List[str]			):
											__lowerCamelCase      :      Optional[Any]         =TFViTModel(__lowercase						,    name='''vision_model'''			)
											__lowerCamelCase      :      Union[str, Any]         =TFBertModel(__lowercase						,    name='''text_model'''			)
											return vision_model, text_model




				def 					__lowercase  (     self  :Union[str, Any]			):
											__lowerCamelCase      :      str         =TFViTModelTester(self			)
											__lowerCamelCase      :      Any         =TFBertModelTester(self			)
											__lowerCamelCase      :      Optional[Any]         =vit_model_tester.prepare_config_and_inputs()
											__lowerCamelCase      :      Tuple         =bert_model_tester.prepare_config_and_inputs()

											__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase      :      Dict         =vision_config_and_inputs

											(
											    (
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						
											)      :      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_tf
class 		SCREAMING_SNAKE_CASE_							(    snake_case__       ,  unittest.TestCase   ):
				"""simple docstring"""




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

											return model, inputs




				def 					__lowercase  (     self  :Any						,    __lowercase  :Tuple						,    __lowercase  :List[Any]						,    __lowercase  :Union[str, Any]						,    __lowercase  :Any						,    __lowercase  :int=None						,    **__lowercase  :List[str]			):
											__lowerCamelCase    ,						__lowerCamelCase      :      Union[str, Any]         =self.get_vision_text_model(__lowercase						,    __lowercase			)
											__lowerCamelCase      :      Optional[int]         =TFVisionTextDualEncoderModel(vision_model=__lowercase						,    text_model=__lowercase			)

											__lowerCamelCase      :      Any         =model(
											    input_ids=__lowercase						,    pixel_values=__lowercase						,    attention_mask=__lowercase						,    output_attentions=__lowercase			)

											__lowerCamelCase      :      str         =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)
											__lowerCamelCase      :      int         =to_atuple(vision_model.config.image_size			)
											__lowerCamelCase      :      Union[str, Any]         =to_atuple(vision_model.config.patch_size			)
											__lowerCamelCase      :      Tuple         =(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
											__lowerCamelCase      :      Any         =num_patches + 2
											self.assertEqual(vision_attentions[0].shape[-3:]						,    (vision_config.num_attention_heads, seq_len, seq_len)			)

											__lowerCamelCase      :      Any         =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 					__lowercase  (     self  :Any						,    __lowercase  :Any						,    __lowercase  :Optional[Any]			):
											__lowerCamelCase      :      str         =TFDeiTModel(__lowercase						,    name='''vision_model'''			)
											__lowerCamelCase      :      List[str]         =TFRobertaModel(__lowercase						,    name='''text_model'''			)
											return vision_model, text_model




				def 					__lowercase  (     self  :Dict			):
											__lowerCamelCase      :      Optional[int]         =TFDeiTModelTester(self			)
											__lowerCamelCase      :      Any         =TFRobertaModelTester(self			)
											__lowerCamelCase      :      Dict         =vit_model_tester.prepare_config_and_inputs()
											__lowerCamelCase      :      Optional[int]         =bert_model_tester.prepare_config_and_inputs()

											__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase      :      Optional[int]         =vision_config_and_inputs

											(
											    (
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						
											)      :      List[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 		SCREAMING_SNAKE_CASE_							(    snake_case__       ,  unittest.TestCase   ):
				"""simple docstring"""




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

											return model, inputs




				def 					__lowercase  (     self  :List[Any]						,    __lowercase  :str						,    __lowercase  :Tuple			):
											__lowerCamelCase      :      int         =TFCLIPVisionModel(__lowercase						,    name='''vision_model'''			)
											__lowerCamelCase      :      Optional[Any]         =TFBertModel(__lowercase						,    name='''text_model'''			)
											return vision_model, text_model




				def 					__lowercase  (     self  :Tuple			):
											__lowerCamelCase      :      List[str]         =TFCLIPVisionModelTester(self			)
											__lowerCamelCase      :      Union[str, Any]         =TFBertModelTester(self			)
											__lowerCamelCase      :      str         =clip_model_tester.prepare_config_and_inputs()
											__lowerCamelCase      :      Tuple         =bert_model_tester.prepare_config_and_inputs()

											__lowerCamelCase    ,						__lowerCamelCase      :      Any         =vision_config_and_inputs

											(
											    (
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						(
											    __lowerCamelCase
											)    ,						
											)      :      str         =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 		SCREAMING_SNAKE_CASE_							(    unittest.TestCase   ):
				"""simple docstring"""




				@slow
				def 					__lowercase  (     self  :List[Any]			):
											__lowerCamelCase      :      Union[str, Any]         =TFVisionTextDualEncoderModel.from_pretrained(
											    '''clip-italian/clip-italian'''						,    logit_scale_init_value=1.0						,    from_pt=__lowercase			)
											__lowerCamelCase      :      str         =VisionTextDualEncoderProcessor.from_pretrained('''clip-italian/clip-italian'''			)

											__lowerCamelCase      :      Optional[Any]         =Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''			)
											__lowerCamelCase      :      int         =processor(
											    text=['''una foto di un gatto''', '''una foto di un cane''']						,    images=__lowercase						,    padding=__lowercase						,    return_tensors='''np'''			)

											__lowerCamelCase      :      str         =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])						,    )

											__lowerCamelCase      :      int         =np.array([[1.2284727, 0.3104122]]			)

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




