






from typing import Optional, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
from configuration_hybrid_clip import HybridCLIPConfig
from flax.core.frozen_dict import FrozenDict

from transformers import FLAX_MODEL_MAPPING, FlaxCLIPVisionModel
from transformers.modeling_flax_utils import FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPOutput
from transformers.utils import logging


_a				     =  logging.get_logger(__name__)





class __lowerCamelCase			(				nn.Module):





							"""simple docstring"""






							UpperCamelCase__								=				42
							UpperCamelCase__								=				jnp.floataa

							def 	UpperCamelCase   (							self   ):
													"""simple docstring"""







													_UpperCAmelCase     =      self.config.text_config
													_UpperCAmelCase     =      self.config.vision_config

													_UpperCAmelCase     =      self.config.projection_dim
													_UpperCAmelCase     =      text_config.hidden_size
													_UpperCAmelCase     =      vision_config.hidden_size

													_UpperCAmelCase     =      FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
													_UpperCAmelCase     =      FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__	,   __UpperCamelCase   ).module_class

													_UpperCAmelCase     =      text_module(__UpperCamelCase	,   dtype=self.dtype   )
													_UpperCAmelCase     =      vision_module(__UpperCamelCase	,   dtype=self.dtype   )

													_UpperCAmelCase     =      nn.Dense(
													    self.projection_dim	,   dtype=self.dtype	,   kernel_init=jax.nn.initializers.normal(0.02   )	,   use_bias=__UpperCamelCase	,   )
													_UpperCAmelCase     =      nn.Dense(
													    self.projection_dim	,   dtype=self.dtype	,   kernel_init=jax.nn.initializers.normal(0.02   )	,   use_bias=__UpperCamelCase	,   )
													_UpperCAmelCase     =      self.param('logit_scale'	,   jax.nn.initializers.ones	,   []   )

							def __call__(							self	,   UpperCAmelCase=None	,   UpperCAmelCase=None	,   UpperCAmelCase=None	,   UpperCAmelCase=None	,   UpperCAmelCase=None	,   UpperCAmelCase = True	,   UpperCAmelCase=None	,   UpperCAmelCase=None	,   UpperCAmelCase=None	,   ):
													"""simple docstring"""







													_UpperCAmelCase     =      return_dict if return_dict is not None else self.config.return_dict

													_UpperCAmelCase     =      self.vision_model(
													    pixel_values=__UpperCamelCase	,   deterministic=__UpperCamelCase	,   output_attentions=__UpperCamelCase	,   output_hidden_states=__UpperCamelCase	,   return_dict=__UpperCamelCase	,   )

													_UpperCAmelCase     =      self.text_model(
													    input_ids=__UpperCamelCase	,   attention_mask=__UpperCamelCase	,   token_type_ids=__UpperCamelCase	,   position_ids=__UpperCamelCase	,   deterministic=__UpperCamelCase	,   output_attentions=__UpperCamelCase	,   output_hidden_states=__UpperCamelCase	,   return_dict=__UpperCamelCase	,   )

													_UpperCAmelCase     =      vision_outputs[1]
													_UpperCAmelCase     =      self.visual_projection(__UpperCamelCase   )

													_UpperCAmelCase     =      text_outputs[1]
													_UpperCAmelCase     =      self.text_projection(__UpperCamelCase   )

													# normalized features
													_UpperCAmelCase     =      image_embeds / jnp.linalg.norm(__UpperCamelCase	,   axis=-1	,   keepdims=__UpperCamelCase   )
													_UpperCAmelCase     =      text_embeds / jnp.linalg.norm(__UpperCamelCase	,   axis=-1	,   keepdims=__UpperCamelCase   )

													# cosine similarity as logits
													_UpperCAmelCase     =      jnp.exp(self.logit_scale   )
													_UpperCAmelCase     =      jnp.matmul(__UpperCamelCase	,   image_embeds.T   ) * logit_scale
													_UpperCAmelCase     =      logits_per_text.T

													if not return_dict:
																			return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)

													return FlaxCLIPOutput(
													    logits_per_image=__UpperCamelCase	,   logits_per_text=__UpperCamelCase	,   text_embeds=__UpperCamelCase	,   image_embeds=__UpperCamelCase	,   text_model_output=__UpperCamelCase	,   vision_model_output=__UpperCamelCase	,   )





