
import argparse
import glob
import logging
import os
import time
from argparse import Namespace

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from torch.utils.data import DataLoader, TensorDataset

from transformers import glue_compute_metrics as compute_metrics
from transformers import glue_convert_examples_to_features as convert_examples_to_features
from transformers import glue_output_modes, glue_tasks_num_labels
from transformers import glue_processors as processors


__lowerCamelCase							=       logging.getLogger(__name__)






class   UpperCAmelCase		(						A_ ):
   A__					:		int						     =     "sequence-classification"



   def __init__(self		:					List[str]    ,						snake_case__		:					str )		->	Dict:


          '''simple docstring'''


          if type(snake_case__ ) == dict:
                 snake_case       :   Dict							=      Namespace(**snake_case__ )
          snake_case       :   Dict							=      glue_output_modes[hparams.task]
          snake_case       :   Union[str, Any]							=      glue_tasks_num_labels[hparams.task]

          super().__init__(snake_case__    ,						snake_case__    ,						self.mode )



   def 				_SCREAMING_SNAKE_CASE   (self		:					Tuple    ,						**snake_case__		:					Optional[Any] )		->	Union[str, Any]:


          '''simple docstring'''


          return self.model(**snake_case__ )



   def 				_SCREAMING_SNAKE_CASE   (self		:					Union[str, Any]    ,						snake_case__		:					List[str]    ,						snake_case__		:					List[Any] )		->	str:


          '''simple docstring'''


          snake_case       :   Union[str, Any]							=      {"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}

          if self.config.model_type not in ["distilbert", "bart"]:
                 snake_case       :   Union[str, Any]							=      batch[2] if self.config.model_type in ["bert", "xlnet", "albert"] else None

          snake_case       :   str							=      self(**snake_case__ )
          snake_case       :   List[str]							=      outputs[0]

          snake_case       :   List[str]							=      self.trainer.lr_schedulers[0]["scheduler"]
          snake_case       :   Union[str, Any]							=      {"loss": loss, "rate": lr_scheduler.get_last_lr()[-1]}
          return {"loss": loss, "log": tensorboard_logs}



   def 				_SCREAMING_SNAKE_CASE   (self		:					Dict )		->	Optional[int]:


          '''simple docstring'''


          snake_case       :   Optional[Any]							=      self.hparams
          snake_case       :   Union[str, Any]							=      processors[args.task]()
          snake_case       :   Union[str, Any]							=      processor.get_labels()

          for mode in ["train", "dev"]:
                 snake_case       :   str							=      self._feature_file(snake_case__ )
                 if os.path.exists(snake_case__ ) and not args.overwrite_cache:
                        logger.info("Loading features from cached file %s"    ,						snake_case__ )
                 else:
                        logger.info("Creating features from dataset file at %s"    ,						args.data_dir )
                        snake_case       :   Union[str, Any]							=      (
                            processor.get_dev_examples(args.data_dir )
                            if mode == "dev"
                            else processor.get_train_examples(args.data_dir )
                        )
                        snake_case       :   Optional[int]							=      convert_examples_to_features(
                            snake_case__    ,						self.tokenizer    ,						max_length=args.max_seq_length    ,						label_list=self.labels    ,						output_mode=args.glue_output_mode    ,						)
                        logger.info("Saving features into cached file %s"    ,						snake_case__ )
                        torch.save(snake_case__    ,						snake_case__ )



   def 				_SCREAMING_SNAKE_CASE   (self		:					Dict    ,						snake_case__		:					str    ,						snake_case__		:					int    ,						snake_case__		:					bool = False )		->	DataLoader:


          '''simple docstring'''


          snake_case       :   Dict							=      "dev" if mode == "test" else mode

          snake_case       :   Any							=      self._feature_file(snake_case__ )
          logger.info("Loading features from cached file %s"    ,						snake_case__ )
          snake_case       :   Any							=      torch.load(snake_case__ )
          snake_case       :   Any							=      torch.tensor([f.input_ids for f in features]    ,						dtype=torch.long )
          snake_case       :   str							=      torch.tensor([f.attention_mask for f in features]    ,						dtype=torch.long )
          snake_case       :   int							=      torch.tensor([f.token_type_ids for f in features]    ,						dtype=torch.long )
          if self.hparams.glue_output_mode == "classification":
                 snake_case       :   List[str]							=      torch.tensor([f.label for f in features]    ,						dtype=torch.long )
          elif self.hparams.glue_output_mode == "regression":
                 snake_case       :   Dict							=      torch.tensor([f.label for f in features]    ,						dtype=torch.float )

          return DataLoader(
              TensorDataset(snake_case__    ,						snake_case__    ,						snake_case__    ,						snake_case__ )    ,						batch_size=snake_case__    ,						shuffle=snake_case__    ,						)



   def 				_SCREAMING_SNAKE_CASE   (self		:					str    ,						snake_case__		:					int    ,						snake_case__		:					Tuple )		->	str:


          '''simple docstring'''


          snake_case       :   str							=      {"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}

          if self.config.model_type not in ["distilbert", "bart"]:
                 snake_case       :   List[Any]							=      batch[2] if self.config.model_type in ["bert", "xlnet", "albert"] else None

          snake_case       :   Dict							=      self(**snake_case__ )
          snake_case				,					snake_case       :   Union[str, Any]							=      outputs[:2]
          snake_case       :   Union[str, Any]							=      logits.detach().cpu().numpy()
          snake_case       :   Dict							=      inputs["labels"].detach().cpu().numpy()

          return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}



   def 				_SCREAMING_SNAKE_CASE   (self		:					List[Any]    ,						snake_case__		:					Optional[int] )		->	tuple:


          '''simple docstring'''


          snake_case       :   int							=      torch.stack([x["val_loss"] for x in outputs] ).mean().detach().cpu().item()
          snake_case       :   Optional[int]							=      np.concatenate([x["pred"] for x in outputs]    ,						axis=0 )

          if self.hparams.glue_output_mode == "classification":
                 snake_case       :   int							=      np.argmax(snake_case__    ,						axis=1 )
          elif self.hparams.glue_output_mode == "regression":
                 snake_case       :   int							=      np.squeeze(snake_case__ )

          snake_case       :   Dict							=      np.concatenate([x["target"] for x in outputs]    ,						axis=0 )
          snake_case       :   List[str]							=      [[] for _ in range(out_label_ids.shape[0] )]
          snake_case       :   Optional[int]							=      [[] for _ in range(out_label_ids.shape[0] )]

          snake_case       :   str							=      {**{"val_loss": val_loss_mean}, **compute_metrics(self.hparams.task    ,						snake_case__    ,						snake_case__ )}

          snake_case       :   Any							=      dict(results.items() )
          snake_case       :   List[Any]							=      results
          return ret, preds_list, out_label_list



   def 				_SCREAMING_SNAKE_CASE   (self		:					Union[str, Any]    ,						snake_case__		:					list )		->	dict:


          '''simple docstring'''


          snake_case				,					snake_case				,					snake_case       :   Union[str, Any]							=      self._eval_end(snake_case__ )
          snake_case       :   List[str]							=      ret["log"]
          return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}



   def 				_SCREAMING_SNAKE_CASE   (self		:					int    ,						snake_case__		:					Any )		->	dict:


          '''simple docstring'''


          snake_case				,					snake_case				,					snake_case       :   Optional[int]							=      self._eval_end(snake_case__ )
          snake_case       :   List[Any]							=      ret["log"]
          # `val_loss` is the key returned by `self._eval_end()` but actually refers to `test_loss`
          return {"avg_test_loss": logs["val_loss"], "log": logs, "progress_bar": logs}



   @staticmethod
   def 				_SCREAMING_SNAKE_CASE   (snake_case__		:					List[Any]    ,						snake_case__		:					List[str] )		->	str:


          '''simple docstring'''


          BaseTransformer.add_model_specific_args(snake_case__    ,						snake_case__ )
          parser.add_argument(
              "--max_seq_length"    ,						default=1_28    ,						type=snake_case__    ,						help=(
                  "The maximum total input sequence length after tokenization. Sequences longer "
                  "than this will be truncated, sequences shorter will be padded."
              )    ,						)

          parser.add_argument(
              "--task"    ,						default=""    ,						type=snake_case__    ,						required=snake_case__    ,						help="The GLUE task to run"    ,						)
          parser.add_argument(
              "--gpus"    ,						default=0    ,						type=snake_case__    ,						help="The number of GPUs allocated for this, it is by default 0 meaning none"    ,						)

          parser.add_argument(
              "--overwrite_cache"    ,						action="store_true"    ,						help="Overwrite the cached training and evaluation sets" )

          return parser



