

'''simple docstring'''



import argparse
import json
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils.deepspeed import DummyOptim, DummyScheduler


_UpperCamelCase						:  List[str]							       =				16
_UpperCamelCase						:  int							       =				32


def       snake_case   (	snake_case      : List[str]					,		snake_case      : List[Any] = 16					,		snake_case      : Union[str, Any] = "bert-base-cased"			) ->  Any:


    """simple docstring"""


    lowerCAmelCase 						=							AutoTokenizer.from_pretrained(A__			)
    lowerCAmelCase 						=							load_dataset('glue'					,		'mrpc'			)

    def tokenize_function(snake_case      : int			):
        # max_length=None => use the model max length (it's actually the default)
        lowerCAmelCase 						=							tokenizer(examples['sentence1']					,		examples['sentence2']					,		truncation=A__					,		max_length=A__			)
        return outputs

    # Apply the method we just defined to all the examples in all the splits of the dataset
    lowerCAmelCase 						=							datasets.map(
        A__					,		batched=A__					,		remove_columns=['idx', 'sentence1', 'sentence2']					,		load_from_cache_file=A__			)

    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
    # transformers library
    lowerCAmelCase 						=							tokenized_datasets.rename_column('label'					,		'labels'			)

    def collate_fn(snake_case      : Union[str, Any]			):
        # On TPU it's best to pad everything to the same length or training will be very slow.
        if accelerator.distributed_type == DistributedType.TPU:
            return tokenizer.pad(A__					,		padding='max_length'					,		max_length=128					,		return_tensors='pt'			)
        return tokenizer.pad(A__					,		padding='longest'					,		return_tensors='pt'			)

    # Instantiate dataloaders.
    lowerCAmelCase 						=							DataLoader(
        tokenized_datasets['train']					,		shuffle=A__					,		collate_fn=A__					,		batch_size=A__			)
    lowerCAmelCase 						=							DataLoader(
        tokenized_datasets['validation']					,		shuffle=A__					,		collate_fn=A__					,		batch_size=A__			)

    return train_dataloader, eval_dataloader


