



'''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        lowerCAmelCase_(				lowerCamelCase_     ):
  '''simple docstring'''




  __lowercase							:      Optional[int]				=      '''token-classification'''




  def __init__(    self  ,__UpperCAmelCase      )    ->		List[Any]:






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




  def 							UpperCAmelCase_	(    self  ,**__UpperCAmelCase      )    ->		Tuple:
      return self.model(**_lowercase      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase      )    ->		Tuple:
      lowerCAmelCase__ :  List[Any]							      =					{"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}
      if self.config.model_type != "distilbert":
          lowerCAmelCase__ :  str							      =					(
              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(**_lowercase      )
      lowerCAmelCase__ :  str							      =					outputs[0]
      # tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
      return {"loss": loss}




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




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase = False      )    ->		Optional[Any]:
      lowerCAmelCase__ :  Dict							      =					self._feature_file(_lowercase      )
      logger.info("""Loading features from cached file %s"""  ,_lowercase      )
      lowerCAmelCase__ :  int							      =					torch.load(_lowercase      )
      lowerCAmelCase__ :  List[Any]							      =					torch.tensor([f.input_ids for f in features]  ,dtype=torch.long      )
      lowerCAmelCase__ :  List[Any]							      =					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__ :  Dict							      =					torch.tensor([0 for f in features]  ,dtype=torch.long      )
          # HACK(we will not use this anymore soon)
      lowerCAmelCase__ :  List[Any]							      =					torch.tensor([f.label_ids for f in features]  ,dtype=torch.long      )
      return DataLoader(
          TensorDataset(_lowercase  ,_lowercase  ,_lowercase  ,_lowercase      )  ,batch_size=_lowercase      )




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




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

      lowerCAmelCase__ :  int							      =					dict(enumerate(self.labels      )      )
      lowerCAmelCase__ :  Any							      =					[[] for _ in range(out_label_ids.shape[0]      )]
      lowerCAmelCase__ :  Tuple							      =					[[] 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(_lowercase  ,_lowercase      ),
          """precision""": precision_score(_lowercase  ,_lowercase      ),
          """recall""": recall_score(_lowercase  ,_lowercase      ),
          """f1""": fa_score(_lowercase  ,_lowercase      ),
      }

      lowerCAmelCase__ :  Any							      =					dict(results.items()      )
      lowerCAmelCase__ :  Optional[Any]							      =					results
      return ret, preds_list, out_label_list




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		int:
      lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__ :  List[Any]							      =					self._eval_end(_lowercase      )
      lowerCAmelCase__ :  Tuple							      =					ret["""log"""]
      return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		List[Any]:
      lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__ :  int							      =					self._eval_end(_lowercase      )

      # Converting to the dict required by pl
      # https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
      # pytorch_lightning/trainer/logging.py#L139
      lowerCAmelCase__ :  Tuple							      =					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 							UpperCAmelCase_	(    __UpperCAmelCase  ,__UpperCAmelCase      )    ->		Tuple:
      BaseTransformer.add_model_specific_args(_lowercase  ,_lowercase      )
      parser.add_argument(
          """--task_type"""  ,default="""NER"""  ,type=_lowercase  ,help="""Task type to fine tune in training (e.g. NER, POS, etc)"""      )
      parser.add_argument(
          """--max_seq_length"""  ,default=128  ,type=_lowercase  ,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=_lowercase  ,help="""Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."""  ,)
      parser.add_argument(
          """--gpus"""  ,default=0  ,type=_lowercase  ,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)



def 	__SCREAMING_SNAKE_CASE  (					__UpperCamelCase    :       int ,     __UpperCamelCase    :       int	)							->					str:







  """simple docstring"""





  if a < 0 or b < 0:
    raise ValueError("""the value of both inputs must be positive"""	)

  SCREAMING_SNAKE_CASE__										=      str(bin(__UpperCamelCase	)	)[2:]  # remove the leading "0b"
  SCREAMING_SNAKE_CASE__										=      str(bin(__UpperCamelCase	)	)[2:]  # remove the leading "0b"

  SCREAMING_SNAKE_CASE__										=      max(len(__UpperCamelCase	) ,     len(__UpperCamelCase	)	)

  return "0b" + "".join(
      str(int(char_a != char_b	)	)
      for char_a, char_b in zip(a_binary.zfill(__UpperCamelCase	) ,     b_binary.zfill(__UpperCamelCase	)	)	)


if __name__ == "__main__":
   import doctest

   doctest.testmod()
