
"""simple docstring"""







from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings







@maybe_allow_in_graph
class 							__magic_name__     (		nn.Module   ):





					def __init__(				self   :							str							,						snake_case_   :							int							,						snake_case_   :							int							,						snake_case_   :							int							,						snake_case_   :							Any=0.0							,						snake_case_   :							Optional[int] = None							,						snake_case_   :							str = "geglu"							,						snake_case_   :							Optional[int] = None							,						snake_case_   :							bool = False							,						snake_case_   :							bool = False							,						snake_case_   :							bool = False							,						snake_case_   :							bool = False							,						snake_case_   :							bool = True							,						snake_case_   :							str = "layer_norm"							,						snake_case_   :							bool = False							,						):
												super().__init__()
												__snake_case     			=       only_cross_attention

												__snake_case     			=       (num_embeds_ada_norm is not None) and norm_type == '''ada_norm_zero'''
												__snake_case     			=       (num_embeds_ada_norm is not None) and norm_type == '''ada_norm'''

												if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
																			raise ValueError(
																			    F'''`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to'''
																			    F''' define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.'''    )

												# Define 3 blocks. Each block has its own normalization layer.
												# 1. Self-Attn
												if self.use_ada_layer_norm:
																			__snake_case     			=       AdaLayerNorm(UpperCAmelCase__							,						UpperCAmelCase__    )
												elif self.use_ada_layer_norm_zero:
																			__snake_case     			=       AdaLayerNormZero(UpperCAmelCase__							,						UpperCAmelCase__    )
												else:
																			__snake_case     			=       nn.LayerNorm(UpperCAmelCase__							,						elementwise_affine=UpperCAmelCase__    )
												__snake_case     			=       Attention(
												    query_dim=UpperCAmelCase__							,						heads=UpperCAmelCase__							,						dim_head=UpperCAmelCase__							,						dropout=UpperCAmelCase__							,						bias=UpperCAmelCase__							,						cross_attention_dim=cross_attention_dim if only_cross_attention else None							,						upcast_attention=UpperCAmelCase__							,						)

												# 2. Cross-Attn
												if cross_attention_dim is not None or double_self_attention:
																			# We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
																			# I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
																			# the second cross attention block.
																			__snake_case     			=       (
																			    AdaLayerNorm(UpperCAmelCase__							,						UpperCAmelCase__    )
																			    if self.use_ada_layer_norm
																			    else nn.LayerNorm(UpperCAmelCase__							,						elementwise_affine=UpperCAmelCase__    )
																			)
																			__snake_case     			=       Attention(
																			    query_dim=UpperCAmelCase__							,						cross_attention_dim=cross_attention_dim if not double_self_attention else None							,						heads=UpperCAmelCase__							,						dim_head=UpperCAmelCase__							,						dropout=UpperCAmelCase__							,						bias=UpperCAmelCase__							,						upcast_attention=UpperCAmelCase__							,						)  # is self-attn if encoder_hidden_states is none
												else:
																			__snake_case     			=       None
																			__snake_case     			=       None

												# 3. Feed-forward
												__snake_case     			=       nn.LayerNorm(UpperCAmelCase__							,						elementwise_affine=UpperCAmelCase__    )
												__snake_case     			=       FeedForward(UpperCAmelCase__							,						dropout=UpperCAmelCase__							,						activation_fn=UpperCAmelCase__							,						final_dropout=UpperCAmelCase__    )

												# let chunk size default to None
												__snake_case     			=       None
												__snake_case     			=       0





					def 							lowerCAmelCase      (				self   :							Optional[int]							,						snake_case_   :							Optional[int]							,						snake_case_   :							int    ):
												__snake_case     			=       chunk_size
												__snake_case     			=       dim





					def 							lowerCAmelCase      (				self   :							Dict							,						snake_case_   :							torch.FloatTensor							,						snake_case_   :							Optional[torch.FloatTensor] = None							,						snake_case_   :							Optional[torch.FloatTensor] = None							,						snake_case_   :							Optional[torch.FloatTensor] = None							,						snake_case_   :							Optional[torch.LongTensor] = None							,						snake_case_   :							Dict[str, Any] = None							,						snake_case_   :							Optional[torch.LongTensor] = None							,						):
												# Notice that normalization is always applied before the real computation in the following blocks.
												# 1. Self-Attention
												if self.use_ada_layer_norm:
																			__snake_case     			=       self.norma(UpperCAmelCase__							,						UpperCAmelCase__    )
												elif self.use_ada_layer_norm_zero:
																			__snake_case     			=       self.norma(
																			    UpperCAmelCase__							,						UpperCAmelCase__							,						UpperCAmelCase__							,						hidden_dtype=hidden_states.dtype    )
												else:
																			__snake_case     			=       self.norma(UpperCAmelCase__    )

												__snake_case     			=       cross_attention_kwargs if cross_attention_kwargs is not None else {}

												__snake_case     			=       self.attna(
												    UpperCAmelCase__							,						encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None							,						attention_mask=UpperCAmelCase__							,						**UpperCAmelCase__							,						)
												if self.use_ada_layer_norm_zero:
																			__snake_case     			=       gate_msa.unsqueeze(1    ) * attn_output
												__snake_case     			=       attn_output + hidden_states

												# 2. Cross-Attention
												if self.attna is not None:
																			__snake_case     			=       (
																			    self.norma(UpperCAmelCase__							,						UpperCAmelCase__    ) if self.use_ada_layer_norm else self.norma(UpperCAmelCase__    )
																			)

																			__snake_case     			=       self.attna(
																			    UpperCAmelCase__							,						encoder_hidden_states=UpperCAmelCase__							,						attention_mask=UpperCAmelCase__							,						**UpperCAmelCase__							,						)
																			__snake_case     			=       attn_output + hidden_states

												# 3. Feed-forward
												__snake_case     			=       self.norma(UpperCAmelCase__    )

												if self.use_ada_layer_norm_zero:
																			__snake_case     			=       norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

												if self._chunk_size is not None:
																			# "feed_forward_chunk_size" can be used to save memory
																			if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
																										raise ValueError(
																										    F'''`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`.'''    )

																			__snake_case     			=       norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
																			__snake_case     			=       torch.cat(
																			    [self.ff(UpperCAmelCase__    ) for hid_slice in norm_hidden_states.chunk(UpperCAmelCase__							,						dim=self._chunk_dim    )]							,						dim=self._chunk_dim							,						)
												else:
																			__snake_case     			=       self.ff(UpperCAmelCase__    )

												if self.use_ada_layer_norm_zero:
																			__snake_case     			=       gate_mlp.unsqueeze(1    ) * ff_output

												__snake_case     			=       ff_output + hidden_states

												return hidden_states