def       snake_case   (	snake_case      : int					,		snake_case      : str			) ->  int:


    """simple docstring"""


    lowerCAmelCase 						=							Accelerator()

    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
    lowerCAmelCase 						=							config['lr']
    lowerCAmelCase 						=							int(config['num_epochs']			)
    lowerCAmelCase 						=							int(config['seed']			)
    lowerCAmelCase 						=							int(config['batch_size']			)
    lowerCAmelCase 						=							args.model_name_or_path

    set_seed(A__			)
    lowerCAmelCase 						=							get_dataloaders(A__					,		A__					,		A__			)

    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
    lowerCAmelCase 						=							AutoModelForSequenceClassification.from_pretrained(A__					,		return_dict=A__			)

    # Instantiate optimizer
    lowerCAmelCase 						=							(
        AdamW
        if accelerator.state.deepspeed_plugin is None
        or 'optimizer' not in accelerator.state.deepspeed_plugin.deepspeed_config
        else DummyOptim
    )
    lowerCAmelCase 						=							optimizer_cls(params=model.parameters()					,		lr=A__			)

    if accelerator.state.deepspeed_plugin is not None:
        lowerCAmelCase 						=							accelerator.state.deepspeed_plugin.deepspeed_config[
            'gradient_accumulation_steps'
        ]
    else:
        lowerCAmelCase 						=							1
    lowerCAmelCase 						=							(len(A__			) * num_epochs) // gradient_accumulation_steps

    # Instantiate scheduler
    if (
        accelerator.state.deepspeed_plugin is None
        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
    ):
        lowerCAmelCase 						=							get_linear_schedule_with_warmup(
            optimizer=A__					,		num_warmup_steps=0					,		num_training_steps=A__					,		)
    else:
        lowerCAmelCase 						=							DummyScheduler(A__					,		total_num_steps=A__					,		warmup_num_steps=0			)

    # Prepare everything
    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
    # prepare method.
    lowerCAmelCase 						=							accelerator.prepare(
        A__					,		A__					,		A__					,		A__					,		A__			)

    # We need to keep track of how many total steps we have iterated over
    lowerCAmelCase 						=							0
    # We also need to keep track of the stating epoch so files are named properly
    lowerCAmelCase 						=							0

    # Now we train the model
    lowerCAmelCase 						=							evaluate.load('glue'					,		'mrpc'			)
    lowerCAmelCase 						=							0
    lowerCAmelCase 						=							{}
    for epoch in range(A__					,		A__			):
        model.train()
        for step, batch in enumerate(A__			):
            lowerCAmelCase 						=							model(**A__			)
            lowerCAmelCase 						=							outputs.loss
            lowerCAmelCase 						=							loss / gradient_accumulation_steps
            accelerator.backward(A__			)
            if step % gradient_accumulation_steps == 0:
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            overall_step += 1

        model.eval()
        lowerCAmelCase 						=							0
        for step, batch in enumerate(A__			):
            # We could avoid this line since we set the accelerator with `device_placement=True`.
            batch.to(accelerator.device			)
            with torch.no_grad():
                lowerCAmelCase 						=							model(**A__			)
            lowerCAmelCase 						=							outputs.logits.argmax(dim=-1			)
            # It is slightly faster to call this once, than multiple times
            lowerCAmelCase 						=							accelerator.gather(
                (predictions, batch['labels'])			)  # If we are in a multiprocess environment, the last batch has duplicates
            if accelerator.use_distributed:
                if step == len(A__			) - 1:
                    lowerCAmelCase 						=							predictions[: len(eval_dataloader.dataset			) - samples_seen]
                    lowerCAmelCase 						=							references[: len(eval_dataloader.dataset			) - samples_seen]
                else:
                    samples_seen += references.shape[0]
            metric.add_batch(
                predictions=A__					,		references=A__					,		)

        lowerCAmelCase 						=							metric.compute()
        # Use accelerator.print to print only on the main process.
        accelerator.print(F'epoch {epoch}:'					,		A__			)
        lowerCAmelCase 						=							eval_metric['accuracy']

        if best_performance < eval_metric["accuracy"]:
            lowerCAmelCase 						=							eval_metric['accuracy']

    if args.performance_lower_bound is not None:
        assert (
            args.performance_lower_bound <= best_performance
        ), F'Best performance metric {best_performance} is lower than the lower bound {args.performance_lower_bound}'

    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        with open(os.path.join(args.output_dir					,		'all_results.json'			)					,		'w'			) as f:
            json.dump(A__					,		A__			)


def       snake_case   (	) ->  Tuple:


    """simple docstring"""


    lowerCAmelCase 						=							argparse.ArgumentParser(description='Simple example of training script tracking peak GPU memory usage.'			)
    parser.add_argument(
        '--model_name_or_path'					,		type=A__					,		default='bert-base-cased'					,		help='Path to pretrained model or model identifier from huggingface.co/models.'					,		required=A__					,		)
    parser.add_argument(
        '--output_dir'					,		type=A__					,		default='.'					,		help='Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'					,		)
    parser.add_argument(
        '--performance_lower_bound'					,		type=A__					,		default=A__					,		help='Optional lower bound for the performance metric. If set, the training will throw error when the performance metric drops below this value.'					,		)
    parser.add_argument(
        '--num_epochs'					,		type=A__					,		default=3					,		help='Number of train epochs.'					,		)
    lowerCAmelCase 						=							parser.parse_args()
    lowerCAmelCase 						=							{'lr': 2e-5, 'num_epochs': args.num_epochs, 'seed': 42, 'batch_size': 16}
    training_function(A__					,		A__			)


if __name__ == "__main__":
     main()









import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_flax_cross_test,
    require_flax,
    require_torch,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import is_flax_available, is_torch_available, is_vision_available

