





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


__lowercase							:Tuple							=	logging.get_logger(__name__)



class  _a  (					nn.Module				):






							"""simple docstring"""

							snake_case_							=    42
							snake_case_							=    jnp.floataa

							def   A_	(   self					:  Tuple     )						->Tuple:
													SCREAMING_SNAKE_CASE__  :					str =       self.config.text_config
													SCREAMING_SNAKE_CASE__  :					Any =       self.config.vision_config

													SCREAMING_SNAKE_CASE__  :					List[str] =       self.config.projection_dim
													SCREAMING_SNAKE_CASE__  :					Optional[int] =       text_config.hidden_size
													SCREAMING_SNAKE_CASE__  :					Optional[int] =       vision_config.hidden_size

													SCREAMING_SNAKE_CASE__  :					List[str] =       FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
													SCREAMING_SNAKE_CASE__  :					Dict =       FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__				,						a     ).module_class

													SCREAMING_SNAKE_CASE__  :					Any =       text_module(a				,						dtype=self.dtype     )
													SCREAMING_SNAKE_CASE__  :					Union[str, Any] =       vision_module(a				,						dtype=self.dtype     )

													SCREAMING_SNAKE_CASE__  :					Tuple =       nn.Dense(
													    self.projection_dim				,						dtype=self.dtype				,						kernel_init=jax.nn.initializers.normal(0.02     )				,						use_bias=a				,						)
													SCREAMING_SNAKE_CASE__  :					Optional[Any] =       nn.Dense(
													    self.projection_dim				,						dtype=self.dtype				,						kernel_init=jax.nn.initializers.normal(0.02     )				,						use_bias=a				,						)
													SCREAMING_SNAKE_CASE__  :					Union[str, Any] =       self.param("logit_scale"				,						jax.nn.initializers.ones				,						[]     )






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

													SCREAMING_SNAKE_CASE__  :					str =       self.vision_model(
													    pixel_values=a				,						deterministic=a				,						output_attentions=a				,						output_hidden_states=a				,						return_dict=a				,						)

													SCREAMING_SNAKE_CASE__  :					List[str] =       self.text_model(
													    input_ids=a				,						attention_mask=a				,						token_type_ids=a				,						position_ids=a				,						deterministic=a				,						output_attentions=a				,						output_hidden_states=a				,						return_dict=a				,						)

													SCREAMING_SNAKE_CASE__  :					Tuple =       vision_outputs[1]
													SCREAMING_SNAKE_CASE__  :					Tuple =       self.visual_projection(a     )

													SCREAMING_SNAKE_CASE__  :					Tuple =       text_outputs[1]
													SCREAMING_SNAKE_CASE__  :					List[Any] =       self.text_projection(a     )

													# normalized features
													SCREAMING_SNAKE_CASE__  :					Optional[Any] =       image_embeds / jnp.linalg.norm(a				,						axis=-1				,						keepdims=a     )
													SCREAMING_SNAKE_CASE__  :					List[str] =       text_embeds / jnp.linalg.norm(a				,						axis=-1				,						keepdims=a     )

													# cosine similarity as logits
													SCREAMING_SNAKE_CASE__  :					str =       jnp.exp(self.logit_scale     )
													SCREAMING_SNAKE_CASE__  :					str =       jnp.matmul(a				,						image_embeds.T     ) * logit_scale
													SCREAMING_SNAKE_CASE__  :					List[Any] =       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=a				,						logits_per_text=a				,						text_embeds=a				,						image_embeds=a				,						text_model_output=a				,						vision_model_output=a				,						)







