

import dataclasses
import re
import string
from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple

import numpy as np

from . import residue_constants


_SCREAMING_SNAKE_CASE  :						Dict         =    Mapping[str, np.ndarray]
_SCREAMING_SNAKE_CASE  :						Any         =    Mapping[str, Any]  # Is a nested dict.
_SCREAMING_SNAKE_CASE  :						str         =    0.0_1






@dataclasses.dataclass(frozen=A__ )
class  UpperCAmelCase__	:






  """simple docstring"""


  a     =      42  # [num_res, num_atom_type, 3]

  # Amino-acid type for each residue represented as an integer between 0 and
  # 20, where 20 is 'X'.
  a     =      42  # [num_res]

  # Binary float mask to indicate presence of a particular atom. 1.0 if an atom
  # is present and 0.0 if not. This should be used for loss masking.
  a     =      42  # [num_res, num_atom_type]

  # Residue index as used in PDB. It is not necessarily continuous or 0-indexed.
  a     =      42  # [num_res]

  # B-factors, or temperature factors, of each residue (in sq. angstroms units),
  # representing the displacement of the residue from its ground truth mean
  # value.
  a     =      42  # [num_res, num_atom_type]

  # Chain indices for multi-chain predictions
  a     =      None

  # Optional remark about the protein. Included as a comment in output PDB
  # files
  a     =      None

  # Templates used to generate this protein (prediction-only)
  a     =      None

  # Chain corresponding to each parent
  a     =      None



