




'''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 		lowercase_   (			A       ):


   """simple docstring"""

   lowerCamelCase_      =						'''token-classification'''



   def __init__(   self :     List[str]  ,				__lowerCamelCase :     Union[str, Any]   ):

      """simple docstring"""
      if type(__lowerCamelCase   ) == dict:
         _SCREAMING_SNAKE_CASE	=				Namespace(**__lowerCamelCase   )
      _SCREAMING_SNAKE_CASE	=				import_module("tasks"   )
      try:
         _SCREAMING_SNAKE_CASE	=				getattr(__lowerCamelCase  ,				hparams.task_type   )
         _SCREAMING_SNAKE_CASE	=				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__()}"""   )
      _SCREAMING_SNAKE_CASE	=				self.token_classification_task.get_labels(hparams.labels   )
      _SCREAMING_SNAKE_CASE	=				CrossEntropyLoss().ignore_index
      super().__init__(__lowerCamelCase  ,				len(self.labels   )  ,				self.mode   )



   def 							lowerCAmelCase_      (   self :     Tuple  ,				**__lowerCamelCase :     str   ):

      """simple docstring"""
      return self.model(**__lowerCamelCase   )



   def 							lowerCAmelCase_      (   self :     Dict  ,				__lowerCamelCase :     List[str]  ,				__lowerCamelCase :     str   ):

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

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



   def 							lowerCAmelCase_      (   self :     Union[str, Any]   ):

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



   def 							lowerCAmelCase_      (   self :     Tuple  ,				__lowerCamelCase :     int  ,				__lowerCamelCase :     int  ,				__lowerCamelCase :     bool = False   ):

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



   def 							lowerCAmelCase_      (   self :     Dict  ,				__lowerCamelCase :     int  ,				__lowerCamelCase :     int   ):

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



   def 							lowerCAmelCase_      (   self :     List[str]  ,				__lowerCamelCase :     List[Any]   ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE	=				torch.stack([x["val_loss"] for x in outputs]   ).mean()
      _SCREAMING_SNAKE_CASE	=				np.concatenate([x["pred"] for x in outputs]  ,				axis=0   )
      _SCREAMING_SNAKE_CASE	=				np.argmax(__lowerCamelCase  ,				axis=2   )
      _SCREAMING_SNAKE_CASE	=				np.concatenate([x["target"] for x in outputs]  ,				axis=0   )

      _SCREAMING_SNAKE_CASE	=				dict(enumerate(self.labels   )   )
      _SCREAMING_SNAKE_CASE	=				[[] for _ in range(out_label_ids.shape[0]   )]
      _SCREAMING_SNAKE_CASE	=				[[] 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]]   )

      _SCREAMING_SNAKE_CASE	=				{
          "val_loss": val_loss_mean,
          "accuracy_score": accuracy_score(__lowerCamelCase  ,				__lowerCamelCase   ),
          "precision": precision_score(__lowerCamelCase  ,				__lowerCamelCase   ),
          "recall": recall_score(__lowerCamelCase  ,				__lowerCamelCase   ),
          "f1": fa_score(__lowerCamelCase  ,				__lowerCamelCase   ),
      }

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



   def 							lowerCAmelCase_      (   self :     Optional[int]  ,				__lowerCamelCase :     Optional[Any]   ):

      """simple docstring"""
      # when stable
      _SCREAMING_SNAKE_CASE,	_SCREAMING_SNAKE_CASE,	_SCREAMING_SNAKE_CASE	=				self._eval_end(__lowerCamelCase   )
      _SCREAMING_SNAKE_CASE	=				ret["log"]
      return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}



   def 							lowerCAmelCase_      (   self :     Optional[int]  ,				__lowerCamelCase :     Dict   ):

      """simple docstring"""
      # updating to test_epoch_end instead of deprecated test_end
      _SCREAMING_SNAKE_CASE,	_SCREAMING_SNAKE_CASE,	_SCREAMING_SNAKE_CASE	=				self._eval_end(__lowerCamelCase   )

      # Converting to the dict required by pl
      # https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
      # pytorch_lightning/trainer/logging.py#L139
      _SCREAMING_SNAKE_CASE	=				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_      (   __lowerCamelCase :     List[str]  ,				__lowerCamelCase :     Optional[Any]   ):

      """simple docstring"""
      # Add NER specific options
      BaseTransformer.add_model_specific_args(__lowerCamelCase  ,				__lowerCamelCase   )
      parser.add_argument(
          "--task_type"  ,				default="NER"  ,				type=__lowerCamelCase  ,				help="Task type to fine tune in training (e.g. NER, POS, etc)"   )
      parser.add_argument(
          "--max_seq_length"  ,				default=1_2_8  ,				type=__lowerCamelCase  ,				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=__lowerCamelCase  ,				help="Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."  ,				)
      parser.add_argument(
          "--gpus"  ,				default=0  ,				type=__lowerCamelCase  ,				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'''



import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible







