


"""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


lowerCAmelCase__      				=   logging.getLogger(__name__)


class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       a  :				Tuple   ="token-classification"

       def __init__(						self		,     snake_case__						):






         """simple docstring"""





         if type(snake_case__						) == dict:
           lowerCAmelCase     :     List[str]								=							Namespace(**snake_case__						)
         lowerCAmelCase     :     Optional[Any]								=							import_module("tasks"						)
         try:
           lowerCAmelCase     :     Dict								=							getattr(snake_case__		,     hparams.task_type						)
           lowerCAmelCase     :     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__()}"""						)
         lowerCAmelCase     :     Dict								=							self.token_classification_task.get_labels(hparams.labels						)
         lowerCAmelCase     :     str								=							CrossEntropyLoss().ignore_index
         super().__init__(snake_case__		,     len(self.labels						)		,     self.mode						)

       def        lowercase__							(						self		,     **snake_case__						):






         """simple docstring"""





         return self.model(**snake_case__						)

       def        lowercase__							(						self		,     snake_case__		,     snake_case__						):






         """simple docstring"""





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

         lowerCAmelCase     :     str								=							self(**snake_case__						)
         lowerCAmelCase     :     Tuple								=							outputs[0]
         # tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
         return {"loss": loss}

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							self.hparams
         for mode in ["train", "dev", "test"]:
           lowerCAmelCase     :     Optional[int]								=							self._feature_file(snake_case__						)
           if os.path.exists(snake_case__						) and not args.overwrite_cache:
             logger.info("Loading features from cached file %s"		,     snake_case__						)
             lowerCAmelCase     :     Dict								=							torch.load(snake_case__						)
           else:
             logger.info("Creating features from dataset file at %s"		,     args.data_dir						)
             lowerCAmelCase     :     Optional[int]								=							self.token_classification_task.read_examples_from_file(args.data_dir		,     snake_case__						)
             lowerCAmelCase     :     Optional[int]								=							self.token_classification_task.convert_examples_to_features(
                 snake_case__		,     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=snake_case__		,     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"		,     snake_case__						)
             torch.save(snake_case__		,     snake_case__						)

       def        lowercase__							(						self		,     snake_case__		,     snake_case__		,     snake_case__ = False						):






         """simple docstring"""





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

       def        lowercase__							(						self		,     snake_case__		,     snake_case__						):






         """simple docstring"""





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

       def        lowercase__							(						self		,     snake_case__						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							torch.stack([x["val_loss"] for x in outputs]						).mean()
         lowerCAmelCase     :     Optional[Any]								=							np.concatenate([x["pred"] for x in outputs]		,     axis=0						)
         lowerCAmelCase     :     List[Any]								=							np.argmax(snake_case__		,     axis=2						)
         lowerCAmelCase     :     Optional[int]								=							np.concatenate([x["target"] for x in outputs]		,     axis=0						)

         lowerCAmelCase     :     str								=							dict(enumerate(self.labels						)						)
         lowerCAmelCase     :     Optional[int]								=							[[] for _ in range(out_label_ids.shape[0]						)]
         lowerCAmelCase     :     Dict								=							[[] 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]]						)

         lowerCAmelCase     :     Dict								=							{
             "val_loss": val_loss_mean,
             "accuracy_score": accuracy_score(snake_case__		,     snake_case__						),
             "precision": precision_score(snake_case__		,     snake_case__						),
             "recall": recall_score(snake_case__		,     snake_case__						),
             "f1": fa_score(snake_case__		,     snake_case__						),
         }

         lowerCAmelCase     :     int								=							dict(results.items()						)
         lowerCAmelCase     :     Tuple								=							results
         return ret, preds_list, out_label_list

       def        lowercase__							(						self		,     snake_case__						):






         """simple docstring"""





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

       def        lowercase__							(						self		,     snake_case__						):






         """simple docstring"""





         lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase     :     int								=							self._eval_end(snake_case__						)

         # Converting to the dict required by pl
         # https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
         # pytorch_lightning/trainer/logging.py#L139
         lowerCAmelCase     :     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        lowercase__							(						snake_case__		,     snake_case__						):






         """simple docstring"""





         BaseTransformer.add_model_specific_args(snake_case__		,     snake_case__						)
         parser.add_argument(
             "--task_type"		,     default="NER"		,     type=snake_case__		,     help="Task type to fine tune in training (e.g. NER, POS, etc)"						)
         parser.add_argument(
             "--max_seq_length"		,     default=128		,     type=snake_case__		,     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=snake_case__		,     help="Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."		,     )
         parser.add_argument(
             "--gpus"		,     default=0		,     type=snake_case__		,     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__":
     lowerCAmelCase__      				=   argparse.ArgumentParser()
     add_generic_args(parser, os.getcwd())
     lowerCAmelCase__      				=   NERTransformer.add_model_specific_args(parser, os.getcwd())
     lowerCAmelCase__      				=   parser.parse_args()
     lowerCAmelCase__      				=   NERTransformer(args)
     lowerCAmelCase__      				=   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
          lowerCAmelCase__      				=   sorted(glob.glob(os.path.join(args.output_dir, '''checkpoint-epoch=*.ckpt'''), recursive=True))
          lowerCAmelCase__      				=   model.load_from_checkpoint(checkpoints[-1])
          trainer.test(model)









"""simple docstring"""







# Lint as: python3
import os
import re
import urllib.parse
from pathlib import Path
from typing import Callable, List, Optional, Union
from zipfile import ZipFile

from ..utils.file_utils import cached_path, hf_github_url
from ..utils.logging import get_logger
from ..utils.version import Version


lowerCAmelCase__      				=   get_logger(__name__)


class       SCREAMING_SNAKE_CASE__				:
       """simple docstring"""


       a  :				Optional[Any]   ="dummy_data"
       a  :				int   ="datasets"
       a  :				Tuple   =False

       def __init__(						self		,     snake_case__		,     snake_case__		,     snake_case__		,     snake_case__ = None		,     snake_case__ = False		,     snake_case__ = True		,     snake_case__ = None		,     ):






         """simple docstring"""





         lowerCAmelCase     :     Tuple								=							0
         lowerCAmelCase     :     int								=							dataset_name
         lowerCAmelCase     :     List[Any]								=							cache_dir
         lowerCAmelCase     :     List[str]								=							use_local_dummy_data
         lowerCAmelCase     :     List[str]								=							config
         # download_callbacks take a single url as input
         lowerCAmelCase     :     List[Callable]								=							download_callbacks or []
         # if False, it doesn't load existing files and it returns the paths of the dummy files relative
         # to the dummy_data zip file root
         lowerCAmelCase     :     Tuple								=							load_existing_dummy_data

         # TODO(PVP, QL) might need to make this more general
         lowerCAmelCase     :     Union[str, Any]								=							str(snake_case__						)
         # to be downloaded
         lowerCAmelCase     :     List[Any]								=							None
         lowerCAmelCase     :     List[Any]								=							None

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         if self._dummy_file is None:
           lowerCAmelCase     :     Any								=							self.download_dummy_data()
         return self._dummy_file

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         if self.config is not None:
           # structure is dummy / config_name / version_name
           return os.path.join("dummy"		,     self.config.name		,     self.version_name						)
         # structure is dummy / version_name
         return os.path.join("dummy"		,     self.version_name						)

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         return os.path.join(self.dummy_data_folder		,     "dummy_data.zip"						)

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     Optional[Any]								=							(
             self.local_path_to_dummy_data if self.use_local_dummy_data is True else self.github_path_to_dummy_data
         )

         lowerCAmelCase     :     str								=							cached_path(
             snake_case__		,     cache_dir=self.cache_dir		,     extract_compressed_file=snake_case__		,     force_extract=snake_case__						)

         return os.path.join(snake_case__		,     self.dummy_file_name						)

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         return os.path.join(self.datasets_scripts_dir		,     self.dataset_name		,     self.dummy_zip_file						)

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         if self._bucket_url is None:
           lowerCAmelCase     :     Union[str, Any]								=							hf_github_url(self.dataset_name		,     self.dummy_zip_file.replace(os.sep		,     "/"						)						)
         return self._bucket_url

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         if os.path.isdir(self.dummy_file						):
           return self.dummy_file
         # else cut off path to file -> example `xsum`.
         return "/".join(self.dummy_file.replace(os.sep		,     "/"						).split("/"						)[:-1]						)

       def        lowercase__							(						self		,     snake_case__		,     *snake_case__						):






         """simple docstring"""





         if self.load_existing_dummy_data:
           # dummy data is downloaded and tested
           lowerCAmelCase     :     int								=							self.dummy_file
         else:
           # dummy data cannot be downloaded and only the path to dummy file is returned
           lowerCAmelCase     :     List[Any]								=							self.dummy_file_name

         # special case when data_url is a dict
         if isinstance(snake_case__		,     snake_case__						):
           return self.create_dummy_data_dict(snake_case__		,     snake_case__						)
         elif isinstance(snake_case__		,     (list, tuple)						):
           return self.create_dummy_data_list(snake_case__		,     snake_case__						)
         else:
           return self.create_dummy_data_single(snake_case__		,     snake_case__						)

       def        lowercase__							(						self		,     snake_case__		,     *snake_case__						):






         """simple docstring"""





         return self.download_and_extract(snake_case__						)

       def        lowercase__							(						self		,     snake_case__		,     snake_case__						):






         """simple docstring"""





         return self.download_and_extract(snake_case__						)

       def        lowercase__							(						self		,     snake_case__		,     *snake_case__		,     **snake_case__						):






         """simple docstring"""





         return path

       def        lowercase__							(						self						):






         """simple docstring"""





         return {}

       def        lowercase__							(						self		,     snake_case__		,     snake_case__						):






         """simple docstring"""





         lowerCAmelCase     :     List[Any]								=							{}
         for key, single_urls in data_url.items():
           for download_callback in self.download_callbacks:
             if isinstance(snake_case__		,     snake_case__						):
               for single_url in single_urls:
                 download_callback(snake_case__						)
             else:
               lowerCAmelCase     :     List[str]								=							single_urls
               download_callback(snake_case__						)
            # we force the name of each key to be the last file / folder name of the url path
            # if the url has arguments, we need to encode them with urllib.parse.quote_plus
           if isinstance(snake_case__		,     snake_case__						):
             lowerCAmelCase     :     Tuple								=							[os.path.join(snake_case__		,     urllib.parse.quote_plus(Path(snake_case__						).name						)						) for x in single_urls]
           else:
             lowerCAmelCase     :     int								=							single_urls
             lowerCAmelCase     :     Any								=							os.path.join(snake_case__		,     urllib.parse.quote_plus(Path(snake_case__						).name						)						)
           lowerCAmelCase     :     Union[str, Any]								=							value

         # make sure that values are unique
         if all(isinstance(snake_case__		,     snake_case__						) for i in dummy_data_dict.values()						) and len(set(dummy_data_dict.values()						)						) < len(
             dummy_data_dict.values()						):
           # append key to value to make its name unique
           lowerCAmelCase     :     Union[str, Any]								=							{key: value + key for key, value in dummy_data_dict.items()}

         return dummy_data_dict

       def        lowercase__							(						self		,     snake_case__		,     snake_case__						):






         """simple docstring"""





         lowerCAmelCase     :     Dict								=							[]
         # trick: if there are many shards named like `data.txt-000001-of-00300`, only use the first one
         lowerCAmelCase     :     Optional[Any]								=							all(bool(re.findall("[0-9]{3,}-of-[0-9]{3,}"		,     snake_case__						)						) for url in data_url						)
         lowerCAmelCase     :     Any								=							all(
             url.startswith("https://ftp.ncbi.nlm.nih.gov/pubmed/baseline/pubmed"						) for url in data_url						)
         if data_url and (is_tf_records or is_pubmed_records):
           lowerCAmelCase     :     int								=							[data_url[0]] * len(snake_case__						)
         for single_url in data_url:
           for download_callback in self.download_callbacks:
             download_callback(snake_case__						)
           # we force the name of each key to be the last file / folder name of the url path
           # if the url has arguments, we need to encode them with urllib.parse.quote_plus
           lowerCAmelCase     :     Dict								=							os.path.join(snake_case__		,     urllib.parse.quote_plus(single_url.split("/"						)[-1]						)						)
           dummy_data_list.append(snake_case__						)
         return dummy_data_list

       def        lowercase__							(						self		,     snake_case__		,     snake_case__						):






         """simple docstring"""





         for download_callback in self.download_callbacks:
           download_callback(snake_case__						)
         # we force the name of each key to be the last file / folder name of the url path
         # if the url has arguments, we need to encode them with urllib.parse.quote_plus
         lowerCAmelCase     :     Tuple								=							os.path.join(snake_case__		,     urllib.parse.quote_plus(data_url.split("/"						)[-1]						)						)
         if os.path.exists(snake_case__						) or not self.load_existing_dummy_data:
           return value
         else:
           # Backward compatibility, maybe deprecate at one point.
           # For many datasets with single url calls to dl_manager.download_and_extract,
           # the dummy_data.zip file is actually the zipped downloaded file
           # while now we expected the dummy_data.zip file to be a directory containing
           # the downloaded file.
           return path_to_dummy_data

       def        lowercase__							(						self						):






         """simple docstring"""





         pass

       def        lowercase__							(						self						):






         """simple docstring"""





         pass

       def        lowercase__							(						self		,     snake_case__						):






         """simple docstring"""





         def _iter_archive_members(snake_case__						):
           # this preserves the order of the members inside the ZIP archive
           lowerCAmelCase     :     str								=							Path(self.dummy_file						).parent
           lowerCAmelCase     :     Optional[Any]								=							path.relative_to(snake_case__						)
           with ZipFile(self.local_path_to_dummy_data						) as zip_file:
             lowerCAmelCase     :     List[Any]								=							zip_file.namelist()
           for member in members:
             if member.startswith(relative_path.as_posix()						):
               yield dummy_parent_path.joinpath(snake_case__						)

         lowerCAmelCase     :     List[Any]								=							Path(snake_case__						)
         lowerCAmelCase     :     str								=							_iter_archive_members(snake_case__						) if self.use_local_dummy_data else path.rglob("*"						)
         for file_path in file_paths:
           if file_path.is_file() and not file_path.name.startswith((".", "__")						):
             yield file_path.relative_to(snake_case__						).as_posix(), file_path.open("rb"						)







       def        lowercase__							(						self		,     snake_case__						):






         """simple docstring"""





         if not isinstance(snake_case__		,     snake_case__						):
           lowerCAmelCase     :     List[Any]								=							[paths]
         for path in paths:
           if os.path.isfile(snake_case__						):
             if os.path.basename(snake_case__						).startswith((".", "__")						):
               return
             yield path
           else:
             for dirpath, dirnames, filenames in os.walk(snake_case__						):
               if os.path.basename(snake_case__						).startswith((".", "__")						):
                 continue
               dirnames.sort()
               for filename in sorted(snake_case__						):
                 if filename.startswith((".", "__")						):
                   continue
                 yield os.path.join(snake_case__		,     snake_case__						)