class 							__magic_name__     (		nn.Module   ):





					def __init__(				self   :							Optional[int]							,						snake_case_   :							int							,						snake_case_   :							Optional[int] = None							,						snake_case_   :							int = 4							,						snake_case_   :							float = 0.0							,						snake_case_   :							str = "geglu"							,						snake_case_   :							bool = False							,						):
												super().__init__()
												__snake_case     			=       int(dim * mult    )
												__snake_case     			=       dim_out if dim_out is not None else dim

												if activation_fn == "gelu":
																			__snake_case     			=       GELU(UpperCAmelCase__							,						UpperCAmelCase__    )
												if activation_fn == "gelu-approximate":
																			__snake_case     			=       GELU(UpperCAmelCase__							,						UpperCAmelCase__							,						approximate="tanh"    )
												elif activation_fn == "geglu":
																			__snake_case     			=       GEGLU(UpperCAmelCase__							,						UpperCAmelCase__    )
												elif activation_fn == "geglu-approximate":
																			__snake_case     			=       ApproximateGELU(UpperCAmelCase__							,						UpperCAmelCase__    )

												__snake_case     			=       nn.ModuleList([]    )
												# project in
												self.net.append(UpperCAmelCase__    )
												# project dropout
												self.net.append(nn.Dropout(UpperCAmelCase__    )    )
												# project out
												self.net.append(nn.Linear(UpperCAmelCase__							,						UpperCAmelCase__    )    )
												# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
												if final_dropout:
																			self.net.append(nn.Dropout(UpperCAmelCase__    )    )





					def 							lowerCAmelCase      (				self   :							int							,						snake_case_   :							int    ):
												for module in self.net:
																			__snake_case     			=       module(UpperCAmelCase__    )
												return hidden_states







