

'''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        A__   (   _snake_case      ):
    lowercase						       =		"sequence-classification"



    def __init__(       self		,					UpperCamelCase__							)    ->   int:





           '''simple docstring'''
           if type(UpperCamelCase__							) == dict:
                  A_						       =   Namespace(**UpperCamelCase__							)
           A_						       =   glue_output_modes[hparams.task]
           A_						       =   glue_tasks_num_labels[hparams.task]

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



    def        snake_case_      (       self		,					**UpperCamelCase__							)    ->   Dict:





           '''simple docstring'''
           return self.model(**UpperCamelCase__							)



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__							)    ->   str:





           '''simple docstring'''
           A_						       =   {"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}

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

           A_						       =   self(**UpperCamelCase__							)
           A_						       =   outputs[0]

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



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





           '''simple docstring'''
           A_						       =   self.hparams
           A_						       =   processors[args.task]()
           A_						       =   processor.get_labels()

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



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__ = False							)    ->   DataLoader:





           '''simple docstring'''
           A_						       =   """dev""" if mode == """test""" else mode

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

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



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__							)    ->   Any:





           '''simple docstring'''
           A_						       =   {"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}

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

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

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



    def        snake_case_      (       self		,					UpperCamelCase__							)    ->   tuple:





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

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

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

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

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



    def        snake_case_      (       self		,					UpperCamelCase__							)    ->   dict:





           '''simple docstring'''
           A_       ,      A_       ,      A_						       =   self._eval_end(UpperCamelCase__							)
           A_						       =   ret["""log"""]
           return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}



    def        snake_case_      (       self		,					UpperCamelCase__							)    ->   dict:





           '''simple docstring'''
           A_       ,      A_       ,      A_						       =   self._eval_end(UpperCamelCase__							)
           A_						       =   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        snake_case_      (       UpperCamelCase__		,					UpperCamelCase__							)    ->   Union[str, Any]:





           '''simple docstring'''
           BaseTransformer.add_model_specific_args(UpperCamelCase__		,					UpperCamelCase__							)
           parser.add_argument(
               """--max_seq_length"""		,					default=128		,					type=UpperCamelCase__		,					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=UpperCamelCase__		,					required=UpperCamelCase__		,					help="""The GLUE task to run"""		,					)
           parser.add_argument(
               """--gpus"""		,					default=0		,					type=UpperCamelCase__		,					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__     (   )      ->      List[str]:
       A_						       =   argparse.ArgumentParser()
       add_generic_args(UpperCAmelCase__,							os.getcwd()					)
       A_						       =   GLUETransformer.add_model_specific_args(UpperCAmelCase__,							os.getcwd()					)
       A_						       =   parser.parse_args()

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

       A_						       =   GLUETransformer(UpperCAmelCase__					)
       A_						       =   generic_train(UpperCAmelCase__,							UpperCAmelCase__					)

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


if __name__ == "__main__":
      main()



'''simple docstring'''







from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


__lowerCamelCase									=     logging.get_logger(__name__)


__lowerCamelCase									=     {
    '''Salesforce/codegen-350M-nl''': '''https://huggingface.co/Salesforce/codegen-350M-nl/resolve/main/config.json''',
    '''Salesforce/codegen-350M-multi''': '''https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/config.json''',
    '''Salesforce/codegen-350M-mono''': '''https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/config.json''',
    '''Salesforce/codegen-2B-nl''': '''https://huggingface.co/Salesforce/codegen-2B-nl/resolve/main/config.json''',
    '''Salesforce/codegen-2B-multi''': '''https://huggingface.co/Salesforce/codegen-2B-multi/resolve/main/config.json''',
    '''Salesforce/codegen-2B-mono''': '''https://huggingface.co/Salesforce/codegen-2B-mono/resolve/main/config.json''',
    '''Salesforce/codegen-6B-nl''': '''https://huggingface.co/Salesforce/codegen-6B-nl/resolve/main/config.json''',
    '''Salesforce/codegen-6B-multi''': '''https://huggingface.co/Salesforce/codegen-6B-multi/resolve/main/config.json''',
    '''Salesforce/codegen-6B-mono''': '''https://huggingface.co/Salesforce/codegen-6B-mono/resolve/main/config.json''',
    '''Salesforce/codegen-16B-nl''': '''https://huggingface.co/Salesforce/codegen-16B-nl/resolve/main/config.json''',
    '''Salesforce/codegen-16B-multi''': '''https://huggingface.co/Salesforce/codegen-16B-multi/resolve/main/config.json''',
    '''Salesforce/codegen-16B-mono''': '''https://huggingface.co/Salesforce/codegen-16B-mono/resolve/main/config.json''',
}