class 		lowercase_   (			unittest.TestCase       ):


   """simple docstring"""

   def 							lowerCAmelCase_      (   self :     List[Any]   ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE	=				[
          "safety_checker/pytorch_model.bin",
          "safety_checker/model.safetensors",
          "vae/diffusion_pytorch_model.bin",
          "vae/diffusion_pytorch_model.safetensors",
          "text_encoder/pytorch_model.bin",
          "text_encoder/model.safetensors",
          "unet/diffusion_pytorch_model.bin",
          "unet/diffusion_pytorch_model.safetensors",
      ]
      self.assertTrue(is_safetensors_compatible(__lowerCamelCase   )   )



   def 							lowerCAmelCase_      (   self :     Dict   ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE	=				[
          "unet/diffusion_pytorch_model.bin",
          "unet/diffusion_pytorch_model.safetensors",
      ]
      self.assertTrue(is_safetensors_compatible(__lowerCamelCase   )   )



   def 							lowerCAmelCase_      (   self :     Tuple   ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE	=				[
          "safety_checker/pytorch_model.bin",
          "safety_checker/model.safetensors",
          "vae/diffusion_pytorch_model.bin",
          "vae/diffusion_pytorch_model.safetensors",
          "text_encoder/pytorch_model.bin",
          "text_encoder/model.safetensors",
          "unet/diffusion_pytorch_model.bin",
          # Removed: 'unet/diffusion_pytorch_model.safetensors',
      ]
      self.assertFalse(is_safetensors_compatible(__lowerCamelCase   )   )



   def 							lowerCAmelCase_      (   self :     List[str]   ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE	=				[
          "text_encoder/pytorch_model.bin",
          "text_encoder/model.safetensors",
      ]
      self.assertTrue(is_safetensors_compatible(__lowerCamelCase   )   )



   def 							lowerCAmelCase_      (   self :     int   ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE	=				[
          "safety_checker/pytorch_model.bin",
          "safety_checker/model.safetensors",
          "vae/diffusion_pytorch_model.bin",
          "vae/diffusion_pytorch_model.safetensors",
          "text_encoder/pytorch_model.bin",
          # Removed: 'text_encoder/model.safetensors',
          "unet/diffusion_pytorch_model.bin",
          "unet/diffusion_pytorch_model.safetensors",
      ]
      self.assertFalse(is_safetensors_compatible(__lowerCamelCase   )   )



   def 							lowerCAmelCase_      (   self :     Optional[int]   ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE	=				[
          "safety_checker/pytorch_model.fp16.bin",
          "safety_checker/model.fp16.safetensors",
          "vae/diffusion_pytorch_model.fp16.bin",
          "vae/diffusion_pytorch_model.fp16.safetensors",
          "text_encoder/pytorch_model.fp16.bin",
          "text_encoder/model.fp16.safetensors",
          "unet/diffusion_pytorch_model.fp16.bin",
          "unet/diffusion_pytorch_model.fp16.safetensors",
      ]
      _SCREAMING_SNAKE_CASE	=				"fp16"
      self.assertTrue(is_safetensors_compatible(__lowerCamelCase  ,				variant=__lowerCamelCase   )   )



   def 							lowerCAmelCase_      (   self :     Any   ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE	=				[
          "unet/diffusion_pytorch_model.fp16.bin",
          "unet/diffusion_pytorch_model.fp16.safetensors",
      ]
      _SCREAMING_SNAKE_CASE	=				"fp16"
      self.assertTrue(is_safetensors_compatible(__lowerCamelCase  ,				variant=__lowerCamelCase   )   )



   def 							lowerCAmelCase_      (   self :     Tuple   ):

      """simple docstring"""
      # pass variant but use the non-variant filenames
      _SCREAMING_SNAKE_CASE	=				[
          "unet/diffusion_pytorch_model.bin",
          "unet/diffusion_pytorch_model.safetensors",
      ]
      _SCREAMING_SNAKE_CASE	=				"fp16"
      self.assertTrue(is_safetensors_compatible(__lowerCamelCase  ,				variant=__lowerCamelCase   )   )



   def 							lowerCAmelCase_      (   self :     int   ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE	=				[
          "safety_checker/pytorch_model.fp16.bin",
          "safety_checker/model.fp16.safetensors",
          "vae/diffusion_pytorch_model.fp16.bin",
          "vae/diffusion_pytorch_model.fp16.safetensors",
          "text_encoder/pytorch_model.fp16.bin",
          "text_encoder/model.fp16.safetensors",
          "unet/diffusion_pytorch_model.fp16.bin",
          # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
      ]
      _SCREAMING_SNAKE_CASE	=				"fp16"
      self.assertFalse(is_safetensors_compatible(__lowerCamelCase  ,				variant=__lowerCamelCase   )   )



   def 							lowerCAmelCase_      (   self :     Dict   ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE	=				[
          "text_encoder/pytorch_model.fp16.bin",
          "text_encoder/model.fp16.safetensors",
      ]
      _SCREAMING_SNAKE_CASE	=				"fp16"
      self.assertTrue(is_safetensors_compatible(__lowerCamelCase  ,				variant=__lowerCamelCase   )   )



   def 							lowerCAmelCase_      (   self :     str   ):

      """simple docstring"""
      # pass variant but use the non-variant filenames
      _SCREAMING_SNAKE_CASE	=				[
          "text_encoder/pytorch_model.bin",
          "text_encoder/model.safetensors",
      ]
      _SCREAMING_SNAKE_CASE	=				"fp16"
      self.assertTrue(is_safetensors_compatible(__lowerCamelCase  ,				variant=__lowerCamelCase   )   )



   def 							lowerCAmelCase_      (   self :     Optional[Any]   ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE	=				[
          "safety_checker/pytorch_model.fp16.bin",
          "safety_checker/model.fp16.safetensors",
          "vae/diffusion_pytorch_model.fp16.bin",
          "vae/diffusion_pytorch_model.fp16.safetensors",
          "text_encoder/pytorch_model.fp16.bin",
          # 'text_encoder/model.fp16.safetensors',
          "unet/diffusion_pytorch_model.fp16.bin",
          "unet/diffusion_pytorch_model.fp16.safetensors",
      ]
      _SCREAMING_SNAKE_CASE	=				"fp16"
      self.assertFalse(is_safetensors_compatible(__lowerCamelCase  ,				variant=__lowerCamelCase   )   )