from ...test_modeling_flax_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_flax_bert import FlaxBertModelTester
from ..clip.test_modeling_flax_clip import FlaxCLIPVisionModelTester
from ..vit.test_modeling_flax_vit import FlaxViTModelTester


if is_flax_available():
				from transformers import (
				    FlaxBertModel,
				    FlaxCLIPVisionModel,
				    FlaxVisionTextDualEncoderModel,
				    FlaxViTModel,
				    VisionTextDualEncoderConfig,
				    VisionTextDualEncoderProcessor,
				)
				from transformers.modeling_flax_pytorch_utils import (
				    convert_pytorch_state_dict_to_flax,
				    load_flax_weights_in_pytorch_model,
				)


if is_torch_available():
				import torch

				from transformers import VisionTextDualEncoderModel

if is_vision_available():
				from PIL import Image






def 						a__			(				A__				):
			if isinstance(A__,      collections.abc.Iterable				):
						return x
			return (x, x)







@require_flax
class 	__lowercase    :





							"""simple docstring"""







							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										pass






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										pass






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										pass






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					np.abs((a - b)       ).max()
										self.assertLessEqual(lowerCAmelCase__			,						lowerCAmelCase__			,						F'''Difference between torch and flax is {diff} (>= {tol}).'''       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__=None			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Tuple					=					VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase__			,						lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     List[str]					=					FlaxVisionTextDualEncoderModel(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Tuple					=					model(input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )

										self.assertEqual(output['text_embeds'].shape			,						(input_ids.shape[0], config.projection_dim)       )
										self.assertEqual(output['image_embeds'].shape			,						(pixel_values.shape[0], config.projection_dim)       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__=None			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.get_vision_text_model(lowerCAmelCase__			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     str					=					{'vision_model': vision_model, 'text_model': text_model}
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCAmelCase__       )

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

										self.assertEqual(output['text_embeds'].shape			,						(input_ids.shape[0], model.config.projection_dim)       )
										self.assertEqual(output['image_embeds'].shape			,						(pixel_values.shape[0], model.config.projection_dim)       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__=None			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     int					=					self.get_vision_text_model(lowerCAmelCase__			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Dict					=					{'vision_model': vision_model, 'text_model': text_model}
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					model(input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					output[0]

										with tempfile.TemporaryDirectory() as tmpdirname:
													model.save_pretrained(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					FlaxVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase__       )

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






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__=None			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.get_vision_text_model(lowerCAmelCase__			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     List[str]					=					{'vision_model': vision_model, 'text_model': text_model}
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     List[str]					=					model(
										    input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__			,						output_attentions=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     List[str]					=					output.vision_model_output.attentions
										self.assertEqual(len(lowerCAmelCase__       )			,						vision_config.num_hidden_layers       )

										# in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
										SCREAMING_SNAKE_CASE_		:     str					=					to_atuple(vision_model.config.image_size       )
										SCREAMING_SNAKE_CASE_		:     Any					=					to_atuple(vision_model.config.patch_size       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
										SCREAMING_SNAKE_CASE_		:     str					=					num_patches + 1
										self.assertEqual(vision_attentions[0].shape[-3:]			,						(vision_config.num_attention_heads, seq_len, seq_len)       )

										SCREAMING_SNAKE_CASE_		:     Tuple					=					output.text_model_output.attentions
										self.assertEqual(len(lowerCAmelCase__       )			,						text_config.num_hidden_layers       )

										self.assertEqual(
										    text_attentions[0].shape[-3:]			,						(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])			,						)






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										pt_model.to(lowerCAmelCase__       )
										pt_model.eval()

										# prepare inputs
										SCREAMING_SNAKE_CASE_		:     str					=					inputs_dict
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					{k: torch.tensor(v.tolist()       ) for k, v in flax_inputs.items()}

										with torch.no_grad():
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					pt_model(**lowerCAmelCase__       ).to_tuple()

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

										# PT -> Flax
										with tempfile.TemporaryDirectory() as tmpdirname:
													pt_model.save_pretrained(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					FlaxVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase__			,						from_pt=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Dict					=					fx_model_loaded(**lowerCAmelCase__       ).to_tuple()
										self.assertEqual(len(lowerCAmelCase__       )			,						len(lowerCAmelCase__       )			,						'Output lengths differ between Flax and PyTorch'       )
										for fx_output_loaded, pt_output in zip(fx_outputs_loaded[:4]			,						pt_outputs[:4]       ):
													self.assert_almost_equals(lowerCAmelCase__			,						pt_output.numpy()			,						4E-2       )

										# Flax -> PT
										with tempfile.TemporaryDirectory() as tmpdirname:
													fx_model.save_pretrained(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					VisionTextDualEncoderModel.from_pretrained(lowerCAmelCase__			,						from_flax=lowerCAmelCase__       )

										pt_model_loaded.to(lowerCAmelCase__       )
										pt_model_loaded.eval()

										with torch.no_grad():
													SCREAMING_SNAKE_CASE_		:     Optional[int]					=					pt_model_loaded(**lowerCAmelCase__       ).to_tuple()

										self.assertEqual(len(lowerCAmelCase__       )			,						len(lowerCAmelCase__       )			,						'Output lengths differ between Flax and PyTorch'       )
										for fx_output, pt_output_loaded in zip(fx_outputs[:4]			,						pt_outputs_loaded[:4]       ):
													self.assert_almost_equals(lowerCAmelCase__			,						pt_output_loaded.numpy()			,						4E-2       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[Any]					=					VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase__			,						lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					VisionTextDualEncoderModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					FlaxVisionTextDualEncoderModel(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     str					=					convert_pytorch_state_dict_to_flax(pt_model.state_dict()			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					fx_state

										self.check_pt_flax_equivalence(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase__			,						lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					VisionTextDualEncoderModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     str					=					FlaxVisionTextDualEncoderModel(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     int					=					load_flax_weights_in_pytorch_model(lowerCAmelCase__			,						fx_model.params       )

										self.check_pt_flax_equivalence(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     str					=					self.prepare_config_and_inputs()
										self.check_model_from_pretrained_configs(**lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					self.prepare_config_and_inputs()
										self.check_vision_text_dual_encoder_from_pretrained(**lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Tuple					=					self.prepare_config_and_inputs()
										self.check_save_load(**lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					self.prepare_config_and_inputs()
										self.check_vision_text_output_attention(**lowerCAmelCase__       )






							@is_pt_flax_cross_test
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					self.prepare_config_and_inputs()
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					config_inputs_dict.pop('vision_config'       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					config_inputs_dict.pop('text_config'       )

										SCREAMING_SNAKE_CASE_		:     Any					=					config_inputs_dict

										self.check_equivalence_pt_to_flax(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )
										self.check_equivalence_flax_to_pt(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )






							@slow
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Dict					=					self.get_pretrained_model_and_inputs()

										SCREAMING_SNAKE_CASE_		:     int					=					model_a(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					outputs[0]

										with tempfile.TemporaryDirectory() as tmp_dirname:
													model_a.save_pretrained(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     int					=					FlaxVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase__       )

													SCREAMING_SNAKE_CASE_		:     str					=					model_a(**lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     Any					=					after_outputs[0]
													SCREAMING_SNAKE_CASE_		:     str					=					np.amax(np.abs(out_a - out_a       )       )
													self.assertLessEqual(lowerCAmelCase__			,						1E-5       )







@require_flax
class 	__lowercase    (__SCREAMING_SNAKE_CASE	,       unittest.TestCase  ):





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
										    'hf-internal-testing/tiny-random-vit'			,						'hf-internal-testing/tiny-bert'			,						vision_from_pt=lowerCAmelCase__			,						text_from_pt=lowerCAmelCase__			,						)
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					1_3
										SCREAMING_SNAKE_CASE_		:     str					=					floats_tensor(
										    [
										        batch_size,
										        model.config.vision_config.num_channels,
										        model.config.vision_config.image_size,
										        model.config.vision_config.image_size,
										    ]       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					ids_tensor([batch_size, 4]			,						model.config.text_config.vocab_size       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					random_attention_mask([batch_size, 4]       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					{'pixel_values': pixel_values, 'input_ids': input_ids, 'attention_mask': attention_mask}

										return model, inputs






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     int					=					FlaxViTModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					FlaxBertModel(lowerCAmelCase__       )
										return vision_model, text_model






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     int					=					FlaxViTModelTester(self       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					FlaxBertModelTester(self       )
										SCREAMING_SNAKE_CASE_		:     str					=					vit_model_tester.prepare_config_and_inputs()
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					bert_model_tester.prepare_config_and_inputs()

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Any					=					vision_config_and_inputs

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     List[Any]					=					text_config_and_inputs

										# make sure that cross attention layers are added
										return {
										    "text_config": text_config,
										    "vision_config": vision_config,
										    "pixel_values": pixel_values,
										    "attention_mask": attention_mask,
										    "input_ids": input_ids,
										    "token_type_ids": token_type_ids,
										}







@require_torch
class 	__lowercase    (__SCREAMING_SNAKE_CASE	,       unittest.TestCase  ):





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     str					=					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
										    'hf-internal-testing/tiny-random-clip'			,						'hf-internal-testing/tiny-bert'			,						vision_from_pt=lowerCAmelCase__			,						text_from_pt=lowerCAmelCase__			,						)
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					1_3
										SCREAMING_SNAKE_CASE_		:     Dict					=					floats_tensor(
										    [
										        batch_size,
										        model.config.vision_config.num_channels,
										        model.config.vision_config.image_size,
										        model.config.vision_config.image_size,
										    ]       )
										SCREAMING_SNAKE_CASE_		:     int					=					ids_tensor([batch_size, 4]			,						model.config.text_config.vocab_size       )
										SCREAMING_SNAKE_CASE_		:     Dict					=					random_attention_mask([batch_size, 4]       )
										SCREAMING_SNAKE_CASE_		:     int					=					{'pixel_values': pixel_values, 'input_ids': input_ids, 'attention_mask': attention_mask}

										return model, inputs






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					FlaxCLIPVisionModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Any					=					FlaxBertModel(lowerCAmelCase__       )
										return vision_model, text_model






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					FlaxCLIPVisionModelTester(self       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					FlaxBertModelTester(self       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					clip_model_tester.prepare_config_and_inputs()
										SCREAMING_SNAKE_CASE_		:     int					=					bert_model_tester.prepare_config_and_inputs()

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     int					=					vision_config_and_inputs

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     str					=					text_config_and_inputs

										# make sure that cross attention layers are added
										return {
										    "text_config": text_config,
										    "vision_config": vision_config,
										    "pixel_values": pixel_values,
										    "attention_mask": attention_mask,
										    "input_ids": input_ids,
										    "token_type_ids": token_type_ids,
										}







@require_flax
@require_vision
class 	__lowercase    (unittest.TestCase  ):





							"""simple docstring"""







							@slow
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[str]					=					FlaxVisionTextDualEncoderModel.from_pretrained('clip-italian/clip-italian'			,						logit_scale_init_value=1.0       )
										SCREAMING_SNAKE_CASE_		:     str					=					VisionTextDualEncoderProcessor.from_pretrained('clip-italian/clip-italian'       )

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

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					model(**lowerCAmelCase__       )

										# verify the logits
										self.assertEqual(outputs.logits_per_image.shape			,						(inputs.pixel_values.shape[0], inputs.input_ids.shape[0])       )
										self.assertEqual(
										    outputs.logits_per_text.shape			,						(inputs.input_ids.shape[0], inputs.pixel_values.shape[0])			,						)

										SCREAMING_SNAKE_CASE_		:     Tuple					=					np.array([[1.2_284_727, 0.3_104_122]]       )

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