class __lowerCamelCase			(				UpperCAmelCase_):





							"""simple docstring"""






							UpperCamelCase__								=				HybridCLIPConfig
							UpperCamelCase__								=				FlaxHybridCLIPModule

							def __init__(							self	,   UpperCAmelCase	,   UpperCAmelCase = None	,   UpperCAmelCase = 0	,   UpperCAmelCase = jnp.floataa	,   **UpperCAmelCase	,   ):
													"""simple docstring"""







													if input_shape is None:
																			_UpperCAmelCase     =      ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

													_UpperCAmelCase     =      self.module_class(config=__UpperCamelCase	,   dtype=__UpperCamelCase	,   **__UpperCamelCase   )
													super().__init__(__UpperCamelCase	,   __UpperCamelCase	,   input_shape=__UpperCamelCase	,   seed=__UpperCamelCase	,   dtype=__UpperCamelCase   )

							def 	UpperCamelCase   (							self	,   UpperCAmelCase	,   UpperCAmelCase	,   UpperCAmelCase = None   ):
													"""simple docstring"""







													_UpperCAmelCase     =      jnp.zeros(input_shape[0]	,   dtype='i4'   )
													_UpperCAmelCase     =      jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase   ).shape[-1]   )	,   input_shape[0]   )
													_UpperCAmelCase     =      jnp.ones_like(__UpperCamelCase   )
													_UpperCAmelCase     =      jnp.ones_like(__UpperCamelCase   )

													_UpperCAmelCase     =      jax.random.normal(__UpperCamelCase	,   input_shape[1]   )

													_UpperCAmelCase ,			_UpperCAmelCase     =      jax.random.split(__UpperCamelCase   )
													_UpperCAmelCase     =      {'params': params_rng, 'dropout': dropout_rng}

													return self.module.init(__UpperCamelCase	,   __UpperCamelCase	,   __UpperCamelCase	,   __UpperCamelCase	,   __UpperCamelCase	,   __UpperCamelCase   )["params"]

							def __call__(							self	,   UpperCAmelCase	,   UpperCAmelCase	,   UpperCAmelCase=None	,   UpperCAmelCase=None	,   UpperCAmelCase=None	,   UpperCAmelCase = None	,   UpperCAmelCase = None	,   UpperCAmelCase = False	,   UpperCAmelCase = None	,   UpperCAmelCase = None	,   UpperCAmelCase = None	,   ):
													"""simple docstring"""







													_UpperCAmelCase     =      output_attentions if output_attentions is not None else self.config.output_attentions
													_UpperCAmelCase     =      (
													    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
													)
													_UpperCAmelCase     =      return_dict if return_dict is not None else self.config.return_dict

													if position_ids is None:
																			_UpperCAmelCase     =      jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase   ).shape[-1]   )	,   input_ids.shape   )

													if token_type_ids is None:
																			_UpperCAmelCase     =      jnp.zeros_like(__UpperCamelCase   )

													if attention_mask is None:
																			_UpperCAmelCase     =      jnp.ones_like(__UpperCamelCase   )

													# Handle any PRNG if needed
													_UpperCAmelCase     =      {}
													if dropout_rng is not None:
																			_UpperCAmelCase     =      dropout_rng

													return self.module.apply(
													    {'params': params or self.params}	,   jnp.array(__UpperCamelCase	,   dtype='i4'   )	,   jnp.array(__UpperCamelCase	,   dtype=jnp.floataa   )	,   jnp.array(__UpperCamelCase	,   dtype='i4'   )	,   jnp.array(__UpperCamelCase	,   dtype='i4'   )	,   jnp.array(__UpperCamelCase	,   dtype='i4'   )	,   not train	,   __UpperCamelCase	,   __UpperCamelCase	,   __UpperCamelCase	,   rngs=__UpperCamelCase	,   )

							def 	UpperCamelCase   (							self	,   UpperCAmelCase	,   UpperCAmelCase=None	,   UpperCAmelCase=None	,   UpperCAmelCase=None	,   UpperCAmelCase = None	,   UpperCAmelCase = None	,   UpperCAmelCase=False	,   ):
													"""simple docstring"""







													if position_ids is None:
																			_UpperCAmelCase     =      jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase   ).shape[-1]   )	,   input_ids.shape   )

													if token_type_ids is None:
																			_UpperCAmelCase     =      jnp.zeros_like(__UpperCamelCase   )

													if attention_mask is None:
																			_UpperCAmelCase     =      jnp.ones_like(__UpperCamelCase   )

													# Handle any PRNG if needed
													_UpperCAmelCase     =      {}
													if dropout_rng is not None:
																			_UpperCAmelCase     =      dropout_rng

													def _get_features(UpperCAmelCase	,   UpperCAmelCase	,   UpperCAmelCase	,   UpperCAmelCase	,   UpperCAmelCase	,   UpperCAmelCase   ):
																			_UpperCAmelCase     =      module.text_model(
																			    input_ids=__UpperCamelCase	,   attention_mask=__UpperCamelCase	,   position_ids=__UpperCamelCase	,   token_type_ids=__UpperCamelCase	,   deterministic=__UpperCamelCase	,   )
																			_UpperCAmelCase     =      text_outputs[1]
																			_UpperCAmelCase     =      module.text_projection(__UpperCamelCase   )
																			return text_features

													return self.module.apply(
													    {'params': params or self.params}	,   jnp.array(__UpperCamelCase	,   dtype='i4'   )	,   jnp.array(__UpperCamelCase	,   dtype='i4'   )	,   jnp.array(__UpperCamelCase	,   dtype='i4'   )	,   jnp.array(__UpperCamelCase	,   dtype='i4'   )	,   not train	,   method=_get_features	,   rngs=__UpperCamelCase	,   )

							def 	UpperCamelCase   (							self	,   UpperCAmelCase	,   UpperCAmelCase = None	,   UpperCAmelCase = None	,   UpperCAmelCase=False   ):
													"""simple docstring"""







													_UpperCAmelCase     =      {}
													if dropout_rng is not None:
																			_UpperCAmelCase     =      dropout_rng

													def _get_features(UpperCAmelCase	,   UpperCAmelCase	,   UpperCAmelCase   ):
																			_UpperCAmelCase     =      module.vision_model(pixel_values=__UpperCamelCase	,   deterministic=__UpperCamelCase   )
																			_UpperCAmelCase     =      vision_outputs[1]  # pooled_output
																			_UpperCAmelCase     =      module.visual_projection(__UpperCamelCase   )
																			return image_features

													return self.module.apply(
													    {'params': params or self.params}	,   jnp.array(__UpperCamelCase	,   dtype=jnp.floataa   )	,   not train	,   method=_get_features	,   rngs=__UpperCamelCase	,   )

							@classmethod
							def 	UpperCamelCase   (							cls	,   UpperCAmelCase = None	,   UpperCAmelCase = None	,   *UpperCAmelCase	,   **UpperCAmelCase	,   ):
													"""simple docstring"""







													_UpperCAmelCase     =      {
													    argument[len('text_'   ) :]: value for argument, value in kwargs.items() if argument.startswith('text_'   )
													}

													_UpperCAmelCase     =      {
													    argument[len('vision_'   ) :]: value for argument, value in kwargs.items() if argument.startswith('vision_'   )
													}

													# remove text, vision kwargs from kwargs
													for key in kwargs_text.keys():
																			del kwargs["text_" + key]
													for key in kwargs_vision.keys():
																			del kwargs["vision_" + key]

													# Load and initialize the text and vision model
													_UpperCAmelCase     =      kwargs_text.pop('model'	,   __UpperCamelCase   )
													if text_model is None:
																			assert (
																			    text_model_name_or_path is not None
																			), "If `model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
																			from transformers import FlaxAutoModel

																			if "config" not in kwargs_text:
																									from transformers import AutoConfig

																									_UpperCAmelCase     =      AutoConfig.from_pretrained(__UpperCamelCase   )
																									_UpperCAmelCase     =      text_config

																			_UpperCAmelCase     =      FlaxAutoModel.from_pretrained(__UpperCamelCase	,   *__UpperCamelCase	,   **__UpperCamelCase   )

													_UpperCAmelCase     =      kwargs_vision.pop('model'	,   __UpperCamelCase   )
													if vision_model is None:
																			assert (
																			    vision_model_name_or_path is not None
																			), "If `model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
																			from transformers import FlaxAutoModel

																			if "config" not in kwargs_vision:
																									from transformers import AutoConfig

																									_UpperCAmelCase     =      AutoConfig.from_pretrained(__UpperCamelCase   )
																									_UpperCAmelCase     =      vision_config

																			_UpperCAmelCase     =      FlaxAutoModel.from_pretrained(__UpperCamelCase	,   *__UpperCamelCase	,   **__UpperCamelCase   )

													# instantiate config with corresponding kwargs
													_UpperCAmelCase     =      kwargs.pop('dtype'	,   jnp.floataa   )
													_UpperCAmelCase     =      HybridCLIPConfig.from_text_vision_configs(text_model.config	,   vision_model.config	,   **__UpperCamelCase   )

													# init model
													_UpperCAmelCase     =      cls(__UpperCamelCase	,   *__UpperCamelCase	,   dtype=__UpperCamelCase	,   **__UpperCamelCase   )

													if vision_config.model_type == "clip":
																			_UpperCAmelCase     =      vision_model.params['vision_model']
																			_UpperCAmelCase     =      vision_model.params['visual_projection']['kernel']
													else:
																			_UpperCAmelCase     =      vision_model.params

													_UpperCAmelCase     =      text_model.params

													return model








