




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



@require_tf
class       a					:

					'''simple docstring'''





					def   lowerCamelCase_       (     self       :      Optional[int] ,   __snake_case       :      List[str] ,   __snake_case       :      Union[str, Any]    ):
									pass
					def   lowerCamelCase_       (     self       :      Dict    ):
									pass
					def   lowerCamelCase_       (     self       :      int    ):
									pass
					def   lowerCamelCase_       (     self       :      Any ,   __snake_case       :      Any ,   __snake_case       :      Tuple ,   __snake_case       :      Optional[Any] ,   __snake_case       :      Any ,   __snake_case       :      List[Any]=None ,   **__snake_case       :      Any    ):
									UpperCAmelCase_										=					VisionTextDualEncoderConfig.from_vision_text_configs(lowercase_ ,   lowercase_    )

									UpperCAmelCase_										=					TFVisionTextDualEncoderModel(lowercase_    )

									UpperCAmelCase_										=					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   lowerCamelCase_       (     self       :      Tuple ,   __snake_case       :      str ,   __snake_case       :      Union[str, Any] ,   __snake_case       :      str ,   __snake_case       :      List[Any] ,   __snake_case       :      List[Any]=None ,   **__snake_case       :      Tuple    ):
									UpperCAmelCase_										=					self.get_vision_text_model(lowercase_ ,   lowercase_    )
									UpperCAmelCase_										=					TFVisionTextDualEncoderModel(vision_model=lowercase_ ,   text_model=lowercase_    )

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

									UpperCAmelCase_										=					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   lowerCamelCase_       (     self       :      Dict ,   __snake_case       :      Any ,   __snake_case       :      List[str] ,   __snake_case       :      Any ,   __snake_case       :      Union[str, Any] ,   __snake_case       :      str=None ,   **__snake_case       :      Any    ):
									UpperCAmelCase_										=					self.get_vision_text_model(lowercase_ ,   lowercase_    )
									UpperCAmelCase_										=					TFVisionTextDualEncoderModel(vision_model=lowercase_ ,   text_model=lowercase_    )

									UpperCAmelCase_										=					model(input_ids=lowercase_ ,   pixel_values=lowercase_ ,   attention_mask=lowercase_    )
									UpperCAmelCase_										=					output[0].numpy()

									with tempfile.TemporaryDirectory() as tmpdirname:
													model.save_pretrained(lowercase_    )
													UpperCAmelCase_										=					TFVisionTextDualEncoderModel.from_pretrained(lowercase_    )

													UpperCAmelCase_										=					model(input_ids=lowercase_ ,   pixel_values=lowercase_ ,   attention_mask=lowercase_    )
													UpperCAmelCase_										=					after_output[0].numpy()
													UpperCAmelCase_										=					np.amax(np.abs(out_a - out_a    )    )
													self.assertLessEqual(lowercase_ ,   1E-5    )
					def   lowerCamelCase_       (     self       :      Tuple ,   __snake_case       :      Any ,   __snake_case       :      Tuple ,   __snake_case       :      Optional[int] ,   __snake_case       :      int ,   __snake_case       :      Any=None ,   **__snake_case       :      Union[str, Any]    ):
									UpperCAmelCase_										=					self.get_vision_text_model(lowercase_ ,   lowercase_    )
									UpperCAmelCase_										=					TFVisionTextDualEncoderModel(vision_model=lowercase_ ,   text_model=lowercase_    )

									UpperCAmelCase_										=					model(
									    input_ids=lowercase_ ,   pixel_values=lowercase_ ,   attention_mask=lowercase_ ,   output_attentions=lowercase_    )

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

									UpperCAmelCase_										=					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   lowerCamelCase_       (     self       :      Union[str, Any] ,   __snake_case       :      np.ndarray ,   __snake_case       :      np.ndarray ,   __snake_case       :      float    ):
									UpperCAmelCase_										=					np.abs((a - b)    ).max()
									self.assertLessEqual(lowercase_ ,   lowercase_ ,   F'Difference between torch and flax is {diff} (>= {tol}).'    )
					def   lowerCamelCase_       (     self       :      List[str]    ):
									UpperCAmelCase_										=					self.prepare_config_and_inputs()
									self.check_vision_text_dual_encoder_model(**lowercase_    )
					def   lowerCamelCase_       (     self       :      Dict    ):
									UpperCAmelCase_										=					self.prepare_config_and_inputs()
									self.check_model_from_pretrained_configs(**lowercase_    )
					def   lowerCamelCase_       (     self       :      Optional[int]    ):
									UpperCAmelCase_										=					self.prepare_config_and_inputs()
									self.check_vision_text_dual_encoder_from_pretrained(**lowercase_    )
					def   lowerCamelCase_       (     self       :      Union[str, Any]    ):
									UpperCAmelCase_										=					self.prepare_config_and_inputs()
									self.check_save_load(**lowercase_    )
					def   lowerCamelCase_       (     self       :      List[str]    ):
									UpperCAmelCase_										=					self.prepare_config_and_inputs()
									self.check_vision_text_output_attention(**lowercase_    )




					@slow
					def   lowerCamelCase_       (     self       :      int    ):
									UpperCAmelCase_										=					self.get_pretrained_model_and_inputs()

									UpperCAmelCase_										=					model_a(**lowercase_    )
									UpperCAmelCase_										=					outputs[0].numpy()

									with tempfile.TemporaryDirectory() as tmp_dirname:
													model_a.save_pretrained(lowercase_    )
													UpperCAmelCase_										=					TFVisionTextDualEncoderModel.from_pretrained(lowercase_    )

													UpperCAmelCase_										=					model_a(**lowercase_    )
													UpperCAmelCase_										=					after_outputs[0].numpy()
													UpperCAmelCase_										=					np.amax(np.abs(out_a - out_a    )    )
													self.assertLessEqual(lowercase_ ,   1E-5    )



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

					'''simple docstring'''





					def   lowerCamelCase_       (     self       :      Tuple    ):
									UpperCAmelCase_										=					TFVisionTextDualEncoderModel.from_vision_text_pretrained(
									    '''hf-internal-testing/tiny-random-vit''' ,   '''hf-internal-testing/tiny-random-bert'''    )
									UpperCAmelCase_										=					13
									UpperCAmelCase_										=					floats_tensor(
									    [
									        batch_size,
									        model.vision_model.config.num_channels,
									        model.vision_model.config.image_size,
									        model.vision_model.config.image_size,
									    ]    )
									UpperCAmelCase_										=					ids_tensor([batch_size, 4] ,   model.text_model.config.vocab_size    )
									UpperCAmelCase_										=					random_attention_mask([batch_size, 4]    )
									UpperCAmelCase_										=					{"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

									return model, inputs
					def   lowerCamelCase_       (     self       :      Any ,   __snake_case       :      Any ,   __snake_case       :      Union[str, Any]    ):
									UpperCAmelCase_										=					TFViTModel(lowercase_ ,   name='''vision_model'''    )
									UpperCAmelCase_										=					TFBertModel(lowercase_ ,   name='''text_model'''    )
									return vision_model, text_model




					def   lowerCamelCase_       (     self       :      str    ):
									UpperCAmelCase_										=					TFViTModelTester(self    )
									UpperCAmelCase_										=					TFBertModelTester(self    )
									UpperCAmelCase_										=					vit_model_tester.prepare_config_and_inputs()
									UpperCAmelCase_										=					bert_model_tester.prepare_config_and_inputs()

									UpperCAmelCase_										=					vision_config_and_inputs

									(
									    UpperCAmelCase_						
									)										=					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      ):

					'''simple docstring'''





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

									return model, inputs
					def   lowerCamelCase_       (     self       :      int ,   __snake_case       :      List[str] ,   __snake_case       :      Dict ,   __snake_case       :      List[str] ,   __snake_case       :      Optional[int] ,   __snake_case       :      Optional[Any]=None ,   **__snake_case       :      Union[str, Any]    ):
									UpperCAmelCase_										=					self.get_vision_text_model(lowercase_ ,   lowercase_    )
									UpperCAmelCase_										=					TFVisionTextDualEncoderModel(vision_model=lowercase_ ,   text_model=lowercase_    )

									UpperCAmelCase_										=					model(
									    input_ids=lowercase_ ,   pixel_values=lowercase_ ,   attention_mask=lowercase_ ,   output_attentions=lowercase_    )

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

									UpperCAmelCase_										=					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   lowerCamelCase_       (     self       :      List[str] ,   __snake_case       :      List[str] ,   __snake_case       :      Dict    ):
									UpperCAmelCase_										=					TFDeiTModel(lowercase_ ,   name='''vision_model'''    )
									UpperCAmelCase_										=					TFRobertaModel(lowercase_ ,   name='''text_model'''    )
									return vision_model, text_model




					def   lowerCamelCase_       (     self       :      Optional[Any]    ):
									UpperCAmelCase_										=					TFDeiTModelTester(self    )
									UpperCAmelCase_										=					TFRobertaModelTester(self    )
									UpperCAmelCase_										=					vit_model_tester.prepare_config_and_inputs()
									UpperCAmelCase_										=					bert_model_tester.prepare_config_and_inputs()

									UpperCAmelCase_										=					vision_config_and_inputs

									(
									    UpperCAmelCase_						
									)										=					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      ):

					'''simple docstring'''





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

									return model, inputs
					def   lowerCamelCase_       (     self       :      Dict ,   __snake_case       :      str ,   __snake_case       :      List[Any]    ):
									UpperCAmelCase_										=					TFCLIPVisionModel(lowercase_ ,   name='''vision_model'''    )
									UpperCAmelCase_										=					TFBertModel(lowercase_ ,   name='''text_model'''    )
									return vision_model, text_model




					def   lowerCamelCase_       (     self       :      int    ):
									UpperCAmelCase_										=					TFCLIPVisionModelTester(self    )
									UpperCAmelCase_										=					TFBertModelTester(self    )
									UpperCAmelCase_										=					clip_model_tester.prepare_config_and_inputs()
									UpperCAmelCase_										=					bert_model_tester.prepare_config_and_inputs()

									UpperCAmelCase_										=					vision_config_and_inputs

									(
									    UpperCAmelCase_						
									)										=					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      ):

					'''simple docstring'''





					@slow
					def   lowerCamelCase_       (     self       :      Tuple    ):
									UpperCAmelCase_										=					TFVisionTextDualEncoderModel.from_pretrained(
									    '''clip-italian/clip-italian''' ,   logit_scale_init_value=1.0 ,   from_pt=lowercase_    )
									UpperCAmelCase_										=					VisionTextDualEncoderProcessor.from_pretrained('''clip-italian/clip-italian'''    )

									UpperCAmelCase_										=					Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''    )
									UpperCAmelCase_										=					processor(
									    text=['''una foto di un gatto''', '''una foto di un cane'''] ,   images=lowercase_ ,   padding=lowercase_ ,   return_tensors='''np'''    )

									UpperCAmelCase_										=					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]) ,   )

									UpperCAmelCase_										=					np.array([[1.2_284_727, 0.3_104_122]]    )

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