class  _a  (					lowercase__				):






							"""simple docstring"""

							snake_case_							=    HybridCLIPConfig
							snake_case_							=    FlaxHybridCLIPModule

							def __init__(   self					:  Optional[int]				,						a					:  HybridCLIPConfig				,						a					:  Optional[Tuple] = None				,						a					:  int = 0				,						a					:  jnp.dtype = jnp.floataa				,						**a					:  Dict				,						)						->int:
													if input_shape is None:
																			SCREAMING_SNAKE_CASE__  :					str =       ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

													SCREAMING_SNAKE_CASE__  :					Union[str, Any] =       self.module_class(config=a				,						dtype=a				,						**a     )
													super().__init__(a				,						a				,						input_shape=a				,						seed=a				,						dtype=a     )

							def   A_	(   self					:  Union[str, Any]				,						a					:  jax.random.PRNGKey				,						a					:  Tuple				,						a					:  FrozenDict = None     )						->FrozenDict:
													# init input tensor
													SCREAMING_SNAKE_CASE__  :					Dict =       jnp.zeros(input_shape[0]				,						dtype="i4"     )
													SCREAMING_SNAKE_CASE__  :					Union[str, Any] =       jnp.broadcast_to(jnp.arange(jnp.atleast_ad(a     ).shape[-1]     )				,						input_shape[0]     )
													SCREAMING_SNAKE_CASE__  :					Dict =       jnp.ones_like(a     )
													SCREAMING_SNAKE_CASE__  :					Dict =       jnp.ones_like(a     )

													SCREAMING_SNAKE_CASE__  :					Union[str, Any] =       jax.random.normal(a				,						input_shape[1]     )

													SCREAMING_SNAKE_CASE__,				SCREAMING_SNAKE_CASE__  :					Tuple =       jax.random.split(a     )
													SCREAMING_SNAKE_CASE__  :					int =       {"params": params_rng, "dropout": dropout_rng}

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

							def __call__(   self					:  Optional[Any]				,						a					:  int				,						a					:  Optional[int]				,						a					:  int=None				,						a					:  List[str]=None				,						a					:  str=None				,						a					:  dict = None				,						a					:  jax.random.PRNGKey = None				,						a					:  bool = False				,						a					:  Optional[bool] = None				,						a					:  Optional[bool] = None				,						a					:  Optional[bool] = None				,						)						->List[Any]:
													SCREAMING_SNAKE_CASE__  :					str =       output_attentions if output_attentions is not None else self.config.output_attentions
													SCREAMING_SNAKE_CASE__  :					Dict =       (
													    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
													)
													SCREAMING_SNAKE_CASE__  :					Tuple =       return_dict if return_dict is not None else self.config.return_dict

													if position_ids is None:
																			SCREAMING_SNAKE_CASE__  :					Optional[int] =       jnp.broadcast_to(jnp.arange(jnp.atleast_ad(a     ).shape[-1]     )				,						input_ids.shape     )

													if token_type_ids is None:
																			SCREAMING_SNAKE_CASE__  :					Tuple =       jnp.zeros_like(a     )

													if attention_mask is None:
																			SCREAMING_SNAKE_CASE__  :					Optional[int] =       jnp.ones_like(a     )

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

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

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

													if token_type_ids is None:
																			SCREAMING_SNAKE_CASE__  :					List[str] =       jnp.zeros_like(a     )

													if attention_mask is None:
																			SCREAMING_SNAKE_CASE__  :					Any =       jnp.ones_like(a     )

													# Handle any PRNG if needed
													SCREAMING_SNAKE_CASE__  :					Any =       {}
													if dropout_rng is not None:
																			SCREAMING_SNAKE_CASE__  :					int =       dropout_rng

													def _get_features(a					:  List[Any]				,						a					:  List[Any]				,						a					:  Optional[Any]				,						a					:  Dict				,						a					:  Dict				,						a					:  Dict     ):
																			SCREAMING_SNAKE_CASE__  :					Optional[Any] =       module.text_model(
																			    input_ids=a				,						attention_mask=a				,						position_ids=a				,						token_type_ids=a				,						deterministic=a				,						)
																			SCREAMING_SNAKE_CASE__  :					List[Any] =       text_outputs[1]
																			SCREAMING_SNAKE_CASE__  :					Optional[Any] =       module.text_projection(a     )
																			return text_features

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

							def   A_	(   self					:  Union[str, Any]				,						a					:  Optional[int]				,						a					:  dict = None				,						a					:  jax.random.PRNGKey = None				,						a					:  Dict=False     )						->int:
													SCREAMING_SNAKE_CASE__  :					Optional[int] =       {}
													if dropout_rng is not None:
																			SCREAMING_SNAKE_CASE__  :					int =       dropout_rng

													def _get_features(a					:  Optional[Any]				,						a					:  List[Any]				,						a					:  List[Any]     ):
																			SCREAMING_SNAKE_CASE__  :					Tuple =       module.vision_model(pixel_values=a				,						deterministic=a     )
																			SCREAMING_SNAKE_CASE__  :					Optional[Any] =       vision_outputs[1]  # pooled_output
																			SCREAMING_SNAKE_CASE__  :					List[Any] =       module.visual_projection(a     )
																			return image_features

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






							@classmethod
							def   A_	(   cls					:  Optional[Any]				,						a					:  str = None				,						a					:  str = None				,						*a					:  Optional[int]				,						**a					:  int				,						)						->FlaxPreTrainedModel:
													SCREAMING_SNAKE_CASE__  :					List[Any] =       {
													    argument[len("text_"     ) :]: value for argument, value in kwargs.items() if argument.startswith("text_"     )
													}

													SCREAMING_SNAKE_CASE__  :					Optional[int] =       {
													    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
													SCREAMING_SNAKE_CASE__  :					Dict =       kwargs_text.pop("model"				,						a     )
													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

																									SCREAMING_SNAKE_CASE__  :					List[str] =       AutoConfig.from_pretrained(a     )
																									SCREAMING_SNAKE_CASE__  :					List[Any] =       text_config

																			SCREAMING_SNAKE_CASE__  :					Optional[Any] =       FlaxAutoModel.from_pretrained(a				,						*a				,						**a     )

													SCREAMING_SNAKE_CASE__  :					Any =       kwargs_vision.pop("model"				,						a     )
													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

																									SCREAMING_SNAKE_CASE__  :					List[Any] =       AutoConfig.from_pretrained(a     )
																									SCREAMING_SNAKE_CASE__  :					Optional[int] =       vision_config

																			SCREAMING_SNAKE_CASE__  :					Dict =       FlaxAutoModel.from_pretrained(a				,						*a				,						**a     )

													# instantiate config with corresponding kwargs
													SCREAMING_SNAKE_CASE__  :					int =       kwargs.pop("dtype"				,						jnp.floataa     )
													SCREAMING_SNAKE_CASE__  :					List[str] =       HybridCLIPConfig.from_text_vision_configs(text_model.config				,						vision_model.config				,						**a     )

													# init model
													SCREAMING_SNAKE_CASE__  :					int =       cls(a				,						*a				,						dtype=a				,						**a     )

													if vision_config.model_type == "clip":
																			SCREAMING_SNAKE_CASE__  :					str =       vision_model.params["vision_model"]
																			SCREAMING_SNAKE_CASE__  :					List[str] =       vision_model.params["visual_projection"]["kernel"]
													else:
																			SCREAMING_SNAKE_CASE__  :					int =       vision_model.params

													SCREAMING_SNAKE_CASE__  :					str =       text_model.params

													return model





from __future__ import annotations

import unittest

from transformers import AlbertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import TF_MODEL_FOR_PRETRAINING_MAPPING
  from transformers.models.albert.modeling_tf_albert import (
      TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFAlbertForMaskedLM,
      TFAlbertForMultipleChoice,
      TFAlbertForPreTraining,
      TFAlbertForQuestionAnswering,
      TFAlbertForSequenceClassification,
      TFAlbertForTokenClassification,
      TFAlbertModel,
  )
class       _UpperCamelCase   :
      def __init__(      self :Optional[int]			,   lowerCamelCase :Tuple			,   lowerCamelCase :Union[str, Any]=13			,   lowerCamelCase :List[str]=7			,   lowerCamelCase :List[Any]=True			,   lowerCamelCase :Any=True			,   lowerCamelCase :Any=True			,   lowerCamelCase :Optional[Any]=True			,   lowerCamelCase :Optional[int]=99			,   lowerCamelCase :Tuple=16			,   lowerCamelCase :Optional[Any]=32			,   lowerCamelCase :Optional[Any]=2			,   lowerCamelCase :List[Any]=4			,   lowerCamelCase :Dict=37			,   lowerCamelCase :Optional[Any]="gelu"			,   lowerCamelCase :Dict=0.1			,   lowerCamelCase :Optional[Any]=0.1			,   lowerCamelCase :Tuple=512			,   lowerCamelCase :Optional[int]=16			,   lowerCamelCase :Any=2			,   lowerCamelCase :Tuple=0.02			,   lowerCamelCase :Union[str, Any]=3			,   lowerCamelCase :str=4			,   lowerCamelCase :Dict=None			,   )  ->      List[Any]:
       UpperCAmelCase__							=    parent
       UpperCAmelCase__							=    13
       UpperCAmelCase__							=    7
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    99
       UpperCAmelCase__							=    16
       UpperCAmelCase__							=    32
       UpperCAmelCase__							=    2
       UpperCAmelCase__							=    4
       UpperCAmelCase__							=    37
       UpperCAmelCase__							=    "gelu"
       UpperCAmelCase__							=    0.1
       UpperCAmelCase__							=    0.1
       UpperCAmelCase__							=    512
       UpperCAmelCase__							=    16
       UpperCAmelCase__							=    2
       UpperCAmelCase__							=    0.02
       UpperCAmelCase__							=    3
       UpperCAmelCase__							=    4
       UpperCAmelCase__							=    None
      def 						UpperCAmelCase_    (      self :Optional[int]   )  ->      Dict:
       UpperCAmelCase__							=    ids_tensor([self.batch_size, self.seq_length]			,   self.vocab_size   )

       UpperCAmelCase__							=    None
       if self.use_input_mask:
        UpperCAmelCase__							=    random_attention_mask([self.batch_size, self.seq_length]   )

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

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

       UpperCAmelCase__							=    AlbertConfig(
           vocab_size=self.vocab_size			,   hidden_size=self.hidden_size			,   embedding_size=self.embedding_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			,   )

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
      def 						UpperCAmelCase_    (      self :Optional[int]			,   lowerCamelCase :str			,   lowerCamelCase :Tuple			,   lowerCamelCase :Tuple			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :int			,   lowerCamelCase :str   )  ->      Optional[Any]:
       UpperCAmelCase__							=    TFAlbertModel(config=lowerCamelCase   )
       # inputs = {'input_ids': input_ids,
       #           'attention_mask': input_mask,
       #           'token_type_ids': token_type_ids}
       # sequence_output, pooled_output = model(**inputs)
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )

       UpperCAmelCase__							=    [input_ids, input_mask]
       UpperCAmelCase__							=    model(lowerCamelCase   )

       UpperCAmelCase__							=    model(lowerCamelCase   )

       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 :Optional[Any]			,   lowerCamelCase :Any			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Any			,   lowerCamelCase :List[str]			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Tuple   )  ->      Optional[int]:
       UpperCAmelCase__							=    self.num_labels
       UpperCAmelCase__							=    TFAlbertForPreTraining(config=lowerCamelCase   )
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(result.prediction_logits.shape			,   (self.batch_size, self.seq_length, self.vocab_size)   )
       self.parent.assertEqual(result.sop_logits.shape			,   (self.batch_size, self.num_labels)   )
      def 						UpperCAmelCase_    (      self :int			,   lowerCamelCase :Dict			,   lowerCamelCase :List[Any]			,   lowerCamelCase :List[str]			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :Tuple			,   lowerCamelCase :Any   )  ->      Optional[Any]:
       UpperCAmelCase__							=    TFAlbertForMaskedLM(config=lowerCamelCase   )
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.seq_length, self.vocab_size)   )
      def 						UpperCAmelCase_    (      self :Tuple			,   lowerCamelCase :str			,   lowerCamelCase :Dict			,   lowerCamelCase :Dict			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Optional[Any]   )  ->      str:
       UpperCAmelCase__							=    self.num_labels
       UpperCAmelCase__							=    TFAlbertForSequenceClassification(config=lowerCamelCase   )
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.num_labels)   )
      def 						UpperCAmelCase_    (      self :Dict			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :str			,   lowerCamelCase :Dict			,   lowerCamelCase :List[Any]			,   lowerCamelCase :List[str]			,   lowerCamelCase :int			,   lowerCamelCase :Optional[Any]   )  ->      Dict:
       UpperCAmelCase__							=    TFAlbertForQuestionAnswering(config=lowerCamelCase   )
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )
       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 :Optional[int]			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :Dict			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Tuple			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Union[str, Any]   )  ->      int:
       UpperCAmelCase__							=    self.num_choices
       UpperCAmelCase__							=    TFAlbertForMultipleChoice(config=lowerCamelCase   )
       UpperCAmelCase__							=    tf.tile(tf.expand_dims(lowerCamelCase			,   1   )			,   (1, self.num_choices, 1)   )
       UpperCAmelCase__							=    tf.tile(tf.expand_dims(lowerCamelCase			,   1   )			,   (1, self.num_choices, 1)   )
       UpperCAmelCase__							=    tf.tile(tf.expand_dims(lowerCamelCase			,   1   )			,   (1, self.num_choices, 1)   )
       UpperCAmelCase__							=    {
           "input_ids": multiple_choice_inputs_ids,
           "attention_mask": multiple_choice_input_mask,
           "token_type_ids": multiple_choice_token_type_ids,
       }
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertListEqual(list(result["logits"].shape   )			,   [self.batch_size, self.num_choices]   )
      def 						UpperCAmelCase_    (      self :int			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :List[str]			,   lowerCamelCase :int			,   lowerCamelCase :int			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Dict			,   lowerCamelCase :Dict   )  ->      List[str]:
       UpperCAmelCase__							=    self.num_labels
       UpperCAmelCase__							=    TFAlbertForTokenClassification(config=lowerCamelCase   )
       UpperCAmelCase__							=    {
           "input_ids": input_ids,
           "attention_mask": input_mask,
           "token_type_ids": token_type_ids,
       }
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertListEqual(list(result["logits"].shape   )			,   [self.batch_size, self.seq_length, self.num_labels]   )







      def 						UpperCAmelCase_    (      self :str   )  ->      Any:
       UpperCAmelCase__							=    self.prepare_config_and_inputs()
       (
           (
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			
       )							=    config_and_inputs
       UpperCAmelCase__							=    {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
       return config, inputs_dict



@require_tf
class       _UpperCamelCase   (					lowerCAmelCase						,							lowerCAmelCase						,							unittest.TestCase		):
      UpperCAmelCase_          =		(
          (
              TFAlbertModel,
              TFAlbertForPreTraining,
              TFAlbertForMaskedLM,
              TFAlbertForSequenceClassification,
              TFAlbertForQuestionAnswering,
              TFAlbertForTokenClassification,
              TFAlbertForMultipleChoice,
          )
          if is_tf_available()
          else ()
      )
      UpperCAmelCase_          =		(
          {
              """feature-extraction""": TFAlbertModel,
              """fill-mask""": TFAlbertForMaskedLM,
              """question-answering""": TFAlbertForQuestionAnswering,
              """text-classification""": TFAlbertForSequenceClassification,
              """token-classification""": TFAlbertForTokenClassification,
              """zero-shot""": TFAlbertForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      UpperCAmelCase_          =		False
      UpperCAmelCase_          =		False
      def 						UpperCAmelCase_    (      self :str			,   lowerCamelCase :Any			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :Optional[int]=False   )  ->      str:
       UpperCAmelCase__							=    super()._prepare_for_class(lowerCamelCase			,   lowerCamelCase			,   return_labels=lowerCamelCase   )

       if return_labels:
        if model_class in get_values(lowerCamelCase   ):
         UpperCAmelCase__							=    tf.zeros(self.model_tester.batch_size			,   dtype=tf.intaa   )

       return inputs_dict
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      List[Any]:
       UpperCAmelCase__							=    TFAlbertModelTester(self   )
       UpperCAmelCase__							=    ConfigTester(self			,   config_class=lowerCamelCase			,   hidden_size=37   )
      def 						UpperCAmelCase_    (      self :Dict   )  ->      Union[str, Any]:
       self.config_tester.run_common_tests()
      def 						UpperCAmelCase_    (      self :Optional[Any]   )  ->      Tuple:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_model(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      Dict:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_pretraining(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      str:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_masked_lm(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Dict   )  ->      int:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_multiple_choice(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Union[str, Any]   )  ->      Dict:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_sequence_classification(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :str   )  ->      int:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_question_answering(*lowerCamelCase   )







      @slow
      def 						UpperCAmelCase_    (      self :int   )  ->      str:
       for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase__							=    TFAlbertModel.from_pretrained(lowerCamelCase   )
        self.assertIsNotNone(lowerCamelCase   )



@require_tf
class       _UpperCamelCase   (					unittest.TestCase		):
      @slow
      def 						UpperCAmelCase_    (      self :List[str]   )  ->      List[str]:
       UpperCAmelCase__							=    TFAlbertForPreTraining.from_pretrained("albert-base-v2"   )
       UpperCAmelCase__							=    tf.constant([[0, 1, 2, 3, 4, 5]]   )
       UpperCAmelCase__							=    model(lowerCamelCase   )[0]

       UpperCAmelCase__							=    [1, 6, 3_0000]
       self.assertEqual(output.shape			,   lowerCamelCase   )

       UpperCAmelCase__							=    tf.constant(
           [
               [
                   [4.59_56_68, 0.74_46_27_54, -1.81_81_47],
                   [4.5_95_43_47, 0.7_45_41_84, -1.8_18_82_58],
                   [4.5_95_49_05, 0.7_44_82_35, -1.8_18_23_16],
               ]
           ]   )
       tf.debugging.assert_near(output[:, :3, :3]			,   lowerCamelCase			,   atol=1e-4   )