import unittest

import pytest

from transformers import (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    MBartaaTokenizerFast,
    MBartConfig,
    MBartForConditionalGeneration,
    TranslationPipeline,
    pipeline,
)
from transformers.testing_utils import is_pipeline_test, require_tf, require_torch, slow

from .test_pipelines_common import ANY





@is_pipeline_test
class __lowerCamelCase			(				unittest.TestCase):





	"""simple docstring"""






	UpperCamelCase__								=				MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
	UpperCamelCase__								=				TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

	def 	UpperCamelCase   (							self	,   UpperCAmelCase	,   UpperCAmelCase	,   UpperCAmelCase   ):
							"""simple docstring"""







							if isinstance(model.config	,   UpperCAmelCase   ):
													_UpperCAmelCase ,			_UpperCAmelCase     =      list(tokenizer.lang_code_to_id.keys()   )[:2]
													_UpperCAmelCase     =      TranslationPipeline(model=UpperCAmelCase	,   tokenizer=UpperCAmelCase	,   src_lang=UpperCAmelCase	,   tgt_lang=UpperCAmelCase   )
							else:
													_UpperCAmelCase     =      TranslationPipeline(model=UpperCAmelCase	,   tokenizer=UpperCAmelCase   )
							return translator, ["Some string", "Some other text"]

	def 	UpperCamelCase   (							self	,   UpperCAmelCase	,   UpperCAmelCase   ):
							"""simple docstring"""







							_UpperCAmelCase     =      translator('Some string'   )
							self.assertEqual(UpperCAmelCase	,   [{'translation_text': ANY(UpperCAmelCase   )}]   )

							_UpperCAmelCase     =      translator(['Some string']   )
							self.assertEqual(UpperCAmelCase	,   [{'translation_text': ANY(UpperCAmelCase   )}]   )

							_UpperCAmelCase     =      translator(['Some string', 'other string']   )
							self.assertEqual(UpperCAmelCase	,   [{'translation_text': ANY(UpperCAmelCase   )}, {'translation_text': ANY(UpperCAmelCase   )}]   )

	@require_torch
	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      pipeline('translation_en_to_ro'	,   model='patrickvonplaten/t5-tiny-random'	,   framework='pt'   )
							_UpperCAmelCase     =      translator('This is a test string'	,   max_length=20   )
							self.assertEqual(
							    UpperCAmelCase	,   [
							        {
							            'translation_text': (
							                'Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide'
							                ' Beide Beide'
							            )
							        }
							    ]	,   )

	@require_tf
	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      pipeline('translation_en_to_ro'	,   model='patrickvonplaten/t5-tiny-random'	,   framework='tf'   )
							_UpperCAmelCase     =      translator('This is a test string'	,   max_length=20   )
							self.assertEqual(
							    UpperCAmelCase	,   [
							        {
							            'translation_text': (
							                'Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide'
							                ' Beide Beide'
							            )
							        }
							    ]	,   )

	@require_torch
	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      pipeline('translation_en_to_de'	,   model='patrickvonplaten/t5-tiny-random'	,   framework='pt'   )
							_UpperCAmelCase     =      translator('This is a test string'	,   max_length=20   )
							self.assertEqual(
							    UpperCAmelCase	,   [
							        {
							            'translation_text': (
							                'monoton monoton monoton monoton monoton monoton monoton monoton monoton monoton urine urine'
							                ' urine urine urine urine urine urine urine'
							            )
							        }
							    ]	,   )

	@require_tf
	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      pipeline('translation_en_to_de'	,   model='patrickvonplaten/t5-tiny-random'	,   framework='tf'   )
							_UpperCAmelCase     =      translator('This is a test string'	,   max_length=20   )
							self.assertEqual(
							    UpperCAmelCase	,   [
							        {
							            'translation_text': (
							                'monoton monoton monoton monoton monoton monoton monoton monoton monoton monoton urine urine'
							                ' urine urine urine urine urine urine urine'
							            )
							        }
							    ]	,   )





