import json
import os
from dataclasses import dataclass
from functools import partial
from typing import Callable

import flax.linen as nn
import jax
import jax.numpy as jnp
import joblib
import optax
import wandb
from flax import jax_utils, struct, traverse_util
from flax.serialization import from_bytes, to_bytes
from flax.training import train_state
from flax.training.common_utils import shard
from tqdm.auto import tqdm

from transformers import BigBirdConfig, FlaxBigBirdForQuestionAnswering
from transformers.models.big_bird.modeling_flax_big_bird import FlaxBigBirdForQuestionAnsweringModule



class     _A   (    __UpperCamelCase							):
		SCREAMING_SNAKE_CASE_	:	Optional[Any]						=42
		SCREAMING_SNAKE_CASE_	:	Any						=jnp.floataa
		SCREAMING_SNAKE_CASE_	:	Union[str, Any]						=True






		def 		_a					(self	)		->      Optional[int]:





							'''simple docstring'''





							super().setup()
							UpperCamelCase__										=      nn.Dense(5   ,				dtype=self.dtype	)






		def __call__(self   ,				*SCREAMING_SNAKE_CASE_   ,				**SCREAMING_SNAKE_CASE_	)		->      Any:





							'''simple docstring'''





							UpperCamelCase__										=      super().__call__(*SCREAMING_SNAKE_CASE_   ,				**SCREAMING_SNAKE_CASE_	)
							UpperCamelCase__										=      self.cls(outputs[2]	)
							return outputs[:2] + (cls_out,)



class     _A   (    __UpperCamelCase							):
		SCREAMING_SNAKE_CASE_	:	str						=FlaxBigBirdForNaturalQuestionsModule







def  __UpperCamelCase    (						A ,							A ,							A ,							A ,							A ,							A							):






					def cross_entropy(A ,							A ,							A=None							):
										UpperCamelCase__										=      logits.shape[-1]
										UpperCamelCase__										=      (labels[..., None] == jnp.arange(_snake_case							)[None]).astype('''f4'''							)
										UpperCamelCase__										=      jax.nn.log_softmax(_snake_case ,							axis=-1							)
										UpperCamelCase__										=      -jnp.sum(labels * logits ,							axis=-1							)
										if reduction is not None:
															UpperCamelCase__										=      reduction(_snake_case							)
										return loss

					UpperCamelCase__										=      partial(_snake_case ,							reduction=jnp.mean							)
					UpperCamelCase__										=      cross_entropy(_snake_case ,							_snake_case							)
					UpperCamelCase__										=      cross_entropy(_snake_case ,							_snake_case							)
					UpperCamelCase__										=      cross_entropy(_snake_case ,							_snake_case							)
					return (start_loss + end_loss + pooled_loss) / 3



@dataclass
class     _A   :
		SCREAMING_SNAKE_CASE_	:	List[str]						="google/bigbird-roberta-base"
		SCREAMING_SNAKE_CASE_	:	Dict						=30_00
		SCREAMING_SNAKE_CASE_	:	Any						=1_05_00

		SCREAMING_SNAKE_CASE_	:	Dict						=1_28
		SCREAMING_SNAKE_CASE_	:	Tuple						=3

		SCREAMING_SNAKE_CASE_	:	Union[str, Any]						=1
		SCREAMING_SNAKE_CASE_	:	Tuple						=5

		# tx_args
		SCREAMING_SNAKE_CASE_	:	str						=3e-5
		SCREAMING_SNAKE_CASE_	:	Union[str, Any]						=0.0
		SCREAMING_SNAKE_CASE_	:	Dict						=2_00_00
		SCREAMING_SNAKE_CASE_	:	Dict						=0.0_0_9_5

		SCREAMING_SNAKE_CASE_	:	Any						="bigbird-roberta-natural-questions"
		SCREAMING_SNAKE_CASE_	:	Any						="training-expt"
		SCREAMING_SNAKE_CASE_	:	List[Any]						="data/nq-training.jsonl"
		SCREAMING_SNAKE_CASE_	:	str						="data/nq-validation.jsonl"






		def 		_a					(self	)		->      int:





							'''simple docstring'''





							os.makedirs(self.base_dir   ,				exist_ok=SCREAMING_SNAKE_CASE_	)
							UpperCamelCase__										=      os.path.join(self.base_dir   ,				self.save_dir	)
							UpperCamelCase__										=      self.batch_size_per_device * jax.device_count()