def 			UpperCamelCase					(    ):
       snake_case       :   int							=      argparse.ArgumentParser()
       add_generic_args(__lowerCamelCase				,					os.getcwd()    )
       snake_case       :   Tuple							=      GLUETransformer.add_model_specific_args(__lowerCamelCase				,					os.getcwd()    )
       snake_case       :   Dict							=      parser.parse_args()

       # If output_dir not provided, a folder will be generated in pwd
       if args.output_dir is None:
              snake_case       :   int							=      os.path.join(
                  "./results"				,					f"""{args.task}_{time.strftime('%Y%m%d_%H%M%S'    )}"""				,					)
              os.makedirs(args.output_dir    )

       snake_case       :   Optional[int]							=      GLUETransformer(__lowerCamelCase    )
       snake_case       :   str							=      generic_train(__lowerCamelCase				,					__lowerCamelCase    )

       # Optionally, predict on dev set and write to output_dir
       if args.do_predict:
              snake_case       :   str							=      sorted(glob.glob(os.path.join(args.output_dir				,					"checkpoint-epoch=*.ckpt"    )				,					recursive=__lowerCamelCase    )    )
              snake_case       :   List[str]							=      model.load_from_checkpoint(checkpoints[-1]    )
              return trainer.test(__lowerCamelCase    )


