




import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


_lowerCAmelCase    : List[Any]		= logging.getLogger(__name__)

require_version("pytorch_lightning>=1.0.4")

_lowerCAmelCase    : Tuple		= {
    "base": AutoModel,
    "sequence-classification": AutoModelForSequenceClassification,
    "question-answering": AutoModelForQuestionAnswering,
    "pretraining": AutoModelForPreTraining,
    "token-classification": AutoModelForTokenClassification,
    "language-modeling": AutoModelWithLMHead,
    "summarization": AutoModelForSeqaSeqLM,
    "translation": AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
_lowerCAmelCase    : Tuple		= {
    "linear": get_linear_schedule_with_warmup,
    "cosine": get_cosine_schedule_with_warmup,
    "cosine_w_restarts": get_cosine_with_hard_restarts_schedule_with_warmup,
    "polynomial": get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
_lowerCAmelCase    : Dict		= sorted(arg_to_scheduler.keys())
_lowerCAmelCase    : Dict		= "{" + ", ".join(arg_to_scheduler_choices) + "}"
class       _UpperCamelCase   (					pl.LightningModule		):
      def __init__(      self :Optional[Any]			,   lowerCamelCase :argparse.Namespace			,   lowerCamelCase :Dict=None			,   lowerCamelCase :Optional[Any]="base"			,   lowerCamelCase :List[Any]=None			,   lowerCamelCase :Optional[int]=None			,   lowerCamelCase :Optional[int]=None			,   **lowerCamelCase :List[str]			,   )  ->      Any:
       super().__init__()
       # TODO: move to self.save_hyperparameters()
       # self.save_hyperparameters()
       # can also expand arguments into trainer signature for easier reading

       self.save_hyperparameters(lowerCamelCase   )
       UpperCAmelCase__							=    0
       UpperCAmelCase__							=    Path(self.hparams.output_dir   )
       UpperCAmelCase__							=    self.hparams.cache_dir if self.hparams.cache_dir else None
       if config is None:
        UpperCAmelCase__							=    AutoConfig.from_pretrained(
            self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path			,   **({"num_labels": num_labels} if num_labels is not None else {})			,   cache_dir=lowerCamelCase			,   **lowerCamelCase			,   )
       else:
        UpperCAmelCase__							=    config

       UpperCAmelCase__							=    ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
       for p in extra_model_params:
        if getattr(self.hparams			,   lowerCamelCase			,   lowerCamelCase   ):
         assert hasattr(self.config			,   lowerCamelCase   ), f'''model config doesn\'t have a `{p}` attribute'''
         setattr(self.config			,   lowerCamelCase			,   getattr(self.hparams			,   lowerCamelCase   )   )

       if tokenizer is None:
        UpperCAmelCase__							=    AutoTokenizer.from_pretrained(
            self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path			,   cache_dir=lowerCamelCase			,   )
       else:
        UpperCAmelCase__							=    tokenizer
       UpperCAmelCase__							=    MODEL_MODES[mode]
       if model is None:
        UpperCAmelCase__							=    self.model_type.from_pretrained(
            self.hparams.model_name_or_path			,   from_tf=bool(".ckpt" in self.hparams.model_name_or_path   )			,   config=self.config			,   cache_dir=lowerCamelCase			,   )
       else:
        UpperCAmelCase__							=    model
      def 						UpperCAmelCase_    (      self :Optional[int]			,   *lowerCamelCase :List[str]			,   **lowerCamelCase :Tuple   )  ->      str:
       UpperCAmelCase__							=    self.model_type.from_pretrained(*lowerCamelCase			,   **lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :str   )  ->      List[Any]:
       UpperCAmelCase__							=    arg_to_scheduler[self.hparams.lr_scheduler]
       UpperCAmelCase__							=    get_schedule_func(
           self.opt			,   num_warmup_steps=self.hparams.warmup_steps			,   num_training_steps=self.total_steps()   )
       UpperCAmelCase__							=    {"scheduler": scheduler, "interval": "step", "frequency": 1}
       return scheduler
      def 						UpperCAmelCase_    (      self :int   )  ->      Optional[int]:
       UpperCAmelCase__							=    self.model
       UpperCAmelCase__							=    ["bias", "LayerNorm.weight"]
       UpperCAmelCase__							=    [
           {
               "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay   )],
               "weight_decay": self.hparams.weight_decay,
           },
           {
               "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay   )],
               "weight_decay": 0.0,
           },
       ]
       if self.hparams.adafactor:
        UpperCAmelCase__							=    Adafactor(
            lowerCamelCase			,   lr=self.hparams.learning_rate			,   scale_parameter=lowerCamelCase			,   relative_step=lowerCamelCase   )

       else:
        UpperCAmelCase__							=    AdamW(
            lowerCamelCase			,   lr=self.hparams.learning_rate			,   eps=self.hparams.adam_epsilon   )
       UpperCAmelCase__							=    optimizer

       UpperCAmelCase__							=    self.get_lr_scheduler()

       return [optimizer], [scheduler]
      def 						UpperCAmelCase_    (      self :Dict			,   lowerCamelCase :Dict			,   lowerCamelCase :List[str]   )  ->      Dict:
       return self.validation_step(lowerCamelCase			,   lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Any			,   lowerCamelCase :Tuple   )  ->      Optional[Any]:
       return self.validation_end(lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :str   )  ->      int:
       UpperCAmelCase__							=    max(1			,   self.hparams.gpus   )  # TODO: consider num_tpu_cores
       UpperCAmelCase__							=    self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
       return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs
      def 						UpperCAmelCase_    (      self :str			,   lowerCamelCase :Optional[Any]   )  ->      Tuple:
       if mode == "test":
        UpperCAmelCase__							=    len(self.test_dataloader().dataset   )
       else:
        UpperCAmelCase__							=    self.get_dataloader("train"			,   self.hparams.train_batch_size			,   shuffle=lowerCamelCase   )
        UpperCAmelCase__							=    len(self.train_dataloader().dataset   )
      def 						UpperCAmelCase_    (      self :List[Any]			,   lowerCamelCase :str			,   lowerCamelCase :int			,   lowerCamelCase :bool = False   )  ->      Tuple:
       raise NotImplementedError("You must implement this for your task"   )
      def 						UpperCAmelCase_    (      self :str   )  ->      List[str]:
       return self.train_loader
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      Tuple:
       return self.get_dataloader("dev"			,   self.hparams.eval_batch_size			,   shuffle=lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      Tuple:
       return self.get_dataloader("test"			,   self.hparams.eval_batch_size			,   shuffle=lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Tuple			,   lowerCamelCase :Optional[Any]   )  ->      Optional[Any]:
       return os.path.join(
           self.hparams.data_dir			,   "cached_{}_{}_{}".format(
               lowerCamelCase			,   list(filter(lowerCamelCase			,   self.hparams.model_name_or_path.split("/"   )   )   ).pop()			,   str(self.hparams.max_seq_length   )			,   )			,   )
      @pl.utilities.rank_zero_only
      def 						UpperCAmelCase_    (      self :Optional[Any]			,   lowerCamelCase :Dict[str, Any]   )  ->      None:
       UpperCAmelCase__							=    self.output_dir.joinpath("best_tfmr"   )
       UpperCAmelCase__							=    self.step_count
       self.model.save_pretrained(lowerCamelCase   )
       self.tokenizer.save_pretrained(lowerCamelCase   )







      @staticmethod
      def 						UpperCAmelCase_    (      lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Tuple   )  ->      Dict:
       parser.add_argument(
           "--model_name_or_path"			,   default=lowerCamelCase			,   type=lowerCamelCase			,   required=lowerCamelCase			,   help="Path to pretrained model or model identifier from huggingface.co/models"			,   )
       parser.add_argument(
           "--config_name"			,   default=""			,   type=lowerCamelCase			,   help="Pretrained config name or path if not the same as model_name"   )
       parser.add_argument(
           "--tokenizer_name"			,   default=lowerCamelCase			,   type=lowerCamelCase			,   help="Pretrained tokenizer name or path if not the same as model_name"			,   )
       parser.add_argument(
           "--cache_dir"			,   default=""			,   type=lowerCamelCase			,   help="Where do you want to store the pre-trained models downloaded from huggingface.co"			,   )
       parser.add_argument(
           "--encoder_layerdrop"			,   type=lowerCamelCase			,   help="Encoder layer dropout probability (Optional). Goes into model.config"			,   )
       parser.add_argument(
           "--decoder_layerdrop"			,   type=lowerCamelCase			,   help="Decoder layer dropout probability (Optional). Goes into model.config"			,   )
       parser.add_argument(
           "--dropout"			,   type=lowerCamelCase			,   help="Dropout probability (Optional). Goes into model.config"			,   )
       parser.add_argument(
           "--attention_dropout"			,   type=lowerCamelCase			,   help="Attention dropout probability (Optional). Goes into model.config"			,   )
       parser.add_argument("--learning_rate"			,   default=5e-5			,   type=lowerCamelCase			,   help="The initial learning rate for Adam."   )
       parser.add_argument(
           "--lr_scheduler"			,   default="linear"			,   choices=lowerCamelCase			,   metavar=lowerCamelCase			,   type=lowerCamelCase			,   help="Learning rate scheduler"			,   )
       parser.add_argument("--weight_decay"			,   default=0.0			,   type=lowerCamelCase			,   help="Weight decay if we apply some."   )
       parser.add_argument("--adam_epsilon"			,   default=1e-8			,   type=lowerCamelCase			,   help="Epsilon for Adam optimizer."   )
       parser.add_argument("--warmup_steps"			,   default=0			,   type=lowerCamelCase			,   help="Linear warmup over warmup_steps."   )
       parser.add_argument("--num_workers"			,   default=4			,   type=lowerCamelCase			,   help="kwarg passed to DataLoader"   )
       parser.add_argument("--num_train_epochs"			,   dest="max_epochs"			,   default=3			,   type=lowerCamelCase   )
       parser.add_argument("--train_batch_size"			,   default=32			,   type=lowerCamelCase   )
       parser.add_argument("--eval_batch_size"			,   default=32			,   type=lowerCamelCase   )
       parser.add_argument("--adafactor"			,   action="store_true"   )



class       _UpperCamelCase   (					pl.Callback		):
      def 						UpperCAmelCase_    (      self :int			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :Tuple   )  ->      Optional[Any]:
       UpperCAmelCase__							=    trainer.lr_schedulers[0]["scheduler"]
       UpperCAmelCase__							=    {f'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()   )}
       pl_module.logger.log_metrics(lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[str]			,   lowerCamelCase :pl.Trainer			,   lowerCamelCase :pl.LightningModule   )  ->      Optional[int]:
       rank_zero_info("***** Validation results *****"   )
       UpperCAmelCase__							=    trainer.callback_metrics
       # Log results
       for key in sorted(lowerCamelCase   ):
        if key not in ["log", "progress_bar"]:
         rank_zero_info("{} = {}\n".format(lowerCamelCase			,   str(metrics[key]   )   )   )







      def 						UpperCAmelCase_    (      self :int			,   lowerCamelCase :pl.Trainer			,   lowerCamelCase :pl.LightningModule   )  ->      int:
       rank_zero_info("***** Test results *****"   )
       UpperCAmelCase__							=    trainer.callback_metrics
       # Log and save results to file
       UpperCAmelCase__							=    os.path.join(pl_module.hparams.output_dir			,   "test_results.txt"   )
       with open(lowerCamelCase			,   "w"   ) as writer:
        for key in sorted(lowerCamelCase   ):
         if key not in ["log", "progress_bar"]:
          rank_zero_info("{} = {}\n".format(lowerCamelCase			,   str(metrics[key]   )   )   )
          writer.write("{} = {}\n".format(lowerCamelCase			,   str(metrics[key]   )   )   )



def     lowerCAmelCase		(		_lowerCAmelCase				: Tuple    ,				_lowerCAmelCase				: Optional[Any]						):




 """simple docstring"""







 parser.add_argument(
     "--output_dir"    ,				default=_lowerCAmelCase    ,				type=_lowerCAmelCase    ,				required=_lowerCAmelCase    ,				help="The output directory where the model predictions and checkpoints will be written."    ,				)
 parser.add_argument(
     "--fp16"    ,				action="store_true"    ,				help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"    ,				)

 parser.add_argument(
     "--fp16_opt_level"    ,				type=_lowerCAmelCase    ,				default="O2"    ,				help=(
         "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
         "See details at https://nvidia.github.io/apex/amp.html"
     )    ,				)
 parser.add_argument("--n_tpu_cores"    ,				dest="tpu_cores"    ,				type=_lowerCAmelCase						)
 parser.add_argument("--max_grad_norm"    ,				dest="gradient_clip_val"    ,				default=1.0    ,				type=_lowerCAmelCase    ,				help="Max gradient norm"						)
 parser.add_argument("--do_train"    ,				action="store_true"    ,				help="Whether to run training."						)
 parser.add_argument("--do_predict"    ,				action="store_true"    ,				help="Whether to run predictions on the test set."						)
 parser.add_argument(
     "--gradient_accumulation_steps"    ,				dest="accumulate_grad_batches"    ,				type=_lowerCAmelCase    ,				default=1    ,				help="Number of updates steps to accumulate before performing a backward/update pass."    ,				)
 parser.add_argument("--seed"    ,				type=_lowerCAmelCase    ,				default=42    ,				help="random seed for initialization"						)
 parser.add_argument(
     "--data_dir"    ,				default=_lowerCAmelCase    ,				type=_lowerCAmelCase    ,				required=_lowerCAmelCase    ,				help="The input data dir. Should contain the training files for the CoNLL-2003 NER task."    ,				)



def     lowerCAmelCase		(		_lowerCAmelCase				: BaseTransformer    ,				_lowerCAmelCase				: argparse.Namespace    ,				_lowerCAmelCase				: Optional[int]=None    ,				_lowerCAmelCase				: str=True    ,				_lowerCAmelCase				: Tuple=[]    ,				_lowerCAmelCase				: Any=None    ,				_lowerCAmelCase				: Any=None    ,				**_lowerCAmelCase				: Union[str, Any]    ,				):




 """simple docstring"""







 pl.seed_everything(args.seed						)

 # init model
 UpperCAmelCase__							=    Path(model.hparams.output_dir						)
 odir.mkdir(exist_ok=_lowerCAmelCase						)

 # add custom checkpoints
 if checkpoint_callback is None:
  UpperCAmelCase__							=    pl.callbacks.ModelCheckpoint(
      filepath=args.output_dir    ,				prefix="checkpoint"    ,				monitor="val_loss"    ,				mode="min"    ,				save_top_k=1						)
 if early_stopping_callback:
  extra_callbacks.append(_lowerCAmelCase						)
 if logging_callback is None:
  UpperCAmelCase__							=    LoggingCallback()

 UpperCAmelCase__							=    {}

 # TODO: remove with PyTorch 1.6 since pl uses native amp
 if args.fpaa:
  UpperCAmelCase__							=    16
  UpperCAmelCase__							=    args.fpaa_opt_level

 if args.gpus > 1:
  UpperCAmelCase__							=    "ddp"

 UpperCAmelCase__							=    args.accumulate_grad_batches
 UpperCAmelCase__							=    extra_train_kwargs.get("accelerator"    ,				_lowerCAmelCase						)
 UpperCAmelCase__							=    extra_train_kwargs.get("profiler"    ,				_lowerCAmelCase						)

 UpperCAmelCase__							=    pl.Trainer.from_argparse_args(
     _lowerCAmelCase    ,				weights_summary=_lowerCAmelCase    ,				callbacks=[logging_callback] + extra_callbacks    ,				logger=_lowerCAmelCase    ,				checkpoint_callback=_lowerCAmelCase    ,				**_lowerCAmelCase    ,				)

 if args.do_train:
  trainer.fit(_lowerCAmelCase						)

 return trainer











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   )