@dataclass
class     _A   :
		SCREAMING_SNAKE_CASE_	:	List[str]						=42
		SCREAMING_SNAKE_CASE_	:	str						=40_96  # no dynamic padding on TPUs






		def __call__(self   ,				SCREAMING_SNAKE_CASE_	)		->      Any:





							'''simple docstring'''





							UpperCamelCase__										=      self.collate_fn(SCREAMING_SNAKE_CASE_	)
							UpperCamelCase__										=      jax.tree_util.tree_map(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)
							return batch






		def 		_a					(self   ,				SCREAMING_SNAKE_CASE_	)		->      Optional[int]:





							'''simple docstring'''





							UpperCamelCase__						,       UpperCamelCase__										=      self.fetch_inputs(features['''input_ids''']	)
							UpperCamelCase__										=      {
							    '''input_ids''': jnp.array(SCREAMING_SNAKE_CASE_   ,				dtype=jnp.intaa	),
							    '''attention_mask''': jnp.array(SCREAMING_SNAKE_CASE_   ,				dtype=jnp.intaa	),
							    '''start_labels''': jnp.array(features['''start_token''']   ,				dtype=jnp.intaa	),
							    '''end_labels''': jnp.array(features['''end_token''']   ,				dtype=jnp.intaa	),
							    '''pooled_labels''': jnp.array(features['''category''']   ,				dtype=jnp.intaa	),
							}
							return batch






		def 		_a					(self   ,				SCREAMING_SNAKE_CASE_	)		->      Optional[Any]:





							'''simple docstring'''





							UpperCamelCase__										=      [self._fetch_inputs(SCREAMING_SNAKE_CASE_	) for ids in input_ids]
							return zip(*SCREAMING_SNAKE_CASE_	)






		def 		_a					(self   ,				SCREAMING_SNAKE_CASE_	)		->      Tuple:





							'''simple docstring'''





							UpperCamelCase__										=      [1 for _ in range(len(SCREAMING_SNAKE_CASE_	)	)]
							while len(SCREAMING_SNAKE_CASE_	) < self.max_length:
												input_ids.append(self.pad_id	)
												attention_mask.append(0	)
							return input_ids, attention_mask







