

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


lowerCamelCase 	=   logging.get_logger(__name__)







class 		__magic_name__     (				nn.Module				):

   '''simple docstring'''




   lowerCamelCase__					:      HybridCLIPConfig
   lowerCamelCase__					:      jnp.dtype        =  jnp.floataa



   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     a__    =self.config.text_config
     a__    =self.config.vision_config

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

     a__    =FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
     a__    =FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__,       lowercase_			).module_class

     a__    =text_module(lowercase_,       dtype=self.dtype			)
     a__    =vision_module(lowercase_,       dtype=self.dtype			)

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



   def __call__(							self,       lowercase_=None,       lowercase_=None,       lowercase_=None,       lowercase_=None,       lowercase_=None,       lowercase_ = True,       lowercase_=None,       lowercase_=None,       lowercase_=None,       )				->       Union[str, Any]:






     """simple docstring"""






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

     a__    =self.vision_model(
         pixel_values=lowercase_,       deterministic=lowercase_,       output_attentions=lowercase_,       output_hidden_states=lowercase_,       return_dict=lowercase_,       )

     a__    =self.text_model(
         input_ids=lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_,       position_ids=lowercase_,       deterministic=lowercase_,       output_attentions=lowercase_,       output_hidden_states=lowercase_,       return_dict=lowercase_,       )

     a__    =vision_outputs[1]
     a__    =self.visual_projection(lowercase_			)

     a__    =text_outputs[1]
     a__    =self.text_projection(lowercase_			)

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

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







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




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



   def __init__(							self,       lowercase_,       lowercase_ = None,       lowercase_ = 0,       lowercase_ = jnp.floataa,       **lowercase_,       )				->       Tuple:






     """simple docstring"""






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

     a__    =self.module_class(config=lowercase_,       dtype=lowercase_,       **lowercase_			)
     super().__init__(lowercase_,       lowercase_,       input_shape=lowercase_,       seed=lowercase_,       dtype=lowercase_			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_ = None			)				->       FrozenDict:






     """simple docstring"""






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

     a__    =jax.random.normal(lowercase_,       input_shape[1]			)

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

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



   def __call__(							self,       lowercase_,       lowercase_,       lowercase_=None,       lowercase_=None,       lowercase_=None,       lowercase_ = None,       lowercase_ = None,       lowercase_ = False,       lowercase_ = None,       lowercase_ = None,       lowercase_ = None,       )				->       List[str]:






     """simple docstring"""






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

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

     if token_type_ids is None:
       a__    =jnp.zeros_like(lowercase_			)

     if attention_mask is None:
       a__    =jnp.ones_like(lowercase_			)

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

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



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_=None,       lowercase_=None,       lowercase_=None,       lowercase_ = None,       lowercase_ = None,       lowercase_=False,       )				->       Optional[int]:






     """simple docstring"""






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

     if token_type_ids is None:
       a__    =jnp.zeros_like(lowercase_			)

     if attention_mask is None:
       a__    =jnp.ones_like(lowercase_			)

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

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

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



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_ = None,       lowercase_ = None,       lowercase_=False			)				->       Union[str, Any]:






     """simple docstring"""






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

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

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



   @classmethod
   def 					_UpperCAmelCase					(							cls,       lowercase_ = None,       lowercase_ = None,       *lowercase_,       **lowercase_,       )				->       FlaxPreTrainedModel:






     """simple docstring"""






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

     a__    ={
         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
     a__    =kwargs_text.pop('''model''',       lowercase_			)
     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

         a__    =AutoConfig.from_pretrained(lowercase_			)
         a__    =text_config

       a__    =FlaxAutoModel.from_pretrained(lowercase_,       *lowercase_,       **lowercase_			)

     a__    =kwargs_vision.pop('''model''',       lowercase_			)
     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

         a__    =AutoConfig.from_pretrained(lowercase_			)
         a__    =vision_config

       a__    =FlaxAutoModel.from_pretrained(lowercase_,       *lowercase_,       **lowercase_			)

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

     # init model
     a__    =cls(lowercase_,       *lowercase_,       dtype=lowercase_,       **lowercase_			)

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

     a__    =text_model.params

     return model



import unittest

from transformers import MarkupLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import cached_property

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 (
           MarkupLMForQuestionAnswering,
           MarkupLMForSequenceClassification,
           MarkupLMForTokenClassification,
           MarkupLMModel,
       )

# TODO check dependencies
from transformers import MarkupLMFeatureExtractor, MarkupLMProcessor, MarkupLMTokenizer







class 		__magic_name__     :

   '''simple docstring'''




   def __init__(							self,       lowercase_,       lowercase_=13,       lowercase_=7,       lowercase_=True,       lowercase_=True,       lowercase_=True,       lowercase_=True,       lowercase_=99,       lowercase_=32,       lowercase_=5,       lowercase_=4,       lowercase_=37,       lowercase_="gelu",       lowercase_=0.1,       lowercase_=0.1,       lowercase_=512,       lowercase_=16,       lowercase_=2,       lowercase_=0.02,       lowercase_=3,       lowercase_=None,       lowercase_=20,       lowercase_=30,       lowercase_=2,       lowercase_=2,       lowercase_=10,       )				->       Dict:






     """simple docstring"""






     a__    =parent
     a__    =batch_size
     a__    =seq_length
     a__    =is_training
     a__    =use_input_mask
     a__    =use_token_type_ids
     a__    =use_labels
     a__    =vocab_size
     a__    =hidden_size
     a__    =num_hidden_layers
     a__    =num_attention_heads
     a__    =intermediate_size
     a__    =hidden_act
     a__    =hidden_dropout_prob
     a__    =attention_probs_dropout_prob
     a__    =max_position_embeddings
     a__    =type_vocab_size
     a__    =type_sequence_label_size
     a__    =initializer_range
     a__    =num_labels
     a__    =scope
     a__    =max_xpath_tag_unit_embeddings
     a__    =max_xpath_subs_unit_embeddings
     a__    =tag_pad_id
     a__    =subs_pad_id
     a__    =max_depth



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






     """simple docstring"""






     a__    =ids_tensor([self.batch_size, self.seq_length],       self.vocab_size			)

     a__    =ids_tensor(
         [self.batch_size, self.seq_length, self.max_depth],       self.max_xpath_tag_unit_embeddings			)

     a__    =ids_tensor(
         [self.batch_size, self.seq_length, self.max_depth],       self.max_xpath_subs_unit_embeddings			)

     a__    =None
     if self.use_input_mask:
       a__    =ids_tensor([self.batch_size, self.seq_length],       vocab_size=2			)

     a__    =None
     if self.use_token_type_ids:
       a__    =ids_tensor([self.batch_size, self.seq_length],       self.type_vocab_size			)

     a__    =None
     a__    =None
     if self.use_labels:
       a__    =ids_tensor([self.batch_size],       self.type_sequence_label_size			)
       a__    =ids_tensor([self.batch_size, self.seq_length],       self.num_labels			)

     a__    =self.get_config()

     return (
         config,
         input_ids,
         xpath_tags_seq,
         xpath_subs_seq,
         token_type_ids,
         input_mask,
         sequence_labels,
         token_labels,
     )



   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     return MarkupLMConfig(
         vocab_size=self.vocab_size,       hidden_size=self.hidden_size,       num_hidden_layers=self.num_hidden_layers,       num_attention_heads=self.num_attention_heads,       intermediate_size=self.intermediate_size,       hidden_act=self.hidden_act,       hidden_dropout_prob=self.hidden_dropout_prob,       attention_probs_dropout_prob=self.attention_probs_dropout_prob,       max_position_embeddings=self.max_position_embeddings,       type_vocab_size=self.type_vocab_size,       initializer_range=self.initializer_range,       max_xpath_tag_unit_embeddings=self.max_xpath_tag_unit_embeddings,       max_xpath_subs_unit_embeddings=self.max_xpath_subs_unit_embeddings,       tag_pad_id=self.tag_pad_id,       subs_pad_id=self.subs_pad_id,       max_depth=self.max_depth,       )



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       )				->       Tuple:






     """simple docstring"""






     a__    =MarkupLMModel(config=lowercase_			)
     model.to(lowercase_			)
     model.eval()
     print('''Configs:''',       model.config.tag_pad_id,       model.config.subs_pad_id			)
     a__    =model(lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_			)
     a__    =model(lowercase_,       token_type_ids=lowercase_			)
     a__    =model(lowercase_			)
     self.parent.assertEqual(result.last_hidden_state.shape,       (self.batch_size, self.seq_length, self.hidden_size)			)
     self.parent.assertEqual(result.pooler_output.shape,       (self.batch_size, self.hidden_size)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       )				->       Tuple:






     """simple docstring"""






     a__    =self.num_labels
     a__    =MarkupLMForSequenceClassification(lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(
         lowercase_,       xpath_tags_seq=lowercase_,       xpath_subs_seq=lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_,       labels=lowercase_,       )
     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.num_labels)			)



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






     """simple docstring"""






     a__    =self.num_labels
     a__    =MarkupLMForTokenClassification(config=lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(
         lowercase_,       xpath_tags_seq=lowercase_,       xpath_subs_seq=lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_,       labels=lowercase_,       )
     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.seq_length, self.num_labels)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       )				->       Any:






     """simple docstring"""






     a__    =MarkupLMForQuestionAnswering(config=lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(
         lowercase_,       xpath_tags_seq=lowercase_,       xpath_subs_seq=lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_,       start_positions=lowercase_,       end_positions=lowercase_,       )
     self.parent.assertEqual(result.start_logits.shape,       (self.batch_size, self.seq_length)			)
     self.parent.assertEqual(result.end_logits.shape,       (self.batch_size, self.seq_length)			)



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






     """simple docstring"""






     a__    =self.prepare_config_and_inputs()
     (
         (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), 
     )    =config_and_inputs
     a__    ={
         '''input_ids''': input_ids,
         '''xpath_tags_seq''': xpath_tags_seq,
         '''xpath_subs_seq''': xpath_subs_seq,
         '''token_type_ids''': token_type_ids,
         '''attention_mask''': input_mask,
     }
     return config, inputs_dict







@require_torch
class 		__magic_name__     (				lowerCamelCase__						,			lowerCamelCase__						,			unittest.TestCase				):

   '''simple docstring'''




   lowerCamelCase__					:      Tuple        =  (
       (
           MarkupLMModel,
           MarkupLMForSequenceClassification,
           MarkupLMForTokenClassification,
           MarkupLMForQuestionAnswering,
       )
       if is_torch_available()
       else None
   )
   lowerCamelCase__					:      Union[str, Any]        =  (
       {
           'feature-extraction': MarkupLMModel,
           'question-answering': MarkupLMForQuestionAnswering,
           'text-classification': MarkupLMForSequenceClassification,
           'token-classification': MarkupLMForTokenClassification,
           'zero-shot': MarkupLMForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )



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






     """simple docstring"""






     return True



   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     a__    =MarkupLMModelTester(self			)
     a__    =ConfigTester(self,       config_class=lowercase_,       hidden_size=37			)



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






     """simple docstring"""






     self.config_tester.run_common_tests()



   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_			)



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






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*lowercase_			)