class 							__magic_name__     (		nn.Module   ):





					def __init__(				self   :							List[str]							,						snake_case_   :							int							,						snake_case_   :							int							,						snake_case_   :							str = "none"    ):
												super().__init__()
												__snake_case     			=       nn.Linear(UpperCAmelCase__							,						UpperCAmelCase__    )
												__snake_case     			=       approximate





					def 							lowerCAmelCase      (				self   :							Optional[int]							,						snake_case_   :							Optional[int]    ):
												if gate.device.type != "mps":
																			return F.gelu(UpperCAmelCase__							,						approximate=self.approximate    )
												# mps: gelu is not implemented for float16
												return F.gelu(gate.to(dtype=torch.floataa    )							,						approximate=self.approximate    ).to(dtype=gate.dtype    )





					def 							lowerCAmelCase      (				self   :							List[Any]							,						snake_case_   :							Optional[int]    ):
												__snake_case     			=       self.proj(UpperCAmelCase__    )
												__snake_case     			=       self.gelu(UpperCAmelCase__    )
												return hidden_states







class 							__magic_name__     (		nn.Module   ):





					def __init__(				self   :							int							,						snake_case_   :							int							,						snake_case_   :							int    ):
												super().__init__()
												__snake_case     			=       nn.Linear(UpperCAmelCase__							,						dim_out * 2    )





					def 							lowerCAmelCase      (				self   :							int							,						snake_case_   :							List[Any]    ):
												if gate.device.type != "mps":
																			return F.gelu(UpperCAmelCase__    )
												# mps: gelu is not implemented for float16
												return F.gelu(gate.to(dtype=torch.floataa    )    ).to(dtype=gate.dtype    )





					def 							lowerCAmelCase      (				self   :							List[str]							,						snake_case_   :							List[str]    ):
												__snake_case     			=       self.proj(UpperCAmelCase__    ).chunk(2							,						dim=-1    )
												return hidden_states * self.gelu(UpperCAmelCase__    )







class 							__magic_name__     (		nn.Module   ):





					def __init__(				self   :							Dict							,						snake_case_   :							int							,						snake_case_   :							int    ):
												super().__init__()
												__snake_case     			=       nn.Linear(UpperCAmelCase__							,						UpperCAmelCase__    )





					def 							lowerCAmelCase      (				self   :							Tuple							,						snake_case_   :							Tuple    ):
												__snake_case     			=       self.proj(UpperCAmelCase__    )
												return x * torch.sigmoid(1.702 * x    )







class 							__magic_name__     (		nn.Module   ):





					def __init__(				self   :							List[Any]							,						snake_case_   :							Dict							,						snake_case_   :							Optional[int]    ):
												super().__init__()
												__snake_case     			=       nn.Embedding(UpperCAmelCase__							,						UpperCAmelCase__    )
												__snake_case     			=       nn.SiLU()
												__snake_case     			=       nn.Linear(UpperCAmelCase__							,						embedding_dim * 2    )
												__snake_case     			=       nn.LayerNorm(UpperCAmelCase__							,						elementwise_affine=UpperCAmelCase__    )





					def 							lowerCAmelCase      (				self   :							Any							,						snake_case_   :							int							,						snake_case_   :							Tuple    ):
												__snake_case     			=       self.linear(self.silu(self.emb(UpperCAmelCase__    )    )    )
												__snake_case     			=       torch.chunk(UpperCAmelCase__							,						2    )
												__snake_case     			=       self.norm(UpperCAmelCase__    ) * (1 + scale) + shift
												return x