class        A__   (   _snake_case      ):
    lowercase						       =		"codegen"
    lowercase						       =		{
        "max_position_embeddings": "n_positions",
        "hidden_size": "n_embd",
        "num_attention_heads": "n_head",
        "num_hidden_layers": "n_layer",
    }



    def __init__(       self		,					UpperCamelCase__=50400		,					UpperCamelCase__=2048		,					UpperCamelCase__=2048		,					UpperCamelCase__=4096		,					UpperCamelCase__=28		,					UpperCamelCase__=16		,					UpperCamelCase__=64		,					UpperCamelCase__=None		,					UpperCamelCase__="gelu_new"		,					UpperCamelCase__=0.0		,					UpperCamelCase__=0.0		,					UpperCamelCase__=0.0		,					UpperCamelCase__=1e-5		,					UpperCamelCase__=0.02		,					UpperCamelCase__=True		,					UpperCamelCase__=50256		,					UpperCamelCase__=50256		,					UpperCamelCase__=False		,					**UpperCamelCase__		,					)    ->   Tuple:





           '''simple docstring'''
           A_						       =   vocab_size
           A_						       =   n_ctx
           A_						       =   n_positions
           A_						       =   n_embd
           A_						       =   n_layer
           A_						       =   n_head
           A_						       =   n_inner
           A_						       =   rotary_dim
           A_						       =   activation_function
           A_						       =   resid_pdrop
           A_						       =   embd_pdrop
           A_						       =   attn_pdrop
           A_						       =   layer_norm_epsilon
           A_						       =   initializer_range
           A_						       =   use_cache

           A_						       =   bos_token_id
           A_						       =   eos_token_id

           super().__init__(
               bos_token_id=UpperCamelCase__		,					eos_token_id=UpperCamelCase__		,					tie_word_embeddings=UpperCamelCase__		,					**UpperCamelCase__							)


class        A__   (   _snake_case      ):



    def __init__(       self		,					UpperCamelCase__		,					UpperCamelCase__ = "default"		,					UpperCamelCase__ = None		,					UpperCamelCase__ = False		,					)    ->   Tuple:





           '''simple docstring'''
           super().__init__(UpperCamelCase__		,					task=UpperCamelCase__		,					patching_specs=UpperCamelCase__		,					use_past=UpperCamelCase__							)
           if not getattr(self._config		,					"""pad_token_id"""		,					UpperCamelCase__							):
                  # TODO: how to do that better?
                  A_						       =   0



    @property
    def        snake_case_      (       self							)    ->   Mapping[str, Mapping[int, str]]:





           '''simple docstring'''
           A_						       =   OrderedDict({"""input_ids""": {0: """batch""", 1: """sequence"""}}							)
           if self.use_past:
                  self.fill_with_past_key_values_(UpperCamelCase__		,					direction="""inputs"""							)
                  A_						       =   {0: """batch""", 1: """past_sequence + sequence"""}
           else:
                  A_						       =   {0: """batch""", 1: """sequence"""}

           return common_inputs



    @property
    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           return self._config.n_layer



    @property
    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           return self._config.n_head



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__ = -1		,					UpperCamelCase__ = -1		,					UpperCamelCase__ = False		,					UpperCamelCase__ = None		,					)    ->   Mapping[str, Any]:





           '''simple docstring'''
           A_						       =   super(UpperCamelCase__		,					self							).generate_dummy_inputs(
               UpperCamelCase__		,					batch_size=UpperCamelCase__		,					seq_length=UpperCamelCase__		,					is_pair=UpperCamelCase__		,					framework=UpperCamelCase__							)

           # We need to order the input in the way they appears in the forward()
           A_						       =   OrderedDict({"""input_ids""": common_inputs["""input_ids"""]}							)

           # Need to add the past_keys
           if self.use_past:
                  if not is_torch_available():
                         raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""							)
                  else:
                         import torch

                         A_       ,      A_						       =   common_inputs["""input_ids"""].shape
                         # Not using the same length for past_key_values
                         A_						       =   seqlen + 2
                         A_						       =   (
                             batch,
                             self.num_attention_heads,
                             past_key_values_length,
                             self._config.hidden_size // self.num_attention_heads,
                         )
                         A_						       =   [
                             (torch.zeros(UpperCamelCase__							), torch.zeros(UpperCamelCase__							)) for _ in range(self.num_layers							)
                         ]

           A_						       =   common_inputs["""attention_mask"""]
           if self.use_past:
                  A_						       =   ordered_inputs["""attention_mask"""].dtype
                  A_						       =   torch.cat(
                      [ordered_inputs["""attention_mask"""], torch.ones(UpperCamelCase__		,					UpperCamelCase__		,					dtype=UpperCamelCase__							)]		,					dim=1							)

           return ordered_inputs







    @property
    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           return 13