class __lowerCamelCase			(				unittest.TestCase):





	"""simple docstring"""






	@require_torch
	@slow
	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							with self.assertRaises(UpperCAmelCase   ):
													pipeline(task='translation_cn_to_ar'   )

							# but we do for this one
							_UpperCAmelCase     =      pipeline(task='translation_en_to_de'   )
							self.assertEqual(translator._preprocess_params['src_lang']	,   'en'   )
							self.assertEqual(translator._preprocess_params['tgt_lang']	,   'de'   )

	@require_torch
	@slow
	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      MBartForConditionalGeneration.from_pretrained('facebook/mbart-large-50-many-to-many-mmt'   )
							_UpperCAmelCase     =      MBartaaTokenizerFast.from_pretrained('facebook/mbart-large-50-many-to-many-mmt'   )

							_UpperCAmelCase     =      pipeline(task='translation'	,   model=UpperCAmelCase	,   tokenizer=UpperCAmelCase   )
							# Missing src_lang, tgt_lang
							with self.assertRaises(UpperCAmelCase   ):
													translator('This is a test'   )

							_UpperCAmelCase     =      translator('This is a test'	,   src_lang='en_XX'	,   tgt_lang='ar_AR'   )
							self.assertEqual(UpperCAmelCase	,   [{'translation_text': 'هذا إختبار'}]   )

							_UpperCAmelCase     =      translator('This is a test'	,   src_lang='en_XX'	,   tgt_lang='hi_IN'   )
							self.assertEqual(UpperCAmelCase	,   [{'translation_text': 'यह एक परीक्षण है'}]   )

							# src_lang, tgt_lang can be defined at pipeline call time
							_UpperCAmelCase     =      pipeline(task='translation'	,   model=UpperCAmelCase	,   tokenizer=UpperCAmelCase	,   src_lang='en_XX'	,   tgt_lang='ar_AR'   )
							_UpperCAmelCase     =      translator('This is a test'   )
							self.assertEqual(UpperCAmelCase	,   [{'translation_text': 'هذا إختبار'}]   )

	@require_torch
	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      'patrickvonplaten/t5-tiny-random'
							_UpperCAmelCase     =      pipeline(task='translation_cn_to_ar'	,   model=UpperCAmelCase   )
							self.assertEqual(translator._preprocess_params['src_lang']	,   'cn'   )
							self.assertEqual(translator._preprocess_params['tgt_lang']	,   'ar'   )

	@require_torch
	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      'patrickvonplaten/t5-tiny-random'
							with pytest.warns(UpperCAmelCase	,   match=R'.*translation_en_to_de.*'   ):
													_UpperCAmelCase     =      pipeline(task='translation'	,   model=UpperCAmelCase   )
							self.assertEqual(translator.task	,   'translation_en_to_de'   )
							self.assertEqual(translator._preprocess_params['src_lang']	,   'en'   )
							self.assertEqual(translator._preprocess_params['tgt_lang']	,   'de'   )

	@require_torch
	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							with self.assertRaises(UpperCAmelCase   ):
													pipeline(task='translation'   )
