



'''simple docstring'''





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 			lowerCAmelCase_(							SCREAMING_SNAKE_CASE_	):
  '''simple docstring'''




  __lowercase				: str           = '''sequence-classification'''




  def __init__(       self    ,__UpperCAmelCase )       ->						List[str]:
      if type(_lowerCAmelCase ) == dict:
          lowerCAmelCase__					:    Optional[int]				 =   Namespace(**_lowerCAmelCase )
      lowerCAmelCase__					:    int				 =   glue_output_modes[hparams.task]
      lowerCAmelCase__					:    Optional[Any]				 =   glue_tasks_num_labels[hparams.task]

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




  def 				UpperCAmelCase_      (       self    ,**__UpperCAmelCase )       ->						Optional[int]:
      return self.model(**_lowerCAmelCase )




  def 				UpperCAmelCase_      (       self    ,__UpperCAmelCase    ,__UpperCAmelCase )       ->						int:
      lowerCAmelCase__					:    Union[str, Any]				 =   {"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}

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

      lowerCAmelCase__					:    Optional[int]				 =   self(**_lowerCAmelCase )
      lowerCAmelCase__					:    int				 =   outputs[0]

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




  def 				UpperCAmelCase_      (       self )       ->						str:
      lowerCAmelCase__					:    List[Any]				 =   self.hparams
      lowerCAmelCase__					:    str				 =   processors[args.task]()
      lowerCAmelCase__					:    List[str]				 =   processor.get_labels()

      for mode in ["train", "dev"]:
          lowerCAmelCase__					:    str				 =   self._feature_file(_lowerCAmelCase )
          if os.path.exists(_lowerCAmelCase ) and not args.overwrite_cache:
              logger.info("""Loading features from cached file %s"""    ,_lowerCAmelCase )
          else:
              logger.info("""Creating features from dataset file at %s"""    ,args.data_dir )
              lowerCAmelCase__					:    int				 =   (
                  processor.get_dev_examples(args.data_dir )
                  if mode == """dev"""
                  else processor.get_train_examples(args.data_dir )
              )
              lowerCAmelCase__					:    List[str]				 =   convert_examples_to_features(
                  _lowerCAmelCase    ,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"""    ,_lowerCAmelCase )
              torch.save(_lowerCAmelCase    ,_lowerCAmelCase )




  def 				UpperCAmelCase_      (       self    ,__UpperCAmelCase    ,__UpperCAmelCase    ,__UpperCAmelCase = False )       ->						DataLoader:
      lowerCAmelCase__					:    str				 =   """dev""" if mode == """test""" else mode

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

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




  def 				UpperCAmelCase_      (       self    ,__UpperCAmelCase    ,__UpperCAmelCase )       ->						Optional[Any]:
      lowerCAmelCase__					:    List[str]				 =   {"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}

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

      lowerCAmelCase__					:    List[str]				 =   self(**_lowerCAmelCase )
      lowerCAmelCase__ ,  lowerCAmelCase__					:    Tuple				 =   outputs[:2]
      lowerCAmelCase__					:    List[str]				 =   logits.detach().cpu().numpy()
      lowerCAmelCase__					:    Tuple				 =   inputs["""labels"""].detach().cpu().numpy()

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




  def 				UpperCAmelCase_      (       self    ,__UpperCAmelCase )       ->						tuple:
      lowerCAmelCase__					:    Optional[int]				 =   torch.stack([x["""val_loss"""] for x in outputs] ).mean().detach().cpu().item()
      lowerCAmelCase__					:    Dict				 =   np.concatenate([x["""pred"""] for x in outputs]    ,axis=0 )

      if self.hparams.glue_output_mode == "classification":
          lowerCAmelCase__					:    Any				 =   np.argmax(_lowerCAmelCase    ,axis=1 )
      elif self.hparams.glue_output_mode == "regression":
          lowerCAmelCase__					:    Union[str, Any]				 =   np.squeeze(_lowerCAmelCase )

      lowerCAmelCase__					:    Tuple				 =   np.concatenate([x["""target"""] for x in outputs]    ,axis=0 )
      lowerCAmelCase__					:    Union[str, Any]				 =   [[] for _ in range(out_label_ids.shape[0] )]
      lowerCAmelCase__					:    Tuple				 =   [[] for _ in range(out_label_ids.shape[0] )]

      lowerCAmelCase__					:    Tuple				 =   {**{"""val_loss""": val_loss_mean}, **compute_metrics(self.hparams.task    ,_lowerCAmelCase    ,_lowerCAmelCase )}

      lowerCAmelCase__					:    Union[str, Any]				 =   dict(results.items() )
      lowerCAmelCase__					:    Tuple				 =   results
      return ret, preds_list, out_label_list




  def 				UpperCAmelCase_      (       self    ,__UpperCAmelCase )       ->						dict:
      lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__					:    Union[str, Any]				 =   self._eval_end(_lowerCAmelCase )
      lowerCAmelCase__					:    int				 =   ret["""log"""]
      return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}




  def 				UpperCAmelCase_      (       self    ,__UpperCAmelCase )       ->						dict:
      lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__					:    Optional[Any]				 =   self._eval_end(_lowerCAmelCase )
      lowerCAmelCase__					:    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 				UpperCAmelCase_      (       __UpperCAmelCase    ,__UpperCAmelCase )       ->						Optional[Any]:
      BaseTransformer.add_model_specific_args(_lowerCAmelCase    ,_lowerCAmelCase )
      parser.add_argument(
          """--max_seq_length"""    ,default=128    ,type=_lowerCAmelCase    ,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=_lowerCAmelCase    ,required=_lowerCAmelCase    ,help="""The GLUE task to run"""    ,)
      parser.add_argument(
          """--gpus"""    ,default=0    ,type=_lowerCAmelCase    ,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   _SCREAMING_SNAKE_CASE    (			):



    """simple docstring"""





    lowerCAmelCase__					:    Union[str, Any]				 =   argparse.ArgumentParser()
    add_generic_args(lowerCAmelCase_				,       os.getcwd()     )
    lowerCAmelCase__					:    Any				 =   GLUETransformer.add_model_specific_args(lowerCAmelCase_				,       os.getcwd()     )
    lowerCAmelCase__					:    Optional[int]				 =   parser.parse_args()

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

    lowerCAmelCase__					:    List[str]				 =   GLUETransformer(lowerCAmelCase_     )
    lowerCAmelCase__					:    Optional[Any]				 =   generic_train(lowerCAmelCase_				,       lowerCAmelCase_     )

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


if __name__ == "__main__":
 main()






from __future__ import annotations



def       _lowerCAmelCase    (				lowerCAmelCase_				:int		,    lowerCAmelCase_				:int				)->list[str]:
   '''simple docstring'''






   if partitions <= 0:
      raise ValueError("partitions must be a positive number!"				)
   if partitions > number_of_bytes:
      raise ValueError("partitions can not > number_of_bytes!"				)
   snake_case_            =       number_of_bytes // partitions
   snake_case_            =       []
   for i in range(lowerCAmelCase_				):
      snake_case_            =       i * bytes_per_partition + 1
      snake_case_            =       (
          number_of_bytes if i == partitions - 1 else (i + 1) * bytes_per_partition
      )
      allocation_list.append(F'''{start_bytes}-{end_bytes}'''				)
   return allocation_list


if __name__ == "__main__":
   import doctest

   doctest.testmod()




