


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



@require_tf
class        a					:





     def  SCREAMING_SNAKE_CASE__       (    self					:				Dict    ,	SCREAMING_SNAKE_CASE_					:				Tuple    ,	SCREAMING_SNAKE_CASE_					:				List[str]							):
        pass





     def  SCREAMING_SNAKE_CASE__       (    self					:				List[Any]							):
        pass





     def  SCREAMING_SNAKE_CASE__       (    self					:				Tuple							):
        pass





     def  SCREAMING_SNAKE_CASE__       (    self					:				List[Any]    ,	SCREAMING_SNAKE_CASE_					:				int    ,	SCREAMING_SNAKE_CASE_					:				str    ,	SCREAMING_SNAKE_CASE_					:				Dict    ,	SCREAMING_SNAKE_CASE_					:				Union[str, Any]    ,	SCREAMING_SNAKE_CASE_					:				List[str]=None    ,	**SCREAMING_SNAKE_CASE_					:				Any							):
        __lowerCamelCase:				Dict     =	VisionTextDualEncoderConfig.from_vision_text_configs(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				Union[str, Any]     =	TFVisionTextDualEncoderModel(SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				int     =	model(input_ids=SCREAMING_SNAKE_CASE_    ,	pixel_values=SCREAMING_SNAKE_CASE_    ,	attention_mask=SCREAMING_SNAKE_CASE_							)

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

        __lowerCamelCase:				Optional[int]     =	model(input_ids=SCREAMING_SNAKE_CASE_    ,	pixel_values=SCREAMING_SNAKE_CASE_    ,	attention_mask=SCREAMING_SNAKE_CASE_							)

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

        __lowerCamelCase:				Union[str, Any]     =	model(input_ids=SCREAMING_SNAKE_CASE_    ,	pixel_values=SCREAMING_SNAKE_CASE_    ,	attention_mask=SCREAMING_SNAKE_CASE_							)

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

        __lowerCamelCase:				List[Any]     =	model(input_ids=SCREAMING_SNAKE_CASE_    ,	pixel_values=SCREAMING_SNAKE_CASE_    ,	attention_mask=SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Tuple     =	output[0].numpy()

        with tempfile.TemporaryDirectory() as tmpdirname:
           model.save_pretrained(SCREAMING_SNAKE_CASE_							)
           __lowerCamelCase:				Union[str, Any]     =	TFVisionTextDualEncoderModel.from_pretrained(SCREAMING_SNAKE_CASE_							)

           __lowerCamelCase:				Any     =	model(input_ids=SCREAMING_SNAKE_CASE_    ,	pixel_values=SCREAMING_SNAKE_CASE_    ,	attention_mask=SCREAMING_SNAKE_CASE_							)
           __lowerCamelCase:				Union[str, Any]     =	after_output[0].numpy()
           __lowerCamelCase:				Optional[Any]     =	np.amax(np.abs(out_a - out_a							)							)
           self.assertLessEqual(SCREAMING_SNAKE_CASE_    ,	1E-5							)





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

        __lowerCamelCase:				Dict     =	model(
            input_ids=SCREAMING_SNAKE_CASE_    ,	pixel_values=SCREAMING_SNAKE_CASE_    ,	attention_mask=SCREAMING_SNAKE_CASE_    ,	output_attentions=SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				int     =	output.vision_model_output.attentions
        self.assertEqual(len(SCREAMING_SNAKE_CASE_							)    ,	vision_config.num_hidden_layers							)

        # in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
        __lowerCamelCase:				Any     =	to_atuple(vision_model.config.image_size							)
        __lowerCamelCase:				List[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:				str     =	output.text_model_output.attentions
        self.assertEqual(len(SCREAMING_SNAKE_CASE_							)    ,	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					:				Union[str, Any]    ,	SCREAMING_SNAKE_CASE_					:				np.ndarray    ,	SCREAMING_SNAKE_CASE_					:				np.ndarray    ,	SCREAMING_SNAKE_CASE_					:				float							):
        __lowerCamelCase:				Optional[int]     =	np.abs((a - b)							).max()
        self.assertLessEqual(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	F'Difference between torch and flax is {diff} (>= {tol}).'							)





     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[Any]							):
        __lowerCamelCase:				Tuple     =	self.prepare_config_and_inputs()
        self.check_vision_text_dual_encoder_model(**SCREAMING_SNAKE_CASE_							)





     def  SCREAMING_SNAKE_CASE__       (    self					:				Dict							):
        __lowerCamelCase:				str     =	self.prepare_config_and_inputs()
        self.check_model_from_pretrained_configs(**SCREAMING_SNAKE_CASE_							)





     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[int]							):
        __lowerCamelCase:				List[Any]     =	self.prepare_config_and_inputs()
        self.check_vision_text_dual_encoder_from_pretrained(**SCREAMING_SNAKE_CASE_							)





     def  SCREAMING_SNAKE_CASE__       (    self					:				List[str]							):
        __lowerCamelCase:				List[str]     =	self.prepare_config_and_inputs()
        self.check_save_load(**SCREAMING_SNAKE_CASE_							)





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







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

        __lowerCamelCase:				List[str]     =	model_a(**SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Optional[int]     =	outputs[0].numpy()

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

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



@require_tf
class        a					(			_UpperCAmelCase     ,unittest.TestCase   ):





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

        return model, inputs





     def  SCREAMING_SNAKE_CASE__       (    self					:				int    ,	SCREAMING_SNAKE_CASE_					:				Any    ,	SCREAMING_SNAKE_CASE_					:				Tuple							):
        __lowerCamelCase:				Union[str, Any]     =	TFViTModel(SCREAMING_SNAKE_CASE_    ,	name="""vision_model"""							)
        __lowerCamelCase:				List[Any]     =	TFBertModel(SCREAMING_SNAKE_CASE_    ,	name="""text_model"""							)
        return vision_model, text_model







     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[int]							):
        __lowerCamelCase:				str     =	TFViTModelTester(self							)
        __lowerCamelCase:				List[str]     =	TFBertModelTester(self							)
        __lowerCamelCase:				Optional[Any]     =	vit_model_tester.prepare_config_and_inputs()
        __lowerCamelCase:				List[str]     =	bert_model_tester.prepare_config_and_inputs()

        __lowerCamelCase   ,     __lowerCamelCase   ,     __lowerCamelCase:				Dict     =	vision_config_and_inputs

        (
            (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     
        ):				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        a					(			_UpperCAmelCase     ,unittest.TestCase   ):





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

        return model, inputs





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

        __lowerCamelCase:				List[Any]     =	model(
            input_ids=SCREAMING_SNAKE_CASE_    ,	pixel_values=SCREAMING_SNAKE_CASE_    ,	attention_mask=SCREAMING_SNAKE_CASE_    ,	output_attentions=SCREAMING_SNAKE_CASE_							)

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

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







     def  SCREAMING_SNAKE_CASE__       (    self					:				Union[str, Any]							):
        __lowerCamelCase:				int     =	TFDeiTModelTester(self							)
        __lowerCamelCase:				int     =	TFRobertaModelTester(self							)
        __lowerCamelCase:				str     =	vit_model_tester.prepare_config_and_inputs()
        __lowerCamelCase:				Tuple     =	bert_model_tester.prepare_config_and_inputs()

        __lowerCamelCase   ,     __lowerCamelCase   ,     __lowerCamelCase:				Any     =	vision_config_and_inputs

        (
            (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     
        ):				Union[str, 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        a					(			_UpperCAmelCase     ,unittest.TestCase   ):





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

        return model, inputs





     def  SCREAMING_SNAKE_CASE__       (    self					:				Union[str, Any]    ,	SCREAMING_SNAKE_CASE_					:				str    ,	SCREAMING_SNAKE_CASE_					:				Tuple							):
        __lowerCamelCase:				List[Any]     =	TFCLIPVisionModel(SCREAMING_SNAKE_CASE_    ,	name="""vision_model"""							)
        __lowerCamelCase:				List[Any]     =	TFBertModel(SCREAMING_SNAKE_CASE_    ,	name="""text_model"""							)
        return vision_model, text_model







     def  SCREAMING_SNAKE_CASE__       (    self					:				List[Any]							):
        __lowerCamelCase:				List[str]     =	TFCLIPVisionModelTester(self							)
        __lowerCamelCase:				List[Any]     =	TFBertModelTester(self							)
        __lowerCamelCase:				Tuple     =	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
        )   ,     (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     (
            __lowerCamelCase
        )   ,     
        ):				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_vision
@require_tf
class        a					(			unittest.TestCase   ):





     @slow
     def  SCREAMING_SNAKE_CASE__       (    self					:				List[str]							):
        __lowerCamelCase:				Tuple     =	TFVisionTextDualEncoderModel.from_pretrained(
            """clip-italian/clip-italian"""    ,	logit_scale_init_value=1.0    ,	from_pt=SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Optional[int]     =	VisionTextDualEncoderProcessor.from_pretrained("""clip-italian/clip-italian"""							)

        __lowerCamelCase:				Any     =	Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""							)
        __lowerCamelCase:				Union[str, Any]     =	processor(
            text=["""una foto di un gatto""", """una foto di un cane"""]    ,	images=SCREAMING_SNAKE_CASE_    ,	padding=SCREAMING_SNAKE_CASE_    ,	return_tensors="""np"""							)

        __lowerCamelCase:				List[Any]     =	model(**SCREAMING_SNAKE_CASE_							)

        # 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:				Tuple     =	np.array([[1.228_4727, 0.310_4122]]							)

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






import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import (
    RobertaTokenizer,
    TrOCRConfig,
    TrOCRForCausalLM,
    TrOCRProcessor,
    VisionEncoderDecoderModel,
    ViTConfig,
    ViTImageProcessor,
    ViTModel,
)
from transformers.utils import logging


logging.set_verbosity_info()
_A		:				List[Any]									=       logging.get_logger(__name__)



def     __lowerCAmelCase						(							snake_case	:    Any						,      snake_case	:    int )       ->							Dict:
   __lowerCamelCase:				Tuple     =	[]
   for i in range(encoder_config.num_hidden_layers ):
      # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
      rename_keys.append(
          (f'encoder.deit.blocks.{i}.norm1.weight', f'encoder.encoder.layer.{i}.layernorm_before.weight') )
      rename_keys.append((f'encoder.deit.blocks.{i}.norm1.bias', f'encoder.encoder.layer.{i}.layernorm_before.bias') )
      rename_keys.append(
          (f'encoder.deit.blocks.{i}.attn.proj.weight', f'encoder.encoder.layer.{i}.attention.output.dense.weight') )
      rename_keys.append(
          (f'encoder.deit.blocks.{i}.attn.proj.bias', f'encoder.encoder.layer.{i}.attention.output.dense.bias') )
      rename_keys.append(
          (f'encoder.deit.blocks.{i}.norm2.weight', f'encoder.encoder.layer.{i}.layernorm_after.weight') )
      rename_keys.append((f'encoder.deit.blocks.{i}.norm2.bias', f'encoder.encoder.layer.{i}.layernorm_after.bias') )
      rename_keys.append(
          (f'encoder.deit.blocks.{i}.mlp.fc1.weight', f'encoder.encoder.layer.{i}.intermediate.dense.weight') )
      rename_keys.append(
          (f'encoder.deit.blocks.{i}.mlp.fc1.bias', f'encoder.encoder.layer.{i}.intermediate.dense.bias') )
      rename_keys.append(
          (f'encoder.deit.blocks.{i}.mlp.fc2.weight', f'encoder.encoder.layer.{i}.output.dense.weight') )
      rename_keys.append((f'encoder.deit.blocks.{i}.mlp.fc2.bias', f'encoder.encoder.layer.{i}.output.dense.bias') )

   # cls token, position embeddings and patch embeddings of encoder
   rename_keys.extend(
       [
           ("""encoder.deit.cls_token""", """encoder.embeddings.cls_token"""),
           ("""encoder.deit.pos_embed""", """encoder.embeddings.position_embeddings"""),
           ("""encoder.deit.patch_embed.proj.weight""", """encoder.embeddings.patch_embeddings.projection.weight"""),
           ("""encoder.deit.patch_embed.proj.bias""", """encoder.embeddings.patch_embeddings.projection.bias"""),
           ("""encoder.deit.norm.weight""", """encoder.layernorm.weight"""),
           ("""encoder.deit.norm.bias""", """encoder.layernorm.bias"""),
       ] )

   return rename_keys



def     __lowerCAmelCase						(							snake_case	:    str						,      snake_case	:    Union[str, Any] )       ->							Tuple:
   for i in range(encoder_config.num_hidden_layers ):
      # queries, keys and values (only weights, no biases)
      __lowerCamelCase:				Union[str, Any]     =	state_dict.pop(f'encoder.deit.blocks.{i}.attn.qkv.weight' )

      __lowerCamelCase:				Any     =	in_proj_weight[
          : encoder_config.hidden_size, :
      ]
      __lowerCamelCase:				Dict     =	in_proj_weight[
          encoder_config.hidden_size : encoder_config.hidden_size * 2, :
      ]
      __lowerCamelCase:				Any     =	in_proj_weight[
          -encoder_config.hidden_size :, :
      ]



def     __lowerCAmelCase						(							snake_case	:    Dict						,      snake_case	:    Union[str, Any]						,      snake_case	:    Optional[int] )       ->							int:
   __lowerCamelCase:				List[str]     =	dct.pop(snake_case )
   __lowerCamelCase:				Tuple     =	val



def     __lowerCAmelCase						(							snake_case	:    Dict )       ->							Any:
   if "handwritten" in checkpoint_url:
      __lowerCamelCase:				Tuple     =	"""https://fki.tic.heia-fr.ch/static/img/a01-122-02-00.jpg"""  # industry
      # url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-12.jpg" # have
      # url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-10.jpg" # let
      # url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg"  #
      # url = "https://fki.tic.heia-fr.ch/static/img/a01-122.jpg"
   elif "printed" in checkpoint_url or "stage1" in checkpoint_url:
      __lowerCamelCase:				Union[str, Any]     =	"""https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg"""
   __lowerCamelCase:				Union[str, Any]     =	Image.open(requests.get(snake_case						,      stream=snake_case ).raw ).convert("""RGB""" )
   return im



@torch.no_grad()
def     __lowerCAmelCase						(							snake_case	:    List[Any]						,      snake_case	:    Any )       ->							List[str]:
   __lowerCamelCase:				Union[str, Any]     =	ViTConfig(image_size=384						,      qkv_bias=snake_case )
   __lowerCamelCase:				Optional[int]     =	TrOCRConfig()

   # size of the architecture
   if "base" in checkpoint_url:
      __lowerCamelCase:				Optional[int]     =	768
   elif "large" in checkpoint_url:
      # use ViT-large encoder
      __lowerCamelCase:				Optional[int]     =	1024
      __lowerCamelCase:				int     =	4096
      __lowerCamelCase:				Tuple     =	24
      __lowerCamelCase:				Optional[Any]     =	16
      __lowerCamelCase:				str     =	1024
   else:
      raise ValueError("""Should either find 'base' or 'large' in checkpoint URL""" )

   # the large-printed + stage1 checkpoints uses sinusoidal position embeddings, no layernorm afterwards
   if "large-printed" in checkpoint_url or "stage1" in checkpoint_url:
      __lowerCamelCase:				Any     =	False
      __lowerCamelCase:				Optional[Any]     =	"""relu"""
      __lowerCamelCase:				List[Any]     =	1024
      __lowerCamelCase:				List[str]     =	True
      __lowerCamelCase:				int     =	False
      __lowerCamelCase:				int     =	False

   # load HuggingFace model
   __lowerCamelCase:				Any     =	ViTModel(snake_case						,      add_pooling_layer=snake_case )
   __lowerCamelCase:				Dict     =	TrOCRForCausalLM(snake_case )
   __lowerCamelCase:				Any     =	VisionEncoderDecoderModel(encoder=snake_case						,      decoder=snake_case )
   model.eval()

   # load state_dict of original model, rename some keys
   __lowerCamelCase:				Any     =	torch.hub.load_state_dict_from_url(snake_case						,      map_location="""cpu"""						,      check_hash=snake_case )["""model"""]

   __lowerCamelCase:				Optional[Any]     =	create_rename_keys(snake_case						,      snake_case )
   for src, dest in rename_keys:
      rename_key(snake_case						,      snake_case						,      snake_case )
   read_in_q_k_v(snake_case						,      snake_case )

   # remove parameters we don't need
   del state_dict["encoder.deit.head.weight"]
   del state_dict["encoder.deit.head.bias"]
   del state_dict["decoder.version"]

   # add prefix to decoder keys
   for key, val in state_dict.copy().items():
      __lowerCamelCase:				int     =	state_dict.pop(snake_case )
      if key.startswith("""decoder""" ) and "output_projection" not in key:
         __lowerCamelCase:				Any     =	val
      else:
         __lowerCamelCase:				List[str]     =	val

    # load state dict
   model.load_state_dict(snake_case )

   # Check outputs on an image
   __lowerCamelCase:				Any     =	ViTImageProcessor(size=encoder_config.image_size )
   __lowerCamelCase:				Tuple     =	RobertaTokenizer.from_pretrained("""roberta-large""" )
   __lowerCamelCase:				List[Any]     =	TrOCRProcessor(snake_case						,      snake_case )

   __lowerCamelCase:				Optional[int]     =	processor(images=prepare_img(snake_case )						,      return_tensors="""pt""" ).pixel_values

   # verify logits
   __lowerCamelCase:				Tuple     =	torch.tensor([[model.config.decoder.decoder_start_token_id]] )
   __lowerCamelCase:				Dict     =	model(pixel_values=snake_case						,      decoder_input_ids=snake_case )
   __lowerCamelCase:				Any     =	outputs.logits

   __lowerCamelCase:				Dict     =	torch.Size([1, 1, 50265] )
   if "trocr-base-handwritten" in checkpoint_url:
      __lowerCamelCase:				Tuple     =	torch.tensor(
          [-1.4_502, -4.6_683, -0.5_347, -2.9_291, 9.1_435, -3.0_571, 8.9_764, 1.7_560, 8.7_358, -1.5_311] )
   elif "trocr-large-handwritten" in checkpoint_url:
      __lowerCamelCase:				Optional[int]     =	torch.tensor(
          [-2.6_437, -1.3_129, -2.2_596, -5.3_455, 6.3_539, 1.7_604, 5.4_991, 1.4_702, 5.6_113, 2.0_170] )
   elif "trocr-base-printed" in checkpoint_url:
      __lowerCamelCase:				Union[str, Any]     =	torch.tensor(
          [-5.6_816, -5.8_388, 1.1_398, -6.9_034, 6.8_505, -2.4_393, 1.2_284, -1.0_232, -1.9_661, -3.9_210] )
   elif "trocr-large-printed" in checkpoint_url:
      __lowerCamelCase:				Union[str, Any]     =	torch.tensor(
          [-6.0_162, -7.0_959, 4.4_155, -5.1_063, 7.0_468, -3.1_631, 2.6_466, -0.3_081, -0.8_106, -1.7_535] )

   if "stage1" not in checkpoint_url:
      assert logits.shape == expected_shape, "Shape of logits not as expected"
      assert torch.allclose(logits[0, 0, :10]						,      snake_case						,      atol=1e-3 ), "First elements of logits not as expected"

   Path(snake_case ).mkdir(exist_ok=snake_case )
   print(f'Saving model to {pytorch_dump_folder_path}' )
   model.save_pretrained(snake_case )
   print(f'Saving processor to {pytorch_dump_folder_path}' )
   processor.save_pretrained(snake_case )


if __name__ == "__main__":
   _A		:				int									=       argparse.ArgumentParser()

   parser.add_argument(
       '''--checkpoint_url''',
       default='''https://layoutlm.blob.core.windows.net/trocr/model_zoo/fairseq/trocr-base-handwritten.pt''',
       type=str,
       help='''URL to the original PyTorch checkpoint (.pth file).''',
   )
   parser.add_argument(
       '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
   )
   _A		:				Optional[int]									=       parser.parse_args()
   convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)


