

'''simple docstring'''

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


a   :			Any         = logging.getLogger(__name__)

require_version("pytorch_lightning>=1.0.4")

a   :			Union[str, Any]         = {
    "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
a   :			str         = {
    "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
}
a   :			Dict         = sorted(arg_to_scheduler.keys())
a   :			Optional[Any]         = "{" + ", ".join(arg_to_scheduler_choices) + "}"







class   UpperCamelCase__    (   pl.LightningModule   ):


		"""simple docstring"""

		def __init__( self    ,			snake_case    ,			snake_case=None    ,			snake_case="base"    ,			snake_case=None    ,			snake_case=None    ,			snake_case=None    ,			**snake_case    ,			):


							'''simple docstring'''

							super().__init__()
							# TODO: move to self.save_hyperparameters()
							# self.save_hyperparameters()
							# can also expand arguments into trainer signature for easier reading

							self.save_hyperparameters(lowercase_			)
							UpperCAmelCase		:							Optional[Any]   					=   0
							UpperCAmelCase		:							List[Any]   					=   Path(self.hparams.output_dir			)
							UpperCAmelCase		:							List[Any]   					=   self.hparams.cache_dir if self.hparams.cache_dir else None
							if config is None:
												UpperCAmelCase		:							Optional[int]   					=   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=lowercase_    ,			**lowercase_    ,			)
							else:
												UpperCAmelCase		:							Union[str, Any]   					=   config

							UpperCAmelCase		:							Union[str, Any]   					=   ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
							for p in extra_model_params:
												if getattr(self.hparams    ,			lowercase_    ,			lowercase_			):
																	assert hasattr(self.config    ,			lowercase_			), f"model config doesn't have a `{p}` attribute"
																	setattr(self.config    ,			lowercase_    ,			getattr(self.hparams    ,			lowercase_			)			)

							if tokenizer is None:
												UpperCAmelCase		:							Tuple   					=   AutoTokenizer.from_pretrained(
												    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path    ,			cache_dir=lowercase_    ,			)
							else:
												UpperCAmelCase		:							Optional[int]   					=   tokenizer
							UpperCAmelCase		:							str   					=   MODEL_MODES[mode]
							if model is None:
												UpperCAmelCase		:							List[Any]   					=   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=lowercase_    ,			)
							else:
												UpperCAmelCase		:							str   					=   model







		def       A_       ( self    ,			*snake_case    ,			**snake_case			):


							'''simple docstring'''

							UpperCAmelCase		:							List[Any]   					=   self.model_type.from_pretrained(*lowercase_    ,			**lowercase_			)







		def       A_       ( self			):


							'''simple docstring'''

							UpperCAmelCase		:							Tuple   					=   arg_to_scheduler[self.hparams.lr_scheduler]
							UpperCAmelCase		:							Optional[int]   					=   get_schedule_func(
							    self.opt    ,			num_warmup_steps=self.hparams.warmup_steps    ,			num_training_steps=self.total_steps()			)
							UpperCAmelCase		:							int   					=   {"scheduler": scheduler, "interval": "step", "frequency": 1}
							return scheduler







		def       A_       ( self			):


							'''simple docstring'''

							UpperCAmelCase		:							Dict   					=   self.model
							UpperCAmelCase		:							Optional[Any]   					=   ["bias", "LayerNorm.weight"]
							UpperCAmelCase		:							Dict   					=   [
							    {
							        "params": [
							            p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay			)
							        ],  # check this named paramters
							        "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		:							Optional[Any]   					=   Adafactor(
												    lowercase_    ,			lr=self.hparams.learning_rate    ,			scale_parameter=lowercase_    ,			relative_step=lowercase_			)

							else:
												UpperCAmelCase		:							List[Any]   					=   AdamW(
												    lowercase_    ,			lr=self.hparams.learning_rate    ,			eps=self.hparams.adam_epsilon			)
							UpperCAmelCase		:							Tuple   					=   optimizer

							UpperCAmelCase		:							Any   					=   self.get_lr_scheduler()

							return [optimizer], [scheduler]







		def       A_       ( self    ,			snake_case    ,			snake_case			):


							'''simple docstring'''

							return self.validation_step(lowercase_    ,			lowercase_			)







		def       A_       ( self    ,			snake_case			):


							'''simple docstring'''

							return self.validation_end(lowercase_			)







		def       A_       ( self			):


							'''simple docstring'''

							UpperCAmelCase		:							str   					=   max(1    ,			self.hparams.gpus			)  # TODO: consider num_tpu_cores
							UpperCAmelCase		:							List[Any]   					=   self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
							return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs







		def       A_       ( self    ,			snake_case			):


							'''simple docstring'''

							if stage == "test":
												UpperCAmelCase		:							int   					=   len(self.test_dataloader().dataset			)
							else:
												UpperCAmelCase		:							List[Any]   					=   self.get_dataloader("train"    ,			self.hparams.train_batch_size    ,			shuffle=lowercase_			)
												UpperCAmelCase		:							Dict   					=   len(self.train_dataloader().dataset			)







		def       A_       ( self    ,			snake_case    ,			snake_case    ,			snake_case = False			):


							'''simple docstring'''

							raise NotImplementedError("You must implement this for your task"			)







		def       A_       ( self			):


							'''simple docstring'''

							return self.train_loader







		def       A_       ( self			):


							'''simple docstring'''

							return self.get_dataloader("dev"    ,			self.hparams.eval_batch_size    ,			shuffle=lowercase_			)







		def       A_       ( self			):


							'''simple docstring'''

							return self.get_dataloader("test"    ,			self.hparams.eval_batch_size    ,			shuffle=lowercase_			)







		def       A_       ( self    ,			snake_case			):


							'''simple docstring'''

							return os.path.join(
							    self.hparams.data_dir    ,			"cached_{}_{}_{}".format(
							        lowercase_    ,			list(filter(lowercase_    ,			self.hparams.model_name_or_path.split("/"			)			)			).pop()    ,			str(self.hparams.max_seq_length			)    ,			)    ,			)







		@pl.utilities.rank_zero_only
		def       A_       ( self    ,			snake_case			):


							'''simple docstring'''

							UpperCAmelCase		:							List[str]   					=   self.output_dir.joinpath("best_tfmr"			)
							UpperCAmelCase		:							Optional[int]   					=   self.step_count
							self.model.save_pretrained(lowercase_			)
							self.tokenizer.save_pretrained(lowercase_			)







		@staticmethod
		def       A_       ( snake_case    ,			snake_case			):


							'''simple docstring'''

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