if __name__ == "__main__":
   main()







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


if is_sentencepiece_available():
   from ..ta.tokenization_ta import TaTokenizer
else:
   from ...utils.dummy_sentencepiece_objects import TaTokenizer

__lowerCamelCase							=       TaTokenizer

if is_tokenizers_available():
   from ..ta.tokenization_ta_fast import TaTokenizerFast
else:
   from ...utils.dummy_tokenizers_objects import TaTokenizerFast

__lowerCamelCase							=       TaTokenizerFast

__lowerCamelCase							=       {"""configuration_mt5""": ["""MT5Config""", """MT5OnnxConfig"""]}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase							=       [
       """MT5EncoderModel""",
       """MT5ForConditionalGeneration""",
       """MT5ForQuestionAnswering""",
       """MT5Model""",
       """MT5PreTrainedModel""",
       """MT5Stack""",
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase							=       ["""TFMT5EncoderModel""", """TFMT5ForConditionalGeneration""", """TFMT5Model"""]

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase							=       ["""FlaxMT5EncoderModel""", """FlaxMT5ForConditionalGeneration""", """FlaxMT5Model"""]


if TYPE_CHECKING:
   from .configuration_mta import MTaConfig, MTaOnnxConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_mta import (
          MTaEncoderModel,
          MTaForConditionalGeneration,
          MTaForQuestionAnswering,
          MTaModel,
          MTaPreTrainedModel,
          MTaStack,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_mta import TFMTaEncoderModel, TFMTaForConditionalGeneration, TFMTaModel

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_mta import FlaxMTaEncoderModel, FlaxMTaForConditionalGeneration, FlaxMTaModel

else:
   import sys

   __lowerCamelCase							=       _LazyModule(
       __name__,
       globals()["""__file__"""],
       _import_structure,
       extra_objects={"""MT5Tokenizer""": MTaTokenizer, """MT5TokenizerFast""": MTaTokenizerFast},
       module_spec=__spec__,
   )