def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					R'''(\[[A-Z]+\]\n)'''
     SCREAMING_SNAKE_CASE__					=					[tag.strip() for tag in re.split(_A   ,	_A       ) if len(_A       ) > 0]
     SCREAMING_SNAKE_CASE__					=					zip(tags[0::2]   ,	[l.split('''\n'''       ) for l in tags[1::2]]       )

     SCREAMING_SNAKE_CASE__					=					["N", "CA", "C"]
     SCREAMING_SNAKE_CASE__					=					None
     SCREAMING_SNAKE_CASE__					=					None
     SCREAMING_SNAKE_CASE__					=					None
     for g in groups:
          if "[PRIMARY]" == g[0]:
               SCREAMING_SNAKE_CASE__					=					g[1][0].strip()
               for i in range(len(_A       )       ):
                    if seq[i] not in residue_constants.restypes:
                         SCREAMING_SNAKE_CASE__					=					'''X'''  # FIXME: strings are immutable
               SCREAMING_SNAKE_CASE__					=					np.array(
                   [residue_constants.restype_order.get(_A   ,	residue_constants.restype_num       ) for res_symbol in seq]       )
          elif "[TERTIARY]" == g[0]:
               SCREAMING_SNAKE_CASE__					=					[]
               for axis in range(3       ):
                    tertiary.append(list(map(_A   ,	g[1][axis].split()       )       )       )
               SCREAMING_SNAKE_CASE__					=					np.array(_A       )
               SCREAMING_SNAKE_CASE__					=					np.zeros((len(tertiary[0]       ) // 3, residue_constants.atom_type_num, 3)       ).astype(np.floataa       )
               for i, atom in enumerate(_A       ):
                    SCREAMING_SNAKE_CASE__					=					np.transpose(tertiary_np[:, i::3]       )
               atom_positions *= PICO_TO_ANGSTROM
          elif "[MASK]" == g[0]:
               SCREAMING_SNAKE_CASE__					=					np.array(list(map({'''-''': 0, '''+''': 1}.get   ,	g[1][0].strip()       )       )       )
               SCREAMING_SNAKE_CASE__					=					np.zeros(
                   (
                       len(_A       ),
                       residue_constants.atom_type_num,
                   )       ).astype(np.floataa       )
               for i, atom in enumerate(_A       ):
                    SCREAMING_SNAKE_CASE__					=					1
               atom_mask *= mask[..., None]

     assert aatype is not None

     return Protein(
         atom_positions=_A   ,	atom_mask=_A   ,	aatype=_A   ,	residue_index=np.arange(len(_A       )       )   ,	b_factors=_A   ,	)
def        UpperCAmelCase_     (		_A   ,	_A = 0       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					[]

     SCREAMING_SNAKE_CASE__					=					prot.remark
     if remark is not None:
          pdb_headers.append(F'''REMARK {remark}'''       )

     SCREAMING_SNAKE_CASE__					=					prot.parents
     SCREAMING_SNAKE_CASE__					=					prot.parents_chain_index
     if parents is not None and parents_chain_index is not None:
          SCREAMING_SNAKE_CASE__					=					[p for i, p in zip(_A   ,	_A       ) if i == chain_id]

     if parents is None or len(_A       ) == 0:
          SCREAMING_SNAKE_CASE__					=					['''N/A''']

     pdb_headers.append(F'''PARENT {" ".join(_A       )}'''       )

     return pdb_headers
def        UpperCAmelCase_     (		_A   ,	_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					[]
     SCREAMING_SNAKE_CASE__					=					pdb_str.split('''\n'''       )

     SCREAMING_SNAKE_CASE__					=					prot.remark
     if remark is not None:
          out_pdb_lines.append(F'''REMARK {remark}'''       )

     SCREAMING_SNAKE_CASE__					=					42
     if prot.parents is not None and len(prot.parents       ) > 0:
          SCREAMING_SNAKE_CASE__					=					[]
          if prot.parents_chain_index is not None:
               SCREAMING_SNAKE_CASE__					=					{}
               for p, i in zip(prot.parents   ,	prot.parents_chain_index       ):
                    parent_dict.setdefault(str(_A       )   ,	[]       )
                    parent_dict[str(_A       )].append(_A       )

               SCREAMING_SNAKE_CASE__					=					max([int(_A       ) for chain_idx in parent_dict]       )
               for i in range(max_idx + 1       ):
                    SCREAMING_SNAKE_CASE__					=					parent_dict.get(str(_A       )   ,	['''N/A''']       )
                    parents_per_chain.append(_A       )
          else:
               parents_per_chain.append(list(prot.parents       )       )
     else:
          SCREAMING_SNAKE_CASE__					=					[['''N/A''']]

     def make_parent_line(_A       ) -> str:
          return F'''PARENT {" ".join(_A       )}'''

     out_pdb_lines.append(make_parent_line(parents_per_chain[0]       )       )

     SCREAMING_SNAKE_CASE__					=					0
     for i, l in enumerate(_A       ):
          if "PARENT" not in l and "REMARK" not in l:
               out_pdb_lines.append(_A       )
          if "TER" in l and "END" not in lines[i + 1]:
               chain_counter += 1
               if not chain_counter >= len(_A       ):
                    SCREAMING_SNAKE_CASE__					=					parents_per_chain[chain_counter]
               else:
                    SCREAMING_SNAKE_CASE__					=					['''N/A''']

               out_pdb_lines.append(make_parent_line(_A       )       )

     return "\n".join(_A       )
def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					residue_constants.restypes + ['''X''']

     def res_atoa(_A       ) -> str:
          return residue_constants.restype_atoa.get(restypes[r]   ,	'''UNK'''       )

     SCREAMING_SNAKE_CASE__					=					residue_constants.atom_types

     SCREAMING_SNAKE_CASE__					=					[]

     SCREAMING_SNAKE_CASE__					=					prot.atom_mask
     SCREAMING_SNAKE_CASE__					=					prot.aatype
     SCREAMING_SNAKE_CASE__					=					prot.atom_positions
     SCREAMING_SNAKE_CASE__					=					prot.residue_index.astype(np.intaa       )
     SCREAMING_SNAKE_CASE__					=					prot.b_factors
     SCREAMING_SNAKE_CASE__					=					prot.chain_index

     if np.any(aatype > residue_constants.restype_num       ):
          raise ValueError('''Invalid aatypes.'''       )

     SCREAMING_SNAKE_CASE__					=					get_pdb_headers(_A       )
     if len(_A       ) > 0:
          pdb_lines.extend(_A       )

     SCREAMING_SNAKE_CASE__					=					aatype.shape[0]
     SCREAMING_SNAKE_CASE__					=					1
     SCREAMING_SNAKE_CASE__					=					0
     SCREAMING_SNAKE_CASE__					=					string.ascii_uppercase
     SCREAMING_SNAKE_CASE__					=					None
     # Add all atom sites.
     for i in range(_A       ):
          SCREAMING_SNAKE_CASE__					=					res_atoa(aatype[i]       )
          for atom_name, pos, mask, b_factor in zip(_A   ,	atom_positions[i]   ,	atom_mask[i]   ,	b_factors[i]       ):
               if mask < 0.5:
                    continue

               SCREAMING_SNAKE_CASE__					=					'''ATOM'''
               SCREAMING_SNAKE_CASE__					=					atom_name if len(_A       ) == 4 else F''' {atom_name}'''
               SCREAMING_SNAKE_CASE__					=					''''''
               SCREAMING_SNAKE_CASE__					=					''''''
               SCREAMING_SNAKE_CASE__					=					1.0_0
               SCREAMING_SNAKE_CASE__					=					atom_name[0]  # Protein supports only C, N, O, S, this works.
               SCREAMING_SNAKE_CASE__					=					''''''

               SCREAMING_SNAKE_CASE__					=					'''A'''
               if chain_index is not None:
                    SCREAMING_SNAKE_CASE__					=					chain_tags[chain_index[i]]

               # PDB is a columnar format, every space matters here!
               SCREAMING_SNAKE_CASE__					=					(
                   F'''{record_type:<6}{atom_index:>5} {name:<4}{alt_loc:>1}'''
                   F'''{res_name_a:>3} {chain_tag:>1}'''
                   F'''{residue_index[i]:>4}{insertion_code:>1}   '''
                   F'''{pos[0]:>8.3f}{pos[1]:>8.3f}{pos[2]:>8.3f}'''
                   F'''{occupancy:>6.2f}{b_factor:>6.2f}          '''
                   F'''{element:>2}{charge:>2}'''
               )
               pdb_lines.append(_A       )
               atom_index += 1

          SCREAMING_SNAKE_CASE__					=					i == n - 1
          if chain_index is not None:
               if i != n - 1 and chain_index[i + 1] != prev_chain_index:
                    SCREAMING_SNAKE_CASE__					=					True
                    SCREAMING_SNAKE_CASE__					=					chain_index[i + 1]

          if should_terminate:
               # Close the chain.
               SCREAMING_SNAKE_CASE__					=					'''TER'''
               SCREAMING_SNAKE_CASE__					=					(
                   F'''{chain_end:<6}{atom_index:>5}      {res_atoa(aatype[i]       ):>3} {chain_tag:>1}{residue_index[i]:>4}'''
               )
               pdb_lines.append(_A       )
               atom_index += 1

               if i != n - 1:
                    # "prev" is a misnomer here. This happens at the beginning of
                    # each new chain.
                    pdb_lines.extend(get_pdb_headers(_A   ,	_A       )       )

     pdb_lines.append('''END'''       )
     pdb_lines.append(''''''       )
     return "\n".join(_A       )
def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     return residue_constants.STANDARD_ATOM_MASK[prot.aatype]







def        UpperCAmelCase_     (		_A   ,	_A   ,	_A = None   ,	_A = None   ,	_A = None   ,	_A = None   ,	_A = None   ,	):





     '''simple docstring'''




     return Protein(
         aatype=features['''aatype''']   ,	atom_positions=result['''final_atom_positions''']   ,	atom_mask=result['''final_atom_mask''']   ,	residue_index=features['''residue_index'''] + 1   ,	b_factors=b_factors if b_factors is not None else np.zeros_like(result['''final_atom_mask''']       )   ,	chain_index=_A   ,	remark=_A   ,	parents=_A   ,	parents_chain_index=_A   ,	)





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  UpperCAmelCase__	:






  """simple docstring"""


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






  def    lowercase_	(							self		:	List[str]      )	->							Optional[int]:
       SCREAMING_SNAKE_CASE__					=					ids_tensor([self.batch_size, self.seq_length]			,     self.vocab_size      )

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

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

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

       SCREAMING_SNAKE_CASE__					=					self.get_config()

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






  def    lowercase_	(							self		:	Dict      )	->							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=__lowerCamelCase			,     initializer_range=self.initializer_range			,     )






  def    lowercase_	(							self		:	List[str]			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	str			,     __lowerCamelCase		:	Tuple			,     __lowerCamelCase		:	Union[str, Any]			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	Optional[Any]			,     __lowerCamelCase		:	List[Any]      )	->							Any:
       SCREAMING_SNAKE_CASE__					=					MobileBertModel(config=__lowerCamelCase      )
       model.to(__lowerCamelCase      )
       model.eval()
       SCREAMING_SNAKE_CASE__					=					model(__lowerCamelCase			,     attention_mask=__lowerCamelCase			,     token_type_ids=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					model(__lowerCamelCase			,     token_type_ids=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					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    lowercase_	(							self		:	Dict			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	List[Any]			,     __lowerCamelCase		:	Dict			,     __lowerCamelCase		:	Dict			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	str			,     __lowerCamelCase		:	Any      )	->							Optional[Any]:
       SCREAMING_SNAKE_CASE__					=					MobileBertForMaskedLM(config=__lowerCamelCase      )
       model.to(__lowerCamelCase      )
       model.eval()
       SCREAMING_SNAKE_CASE__					=					model(__lowerCamelCase			,     attention_mask=__lowerCamelCase			,     token_type_ids=__lowerCamelCase			,     labels=__lowerCamelCase      )
       self.parent.assertEqual(result.logits.shape			,     (self.batch_size, self.seq_length, self.vocab_size)      )






  def    lowercase_	(							self		:	int			,     __lowerCamelCase		:	Any			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	List[str]			,     __lowerCamelCase		:	List[str]			,     __lowerCamelCase		:	List[Any]			,     __lowerCamelCase		:	List[Any]      )	->							int:
       SCREAMING_SNAKE_CASE__					=					MobileBertForNextSentencePrediction(config=__lowerCamelCase      )
       model.to(__lowerCamelCase      )
       model.eval()
       SCREAMING_SNAKE_CASE__					=					model(
           __lowerCamelCase			,     attention_mask=__lowerCamelCase			,     token_type_ids=__lowerCamelCase			,     labels=__lowerCamelCase			,     )
       self.parent.assertEqual(result.logits.shape			,     (self.batch_size, 2)      )






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






  def    lowercase_	(							self		:	int			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	Any			,     __lowerCamelCase		:	str			,     __lowerCamelCase		:	Any			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	str			,     __lowerCamelCase		:	Union[str, Any]      )	->							int:
       SCREAMING_SNAKE_CASE__					=					self.num_labels
       SCREAMING_SNAKE_CASE__					=					MobileBertForTokenClassification(config=__lowerCamelCase      )
       model.to(__lowerCamelCase      )
       model.eval()
       SCREAMING_SNAKE_CASE__					=					model(__lowerCamelCase			,     attention_mask=__lowerCamelCase			,     token_type_ids=__lowerCamelCase			,     labels=__lowerCamelCase      )
       self.parent.assertEqual(result.logits.shape			,     (self.batch_size, self.seq_length, self.num_labels)      )






  def    lowercase_	(							self		:	Union[str, Any]			,     __lowerCamelCase		:	Dict			,     __lowerCamelCase		:	List[str]			,     __lowerCamelCase		:	Any			,     __lowerCamelCase		:	Dict			,     __lowerCamelCase		:	Any			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	Tuple      )	->							List[Any]:
       SCREAMING_SNAKE_CASE__					=					self.num_choices
       SCREAMING_SNAKE_CASE__					=					MobileBertForMultipleChoice(config=__lowerCamelCase      )
       model.to(__lowerCamelCase      )
       model.eval()
       SCREAMING_SNAKE_CASE__					=					input_ids.unsqueeze(1      ).expand(-1			,     self.num_choices			,     -1      ).contiguous()
       SCREAMING_SNAKE_CASE__					=					token_type_ids.unsqueeze(1      ).expand(-1			,     self.num_choices			,     -1      ).contiguous()
       SCREAMING_SNAKE_CASE__					=					input_mask.unsqueeze(1      ).expand(-1			,     self.num_choices			,     -1      ).contiguous()
       SCREAMING_SNAKE_CASE__					=					model(
           __lowerCamelCase			,     attention_mask=__lowerCamelCase			,     token_type_ids=__lowerCamelCase			,     labels=__lowerCamelCase			,     )
       self.parent.assertEqual(result.logits.shape			,     (self.batch_size, self.num_choices)      )







  def    lowercase_	(							self		:	List[str]      )	->							Tuple:
       SCREAMING_SNAKE_CASE__					=					self.prepare_config_and_inputs()
       (
           (
           SCREAMING_SNAKE_CASE__
       ),(
           SCREAMING_SNAKE_CASE__
       ),(
           SCREAMING_SNAKE_CASE__
       ),(
           SCREAMING_SNAKE_CASE__
       ),(
           SCREAMING_SNAKE_CASE__
       ),(
           SCREAMING_SNAKE_CASE__
       ),(
           SCREAMING_SNAKE_CASE__
       ),
       )					=					config_and_inputs
       SCREAMING_SNAKE_CASE__					=					{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
       return config, inputs_dict






@require_torch
class  UpperCAmelCase__	( A__		,  A__		,  unittest.TestCase ):






  """simple docstring"""


  a     =      (
      (
          MobileBertModel,
          MobileBertForMaskedLM,
          MobileBertForMultipleChoice,
          MobileBertForNextSentencePrediction,
          MobileBertForPreTraining,
          MobileBertForQuestionAnswering,
          MobileBertForSequenceClassification,
          MobileBertForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  a     =      (
      {
          "feature-extraction": MobileBertModel,
          "fill-mask": MobileBertForMaskedLM,
          "question-answering": MobileBertForQuestionAnswering,
          "text-classification": MobileBertForSequenceClassification,
          "token-classification": MobileBertForTokenClassification,
          "zero-shot": MobileBertForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  a     =      True






  def    lowercase_	(							self		:	int			,     __lowerCamelCase		:	Optional[Any]			,     __lowerCamelCase		:	List[str]			,     __lowerCamelCase		:	Optional[int]=False      )	->							str:
       SCREAMING_SNAKE_CASE__					=					super()._prepare_for_class(__lowerCamelCase			,     __lowerCamelCase			,     return_labels=__lowerCamelCase      )

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






  def    lowercase_	(							self		:	Optional[Any]      )	->							List[str]:
       SCREAMING_SNAKE_CASE__					=					MobileBertModelTester(self      )
       SCREAMING_SNAKE_CASE__					=					ConfigTester(self			,     config_class=__lowerCamelCase			,     hidden_size=37      )






  def    lowercase_	(							self		:	Any      )	->							Union[str, Any]:
       self.config_tester.run_common_tests()






  def    lowercase_	(							self		:	Optional[int]      )	->							Union[str, Any]:
       SCREAMING_SNAKE_CASE__					=					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_model(*__lowerCamelCase      )






  def    lowercase_	(							self		:	Tuple      )	->							Any:
       SCREAMING_SNAKE_CASE__					=					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_masked_lm(*__lowerCamelCase      )






  def    lowercase_	(							self		:	Dict      )	->							Optional[Any]:
       SCREAMING_SNAKE_CASE__					=					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_multiple_choice(*__lowerCamelCase      )






  def    lowercase_	(							self		:	Union[str, Any]      )	->							Optional[Any]:
       SCREAMING_SNAKE_CASE__					=					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*__lowerCamelCase      )






  def    lowercase_	(							self		:	Any      )	->							str:
       SCREAMING_SNAKE_CASE__					=					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_pretraining(*__lowerCamelCase      )






  def    lowercase_	(							self		:	Dict      )	->							str:
       SCREAMING_SNAKE_CASE__					=					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_question_answering(*__lowerCamelCase      )






  def    lowercase_	(							self		:	List[str]      )	->							str:
       SCREAMING_SNAKE_CASE__					=					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_sequence_classification(*__lowerCamelCase      )







  def    lowercase_	(							self		:	Tuple      )	->							List[Any]:
       SCREAMING_SNAKE_CASE__					=					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_token_classification(*__lowerCamelCase      )



def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




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


_SCREAMING_SNAKE_CASE  :						Optional[Any]         =    1e-3







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






  """simple docstring"""


  @slow
  def    lowercase_	(							self		:	List[str]      )	->							List[Any]:
       SCREAMING_SNAKE_CASE__					=					MobileBertModel.from_pretrained('''google/mobilebert-uncased'''      ).to(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					_long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]      )
       with torch.no_grad():
            SCREAMING_SNAKE_CASE__					=					model(__lowerCamelCase      )[0]
       SCREAMING_SNAKE_CASE__					=					torch.Size((1, 9, 512)      )
       self.assertEqual(output.shape			,     __lowerCamelCase      )
       SCREAMING_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=__lowerCamelCase			,     )

       # 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
       SCREAMING_SNAKE_CASE__					=					torch.all((expected_slice / output[..., :3, :3]) >= 1 - TOLERANCE      )
       SCREAMING_SNAKE_CASE__					=					torch.all((expected_slice / output[..., :3, :3]) <= 1 + TOLERANCE      )

       self.assertTrue(lower_bound and upper_bound      )