class   UpperCamelCase__    (   pl.Callback   ):


		"""simple docstring"""

		def       A_       ( self    ,			snake_case    ,			snake_case			):


							'''simple docstring'''

							if (
							    trainer.is_global_zero and trainer.global_rank == 0
							):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
												pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.







class   UpperCamelCase__    (   pl.Callback   ):


		"""simple docstring"""

		def       A_       ( self    ,			snake_case    ,			snake_case			):


							'''simple docstring'''

							for name, param in pl_module.model.rag.named_parameters():
												if param.grad is None:
																	print(lowercase_			)







class   UpperCamelCase__    (   pl.Callback   ):


		"""simple docstring"""

		def       A_       ( self    ,			snake_case    ,			snake_case			):


							'''simple docstring'''

							UpperCAmelCase		:							Optional[int]   					=   trainer.lr_schedulers[0]["scheduler"]
							UpperCAmelCase		:							List[Any]   					=   {f"lr_group_{i}": lr for i, lr in enumerate(lr_scheduler.get_lr()			)}
							pl_module.logger.log_metrics(lowercase_			)







		def       A_       ( self    ,			snake_case    ,			snake_case			):


							'''simple docstring'''

							rank_zero_info("***** Validation results *****"			)
							UpperCAmelCase		:							List[str]   					=   trainer.callback_metrics
							# Log results
							for key in sorted(lowercase_			):
												if key not in ["log", "progress_bar"]:
																	rank_zero_info("{} = {}\n".format(lowercase_    ,			str(metrics[key]			)			)			)







		def       A_       ( self    ,			snake_case    ,			snake_case			):


							'''simple docstring'''

							rank_zero_info("***** Test results *****"			)
							UpperCAmelCase		:							Tuple   					=   trainer.callback_metrics
							# Log and save results to file
							UpperCAmelCase		:							List[str]   					=   os.path.join(pl_module.hparams.output_dir    ,			"test_results.txt"			)
							with open(lowercase_    ,			"w"			) as writer:
												for key in sorted(lowercase_			):
																	if key not in ["log", "progress_bar"]:
																						rank_zero_info("{} = {}\n".format(lowercase_    ,			str(metrics[key]			)			)			)
																						writer.write("{} = {}\n".format(lowercase_    ,			str(metrics[key]			)			)			)