def  __UpperCamelCase    (						A ,							A ,							A=None							):






					if seed is not None:
										UpperCamelCase__										=      dataset.shuffle(seed=_snake_case							)
					for i in range(len(_snake_case							) // batch_size							):
										UpperCamelCase__										=      dataset[i * batch_size : (i + 1) * batch_size]
										yield dict(_snake_case							)







@partial(jax.pmap ,							axis_name='''batch'''							)
def  __UpperCamelCase    (						A ,							A ,							**A							):






					def loss_fn(A							):
										UpperCamelCase__										=      model_inputs.pop('''start_labels'''							)
										UpperCamelCase__										=      model_inputs.pop('''end_labels'''							)
										UpperCamelCase__										=      model_inputs.pop('''pooled_labels'''							)

										UpperCamelCase__										=      state.apply_fn(**_snake_case ,							params=_snake_case ,							dropout_rng=_snake_case ,							train=_snake_case							)
										UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      outputs

										return state.loss_fn(
										    _snake_case ,							_snake_case ,							_snake_case ,							_snake_case ,							_snake_case ,							_snake_case ,							)

					UpperCamelCase__						,       UpperCamelCase__										=      jax.random.split(_snake_case							)
					UpperCamelCase__										=      jax.value_and_grad(_snake_case							)
					UpperCamelCase__						,       UpperCamelCase__										=      grad_fn(state.params							)
					UpperCamelCase__										=      jax.lax.pmean({'''loss''': loss} ,							axis_name='''batch'''							)
					UpperCamelCase__										=      jax.lax.pmean(_snake_case ,							'''batch'''							)

					UpperCamelCase__										=      state.apply_gradients(grads=_snake_case							)
					return state, metrics, new_drp_rng







@partial(jax.pmap ,							axis_name='''batch'''							)
def  __UpperCamelCase    (						A ,							**A							):
					UpperCamelCase__										=      model_inputs.pop('''start_labels'''							)
					UpperCamelCase__										=      model_inputs.pop('''end_labels'''							)
					UpperCamelCase__										=      model_inputs.pop('''pooled_labels'''							)

					UpperCamelCase__										=      state.apply_fn(**_snake_case ,							params=state.params ,							train=_snake_case							)
					UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      outputs

					UpperCamelCase__										=      state.loss_fn(_snake_case ,							_snake_case ,							_snake_case ,							_snake_case ,							_snake_case ,							_snake_case							)
					UpperCamelCase__										=      jax.lax.pmean({'''loss''': loss} ,							axis_name='''batch'''							)
					return metrics



class     _A   (    train_state.TrainState							):
		SCREAMING_SNAKE_CASE_	:	Tuple						=struct.field(pytree_node=__UpperCamelCase							)



@dataclass
class     _A   :
		SCREAMING_SNAKE_CASE_	:	Tuple						=42
		SCREAMING_SNAKE_CASE_	:	str						=42
		SCREAMING_SNAKE_CASE_	:	int						=42
		SCREAMING_SNAKE_CASE_	:	Any						=42
		SCREAMING_SNAKE_CASE_	:	int						=42
		SCREAMING_SNAKE_CASE_	:	Optional[int]						=42
		SCREAMING_SNAKE_CASE_	:	Union[str, Any]						=None






		def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=None	)		->      Dict:





							'''simple docstring'''





							UpperCamelCase__										=      model.params
							UpperCamelCase__										=      TrainState.create(
							    apply_fn=model.__call__   ,				params=SCREAMING_SNAKE_CASE_   ,				tx=SCREAMING_SNAKE_CASE_   ,				loss_fn=SCREAMING_SNAKE_CASE_   ,				)
							if ckpt_dir is not None:
												UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      restore_checkpoint(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)
												UpperCamelCase__										=      {
												    '''lr''': args.lr,
												    '''init_lr''': args.init_lr,
												    '''warmup_steps''': args.warmup_steps,
												    '''num_train_steps''': num_train_steps,
												    '''weight_decay''': args.weight_decay,
												}
												UpperCamelCase__						,       UpperCamelCase__										=      build_tx(**SCREAMING_SNAKE_CASE_	)
												UpperCamelCase__										=      train_state.TrainState(
												    step=SCREAMING_SNAKE_CASE_   ,				apply_fn=model.__call__   ,				params=SCREAMING_SNAKE_CASE_   ,				tx=SCREAMING_SNAKE_CASE_   ,				opt_state=SCREAMING_SNAKE_CASE_   ,				)
												UpperCamelCase__										=      args
												UpperCamelCase__										=      data_collator
												UpperCamelCase__										=      lr
												UpperCamelCase__										=      params
							UpperCamelCase__										=      jax_utils.replicate(SCREAMING_SNAKE_CASE_	)
							return state






		def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Tuple:





							'''simple docstring'''





							UpperCamelCase__										=      self.args
							UpperCamelCase__										=      len(SCREAMING_SNAKE_CASE_	) // args.batch_size

							UpperCamelCase__										=      jax.random.PRNGKey(0	)
							UpperCamelCase__										=      jax.random.split(SCREAMING_SNAKE_CASE_   ,				jax.device_count()	)
							for epoch in range(args.max_epochs	):
												UpperCamelCase__										=      jnp.array(0   ,				dtype=jnp.floataa	)
												UpperCamelCase__										=      get_batched_dataset(SCREAMING_SNAKE_CASE_   ,				args.batch_size   ,				seed=SCREAMING_SNAKE_CASE_	)
												UpperCamelCase__										=      0
												for batch in tqdm(SCREAMING_SNAKE_CASE_   ,				total=SCREAMING_SNAKE_CASE_   ,				desc=F"Running EPOCH-{epoch}"	):
																	UpperCamelCase__										=      self.data_collator(SCREAMING_SNAKE_CASE_	)
																	UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      self.train_step_fn(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				**SCREAMING_SNAKE_CASE_	)
																	running_loss += jax_utils.unreplicate(metrics['''loss''']	)
																	i += 1
																	if i % args.logging_steps == 0:
																						UpperCamelCase__										=      jax_utils.unreplicate(state.step	)
																						UpperCamelCase__										=      running_loss.item() / i
																						UpperCamelCase__										=      self.scheduler_fn(state_step - 1	)

																						UpperCamelCase__										=      self.evaluate(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)
																						UpperCamelCase__										=      {
																						    '''step''': state_step.item(),
																						    '''eval_loss''': eval_loss.item(),
																						    '''tr_loss''': tr_loss,
																						    '''lr''': lr.item(),
																						}
																						tqdm.write(str(SCREAMING_SNAKE_CASE_	)	)
																						self.logger.log(SCREAMING_SNAKE_CASE_   ,				commit=SCREAMING_SNAKE_CASE_	)

																	if i % args.save_steps == 0:
																						self.save_checkpoint(args.save_dir + F"-e{epoch}-s{i}"   ,				state=SCREAMING_SNAKE_CASE_	)






		def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Tuple:





							'''simple docstring'''





							UpperCamelCase__										=      get_batched_dataset(SCREAMING_SNAKE_CASE_   ,				self.args.batch_size	)
							UpperCamelCase__										=      len(SCREAMING_SNAKE_CASE_	) // self.args.batch_size
							UpperCamelCase__										=      jnp.array(0   ,				dtype=jnp.floataa	)
							UpperCamelCase__										=      0
							for batch in tqdm(SCREAMING_SNAKE_CASE_   ,				total=SCREAMING_SNAKE_CASE_   ,				desc='''Evaluating ... '''	):
												UpperCamelCase__										=      self.data_collator(SCREAMING_SNAKE_CASE_	)
												UpperCamelCase__										=      self.val_step_fn(SCREAMING_SNAKE_CASE_   ,				**SCREAMING_SNAKE_CASE_	)
												running_loss += jax_utils.unreplicate(metrics['''loss''']	)
												i += 1
							return running_loss / i






		def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Dict:





							'''simple docstring'''





							UpperCamelCase__										=      jax_utils.unreplicate(SCREAMING_SNAKE_CASE_	)
							print(F"SAVING CHECKPOINT IN {save_dir}"   ,				end=''' ... '''	)
							self.model_save_fn(SCREAMING_SNAKE_CASE_   ,				params=state.params	)
							with open(os.path.join(SCREAMING_SNAKE_CASE_   ,				'''opt_state.msgpack'''	)   ,				'''wb'''	) as f:
												f.write(to_bytes(state.opt_state	)	)
							joblib.dump(self.args   ,				os.path.join(SCREAMING_SNAKE_CASE_   ,				'''args.joblib'''	)	)
							joblib.dump(self.data_collator   ,				os.path.join(SCREAMING_SNAKE_CASE_   ,				'''data_collator.joblib'''	)	)
							with open(os.path.join(SCREAMING_SNAKE_CASE_   ,				'''training_state.json'''	)   ,				'''w'''	) as f:
												json.dump({'''step''': state.step.item()}   ,				SCREAMING_SNAKE_CASE_	)
							print('''DONE'''	)