import argparse
import logging
import pickle
from collections import Counter


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO
)
_lowerCamelCase      			=						logging.getLogger(__name__)

if __name__ == "__main__":
 _lowerCamelCase      			=						argparse.ArgumentParser(
     description='Token Counts for smoothing the masking probabilities in MLM (cf XLM/word2vec)'
 )
 parser.add_argument(
     '--data_file', type=str, default='data/dump.bert-base-uncased.pickle', help='The binarized dataset.'
 )
 parser.add_argument(
     '--token_counts_dump', type=str, default='data/token_counts.bert-base-uncased.pickle', help='The dump file.'
 )
 parser.add_argument('--vocab_size', default=3_05_22, type=int)
 _lowerCamelCase      			=						parser.parse_args()

 logger.info(F"Loading data from {args.data_file}")
 with open(args.data_file, 'rb') as fp:
  _lowerCamelCase      			=						pickle.load(fp)

 logger.info('Counting occurrences for MLM.')
 _lowerCamelCase      			=						Counter()
 for tk_ids in data:
  counter.update(tk_ids)
 _lowerCamelCase      			=						[0] * args.vocab_size
 for k, v in counter.items():
  _lowerCamelCase      			=						v

 logger.info(F"Dump to {args.token_counts_dump}")
 with open(args.token_counts_dump, 'wb') as handle:
  pickle.dump(counts, handle, protocol=pickle.HIGHEST_PROTOCOL)