def       lowercase					(     __magic_name__ ,    __magic_name__							):



					'''simple docstring'''







					parser.add_argument(
					    "--output_dir" ,    default=str(Path(_A							).parent / "test_run" / "model_checkpoints"							) ,    type=_A ,    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=_A ,    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=_A							)
					parser.add_argument("--max_grad_norm" ,    dest="gradient_clip_val" ,    default=1.0 ,    type=_A ,    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=_A ,    default=1 ,    help="Number of updates steps to accumulate before performing a backward/update pass." ,    )
					parser.add_argument("--seed" ,    type=_A ,    default=42 ,    help="random seed for initialization"							)
					parser.add_argument(
					    "--data_dir" ,    default=str(Path(_A							).parent / "test_run" / "dummy-train-data"							) ,    type=_A ,    help="The input data dir. Should contain the training files for the CoNLL-2003 NER task." ,    )





def       lowercase					(     __magic_name__ ,    __magic_name__ ,    __magic_name__=None ,    __magic_name__=True ,    __magic_name__=[] ,    __magic_name__=None ,    __magic_name__=None ,    **__magic_name__ ,    ):



					'''simple docstring'''







					pl.seed_everything(args.seed							)

					# init model
					UpperCAmelCase		:							Optional[int]   					=   Path(model.hparams.output_dir							)
					odir.mkdir(exist_ok=_A							)

					# add custom checkpoints
					if checkpoint_callback is None:
										UpperCAmelCase		:							List[str]   					=   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(_A							)
					if logging_callback is None:
										UpperCAmelCase		:							Tuple   					=   LoggingCallback()

					UpperCAmelCase		:							List[str]   					=   {}

					if args.fpaa:
										UpperCAmelCase		:							List[str]   					=   16

					if args.gpus > 1:
										UpperCAmelCase		:							int   					=   "auto"
										UpperCAmelCase		:							Optional[Any]   					=   "ddp"

					UpperCAmelCase		:							Dict   					=   args.accumulate_grad_batches
					UpperCAmelCase		:							Dict   					=   None
					UpperCAmelCase		:							Union[str, Any]   					=   "auto"

					UpperCAmelCase		:							Union[str, Any]   					=   pl.Trainer.from_argparse_args(
					    _A ,    weights_summary=_A ,    callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback] ,    logger=_A ,    val_check_interval=1 ,    num_sanity_val_steps=2 ,    **_A ,    )

					if args.do_train:
										trainer.fit(_A							)

					else:
										print("RAG modeling tests with new set functions successfuly executed!"							)
					return trainer



import os
import tempfile
import unittest

from transformers import FlaubertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, 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 (
           FlaubertForMultipleChoice,
           FlaubertForQuestionAnswering,
           FlaubertForQuestionAnsweringSimple,
           FlaubertForSequenceClassification,
           FlaubertForTokenClassification,
           FlaubertModel,
           FlaubertWithLMHeadModel,
       )
       from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST







class 						__magic_name__			(							lowerCamelCase__							):

   '''simple docstring'''




   def __init__(				self,    lowercase_,    lowercase_=13,    lowercase_=7,    lowercase_=True,    lowercase_=True,    lowercase_=True,    lowercase_=True,    lowercase_=True,    lowercase_=False,    lowercase_=False,    lowercase_=False,    lowercase_=2,    lowercase_=99,    lowercase_=0,    lowercase_=32,    lowercase_=5,    lowercase_=4,    lowercase_=0.1,    lowercase_=0.1,    lowercase_=512,    lowercase_=12,    lowercase_=2,    lowercase_=0.02,    lowercase_=3,    lowercase_=4,    lowercase_="last",    lowercase_=None,    lowercase_=None,    )							->    List[Any]:






     """simple docstring"""






     a__  =parent
     a__  =batch_size
     a__  =seq_length
     a__  =is_training
     a__  =use_input_lengths
     a__  =use_token_type_ids
     a__  =use_labels
     a__  =gelu_activation
     a__  =sinusoidal_embeddings
     a__  =causal
     a__  =asm
     a__  =n_langs
     a__  =vocab_size
     a__  =n_special
     a__  =hidden_size
     a__  =num_hidden_layers
     a__  =num_attention_heads
     a__  =hidden_dropout_prob
     a__  =attention_probs_dropout_prob
     a__  =max_position_embeddings
     a__  =type_vocab_size
     a__  =type_sequence_label_size
     a__  =initializer_range
     a__  =num_labels
     a__  =num_choices
     a__  =summary_type
     a__  =use_proj
     a__  =scope



   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






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

     a__  =None
     if self.use_input_lengths:
       a__  =(
           ids_tensor([self.batch_size],    vocab_size=2       ) + self.seq_length - 2
       )  # small variation of seq_length

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

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

     a__  =self.get_config()

     return (
         config,
         input_ids,
         token_type_ids,
         input_lengths,
         sequence_labels,
         token_labels,
         is_impossible_labels,
         choice_labels,
         input_mask,
     )



   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






     return FlaubertConfig(
         vocab_size=self.vocab_size,    n_special=self.n_special,    emb_dim=self.hidden_size,    n_layers=self.num_hidden_layers,    n_heads=self.num_attention_heads,    dropout=self.hidden_dropout_prob,    attention_dropout=self.attention_probs_dropout_prob,    gelu_activation=self.gelu_activation,    sinusoidal_embeddings=self.sinusoidal_embeddings,    asm=self.asm,    causal=self.causal,    n_langs=self.n_langs,    max_position_embeddings=self.max_position_embeddings,    initializer_range=self.initializer_range,    summary_type=self.summary_type,    use_proj=self.use_proj,    )



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






     """simple docstring"""






     a__  =FlaubertModel(config=lowercase_       )
     model.to(lowercase_       )
     model.eval()
     a__  =model(lowercase_,    lengths=lowercase_,    langs=lowercase_       )
     a__  =model(lowercase_,    langs=lowercase_       )
     a__  =model(lowercase_       )
     self.parent.assertEqual(result.last_hidden_state.shape,    (self.batch_size, self.seq_length, self.hidden_size)       )



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






     """simple docstring"""






     a__  =FlaubertWithLMHeadModel(lowercase_       )
     model.to(lowercase_       )
     model.eval()

     a__  =model(lowercase_,    token_type_ids=lowercase_,    labels=lowercase_       )
     self.parent.assertEqual(result.loss.shape,    ()       )
     self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.seq_length, self.vocab_size)       )



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






     """simple docstring"""






     a__  =FlaubertForQuestionAnsweringSimple(lowercase_       )
     model.to(lowercase_       )
     model.eval()

     a__  =model(lowercase_       )

     a__  =model(lowercase_,    start_positions=lowercase_,    end_positions=lowercase_       )
     self.parent.assertEqual(result.start_logits.shape,    (self.batch_size, self.seq_length)       )
     self.parent.assertEqual(result.end_logits.shape,    (self.batch_size, self.seq_length)       )



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






     """simple docstring"""






     a__  =FlaubertForQuestionAnswering(lowercase_       )
     model.to(lowercase_       )
     model.eval()

     a__  =model(lowercase_       )

     a__  =model(
         lowercase_,    start_positions=lowercase_,    end_positions=lowercase_,    cls_index=lowercase_,    is_impossible=lowercase_,    p_mask=lowercase_,    )

     a__  =model(
         lowercase_,    start_positions=lowercase_,    end_positions=lowercase_,    cls_index=lowercase_,    is_impossible=lowercase_,    )

     ((a__),      )  =result_with_labels.to_tuple()

     a__  =model(lowercase_,    start_positions=lowercase_,    end_positions=lowercase_       )

     ((a__),      )  =result_with_labels.to_tuple()

     self.parent.assertEqual(result_with_labels.loss.shape,    ()       )
     self.parent.assertEqual(result.start_top_log_probs.shape,    (self.batch_size, model.config.start_n_top)       )
     self.parent.assertEqual(result.start_top_index.shape,    (self.batch_size, model.config.start_n_top)       )
     self.parent.assertEqual(
         result.end_top_log_probs.shape,    (self.batch_size, model.config.start_n_top * model.config.end_n_top)       )
     self.parent.assertEqual(
         result.end_top_index.shape,    (self.batch_size, model.config.start_n_top * model.config.end_n_top)       )
     self.parent.assertEqual(result.cls_logits.shape,    (self.batch_size,)       )



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






     """simple docstring"""






     a__  =FlaubertForSequenceClassification(lowercase_       )
     model.to(lowercase_       )
     model.eval()

     a__  =model(lowercase_       )
     a__  =model(lowercase_,    labels=lowercase_       )

     self.parent.assertEqual(result.loss.shape,    ()       )
     self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.type_sequence_label_size)       )



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






     """simple docstring"""






     a__  =self.num_labels
     a__  =FlaubertForTokenClassification(lowercase_       )
     model.to(lowercase_       )
     model.eval()

     a__  =model(lowercase_,    attention_mask=lowercase_,    labels=lowercase_       )
     self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.seq_length, self.num_labels)       )



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






     """simple docstring"""






     a__  =self.num_choices
     a__  =FlaubertForMultipleChoice(config=lowercase_       )
     model.to(lowercase_       )
     model.eval()
     a__  =input_ids.unsqueeze(1       ).expand(-1,    self.num_choices,    -1       ).contiguous()
     a__  =token_type_ids.unsqueeze(1       ).expand(-1,    self.num_choices,    -1       ).contiguous()
     a__  =input_mask.unsqueeze(1       ).expand(-1,    self.num_choices,    -1       ).contiguous()
     a__  =model(
         lowercase_,    attention_mask=lowercase_,    token_type_ids=lowercase_,    labels=lowercase_,    )
     self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.num_choices)       )



   def  _UpperCAmelCase (				self       )							->    Dict:






     """simple docstring"""






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







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

   '''simple docstring'''




   lowerCamelCase__ :	str    =   (
       (
           FlaubertModel,
           FlaubertWithLMHeadModel,
           FlaubertForQuestionAnswering,
           FlaubertForQuestionAnsweringSimple,
           FlaubertForSequenceClassification,
           FlaubertForTokenClassification,
           FlaubertForMultipleChoice,
       )
       if is_torch_available()
       else ()
   )
   lowerCamelCase__ :	Dict    =   (
       {
           'feature-extraction': FlaubertModel,
           'fill-mask': FlaubertWithLMHeadModel,
           'question-answering': FlaubertForQuestionAnsweringSimple,
           'text-classification': FlaubertForSequenceClassification,
           'token-classification': FlaubertForTokenClassification,
           'zero-shot': FlaubertForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )



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






     """simple docstring"""






     if (
         pipeline_test_casse_name == "QAPipelineTests"
         and tokenizer_name is not None
         and not tokenizer_name.endswith('''Fast'''       )
     ):
       # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
       # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
       # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
       return True

     return False



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_,    lowercase_=False       )							->    str:






     """simple docstring"""






     a__  =super()._prepare_for_class(lowercase_,    lowercase_,    return_labels=lowercase_       )

     if return_labels:
       if model_class.__name__ == "FlaubertForQuestionAnswering":
         a__  =torch.zeros(
             self.model_tester.batch_size,    dtype=torch.long,    device=lowercase_       )
         a__  =torch.zeros(
             self.model_tester.batch_size,    dtype=torch.long,    device=lowercase_       )

     return inputs_dict



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






     """simple docstring"""






     a__  =FlaubertModelTester(self       )
     a__  =ConfigTester(self,    config_class=lowercase_,    emb_dim=37       )



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






     """simple docstring"""






     self.config_tester.run_common_tests()



   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self       )							->    str:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self       )							->    Dict:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self       )							->    Dict:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self       )							->    Tuple:






     """simple docstring"""






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



   @slow
   def  _UpperCAmelCase (				self       )							->    Tuple:






     """simple docstring"""






     for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       a__  =FlaubertModel.from_pretrained(lowercase_       )
       self.assertIsNotNone(lowercase_       )



   @slow
   @require_torch_gpu
   def  _UpperCAmelCase (				self       )							->    int:






     """simple docstring"""






     a__,      a__  =self.model_tester.prepare_config_and_inputs_for_common()
     for model_class in self.all_model_classes:
       # FlauBertForMultipleChoice behaves incorrectly in JIT environments.
       if model_class == FlaubertForMultipleChoice:
         return

       a__  =True
       a__  =model_class(config=lowercase_       )

       a__  =self._prepare_for_class(lowercase_,    lowercase_       )
       a__  =torch.jit.trace(
           lowercase_,    (inputs_dict['''input_ids'''].to('''cpu'''       ), inputs_dict['''attention_mask'''].to('''cpu'''       ))       )

       with tempfile.TemporaryDirectory() as tmp:
         torch.jit.save(lowercase_,    os.path.join(lowercase_,    '''traced_model.pt'''       )       )
         a__  =torch.jit.load(os.path.join(lowercase_,    '''traced_model.pt'''       ),    map_location=lowercase_       )
         loaded(inputs_dict['''input_ids'''].to(lowercase_       ),    inputs_dict['''attention_mask'''].to(lowercase_       )       )







@require_torch
class 						__magic_name__			(							unittest.TestCase							):

   '''simple docstring'''




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






     """simple docstring"""






     a__  =FlaubertModel.from_pretrained('''flaubert/flaubert_base_cased'''       )
     a__  =torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]       )
     with torch.no_grad():
       a__  =model(lowercase_       )[0]
     a__  =torch.Size((1, 11, 768)       )
     self.assertEqual(output.shape,    lowercase_       )
     a__  =torch.tensor(
         [[[-2.6251, -1.4298, -0.0227], [-2.8510, -1.6387, 0.2258], [-2.8114, -1.1832, -0.3066]]]       )

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