def      UpperCAmelCase__							(     ):





  '''simple docstring'''







  a__    ='''
    <!DOCTYPE html>
    <html>
    <head>
    <title>Page Title</title>
    </head>
    <body>

    <h1>This is a Heading</h1>
    <p>This is a paragraph.</p>

    </body>
    </html>
    '''

  return html_string







@require_torch
class 		__magic_name__     (				unittest.TestCase				):

   '''simple docstring'''




   @cached_property
   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     a__    =MarkupLMFeatureExtractor()
     a__    =MarkupLMTokenizer.from_pretrained('''microsoft/markuplm-base'''			)

     return MarkupLMProcessor(lowercase_,       lowercase_			)



   @slow
   def 					_UpperCAmelCase					(							self			)				->       Optional[Any]:






     """simple docstring"""






     a__    =MarkupLMModel.from_pretrained('''microsoft/markuplm-base'''			).to(lowercase_			)

     a__    =self.default_processor

     a__    =processor(prepare_html_string(),       return_tensors='''pt'''			)
     a__    =inputs.to(lowercase_			)

     # forward pass
     with torch.no_grad():
       a__    =model(**lowercase_			)

     # verify the last hidden states
     a__    =torch.Size([1, 14, 768]			)
     self.assertEqual(outputs.last_hidden_state.shape,       lowercase_			)

     a__    =torch.tensor(
         [[0.0675, -0.0052, 0.5001], [-0.2281, 0.0802, 0.2192], [-0.0583, -0.3311, 0.1185]]			).to(lowercase_			)

     self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3],       lowercase_,       atol=1E-4			)			)