class 							__magic_name__     (		nn.Module   ):





					def __init__(				self   :							int							,						snake_case_   :							List[str]							,						snake_case_   :							List[str]    ):
												super().__init__()

												__snake_case     			=       CombinedTimestepLabelEmbeddings(UpperCAmelCase__							,						UpperCAmelCase__    )

												__snake_case     			=       nn.SiLU()
												__snake_case     			=       nn.Linear(UpperCAmelCase__							,						6 * embedding_dim							,						bias=UpperCAmelCase__    )
												__snake_case     			=       nn.LayerNorm(UpperCAmelCase__							,						elementwise_affine=UpperCAmelCase__							,						eps=1e-6    )





					def 							lowerCAmelCase      (				self   :							Dict							,						snake_case_   :							Dict							,						snake_case_   :							str							,						snake_case_   :							Any							,						snake_case_   :							str=None    ):
												__snake_case     			=       self.linear(self.silu(self.emb(UpperCAmelCase__							,						UpperCAmelCase__							,						hidden_dtype=UpperCAmelCase__    )    )    )
												__snake_case     			=       emb.chunk(6							,						dim=1    )
												__snake_case     			=       self.norm(UpperCAmelCase__    ) * (1 + scale_msa[:, None]) + shift_msa[:, None]
												return x, gate_msa, shift_mlp, scale_mlp, gate_mlp







class 							__magic_name__     (		nn.Module   ):





					def __init__(				self   :							Any							,						snake_case_   :							int							,						snake_case_   :							int							,						snake_case_   :							int							,						snake_case_   :							Optional[str] = None							,						snake_case_   :							float = 1e-5    ):
												super().__init__()
												__snake_case     			=       num_groups
												__snake_case     			=       eps

												if act_fn is None:
																			__snake_case     			=       None
												else:
																			__snake_case     			=       get_activation(UpperCAmelCase__    )

												__snake_case     			=       nn.Linear(UpperCAmelCase__							,						out_dim * 2    )





					def 							lowerCAmelCase      (				self   :							Optional[Any]							,						snake_case_   :							int							,						snake_case_   :							Optional[Any]    ):
												if self.act:
																			__snake_case     			=       self.act(UpperCAmelCase__    )
												__snake_case     			=       self.linear(UpperCAmelCase__    )
												__snake_case     			=       emb[:, :, None, None]
												__snake_case     			=       emb.chunk(2							,						dim=1    )

												__snake_case     			=       F.group_norm(UpperCAmelCase__							,						self.num_groups							,						eps=self.eps    )
												__snake_case     			=       x * (1 + scale) + shift
												return x


"""simple docstring"""







import unittest

from transformers import MobileBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    MODEL_FOR_PRETRAINING_MAPPING,
		    MobileBertForMaskedLM,
		    MobileBertForMultipleChoice,
		    MobileBertForNextSentencePrediction,
		    MobileBertForPreTraining,
		    MobileBertForQuestionAnswering,
		    MobileBertForSequenceClassification,
		    MobileBertForTokenClassification,
		    MobileBertModel,
		)