def  __UpperCamelCase    (						A ,							A							):
					print(f"RESTORING CHECKPOINT FROM {save_dir}" ,							end=''' ... '''							)
					with open(os.path.join(_snake_case ,							'''flax_model.msgpack'''							) ,							'''rb'''							) as f:
										UpperCamelCase__										=      from_bytes(state.params ,							f.read()							)

					with open(os.path.join(_snake_case ,							'''opt_state.msgpack'''							) ,							'''rb'''							) as f:
										UpperCamelCase__										=      from_bytes(state.opt_state ,							f.read()							)

					UpperCamelCase__										=      joblib.load(os.path.join(_snake_case ,							'''args.joblib'''							)							)
					UpperCamelCase__										=      joblib.load(os.path.join(_snake_case ,							'''data_collator.joblib'''							)							)

					with open(os.path.join(_snake_case ,							'''training_state.json'''							) ,							'''r'''							) as f:
										UpperCamelCase__										=      json.load(_snake_case							)
					UpperCamelCase__										=      training_state['''step''']

					print('''DONE'''							)
					return params, opt_state, step, args, data_collator







def  __UpperCamelCase    (						A ,							A ,							A ,							A							):
					UpperCamelCase__										=      num_train_steps - warmup_steps
					UpperCamelCase__										=      optax.linear_schedule(init_value=_snake_case ,							end_value=_snake_case ,							transition_steps=_snake_case							)
					UpperCamelCase__										=      optax.linear_schedule(init_value=_snake_case ,							end_value=1e-7 ,							transition_steps=_snake_case							)
					UpperCamelCase__										=      optax.join_schedules(schedules=[warmup_fn, decay_fn] ,							boundaries=[warmup_steps]							)
					return lr







