




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"""

      lowerCAmelCase__     :	HybridCLIPConfig
      lowerCAmelCase__     :	jnp.dtype      =   jnp.floataa







      def  snake_case		(       self   :	Optional[int]			):
          __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__	,    snake_case			).module_class

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

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







      def __call__(       self   :	Union[str, Any]	,    snake_case   :	Tuple=None	,    snake_case   :	Optional[int]=None	,    snake_case   :	List[str]=None	,    snake_case   :	List[Any]=None	,    snake_case   :	int=None	,    snake_case   :	bool = True	,    snake_case   :	Union[str, Any]=None	,    snake_case   :	Dict=None	,    snake_case   :	Dict=None	,    ):
          __UpperCamelCase											=  return_dict if return_dict is not None else self.config.return_dict

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

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

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

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

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

          # cosine similarity as logits
          __UpperCamelCase											=  jnp.exp(self.logit_scale			)
          __UpperCamelCase											=  jnp.matmul(snake_case	,    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=snake_case	,    logits_per_text=snake_case	,    text_embeds=snake_case	,    image_embeds=snake_case	,    text_model_output=snake_case	,    vision_model_output=snake_case	,    )





class 	_lowerCamelCase  (     _SCREAMING_SNAKE_CASE				):

      """simple docstring"""

      lowerCAmelCase__     :	List[Any]      =   HybridCLIPConfig
      lowerCAmelCase__     :	Union[str, Any]      =   FlaxHybridCLIPModule







      def __init__(       self   :	Any	,    snake_case   :	HybridCLIPConfig	,    snake_case   :	Optional[Tuple] = None	,    snake_case   :	int = 0	,    snake_case   :	jnp.dtype = jnp.floataa	,    **snake_case   :	Dict	,    ):
          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=snake_case	,    dtype=snake_case	,    **snake_case			)
          super().__init__(snake_case	,    snake_case	,    input_shape=snake_case	,    seed=snake_case	,    dtype=snake_case			)







      def  snake_case		(       self   :	Any	,    snake_case   :	jax.random.PRNGKey	,    snake_case   :	Tuple	,    snake_case   :	FrozenDict = None			):
          # init input tensor
          __UpperCamelCase											=  jnp.zeros(input_shape[0]	,    dtype='''i4'''			)
          __UpperCamelCase											=  jnp.broadcast_to(jnp.arange(jnp.atleast_ad(snake_case			).shape[-1]			)	,    input_shape[0]			)
          __UpperCamelCase											=  jnp.ones_like(snake_case			)
          __UpperCamelCase											=  jnp.ones_like(snake_case			)

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

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

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







      def __call__(       self   :	Optional[Any]	,    snake_case   :	List[str]	,    snake_case   :	Optional[Any]	,    snake_case   :	Optional[int]=None	,    snake_case   :	List[Any]=None	,    snake_case   :	List[Any]=None	,    snake_case   :	dict = None	,    snake_case   :	jax.random.PRNGKey = None	,    snake_case   :	bool = False	,    snake_case   :	Optional[bool] = None	,    snake_case   :	Optional[bool] = None	,    snake_case   :	Optional[bool] = None	,    ):
          __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(snake_case			).shape[-1]			)	,    input_ids.shape			)

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

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

          # 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(snake_case	,    dtype='''i4'''			)	,    jnp.array(snake_case	,    dtype=jnp.floataa			)	,    jnp.array(snake_case	,    dtype='''i4'''			)	,    jnp.array(snake_case	,    dtype='''i4'''			)	,    jnp.array(snake_case	,    dtype='''i4'''			)	,    not train	,    snake_case	,    snake_case	,    snake_case	,    rngs=snake_case	,    )







      def  snake_case		(       self   :	Dict	,    snake_case   :	Optional[Any]	,    snake_case   :	Optional[int]=None	,    snake_case   :	Union[str, Any]=None	,    snake_case   :	List[Any]=None	,    snake_case   :	dict = None	,    snake_case   :	jax.random.PRNGKey = None	,    snake_case   :	Union[str, Any]=False	,    ):
          if position_ids is None:
              __UpperCamelCase											=  jnp.broadcast_to(jnp.arange(jnp.atleast_ad(snake_case			).shape[-1]			)	,    input_ids.shape			)

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

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

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

          def _get_features(snake_case   :	Union[str, Any]	,    snake_case   :	Union[str, Any]	,    snake_case   :	List[Any]	,    snake_case   :	int	,    snake_case   :	Dict	,    snake_case   :	Union[str, Any]			):
              __UpperCamelCase											=  module.text_model(
                  input_ids=snake_case	,    attention_mask=snake_case	,    position_ids=snake_case	,    token_type_ids=snake_case	,    deterministic=snake_case	,    )
              __UpperCamelCase											=  text_outputs[1]
              __UpperCamelCase											=  module.text_projection(snake_case			)
              return text_features

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







      def  snake_case		(       self   :	List[str]	,    snake_case   :	Optional[int]	,    snake_case   :	dict = None	,    snake_case   :	jax.random.PRNGKey = None	,    snake_case   :	Any=False			):
          __UpperCamelCase											=  {}
          if dropout_rng is not None:
              __UpperCamelCase											=  dropout_rng

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

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







      @classmethod
      def  snake_case		(       cls   :	Optional[Any]	,    snake_case   :	str = None	,    snake_case   :	str = None	,    *snake_case   :	List[Any]	,    **snake_case   :	Union[str, Any]	,    ):
          __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'''	,    snake_case			)
          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(snake_case			)
                  __UpperCamelCase											=  text_config

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

          __UpperCamelCase											=  kwargs_vision.pop('''model'''	,    snake_case			)
          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(snake_case			)
                  __UpperCamelCase											=  vision_config

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

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

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

          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

from transformers.testing_utils import is_torch_available, require_torch, tooslow

from ...generation.test_utils import torch_device
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     CpmAntConfig,
     CpmAntForCausalLM,
     CpmAntModel,
     CpmAntTokenizer,
 )

@require_torch
class 	_lowerCamelCase  :

      """simple docstring"""

      def __init__(       self   :	str	,    snake_case   :	str	,    snake_case   :	List[str]=2	,    snake_case   :	Optional[Any]=8	,    snake_case   :	List[Any]=True	,    snake_case   :	Optional[Any]=False	,    snake_case   :	str=False	,    snake_case   :	Tuple=False	,    snake_case   :	str=False	,    snake_case   :	int=99	,    snake_case   :	Tuple=32	,    snake_case   :	Union[str, Any]=3	,    snake_case   :	str=4	,    snake_case   :	List[Any]=37	,    snake_case   :	Tuple=32	,    snake_case   :	Tuple=128	,    snake_case   :	Optional[int]=8	,    snake_case   :	Optional[Any]=8	,    snake_case   :	List[Any]=8	,    snake_case   :	List[str]=1.0	,    snake_case   :	List[Any]=True	,    ):
          __UpperCamelCase											=  parent
          __UpperCamelCase											=  batch_size
          __UpperCamelCase											=  seq_length
          __UpperCamelCase											=  is_training
          __UpperCamelCase											=  use_token_type_ids
          __UpperCamelCase											=  use_input_mask
          __UpperCamelCase											=  use_labels
          __UpperCamelCase											=  use_mc_token_ids
          __UpperCamelCase											=  vocab_size
          __UpperCamelCase											=  hidden_size
          __UpperCamelCase											=  num_hidden_layers
          __UpperCamelCase											=  num_attention_heads
          __UpperCamelCase											=  intermediate_size
          __UpperCamelCase											=  num_buckets
          __UpperCamelCase											=  max_distance
          __UpperCamelCase											=  prompt_length
          __UpperCamelCase											=  prompt_types
          __UpperCamelCase											=  segment_types
          __UpperCamelCase											=  init_std
          __UpperCamelCase											=  return_dict







      def  snake_case		(       self   :	Tuple			):
          __UpperCamelCase											=  {}
          __UpperCamelCase											=  ids_tensor([self.batch_size, self.seq_length]	,    self.vocab_size			).type(torch.intaa			)
          __UpperCamelCase											=  False

          __UpperCamelCase											=  self.get_config()

          return (config, input_ids)







      def  snake_case		(       self   :	List[str]			):
          return CpmAntConfig(
              vocab_size=self.vocab_size	,    hidden_size=self.hidden_size	,    num_hidden_layers=self.num_hidden_layers	,    num_attention_heads=self.num_attention_heads	,    dim_ff=self.intermediate_size	,    position_bias_num_buckets=self.num_buckets	,    position_bias_max_distance=self.max_distance	,    prompt_types=self.prompt_types	,    prompt_length=self.prompt_length	,    segment_types=self.segment_types	,    use_cache=snake_case	,    init_std=self.init_std	,    return_dict=self.return_dict	,    )







      def  snake_case		(       self   :	Union[str, Any]	,    snake_case   :	Tuple	,    snake_case   :	List[str]	,    *snake_case   :	str			):
          __UpperCamelCase											=  CpmAntModel(config=snake_case			)
          model.to(snake_case			)
          model.eval()

          __UpperCamelCase											=  model(**snake_case			).last_hidden_state

          self.parent.assertEqual(hidden_states.shape	,    (self.batch_size, self.seq_length, config.hidden_size)			)







      def  snake_case		(       self   :	List[str]	,    snake_case   :	Optional[int]	,    snake_case   :	List[str]	,    *snake_case   :	Any			):
          __UpperCamelCase											=  CpmAntForCausalLM(snake_case			)
          model.to(snake_case			)
          __UpperCamelCase											=  input_ids['''input_ids'''].to(snake_case			)
          model.eval()

          __UpperCamelCase											=  model(**snake_case			)
          self.parent.assertEqual(
              model_output.logits.shape	,    (self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length)	,    )







      def  snake_case		(       self   :	List[str]			):
          __UpperCamelCase						, __UpperCamelCase											=  self.prepare_config_and_inputs()
          return config, inputs_dict







@require_torch
class 	_lowerCamelCase  (     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     unittest.TestCase				):

      """simple docstring"""

      lowerCAmelCase__     :	List[Any]      =   (CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
      lowerCAmelCase__     :	Union[str, Any]      =   (
          {"feature-extraction": CpmAntModel, "text-generation": CpmAntForCausalLM} if is_torch_available() else {}
      )

      lowerCAmelCase__     :	Optional[int]      =   False
      lowerCAmelCase__     :	Union[str, Any]      =   False
      lowerCAmelCase__     :	Optional[int]      =   False
      lowerCAmelCase__     :	Tuple      =   False
      lowerCAmelCase__     :	Union[str, Any]      =   False







      def  snake_case		(       self   :	Any			):
          __UpperCamelCase											=  CpmAntModelTester(self			)
          __UpperCamelCase											=  ConfigTester(self	,    config_class=snake_case			)







      def  snake_case		(       self   :	Optional[Any]			):
          self.config_tester.create_and_test_config_common_properties()
          self.config_tester.create_and_test_config_to_json_string()
          self.config_tester.create_and_test_config_to_json_file()
          self.config_tester.create_and_test_config_from_and_save_pretrained()
          self.config_tester.check_config_can_be_init_without_params()
          self.config_tester.check_config_arguments_init()







      def  snake_case		(       self   :	List[str]			):
          unittest.skip('''CPMAnt doesn\'t support input_embeds.'''			)(self.test_inputs_embeds			)







      def  snake_case		(       self   :	int			):
          unittest.skip(
              '''CPMAnt doesn\'t support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
                 So is attentions. We strongly recommand you use loss to tune model.'''			)(self.test_retain_grad_hidden_states_attentions			)







      def  snake_case		(       self   :	Union[str, Any]			):
          __UpperCamelCase						, __UpperCamelCase											=  self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_cpmant_model(snake_case	,    snake_case			)







      def  snake_case		(       self   :	Any			):
          __UpperCamelCase						, __UpperCamelCase											=  self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_lm_head_model(snake_case	,    snake_case			)







@require_torch
class 	_lowerCamelCase  (     unittest.TestCase				):

      """simple docstring"""

      @tooslow
      def  snake_case		(       self   :	Dict			):
          __UpperCamelCase											=  '''今天天气真好！'''
          __UpperCamelCase											=  '''openbmb/cpm-ant-10b'''
          __UpperCamelCase											=  CpmAntModel.from_pretrained(snake_case			)
          __UpperCamelCase											=  CpmAntTokenizer.from_pretrained(snake_case			)
          __UpperCamelCase											=  tokenizer(snake_case	,    return_tensors='''pt'''			)
          __UpperCamelCase											=  model(**snake_case			).last_hidden_state

          __UpperCamelCase											=  torch.tensor(
              [[[6.17_08, 5.92_44, 1.08_35], [6.52_07, 6.28_93, -11.33_24], [-1.01_07, -0.05_76, -5.95_77]]]	,    )
          self.assertTrue(torch.allclose(hidden_states[:, :3, :3]	,    snake_case	,    atol=1E-2			)			)







@require_torch
class 	_lowerCamelCase  (     unittest.TestCase				):

      """simple docstring"""

      @tooslow
      def  snake_case		(       self   :	Optional[int]			):
          __UpperCamelCase											=  '''今天天气真好！'''
          __UpperCamelCase											=  '''openbmb/cpm-ant-10b'''
          __UpperCamelCase											=  CpmAntForCausalLM.from_pretrained(snake_case			)
          __UpperCamelCase											=  CpmAntTokenizer.from_pretrained(snake_case			)
          __UpperCamelCase											=  tokenizer(snake_case	,    return_tensors='''pt'''			)
          __UpperCamelCase											=  model(**snake_case			).logits

          __UpperCamelCase											=  torch.tensor(
              [[[-6.42_67, -6.40_83, -6.39_58], [-5.88_02, -5.94_47, -5.78_11], [-5.38_96, -5.48_20, -5.42_95]]]	,    )
          self.assertTrue(torch.allclose(hidden_states[:, :3, :3]	,    snake_case	,    atol=1E-2			)			)







      @tooslow
      def  snake_case		(       self   :	List[str]			):
          __UpperCamelCase											=  '''openbmb/cpm-ant-10b'''
          __UpperCamelCase											=  CpmAntForCausalLM.from_pretrained(snake_case			)
          __UpperCamelCase											=  CpmAntTokenizer.from_pretrained(snake_case			)
          __UpperCamelCase											=  '''今天天气不错，'''
          __UpperCamelCase											=  '''今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的'''
          __UpperCamelCase											=  tokenizer(snake_case	,    return_tensors='''pt'''			)
          __UpperCamelCase											=  model.generate(**snake_case			)
          __UpperCamelCase											=  tokenizer.batch_decode(snake_case			)
          self.assertEqual(snake_case	,    snake_case			)







      @tooslow
      def  snake_case		(       self   :	Any			):
          __UpperCamelCase											=  '''openbmb/cpm-ant-10b'''
          __UpperCamelCase											=  CpmAntForCausalLM.from_pretrained(snake_case			)
          __UpperCamelCase											=  CpmAntTokenizer.from_pretrained(snake_case			)
          __UpperCamelCase											=  ['''今天天气不错，''', '''新年快乐，万事如意！''']
          __UpperCamelCase											=  [
              '''今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的''',
              '''新年快乐，万事如意！在这辞旧迎新的美好时刻，我谨代表《农村新技术》杂志社全体同仁，向一直以来关心、支持《农村新技术》杂志发展的各级领导、各界朋友和广大读者致以最诚挚的''',
          ]
          __UpperCamelCase											=  tokenizer(snake_case	,    return_tensors='''pt'''	,    padding=snake_case			)
          __UpperCamelCase											=  model.generate(**snake_case			)
          __UpperCamelCase											=  tokenizer.batch_decode(snake_case			)
          self.assertEqual(snake_case	,    snake_case			)