class 							__magic_name__     :





					def __init__(				self   :							Optional[int]							,						snake_case_   :							List[str]							,						snake_case_   :							str=13							,						snake_case_   :							Optional[int]=7							,						snake_case_   :							Optional[Any]=True							,						snake_case_   :							Union[str, Any]=True							,						snake_case_   :							Tuple=True							,						snake_case_   :							int=True							,						snake_case_   :							int=99							,						snake_case_   :							Optional[int]=64							,						snake_case_   :							Dict=32							,						snake_case_   :							Dict=5							,						snake_case_   :							List[str]=4							,						snake_case_   :							List[Any]=37							,						snake_case_   :							Optional[Any]="gelu"							,						snake_case_   :							Optional[Any]=0.1							,						snake_case_   :							List[Any]=0.1							,						snake_case_   :							Union[str, Any]=512							,						snake_case_   :							int=16							,						snake_case_   :							List[str]=2							,						snake_case_   :							Union[str, Any]=0.02							,						snake_case_   :							Union[str, Any]=3							,						snake_case_   :							str=4							,						snake_case_   :							int=None							,						):
												__snake_case     			=       parent
												__snake_case     			=       batch_size
												__snake_case     			=       seq_length
												__snake_case     			=       is_training
												__snake_case     			=       use_input_mask
												__snake_case     			=       use_token_type_ids
												__snake_case     			=       use_labels
												__snake_case     			=       vocab_size
												__snake_case     			=       hidden_size
												__snake_case     			=       embedding_size
												__snake_case     			=       num_hidden_layers
												__snake_case     			=       num_attention_heads
												__snake_case     			=       intermediate_size
												__snake_case     			=       hidden_act
												__snake_case     			=       hidden_dropout_prob
												__snake_case     			=       attention_probs_dropout_prob
												__snake_case     			=       max_position_embeddings
												__snake_case     			=       type_vocab_size
												__snake_case     			=       type_sequence_label_size
												__snake_case     			=       initializer_range
												__snake_case     			=       num_labels
												__snake_case     			=       num_choices
												__snake_case     			=       scope





					def 							lowerCAmelCase      (				self   :							Union[str, Any]    ):
												__snake_case     			=       ids_tensor([self.batch_size, self.seq_length]							,						self.vocab_size    )

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

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

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

												__snake_case     			=       self.get_config()

												return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





					def 							lowerCAmelCase      (				self   :							Union[str, Any]    ):
												return MobileBertConfig(
												    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							,						embedding_size=self.embedding_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							,						is_decoder=snake_case_							,						initializer_range=self.initializer_range							,						)





					def 							lowerCAmelCase      (				self   :							Optional[Any]							,						snake_case_   :							Any							,						snake_case_   :							Optional[Any]							,						snake_case_   :							Tuple							,						snake_case_   :							Any							,						snake_case_   :							str							,						snake_case_   :							Dict							,						snake_case_   :							List[str]    ):
												__snake_case     			=       MobileBertModel(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_    )
												__snake_case     			=       model(snake_case_							,						token_type_ids=snake_case_    )
												__snake_case     			=       model(snake_case_    )

												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 							lowerCAmelCase      (				self   :							List[str]							,						snake_case_   :							List[Any]							,						snake_case_   :							Any							,						snake_case_   :							Optional[int]							,						snake_case_   :							int							,						snake_case_   :							int							,						snake_case_   :							str							,						snake_case_   :							str    ):
												__snake_case     			=       MobileBertForMaskedLM(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						labels=snake_case_    )
												self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.seq_length, self.vocab_size)    )





					def 							lowerCAmelCase      (				self   :							Optional[Any]							,						snake_case_   :							Optional[int]							,						snake_case_   :							Union[str, Any]							,						snake_case_   :							List[str]							,						snake_case_   :							int							,						snake_case_   :							int							,						snake_case_   :							int							,						snake_case_   :							Any    ):
												__snake_case     			=       MobileBertForNextSentencePrediction(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(
												    snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						labels=snake_case_							,						)
												self.parent.assertEqual(result.logits.shape							,						(self.batch_size, 2)    )





					def 							lowerCAmelCase      (				self   :							Optional[int]							,						snake_case_   :							List[str]							,						snake_case_   :							List[str]							,						snake_case_   :							str							,						snake_case_   :							Optional[Any]							,						snake_case_   :							List[str]							,						snake_case_   :							List[str]							,						snake_case_   :							Dict    ):
												__snake_case     			=       MobileBertForPreTraining(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(
												    snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						labels=snake_case_							,						next_sentence_label=snake_case_							,						)
												self.parent.assertEqual(result.prediction_logits.shape							,						(self.batch_size, self.seq_length, self.vocab_size)    )
												self.parent.assertEqual(result.seq_relationship_logits.shape							,						(self.batch_size, 2)    )





					def 							lowerCAmelCase      (				self   :							Any							,						snake_case_   :							List[str]							,						snake_case_   :							int							,						snake_case_   :							Any							,						snake_case_   :							Union[str, Any]							,						snake_case_   :							Dict							,						snake_case_   :							int							,						snake_case_   :							int    ):
												__snake_case     			=       MobileBertForQuestionAnswering(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(
												    snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						start_positions=snake_case_							,						end_positions=snake_case_							,						)
												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 							lowerCAmelCase      (				self   :							str							,						snake_case_   :							Tuple							,						snake_case_   :							int							,						snake_case_   :							str							,						snake_case_   :							Dict							,						snake_case_   :							Optional[int]							,						snake_case_   :							str							,						snake_case_   :							List[str]    ):
												__snake_case     			=       self.num_labels
												__snake_case     			=       MobileBertForSequenceClassification(snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						labels=snake_case_    )
												self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.num_labels)    )





					def 							lowerCAmelCase      (				self   :							Optional[Any]							,						snake_case_   :							Tuple							,						snake_case_   :							str							,						snake_case_   :							Union[str, Any]							,						snake_case_   :							Tuple							,						snake_case_   :							Optional[int]							,						snake_case_   :							Tuple							,						snake_case_   :							List[Any]    ):
												__snake_case     			=       self.num_labels
												__snake_case     			=       MobileBertForTokenClassification(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						labels=snake_case_    )
												self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.seq_length, self.num_labels)    )





					def 							lowerCAmelCase      (				self   :							Tuple							,						snake_case_   :							Union[str, Any]							,						snake_case_   :							Tuple							,						snake_case_   :							Optional[Any]							,						snake_case_   :							Any							,						snake_case_   :							Any							,						snake_case_   :							Dict							,						snake_case_   :							Union[str, Any]    ):
												__snake_case     			=       self.num_choices
												__snake_case     			=       MobileBertForMultipleChoice(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       input_ids.unsqueeze(1    ).expand(-1							,						self.num_choices							,						-1    ).contiguous()
												__snake_case     			=       token_type_ids.unsqueeze(1    ).expand(-1							,						self.num_choices							,						-1    ).contiguous()
												__snake_case     			=       input_mask.unsqueeze(1    ).expand(-1							,						self.num_choices							,						-1    ).contiguous()
												__snake_case     			=       model(
												    snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						labels=snake_case_							,						)
												self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.num_choices)    )





					def 							lowerCAmelCase      (				self   :							List[Any]    ):
												__snake_case     			=       self.prepare_config_and_inputs()
												(
												    (
												    __snake_case
												)			,   (
												    __snake_case
												)			,   (
												    __snake_case
												)			,   (
												    __snake_case
												)			,   (
												    __snake_case
												)			,   (
												    __snake_case
												)			,   (
												    __snake_case
												)			,   
												)     			=       config_and_inputs
												__snake_case     			=       {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
												return config, inputs_dict







@require_torch
class 							__magic_name__     (		lowercase__      ,		lowercase__      ,		unittest.TestCase   ):
					_SCREAMING_SNAKE_CASE       :					Dict     							=   (
					    (
					        MobileBertModel,
					        MobileBertForMaskedLM,
					        MobileBertForMultipleChoice,
					        MobileBertForNextSentencePrediction,
					        MobileBertForPreTraining,
					        MobileBertForQuestionAnswering,
					        MobileBertForSequenceClassification,
					        MobileBertForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					_SCREAMING_SNAKE_CASE       :					Optional[int]     							=   (
					    {
					        'feature-extraction': MobileBertModel,
					        'fill-mask': MobileBertForMaskedLM,
					        'question-answering': MobileBertForQuestionAnswering,
					        'text-classification': MobileBertForSequenceClassification,
					        'token-classification': MobileBertForTokenClassification,
					        'zero-shot': MobileBertForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					_SCREAMING_SNAKE_CASE       :					int     							=   True





					def 							lowerCAmelCase      (				self   :							List[str]							,						snake_case_   :							List[str]							,						snake_case_   :							str							,						snake_case_   :							str=False    ):
												__snake_case     			=       super()._prepare_for_class(snake_case_							,						snake_case_							,						return_labels=snake_case_    )

												if return_labels:
																			if model_class in get_values(snake_case_    ):
																										__snake_case     			=       torch.zeros(
																										    (self.model_tester.batch_size, self.model_tester.seq_length)							,						dtype=torch.long							,						device=snake_case_    )
																										__snake_case     			=       torch.zeros(
																										    self.model_tester.batch_size							,						dtype=torch.long							,						device=snake_case_    )
												return inputs_dict





					def 							lowerCAmelCase      (				self   :							Optional[Any]    ):
												__snake_case     			=       MobileBertModelTester(self    )
												__snake_case     			=       ConfigTester(self							,						config_class=snake_case_							,						hidden_size=37    )





					def 							lowerCAmelCase      (				self   :							List[Any]    ):
												self.config_tester.run_common_tests()





					def 							lowerCAmelCase      (				self   :							Tuple    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_model(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							Optional[Any]    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_masked_lm(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							Tuple    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_multiple_choice(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							Any    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							Any    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_pretraining(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							List[str]    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_question_answering(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							List[Any]    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_sequence_classification(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							str    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_token_classification(*snake_case_    )




def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE )  ->     List[Any]:







							"""simple docstring"""


							return torch.tensor(
							    SCREAMING_SNAKE_CASE							,    dtype=torch.long							,    device=SCREAMING_SNAKE_CASE							,    )


_SCREAMING_SNAKE_CASE	       =      1E-3







@require_torch
@require_sentencepiece
@require_tokenizers
class 							__magic_name__     (		unittest.TestCase   ):





					@slow
					def 							lowerCAmelCase      (				self   :							Tuple    ):
												__snake_case     			=       MobileBertModel.from_pretrained("google/mobilebert-uncased"    ).to(snake_case_    )
												__snake_case     			=       _long_tensor([[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]]    )
												with torch.no_grad():
																			__snake_case     			=       model(snake_case_    )[0]
												__snake_case     			=       torch.Size((1, 9, 512)    )
												self.assertEqual(output.shape							,						snake_case_    )
												__snake_case     			=       torch.tensor(
												    [
												        [
												            [-2.4_73_65_26e07, 8.2_69_16_56e04, 1.6_52_18_38e05],
												            [-5.7_54_17_04e-01, 3.9_05_60_22e00, 4.4_01_15_07e00],
												            [2.6_04_73_59e00, 1.5_67_76_52e00, -1.7_32_41_88e-01],
												        ]
												    ]							,						device=snake_case_							,						)

												# MobileBERT results range from 10e0 to 10e8. Even a 0.0000001% difference with a value of 10e8 results in a
												# ~1 difference, it's therefore not a good idea to measure using addition.
												# Here, we instead divide the expected result with the result in order to obtain ~1. We then check that the
												# result is held between bounds: 1 - TOLERANCE < expected_result / result < 1 + TOLERANCE
												__snake_case     			=       torch.all((expected_slice / output[..., :3, :3]) >= 1 - TOLERANCE    )
												__snake_case     			=       torch.all((expected_slice / output[..., :3, :3]) <= 1 + TOLERANCE    )

												self.assertTrue(lower_bound and upper_bound    )
