



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


lowercase__		:				Any          =	logging.getLogger(__name__)







class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     """token-classification"""



      def __init__(      self				:      Optional[Any]    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
        if type(SCREAMING_SNAKE_CASE_       ) == dict:
          lowerCAmelCase_   :				Union[str, Any]								= Namespace(**SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				List[Any]								= import_module('tasks'       )
        try:
          lowerCAmelCase_   :				List[Any]								= getattr(SCREAMING_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_   :				int								= CrossEntropyLoss().ignore_index
        super().__init__(SCREAMING_SNAKE_CASE_    ,			len(self.labels       )    ,			self.mode       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[Any]    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        return self.model(**SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      str    ,			SCREAMING_SNAKE_CASE_				:      Tuple       ):
        lowerCAmelCase_   :				Any								= {'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[3]}
        if self.config.model_type != "distilbert":
          lowerCAmelCase_   :				Dict								= (
              batch[2] if self.config.model_type in ['bert', 'xlnet'] else None
          )  # XLM and RoBERTa don"t use token_type_ids

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



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



      def      SCREAMING_SNAKE_CASE__ (      self				:      int    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      bool = False       ):
        lowerCAmelCase_   :				Union[str, Any]								= self._feature_file(SCREAMING_SNAKE_CASE_       )
        logger.info('Loading features from cached file %s'    ,			SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Union[str, Any]								= torch.load(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				int								= torch.tensor([f.input_ids for f in features]    ,			dtype=torch.long       )
        lowerCAmelCase_   :				Tuple								= torch.tensor([f.attention_mask for f in features]    ,			dtype=torch.long       )
        if features[0].token_type_ids is not None:
          lowerCAmelCase_   :				List[str]								= torch.tensor([f.token_type_ids for f in features]    ,			dtype=torch.long       )
        else:
          lowerCAmelCase_   :				Dict								= 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(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )    ,			batch_size=SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      List[Any]       ):
        """Compute validation""" ""
        lowerCAmelCase_   :				Optional[int]								= {'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_   :				List[Any]								= self(**SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_						,lowerCAmelCase_   :				Dict								= outputs[:2]
        lowerCAmelCase_   :				int								= logits.detach().cpu().numpy()
        lowerCAmelCase_   :				str								= inputs['labels'].detach().cpu().numpy()
        return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}



      def      SCREAMING_SNAKE_CASE__ (      self				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      Dict       ):
        lowerCAmelCase_   :				List[Any]								= torch.stack([x['val_loss'] for x in outputs]       ).mean()
        lowerCAmelCase_   :				Any								= np.concatenate([x['pred'] for x in outputs]    ,			axis=0       )
        lowerCAmelCase_   :				Any								= np.argmax(SCREAMING_SNAKE_CASE_    ,			axis=2       )
        lowerCAmelCase_   :				Dict								= np.concatenate([x['target'] for x in outputs]    ,			axis=0       )

        lowerCAmelCase_   :				int								= dict(enumerate(self.labels       )       )
        lowerCAmelCase_   :				Optional[int]								= [[] for _ in range(out_label_ids.shape[0]       )]
        lowerCAmelCase_   :				str								= [[] 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_   :				List[Any]								= {
            'val_loss': val_loss_mean,
            'accuracy_score': accuracy_score(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ),
            'precision': precision_score(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ),
            'recall': recall_score(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ),
            'f1': fa_score(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ),
        }

        lowerCAmelCase_   :				List[Any]								= dict(results.items()       )
        lowerCAmelCase_   :				List[str]								= results
        return ret, preds_list, out_label_list



      def      SCREAMING_SNAKE_CASE__ (      self				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      int       ):
        # when stable
        lowerCAmelCase_						,lowerCAmelCase_						,lowerCAmelCase_   :				Optional[Any]								= self._eval_end(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				List[str]								= ret['log']
        return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}



      def      SCREAMING_SNAKE_CASE__ (      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      str       ):
        # updating to test_epoch_end instead of deprecated test_end
        lowerCAmelCase_						,lowerCAmelCase_						,lowerCAmelCase_   :				Any								= self._eval_end(SCREAMING_SNAKE_CASE_       )

        # Converting to the dict required by pl
        # https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
        # pytorch_lightning/trainer/logging.py#L139
        lowerCAmelCase_   :				List[Any]								= 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      SCREAMING_SNAKE_CASE__ (      SCREAMING_SNAKE_CASE_				:      Optional[Any]    ,			SCREAMING_SNAKE_CASE_				:      Tuple       ):
        # Add NER specific options
        BaseTransformer.add_model_specific_args(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
        parser.add_argument(
            '--task_type'    ,			default='NER'    ,			type=SCREAMING_SNAKE_CASE_    ,			help='Task type to fine tune in training (e.g. NER, POS, etc)'       )
        parser.add_argument(
            '--max_seq_length'    ,			default=1_2_8    ,			type=SCREAMING_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=SCREAMING_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=SCREAMING_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__":
     lowercase__		:				int          =	argparse.ArgumentParser()
     add_generic_args(parser, os.getcwd())
     lowercase__		:				Tuple          =	NERTransformer.add_model_specific_args(parser, os.getcwd())
     lowercase__		:				Tuple          =	parser.parse_args()
     lowercase__		:				int          =	NERTransformer(args)
     lowercase__		:				List[Any]          =	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
          lowercase__		:				Tuple          =	sorted(glob.glob(os.path.join(args.output_dir, """checkpoint-epoch=*.ckpt"""), recursive=True))
          lowercase__		:				Optional[Any]          =	model.load_from_checkpoint(checkpoints[-1])
          trainer.test(model)










"""simple docstring"""







from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase__		:				Tuple          =	logging.get_logger(__name__)

lowercase__		:				Any          =	{
    """facebook/dpr-ctx_encoder-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-question_encoder-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-reader-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-ctx_encoder-multiset-base""": (
        """https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/config.json"""
    ),
    """facebook/dpr-question_encoder-multiset-base""": (
        """https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/config.json"""
    ),
    """facebook/dpr-reader-multiset-base""": (
        """https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/config.json"""
    ),
}







class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     """dpr"""



      def __init__(      self				:      List[Any]    ,			SCREAMING_SNAKE_CASE_				:      Tuple=3_0_5_2_2    ,			SCREAMING_SNAKE_CASE_				:      Any=7_6_8    ,			SCREAMING_SNAKE_CASE_				:      Dict=1_2    ,			SCREAMING_SNAKE_CASE_				:      List[str]=1_2    ,			SCREAMING_SNAKE_CASE_				:      int=3_0_7_2    ,			SCREAMING_SNAKE_CASE_				:      List[Any]="gelu"    ,			SCREAMING_SNAKE_CASE_				:      int=0.1    ,			SCREAMING_SNAKE_CASE_				:      int=0.1    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]=5_1_2    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]=2    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]=0.02    ,			SCREAMING_SNAKE_CASE_				:      List[str]=1E-12    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=0    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]="absolute"    ,			SCREAMING_SNAKE_CASE_				:      int = 0    ,			**SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			):
        super().__init__(pad_token_id=SCREAMING_SNAKE_CASE_    ,			**SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				Optional[int]								= vocab_size
        lowerCAmelCase_   :				Tuple								= hidden_size
        lowerCAmelCase_   :				int								= num_hidden_layers
        lowerCAmelCase_   :				str								= num_attention_heads
        lowerCAmelCase_   :				Any								= hidden_act
        lowerCAmelCase_   :				List[str]								= intermediate_size
        lowerCAmelCase_   :				int								= hidden_dropout_prob
        lowerCAmelCase_   :				Union[str, Any]								= attention_probs_dropout_prob
        lowerCAmelCase_   :				Tuple								= max_position_embeddings
        lowerCAmelCase_   :				List[str]								= type_vocab_size
        lowerCAmelCase_   :				str								= initializer_range
        lowerCAmelCase_   :				str								= layer_norm_eps
        lowerCAmelCase_   :				List[str]								= projection_dim
        lowerCAmelCase_   :				List[str]								= position_embedding_type