def  __UpperCamelCase    (						A ,							A ,							A ,							A ,							A							):






					def weight_decay_mask(A							):
										UpperCamelCase__										=      traverse_util.flatten_dict(_snake_case							)
										UpperCamelCase__										=      {k: (v[-1] != '''bias''' and v[-2:] != ('''LayerNorm''', '''scale''')) for k, v in params.items()}
										return traverse_util.unflatten_dict(_snake_case							)

					UpperCamelCase__										=      scheduler_fn(_snake_case ,							_snake_case ,							_snake_case ,							_snake_case							)

					UpperCamelCase__										=      optax.adamw(learning_rate=_snake_case ,							weight_decay=_snake_case ,							mask=_snake_case							)
					return tx, lr




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



@require_flax
class     _A   :






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Tuple:





       '''simple docstring'''





       pass






  def 		_a					(self	)		->      Optional[int]:





       '''simple docstring'''





       pass






  def 		_a					(self	)		->      Tuple:





       '''simple docstring'''





       pass






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Any:





       '''simple docstring'''





       UpperCamelCase__										=      np.abs((a - b)	).max()
       self.assertLessEqual(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				F"Difference between torch and flax is {diff} (>= {tol})."	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=None   ,				**SCREAMING_SNAKE_CASE_	)		->      Dict:





       '''simple docstring'''





       UpperCamelCase__										=      VisionTextDualEncoderConfig.from_vision_text_configs(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      FlaxVisionTextDualEncoderModel(SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      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 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=None   ,				**SCREAMING_SNAKE_CASE_	)		->      int:





       '''simple docstring'''





       UpperCamelCase__						,       UpperCamelCase__										=      self.get_vision_text_model(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      {'''vision_model''': vision_model, '''text_model''': text_model}
       UpperCamelCase__										=      FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      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 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=None   ,				**SCREAMING_SNAKE_CASE_	)		->      List[Any]:





       '''simple docstring'''





       UpperCamelCase__						,       UpperCamelCase__										=      self.get_vision_text_model(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      {'''vision_model''': vision_model, '''text_model''': text_model}
       UpperCamelCase__										=      FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      model(input_ids=SCREAMING_SNAKE_CASE_   ,				pixel_values=SCREAMING_SNAKE_CASE_   ,				attention_mask=SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      output[0]

       with tempfile.TemporaryDirectory() as tmpdirname:
            model.save_pretrained(SCREAMING_SNAKE_CASE_	)
            UpperCamelCase__										=      FlaxVisionTextDualEncoderModel.from_pretrained(SCREAMING_SNAKE_CASE_	)

            UpperCamelCase__										=      model(input_ids=SCREAMING_SNAKE_CASE_   ,				pixel_values=SCREAMING_SNAKE_CASE_   ,				attention_mask=SCREAMING_SNAKE_CASE_	)
            UpperCamelCase__										=      after_output[0]
            UpperCamelCase__										=      np.amax(np.abs(out_a - out_a	)	)
            self.assertLessEqual(SCREAMING_SNAKE_CASE_   ,				1E-3	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=None   ,				**SCREAMING_SNAKE_CASE_	)		->      Any:





       '''simple docstring'''





       UpperCamelCase__						,       UpperCamelCase__										=      self.get_vision_text_model(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      {'''vision_model''': vision_model, '''text_model''': text_model}
       UpperCamelCase__										=      FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      model(
           input_ids=SCREAMING_SNAKE_CASE_   ,				pixel_values=SCREAMING_SNAKE_CASE_   ,				attention_mask=SCREAMING_SNAKE_CASE_   ,				output_attentions=SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      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)
       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(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 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      List[str]:





       '''simple docstring'''





       pt_model.to(SCREAMING_SNAKE_CASE_	)
       pt_model.eval()

       # prepare inputs
       UpperCamelCase__										=      inputs_dict
       UpperCamelCase__										=      {k: torch.tensor(v.tolist()	) for k, v in flax_inputs.items()}

       with torch.no_grad():
            UpperCamelCase__										=      pt_model(**SCREAMING_SNAKE_CASE_	).to_tuple()

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

       # PT -> Flax
       with tempfile.TemporaryDirectory() as tmpdirname:
            pt_model.save_pretrained(SCREAMING_SNAKE_CASE_	)
            UpperCamelCase__										=      FlaxVisionTextDualEncoderModel.from_pretrained(SCREAMING_SNAKE_CASE_   ,				from_pt=SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      fx_model_loaded(**SCREAMING_SNAKE_CASE_	).to_tuple()
       self.assertEqual(len(SCREAMING_SNAKE_CASE_	)   ,				len(SCREAMING_SNAKE_CASE_	)   ,				'''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(SCREAMING_SNAKE_CASE_   ,				pt_output.numpy()   ,				4E-2	)

       # Flax -> PT
       with tempfile.TemporaryDirectory() as tmpdirname:
            fx_model.save_pretrained(SCREAMING_SNAKE_CASE_	)
            UpperCamelCase__										=      VisionTextDualEncoderModel.from_pretrained(SCREAMING_SNAKE_CASE_   ,				from_flax=SCREAMING_SNAKE_CASE_	)

       pt_model_loaded.to(SCREAMING_SNAKE_CASE_	)
       pt_model_loaded.eval()

       with torch.no_grad():
            UpperCamelCase__										=      pt_model_loaded(**SCREAMING_SNAKE_CASE_	).to_tuple()

       self.assertEqual(len(SCREAMING_SNAKE_CASE_	)   ,				len(SCREAMING_SNAKE_CASE_	)   ,				'''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(SCREAMING_SNAKE_CASE_   ,				pt_output_loaded.numpy()   ,				4E-2	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Tuple:





       '''simple docstring'''





       UpperCamelCase__										=      VisionTextDualEncoderConfig.from_vision_text_configs(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      VisionTextDualEncoderModel(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      FlaxVisionTextDualEncoderModel(SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      convert_pytorch_state_dict_to_flax(pt_model.state_dict()   ,				SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      fx_state

       self.check_pt_flax_equivalence(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      str:





       '''simple docstring'''





       UpperCamelCase__										=      VisionTextDualEncoderConfig.from_vision_text_configs(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      VisionTextDualEncoderModel(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      FlaxVisionTextDualEncoderModel(SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE_   ,				fx_model.params	)

       self.check_pt_flax_equivalence(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      str:





       '''simple docstring'''





       UpperCamelCase__										=      self.prepare_config_and_inputs()
       self.check_model_from_pretrained_configs(**SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      str:





       '''simple docstring'''





       UpperCamelCase__										=      self.prepare_config_and_inputs()
       self.check_vision_text_dual_encoder_from_pretrained(**SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      Dict:





       '''simple docstring'''





       UpperCamelCase__										=      self.prepare_config_and_inputs()
       self.check_save_load(**SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      Union[str, Any]:





       '''simple docstring'''





       UpperCamelCase__										=      self.prepare_config_and_inputs()
       self.check_vision_text_output_attention(**SCREAMING_SNAKE_CASE_	)






  @is_pt_flax_cross_test
  def 		_a					(self	)		->      str:





       '''simple docstring'''





       UpperCamelCase__										=      self.prepare_config_and_inputs()
       UpperCamelCase__										=      config_inputs_dict.pop('''vision_config'''	)
       UpperCamelCase__										=      config_inputs_dict.pop('''text_config'''	)

       UpperCamelCase__										=      config_inputs_dict

       self.check_equivalence_pt_to_flax(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)
       self.check_equivalence_flax_to_pt(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)






  @slow
  def 		_a					(self	)		->      Tuple:





       '''simple docstring'''





       UpperCamelCase__						,       UpperCamelCase__										=      self.get_pretrained_model_and_inputs()

       UpperCamelCase__										=      model_a(**SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      outputs[0]

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

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



@require_flax
class     _A   (    __UpperCamelCase	, unittest.TestCase							):






  def 		_a					(self	)		->      Any:





       '''simple docstring'''





       UpperCamelCase__										=      FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
           '''hf-internal-testing/tiny-random-vit'''   ,				'''hf-internal-testing/tiny-bert'''   ,				vision_from_pt=SCREAMING_SNAKE_CASE_   ,				text_from_pt=SCREAMING_SNAKE_CASE_   ,				)
       UpperCamelCase__										=      13
       UpperCamelCase__										=      floats_tensor(
           [
               batch_size,
               model.config.vision_config.num_channels,
               model.config.vision_config.image_size,
               model.config.vision_config.image_size,
           ]	)
       UpperCamelCase__										=      ids_tensor([batch_size, 4]   ,				model.config.text_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 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Union[str, Any]:





       '''simple docstring'''





       UpperCamelCase__										=      FlaxViTModel(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      FlaxBertModel(SCREAMING_SNAKE_CASE_	)
       return vision_model, text_model






  def 		_a					(self	)		->      Optional[Any]:





       '''simple docstring'''





       UpperCamelCase__										=      FlaxViTModelTester(self	)
       UpperCamelCase__										=      FlaxBertModelTester(self	)
       UpperCamelCase__										=      vit_model_tester.prepare_config_and_inputs()
       UpperCamelCase__										=      bert_model_tester.prepare_config_and_inputs()

       UpperCamelCase__						,       UpperCamelCase__										=      vision_config_and_inputs

       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      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     _A   (    __UpperCamelCase	, unittest.TestCase							):






  def 		_a					(self	)		->      Dict:





       '''simple docstring'''





       UpperCamelCase__										=      FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
           '''hf-internal-testing/tiny-random-clip'''   ,				'''hf-internal-testing/tiny-bert'''   ,				vision_from_pt=SCREAMING_SNAKE_CASE_   ,				text_from_pt=SCREAMING_SNAKE_CASE_   ,				)
       UpperCamelCase__										=      13
       UpperCamelCase__										=      floats_tensor(
           [
               batch_size,
               model.config.vision_config.num_channels,
               model.config.vision_config.image_size,
               model.config.vision_config.image_size,
           ]	)
       UpperCamelCase__										=      ids_tensor([batch_size, 4]   ,				model.config.text_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 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      int:





       '''simple docstring'''





       UpperCamelCase__										=      FlaxCLIPVisionModel(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      FlaxBertModel(SCREAMING_SNAKE_CASE_	)
       return vision_model, text_model






  def 		_a					(self	)		->      List[Any]:





       '''simple docstring'''





       UpperCamelCase__										=      FlaxCLIPVisionModelTester(self	)
       UpperCamelCase__										=      FlaxBertModelTester(self	)
       UpperCamelCase__										=      clip_model_tester.prepare_config_and_inputs()
       UpperCamelCase__										=      bert_model_tester.prepare_config_and_inputs()

       UpperCamelCase__						,       UpperCamelCase__										=      vision_config_and_inputs

       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      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     _A   (    unittest.TestCase							):






  @slow
  def 		_a					(self	)		->      int:





       '''simple docstring'''





       UpperCamelCase__										=      FlaxVisionTextDualEncoderModel.from_pretrained('''clip-italian/clip-italian'''   ,				logit_scale_init_value=1.0	)
       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=SCREAMING_SNAKE_CASE_   ,				padding=SCREAMING_SNAKE_CASE_   ,				return_tensors='''np'''	)

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

       UpperCamelCase__										=      np.array([[1.2284727, 0.3104122]]	)

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



