

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


lowercase_   							=  logging.getLogger(__name__)





class 					__lowerCAmelCase	(   SCREAMING_SNAKE_CASE				):
  _a		     =    """sequence-classification"""




  def __init__(							self				,  lowerCAmelCase   )    ->					Optional[int]:
        '''simple docstring'''






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

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




  def     A__	(							self				,  **lowerCAmelCase   )    ->					List[str]:
        '''simple docstring'''






        return self.model(**lowerCAmelCase   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase   )    ->					Optional[int]:
        '''simple docstring'''






        _lowercase      			={'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[3]}

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

        _lowercase      			=self(**lowerCAmelCase   )
        _lowercase      			=outputs[0]

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




  def     A__	(							self   )    ->					Dict:
        '''simple docstring'''






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

        for mode in ["train", "dev"]:
              _lowercase      			=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   )
                    _lowercase      			=(
                        processor.get_dev_examples(args.data_dir   )
                        if mode == 'dev'
                        else processor.get_train_examples(args.data_dir   )
                    )
                    _lowercase      			=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     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase = False   )    ->					DataLoader:
        '''simple docstring'''






        _lowercase      			='dev' if mode == 'test' else mode

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

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




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase   )    ->					int:
        '''simple docstring'''






        _lowercase      			={'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[3]}

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

        _lowercase      			=self(**lowerCAmelCase   )
        _lowercase			,			_lowercase      			=outputs[:2]
        _lowercase      			=logits.detach().cpu().numpy()
        _lowercase      			=inputs['labels'].detach().cpu().numpy()

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




  def     A__	(							self				,  lowerCAmelCase   )    ->					tuple:
        '''simple docstring'''






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

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

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

        _lowercase      			={**{'val_loss': val_loss_mean}, **compute_metrics(self.hparams.task				,  lowerCAmelCase				,  lowerCAmelCase   )}

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




  def     A__	(							self				,  lowerCAmelCase   )    ->					dict:
        '''simple docstring'''






        _lowercase			,			_lowercase			,			_lowercase      			=self._eval_end(lowerCAmelCase   )
        _lowercase      			=ret['log']
        return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}




  def     A__	(							self				,  lowerCAmelCase   )    ->					dict:
        '''simple docstring'''






        _lowercase			,			_lowercase			,			_lowercase      			=self._eval_end(lowerCAmelCase   )
        _lowercase      			=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     A__	(							lowerCAmelCase				,  lowerCAmelCase   )    ->					Optional[Any]:
        '''simple docstring'''






        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 	a (     )       -> List[str]:





      """simple docstring"""



      _lowercase      			=argparse.ArgumentParser()
      add_generic_args(A__ ,  os.getcwd()	)
      _lowercase      			=GLUETransformer.add_model_specific_args(A__ ,  os.getcwd()	)
      _lowercase      			=parser.parse_args()

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

      _lowercase      			=GLUETransformer(A__	)
      _lowercase      			=generic_train(A__ ,  A__	)

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


if __name__ == "__main__":
  main()





import json
import os
import tempfile

import transformers

import datasets
from utils import generate_example_dataset, get_duration


lowercase_   							=  5_0_0_0_0_0

lowercase_							,		lowercase_   							=  os.path.split(__file__)
lowercase_   							=  os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))
@get_duration
def 	a (     A__      :  datasets.Dataset ,  **A__      :  List[str]	)       -> Optional[Any]:





      """simple docstring"""



      _lowercase      			=dataset.map(**A__	)
@get_duration
def 	a (     A__      :  datasets.Dataset ,  **A__      :  List[Any]	)       -> List[str]:





      """simple docstring"""



      _lowercase      			=dataset.filter(**A__	)
def 	a (     )       -> Union[str, Any]:





      """simple docstring"""



      _lowercase      			={'num examples': SPEED_TEST_N_EXAMPLES}
      with tempfile.TemporaryDirectory() as tmp_dir:
            _lowercase      			=datasets.Features({'text': datasets.Value('string'	), 'numbers': datasets.Value('float32'	)}	)
            _lowercase      			=generate_example_dataset(
                os.path.join(A__ ,  'dataset.arrow'	) ,  A__ ,  num_examples=A__	)

            _lowercase      			=transformers.AutoTokenizer.from_pretrained('bert-base-cased' ,  use_fast=A__	)

            def tokenize(A__      :  int	):
                  return tokenizer(examples['text']	)

            _lowercase      			=map(A__	)

            _lowercase      			=map(A__ ,  batched=A__	)

            _lowercase      			=map(A__ ,  function=lambda A__	: None ,  batched=A__	)

            with dataset.formatted_as(type='numpy'	):
                  _lowercase      			=map(A__ ,  function=lambda A__	: None ,  batched=A__	)

            with dataset.formatted_as(type='pandas'	):
                  _lowercase      			=map(A__ ,  function=lambda A__	: None ,  batched=A__	)

            with dataset.formatted_as(type='torch' ,  columns='numbers'	):
                  _lowercase      			=map(A__ ,  function=lambda A__	: None ,  batched=A__	)

            with dataset.formatted_as(type='tensorflow' ,  columns='numbers'	):
                  _lowercase      			=map(A__ ,  function=lambda A__	: None ,  batched=A__	)

            _lowercase      			=map(A__ ,  function=A__ ,  batched=A__	)

            _lowercase      			=filter(A__	)

            # Activate later when tokenizer support batched inputs
            # with dataset.formatted_as(type='numpy'):
            #     times[func.__name__ + " fast-tokenizer batched numpy"] = func(dataset, function=tokenize, batched=True)

      with open(A__ ,  'wb'	) as f:
            f.write(json.dumps(A__	).encode('utf-8'	)	)


if __name__ == "__main__":  # useful to run the profiler
  benchmark_map_filter()


