"""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 				lowercase__       (   __SCREAMING_SNAKE_CASE				):







       '''simple docstring'''


       UpperCamelCase						=					'sequence-classification'



       def __init__(							self      :  str		,    _UpperCAmelCase      :  Union[str, Any]   )			->						Optional[Any]:


             '''simple docstring'''




             if type(_a   ) == dict:
                   UpperCAmelCase_      			=							Namespace(**_a   )
             UpperCAmelCase_      			=							glue_output_modes[hparams.task]
             UpperCAmelCase_      			=							glue_tasks_num_labels[hparams.task]

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



       def 					lowercase__    (							self      :  Optional[Any]		,    **_UpperCAmelCase      :  int   )			->						Union[str, Any]:


             '''simple docstring'''




             return self.model(**_a   )



       def 					lowercase__    (							self      :  Tuple		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  List[str]   )			->						List[str]:


             '''simple docstring'''




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

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

             UpperCAmelCase_      			=							self(**_a   )
             UpperCAmelCase_      			=							outputs[0]

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



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


             '''simple docstring'''




             UpperCAmelCase_      			=							self.hparams
             UpperCAmelCase_      			=							processors[args.task]()
             UpperCAmelCase_      			=							processor.get_labels()

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



       def 					lowercase__    (							self      :  List[Any]		,    _UpperCAmelCase      :  Optional[int]		,    _UpperCAmelCase      :  List[str]		,    _UpperCAmelCase      :  List[str] = False   )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							"dev" if mode == "test" else mode

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

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



       def 					lowercase__    (							self      :  List[Any]		,    _UpperCAmelCase      :  Dict		,    _UpperCAmelCase      :  List[str]   )			->						int:


             '''simple docstring'''




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

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

             UpperCAmelCase_      			=							self(**_a   )
             UpperCAmelCase_			,   UpperCAmelCase_      			=							outputs[:2]
             UpperCAmelCase_      			=							logits.detach().cpu().numpy()
             UpperCAmelCase_      			=							inputs["labels"].detach().cpu().numpy()

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



       def 					lowercase__    (							self      :  Dict		,    _UpperCAmelCase      :  Optional[int]   )			->						str:


             '''simple docstring'''




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

             if self.hparams.glue_output_mode == "classification":
                   UpperCAmelCase_      			=							np.argmax(_a		,    axis=1   )
             elif self.hparams.glue_output_mode == "regression":
                   UpperCAmelCase_      			=							np.squeeze(_a   )

             UpperCAmelCase_      			=							np.concatenate([x["target"] for x in outputs]		,    axis=0   )
             UpperCAmelCase_      			=							[[] for _ in range(out_label_ids.shape[0]   )]
             UpperCAmelCase_      			=							[[] for _ in range(out_label_ids.shape[0]   )]

             UpperCAmelCase_      			=							{**{"val_loss": val_loss_mean}, **compute_metrics(self.hparams.task		,    _a		,    _a   )}

             UpperCAmelCase_      			=							dict(results.items()   )
             UpperCAmelCase_      			=							results
             return ret, preds_list, out_label_list



       def 					lowercase__    (							self      :  Any		,    _UpperCAmelCase      :  Tuple   )			->						List[str]:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							self._eval_end(_a   )
             UpperCAmelCase_      			=							ret["log"]
             return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}



       def 					lowercase__    (							self      :  List[Any]		,    _UpperCAmelCase      :  Optional[int]   )			->						int:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							self._eval_end(_a   )
             UpperCAmelCase_      			=							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 					lowercase__    (							_UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  List[Any]   )			->						Union[str, Any]:


             '''simple docstring'''




             BaseTransformer.add_model_specific_args(_a		,    _a   )
             parser.add_argument(
                 "--max_seq_length"		,    default=128		,    type=_a		,    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=_a		,    required=_a		,    help="The GLUE task to run"		,    )
             parser.add_argument(
                 "--gpus"		,    default=0		,    type=_a		,    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        a__     (				):
      UpperCAmelCase_      			=							argparse.ArgumentParser()
      add_generic_args(lowerCAmelCase__     ,      os.getcwd()				)
      UpperCAmelCase_      			=							GLUETransformer.add_model_specific_args(lowerCAmelCase__     ,      os.getcwd()				)
      UpperCAmelCase_      			=							parser.parse_args()

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

      UpperCAmelCase_      			=							GLUETransformer(lowerCAmelCase__				)
      UpperCAmelCase_      			=							generic_train(lowerCAmelCase__     ,      lowerCAmelCase__				)

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


if __name__ == "__main__":
  main()









"""simple docstring"""





import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


lowerCamelCase       	=						datasets.utils.logging.get_logger(__name__)
@dataclass
class 				lowercase__       (   datasets.BuilderConfig				):







       '''simple docstring'''


       UpperCamelCase						=					None
       UpperCamelCase						=					"utf-8"
       UpperCamelCase						=					None
       UpperCamelCase						=					None
       UpperCamelCase						=					True  # deprecated
       UpperCamelCase						=					None  # deprecated
       UpperCamelCase						=					10 << 20  # 10MB
       UpperCamelCase						=					None



class 				lowercase__       (   datasets.ArrowBasedBuilder				):







       '''simple docstring'''


       UpperCamelCase						=					JsonConfig



       def 					lowercase__    (							self      :  List[Any]   )			->						Dict:


             '''simple docstring'''




             if self.config.block_size is not None:
                   logger.warning("The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead"   )
                   UpperCAmelCase_      			=							self.config.block_size
             if self.config.use_threads is not True:
                   logger.warning(
                       "The JSON loader parameter `use_threads` is deprecated and doesn't have any effect anymore."   )
             if self.config.newlines_in_values is not None:
                   raise ValueError("The JSON loader parameter `newlines_in_values` is no longer supported"   )
             return datasets.DatasetInfo(features=self.config.features   )



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


             '''simple docstring'''




             if not self.config.data_files:
                   raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}"""   )
             UpperCAmelCase_      			=							dl_manager.download_and_extract(self.config.data_files   )
             if isinstance(_UpperCAmelCase		,    (str, list, tuple)   ):
                   UpperCAmelCase_      			=							data_files
                   if isinstance(_UpperCAmelCase		,    _UpperCAmelCase   ):
                         UpperCAmelCase_      			=							[files]
                   UpperCAmelCase_      			=							[dl_manager.iter_files(_UpperCAmelCase   ) for file in files]
                   return [datasets.SplitGenerator(name=datasets.Split.TRAIN		,    gen_kwargs={"files": files}   )]
             UpperCAmelCase_      			=							[]
             for split_name, files in data_files.items():
                   if isinstance(_UpperCAmelCase		,    _UpperCAmelCase   ):
                         UpperCAmelCase_      			=							[files]
                   UpperCAmelCase_      			=							[dl_manager.iter_files(_UpperCAmelCase   ) for file in files]
                   splits.append(datasets.SplitGenerator(name=_UpperCAmelCase		,    gen_kwargs={"files": files}   )   )
             return splits



       def 					lowercase__    (							self      :  str		,    _UpperCAmelCase      :  pa.Table   )			->						pa.Table:


             '''simple docstring'''




             if self.config.features is not None:
                   # adding missing columns
                   for column_name in set(self.config.features   ) - set(pa_table.column_names   ):
                         UpperCAmelCase_      			=							self.config.features.arrow_schema.field(_UpperCAmelCase   ).type
                         UpperCAmelCase_      			=							pa_table.append_column(_UpperCAmelCase		,    pa.array([None] * len(_UpperCAmelCase   )		,    type=_UpperCAmelCase   )   )
                   # more expensive cast to support nested structures with keys in a different order
                   # allows str <-> int/float or str to Audio for example
                   UpperCAmelCase_      			=							table_cast(_UpperCAmelCase		,    self.config.features.arrow_schema   )
             return pa_table







       def 					lowercase__    (							self      :  Optional[int]		,    _UpperCAmelCase      :  List[Any]   )			->						str:


             '''simple docstring'''




             for file_idx, file in enumerate(itertools.chain.from_iterable(_UpperCAmelCase   )   ):
                   # If the file is one json object and if we need to look at the list of items in one specific field
                   if self.config.field is not None:
                         with open(_UpperCAmelCase		,    encoding=self.config.encoding		,    errors=self.config.encoding_errors   ) as f:
                               UpperCAmelCase_      			=							json.load(_UpperCAmelCase   )

                         # We keep only the field we are interested in
                         UpperCAmelCase_      			=							dataset[self.config.field]

                         # We accept two format: a list of dicts or a dict of lists
                         if isinstance(_UpperCAmelCase		,    (list, tuple)   ):
                               UpperCAmelCase_      			=							set().union(*[row.keys() for row in dataset]   )
                               UpperCAmelCase_      			=							{col: [row.get(_UpperCAmelCase   ) for row in dataset] for col in keys}
                         else:
                               UpperCAmelCase_      			=							dataset
                         UpperCAmelCase_      			=							pa.Table.from_pydict(_UpperCAmelCase   )
                         yield file_idx, self._cast_table(_UpperCAmelCase   )

                   # If the file has one json object per line
                   else:
                         with open(_UpperCAmelCase		,    "rb"   ) as f:
                               UpperCAmelCase_      			=							0
                               # Use block_size equal to the chunk size divided by 32 to leverage multithreading
                               # Set a default minimum value of 16kB if the chunk size is really small
                               UpperCAmelCase_      			=							max(self.config.chunksize // 32		,    16 << 10   )
                               UpperCAmelCase_      			=							(
                                   self.config.encoding_errors if self.config.encoding_errors is not None else "strict"
                               )
                               while True:
                                     UpperCAmelCase_      			=							f.read(self.config.chunksize   )
                                     if not batch:
                                           break
                                     # Finish current line
                                     try:
                                           batch += f.readline()
                                     except (AttributeError, io.UnsupportedOperation):
                                           batch += readline(_UpperCAmelCase   )
                                     # PyArrow only accepts utf-8 encoded bytes
                                     if self.config.encoding != "utf-8":
                                           UpperCAmelCase_      			=							batch.decode(self.config.encoding		,    errors=_UpperCAmelCase   ).encode("utf-8"   )
                                     try:
                                           while True:
                                                 try:
                                                       UpperCAmelCase_      			=							paj.read_json(
                                                           io.BytesIO(_UpperCAmelCase   )		,    read_options=paj.ReadOptions(block_size=_UpperCAmelCase   )   )
                                                       break
                                                 except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
                                                       if (
                                                           isinstance(_UpperCAmelCase		,    pa.ArrowInvalid   )
                                                           and "straddling" not in str(_UpperCAmelCase   )
                                                           or block_size > len(_UpperCAmelCase   )
                                                       ):
                                                             raise
                                                       else:
                                                             # Increase the block size in case it was too small.
                                                             # The block size will be reset for the next file.
                                                             logger.debug(
                                                                 F"""Batch of {len(_UpperCAmelCase   )} bytes couldn't be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}."""   )
                                                             block_size *= 2
                                     except pa.ArrowInvalid as e:
                                           try:
                                                 with open(
                                                     _UpperCAmelCase		,    encoding=self.config.encoding		,    errors=self.config.encoding_errors   ) as f:
                                                       UpperCAmelCase_      			=							json.load(_UpperCAmelCase   )
                                           except json.JSONDecodeError:
                                                 logger.error(F"""Failed to read file '{file}' with error {type(_UpperCAmelCase   )}: {e}"""   )
                                                 raise e
                                           # If possible, parse the file as a list of json objects and exit the loop
                                           if isinstance(_UpperCAmelCase		,    _UpperCAmelCase   ):  # list is the only sequence type supported in JSON
                                                 try:
                                                       UpperCAmelCase_      			=							set().union(*[row.keys() for row in dataset]   )
                                                       UpperCAmelCase_      			=							{col: [row.get(_UpperCAmelCase   ) for row in dataset] for col in keys}
                                                       UpperCAmelCase_      			=							pa.Table.from_pydict(_UpperCAmelCase   )
                                                 except (pa.ArrowInvalid, AttributeError) as e:
                                                       logger.error(F"""Failed to read file '{file}' with error {type(_UpperCAmelCase   )}: {e}"""   )
                                                       raise ValueError(F"""Not able to read records in the JSON file at {file}."""   ) from None
                                                 yield file_idx, self._cast_table(_UpperCAmelCase   )
                                                 break
                                           else:
                                                 logger.error(F"""Failed to read file '{file}' with error {type(_UpperCAmelCase   )}: {e}"""   )
                                                 raise ValueError(
                                                     F"""Not able to read records in the JSON file at {file}. """
                                                     F"""You should probably indicate the field of the JSON file containing your records. """
                                                     F"""This JSON file contain the following fields: {str(list(dataset.keys()   )   )}. """
                                                     F"""Select the correct one and provide it as `field='XXX'` to the dataset loading method. """   ) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                                     yield (file_idx, batch_idx), self._cast_table(_UpperCAmelCase   )
                                     batch_idx += 1




