
"""simple docstring"""




import argparse
import glob
import logging
import os
from argparse import Namespace
from importlib import import_module

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from seqeval.metrics import accuracy_score, fa_score, precision_score, recall_score
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader, TensorDataset
from utils_ner import TokenClassificationTask


A__      :	int          =    logging.getLogger(__name__)







class   _lowercase						(       lowerCAmelCase_				):


     '''simple docstring'''





     _A								=							'token-classification'



     def __init__(      self ,			__UpperCamelCase  )->					Union[str, Any]:
          if type(__UpperCamelCase  ) == dict:
               UpperCAmelCase__							: int 							=						Namespace(**__UpperCamelCase  )
          UpperCAmelCase__							: Dict 							=						import_module("tasks"  )
          try:
               UpperCAmelCase__							: Optional[int] 							=						getattr(__UpperCamelCase ,			hparams.task_type  )
               UpperCAmelCase__							: TokenClassificationTask 							=						token_classification_task_clazz()
          except AttributeError:
               raise ValueError(
                   F"Task {hparams.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. "
                   F"Available tasks classes are: {TokenClassificationTask.__subclasses__()}"  )
          UpperCAmelCase__							: List[Any] 							=						self.token_classification_task.get_labels(hparams.labels  )
          UpperCAmelCase__							: Optional[int] 							=						CrossEntropyLoss().ignore_index
          super().__init__(__UpperCamelCase ,			len(self.labels  ) ,			self.mode  )



     def  lowerCAmelCase__       (      self ,			**__UpperCamelCase  )->					Dict:
          return self.model(**__UpperCamelCase  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase  )->					List[Any]:
          UpperCAmelCase__							: int 							=						{"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}
          if self.config.model_type != "distilbert":
               UpperCAmelCase__							: Any 							=						(
                   batch[2] if self.config.model_type in ["bert", "xlnet"] else None
               )  # XLM and RoBERTa don"t use token_type_ids

          UpperCAmelCase__							: Union[str, Any] 							=						self(**__UpperCamelCase  )
          UpperCAmelCase__							: List[str] 							=						outputs[0]
          # tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
          return {"loss": loss}



     def  lowerCAmelCase__       (      self  )->					Dict:
          UpperCAmelCase__							: str 							=						self.hparams
          for mode in ["train", "dev", "test"]:
               UpperCAmelCase__							: Any 							=						self._feature_file(__UpperCamelCase  )
               if os.path.exists(__UpperCamelCase  ) and not args.overwrite_cache:
                    logger.info("Loading features from cached file %s" ,			__UpperCamelCase  )
                    UpperCAmelCase__							: str 							=						torch.load(__UpperCamelCase  )
               else:
                    logger.info("Creating features from dataset file at %s" ,			args.data_dir  )
                    UpperCAmelCase__							: List[Any] 							=						self.token_classification_task.read_examples_from_file(args.data_dir ,			__UpperCamelCase  )
                    UpperCAmelCase__							: Optional[int] 							=						self.token_classification_task.convert_examples_to_features(
                        __UpperCamelCase ,			self.labels ,			args.max_seq_length ,			self.tokenizer ,			cls_token_at_end=bool(self.config.model_type in ["xlnet"]  ) ,			cls_token=self.tokenizer.cls_token ,			cls_token_segment_id=2 if self.config.model_type in ["xlnet"] else 0 ,			sep_token=self.tokenizer.sep_token ,			sep_token_extra=__UpperCamelCase ,			pad_on_left=bool(self.config.model_type in ["xlnet"]  ) ,			pad_token=self.tokenizer.pad_token_id ,			pad_token_segment_id=self.tokenizer.pad_token_type_id ,			pad_token_label_id=self.pad_token_label_id ,			)
                    logger.info("Saving features into cached file %s" ,			__UpperCamelCase  )
                    torch.save(__UpperCamelCase ,			__UpperCamelCase  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase = False  )->					DataLoader:
          UpperCAmelCase__							: Optional[Any] 							=						self._feature_file(__UpperCamelCase  )
          logger.info("Loading features from cached file %s" ,			__UpperCamelCase  )
          UpperCAmelCase__							: List[Any] 							=						torch.load(__UpperCamelCase  )
          UpperCAmelCase__							: Tuple 							=						torch.tensor([f.input_ids for f in features] ,			dtype=torch.long  )
          UpperCAmelCase__							: Optional[int] 							=						torch.tensor([f.attention_mask for f in features] ,			dtype=torch.long  )
          if features[0].token_type_ids is not None:
               UpperCAmelCase__							: str 							=						torch.tensor([f.token_type_ids for f in features] ,			dtype=torch.long  )
          else:
               UpperCAmelCase__							: Optional[Any] 							=						torch.tensor([0 for f in features] ,			dtype=torch.long  )
               # HACK(we will not use this anymore soon)
          UpperCAmelCase__							: Optional[int] 							=						torch.tensor([f.label_ids for f in features] ,			dtype=torch.long  )
          return DataLoader(
              TensorDataset(__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase  ) ,			batch_size=__UpperCamelCase  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase  )->					str:
          """Compute validation""" ""
          UpperCAmelCase__							: Tuple 							=						{"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}
          if self.config.model_type != "distilbert":
               UpperCAmelCase__							: List[Any] 							=						(
                   batch[2] if self.config.model_type in ["bert", "xlnet"] else None
               )  # XLM and RoBERTa don"t use token_type_ids
          UpperCAmelCase__							: Dict 							=						self(**__UpperCamelCase  )
          UpperCAmelCase__							,	UpperCAmelCase__							: str 							=						outputs[:2]
          UpperCAmelCase__							: str 							=						logits.detach().cpu().numpy()
          UpperCAmelCase__							: Optional[int] 							=						inputs["labels"].detach().cpu().numpy()
          return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase  )->					Union[str, Any]:
          UpperCAmelCase__							: List[Any] 							=						torch.stack([x["val_loss"] for x in outputs]  ).mean()
          UpperCAmelCase__							: Optional[Any] 							=						np.concatenate([x["pred"] for x in outputs] ,			axis=0  )
          UpperCAmelCase__							: int 							=						np.argmax(__UpperCamelCase ,			axis=2  )
          UpperCAmelCase__							: Any 							=						np.concatenate([x["target"] for x in outputs] ,			axis=0  )

          UpperCAmelCase__							: Dict 							=						dict(enumerate(self.labels  )  )
          UpperCAmelCase__							: List[Any] 							=						[[] for _ in range(out_label_ids.shape[0]  )]
          UpperCAmelCase__							: List[Any] 							=						[[] for _ in range(out_label_ids.shape[0]  )]

          for i in range(out_label_ids.shape[0]  ):
               for j in range(out_label_ids.shape[1]  ):
                    if out_label_ids[i, j] != self.pad_token_label_id:
                         out_label_list[i].append(label_map[out_label_ids[i][j]]  )
                         preds_list[i].append(label_map[preds[i][j]]  )

          UpperCAmelCase__							: List[Any] 							=						{
              "val_loss": val_loss_mean,
              "accuracy_score": accuracy_score(__UpperCamelCase ,			__UpperCamelCase  ),
              "precision": precision_score(__UpperCamelCase ,			__UpperCamelCase  ),
              "recall": recall_score(__UpperCamelCase ,			__UpperCamelCase  ),
              "f1": fa_score(__UpperCamelCase ,			__UpperCamelCase  ),
          }

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



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase  )->					Any:
          # when stable
          UpperCAmelCase__							,	UpperCAmelCase__							,	UpperCAmelCase__							: Optional[Any] 							=						self._eval_end(__UpperCamelCase  )
          UpperCAmelCase__							: Optional[Any] 							=						ret["log"]
          return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase  )->					List[str]:
          # updating to test_epoch_end instead of deprecated test_end
          UpperCAmelCase__							,	UpperCAmelCase__							,	UpperCAmelCase__							: List[str] 							=						self._eval_end(__UpperCamelCase  )

          # Converting to the dict required by pl
          # https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
          # pytorch_lightning/trainer/logging.py#L139
          UpperCAmelCase__							: List[str] 							=						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  lowerCAmelCase__       (      __UpperCamelCase ,			__UpperCamelCase  )->					Optional[int]:
          # Add NER specific options
          BaseTransformer.add_model_specific_args(__UpperCamelCase ,			__UpperCamelCase  )
          parser.add_argument(
              "--task_type" ,			default="NER" ,			type=__UpperCamelCase ,			help="Task type to fine tune in training (e.g. NER, POS, etc)"  )
          parser.add_argument(
              "--max_seq_length" ,			default=1_28 ,			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(
              "--labels" ,			default="" ,			type=__UpperCamelCase ,			help="Path to a file containing all labels. If not specified, CoNLL-2003 labels are used." ,			)
          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


if __name__ == "__main__":
    A__      :	Optional[Any]          =    argparse.ArgumentParser()
    add_generic_args(parser, os.getcwd())
    A__      :	str          =    NERTransformer.add_model_specific_args(parser, os.getcwd())
    A__      :	Dict          =    parser.parse_args()
    A__      :	Union[str, Any]          =    NERTransformer(args)
    A__      :	str          =    generic_train(model, args)

    if args.do_predict:
        # See https://github.com/huggingface/transformers/issues/3159
        # pl use this default format to create a checkpoint:
        # https://github.com/PyTorchLightning/pytorch-lightning/blob/master\
        # /pytorch_lightning/callbacks/model_checkpoint.py#L322
        A__      :	Optional[int]          =    sorted(glob.glob(os.path.join(args.output_dir, """checkpoint-epoch=*.ckpt"""), recursive=True))
        A__      :	str          =    model.load_from_checkpoint(checkpoints[-1])
        trainer.test(model)





"""simple docstring"""




import argparse

import torch

from transformers import GPTaConfig, GPTaModel, load_tf_weights_in_gpta
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()



def 							a__ (				lowerCAmelCase	:       List[Any]       ,    lowerCAmelCase	:       Any       ,    lowerCAmelCase	:       List[Any]		):






     '''simple docstring'''
     # Construct model
     if gpta_config_file == "":
          UpperCAmelCase__							: Optional[int] 							=						GPTaConfig()
     else:
          UpperCAmelCase__							: Dict 							=						GPTaConfig.from_json_file(lowerCAmelCase		)
     UpperCAmelCase__							: Optional[Any] 							=						GPTaModel(lowerCAmelCase		)

     # Load weights from numpy
     load_tf_weights_in_gpta(lowerCAmelCase       ,    lowerCAmelCase       ,    lowerCAmelCase		)

     # Save pytorch-model
     UpperCAmelCase__							: Optional[int] 							=						pytorch_dump_folder_path + "/" + WEIGHTS_NAME
     UpperCAmelCase__							: Any 							=						pytorch_dump_folder_path + "/" + CONFIG_NAME
     print(F"Save PyTorch model to {pytorch_weights_dump_path}"		)
     torch.save(model.state_dict()       ,    lowerCAmelCase		)
     print(F"Save configuration file to {pytorch_config_dump_path}"		)
     with open(lowerCAmelCase       ,    "w"       ,    encoding="utf-8"		) as f:
          f.write(config.to_json_string()		)


if __name__ == "__main__":
    A__      :	Optional[int]          =    argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        """--gpt2_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
    )
    parser.add_argument(
        """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
    )
    parser.add_argument(
        """--gpt2_config_file""",
        default="""""",
        type=str,
        help=(
            """An optional config json file corresponding to the pre-trained OpenAI model. \n"""
            """This specifies the model architecture."""
        ),
    )
    A__      :	Optional[Any]          =    parser.parse_args()
    convert_gpta_checkpoint_to_pytorch(args.gpta_checkpoint_path, args.gpta_config_file, args.pytorch_dump_folder_path)



