



'''simple docstring'''






import copy
import os
from typing import Union

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


A_							: int						 =	logging.get_logger(__name__)

A_							: List[str]						 =	{
    """google/pix2struct-textcaps-base""": (
        """https://huggingface.co/google/pix2struct-textcaps-base/resolve/main/config.json"""
    ),
}



class 							lowercase       (  _lowerCamelCase				):


 """simple docstring"""


 UpperCAmelCase							=  "pix2struct_text_model"
 UpperCAmelCase							=  ["past_key_values"]
 UpperCAmelCase							=  {
     "hidden_size": "hidden_size",
     "num_attention_heads": "num_heads",
     "num_hidden_layers": "num_layers",
 }







 def __init__(							self  ,a_=50_244  ,a_=768  ,a_=64  ,a_=2_048  ,a_=12  ,a_=12  ,a_=32  ,a_=128  ,a_=0.1  ,a_=1E-6  ,a_=1.0  ,a_="gelu_new"  ,a_=0  ,a_=False  ,a_=0  ,a_=1  ,a_=False  ,a_=True  ,**a_  ,)		->   int:
        _UpperCAmelCase  :						List[Any]			     = vocab_size
        _UpperCAmelCase  :						List[str]			     = hidden_size
        _UpperCAmelCase  :						List[Any]			     = d_kv
        _UpperCAmelCase  :						str			     = d_ff
        _UpperCAmelCase  :						Union[str, Any]			     = num_layers
        _UpperCAmelCase  :						Dict			     = num_heads
        _UpperCAmelCase  :						Optional[int]			     = relative_attention_num_buckets
        _UpperCAmelCase  :						Union[str, Any]			     = relative_attention_max_distance
        _UpperCAmelCase  :						str			     = dropout_rate
        _UpperCAmelCase  :						Union[str, Any]			     = layer_norm_epsilon
        _UpperCAmelCase  :						List[Any]			     = initializer_factor
        _UpperCAmelCase  :						int			     = use_cache

        _UpperCAmelCase  :						List[Any]			     = eos_token_id
        _UpperCAmelCase  :						Tuple			     = decoder_start_token_id

        # for backwards compatibility
        _UpperCAmelCase  :						int			     = dense_act_fn

        super().__init__(
            pad_token_id=snake_case__  ,eos_token_id=snake_case__  ,decoder_start_token_id=snake_case__  ,tie_word_embeddings=snake_case__  ,is_decoder=snake_case__  ,**snake_case__  ,)







 @classmethod
 def _snake_case (							cls  ,a_  ,**a_							)		->   List[str]:
        cls._set_token_in_kwargs(snake_case__							)

        _UpperCAmelCase  :						int			     = cls.get_config_dict(snake_case__  ,**snake_case__							)

        # get the text config dict if we are loading from Pix2StructConfig
        if config_dict.get("""model_type"""							) == "pix2struct":
               _UpperCAmelCase  :						List[Any]			     = config_dict["text_config"]

        if "model_type" in config_dict and hasattr(cls  ,"""model_type"""							) and config_dict["model_type"] != cls.model_type:
               logger.warning(
                   f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
                   f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

        return cls.from_dict(snake_case__  ,**snake_case__							)







class 							lowercase       (  _lowerCamelCase				):


 """simple docstring"""


 UpperCAmelCase							=  "pix2struct_vision_model"







 def __init__(							self  ,a_=768  ,a_=768  ,a_=2_048  ,a_=64  ,a_=12  ,a_=12  ,a_="gelu_new"  ,a_=1E-6  ,a_=0.0  ,a_=0.0  ,a_=1E-1_0  ,a_=1.0  ,a_=4_096  ,a_=32  ,a_=128  ,**a_  ,)		->   Optional[int]:
        super().__init__(**snake_case__							)

        _UpperCAmelCase  :						Any			     = hidden_size
        _UpperCAmelCase  :						Tuple			     = patch_embed_hidden_size
        _UpperCAmelCase  :						Tuple			     = d_ff
        _UpperCAmelCase  :						Optional[int]			     = dropout_rate
        _UpperCAmelCase  :						Dict			     = num_hidden_layers
        _UpperCAmelCase  :						int			     = num_attention_heads
        _UpperCAmelCase  :						Optional[int]			     = initializer_range
        _UpperCAmelCase  :						List[str]			     = initializer_factor
        _UpperCAmelCase  :						List[Any]			     = attention_dropout
        _UpperCAmelCase  :						Dict			     = layer_norm_eps
        _UpperCAmelCase  :						Optional[int]			     = dense_act_fn
        _UpperCAmelCase  :						Union[str, Any]			     = seq_len
        _UpperCAmelCase  :						Optional[Any]			     = relative_attention_num_buckets
        _UpperCAmelCase  :						Any			     = relative_attention_max_distance
        _UpperCAmelCase  :						Dict			     = d_kv







 @classmethod
 def _snake_case (							cls  ,a_  ,**a_							)		->   Optional[Any]:
        cls._set_token_in_kwargs(snake_case__							)

        _UpperCAmelCase  :						Dict			     = cls.get_config_dict(snake_case__  ,**snake_case__							)

        # get the vision config dict if we are loading from Pix2StructConfig
        if config_dict.get("""model_type"""							) == "pix2struct":
               _UpperCAmelCase  :						List[Any]			     = config_dict["vision_config"]

        if "model_type" in config_dict and hasattr(cls  ,"""model_type"""							) and config_dict["model_type"] != cls.model_type:
               logger.warning(
                   f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
                   f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

        return cls.from_dict(snake_case__  ,**snake_case__							)







class 							lowercase       (  _lowerCamelCase				):


 """simple docstring"""


 UpperCAmelCase							=  "pix2struct"
 UpperCAmelCase							=  True







 def __init__(							self  ,a_=None  ,a_=None  ,a_=1.0  ,a_=0.02  ,a_=False  ,a_=False  ,a_=True  ,**a_  ,)		->   Union[str, Any]:
        super().__init__(tie_word_embeddings=snake_case__  ,is_encoder_decoder=snake_case__  ,**snake_case__							)

        if text_config is None:
               _UpperCAmelCase  :						Any			     = {}
               logger.info("""text_config is None. Initializing the Pix2StructTextConfig with default values."""							)

        if vision_config is None:
               _UpperCAmelCase  :						List[Any]			     = {}
               logger.info("""vision_config is None. Initializing the Pix2StructVisionConfig with default values."""							)

        _UpperCAmelCase  :						Union[str, Any]			     = PixaStructTextConfig(**snake_case__							)
        _UpperCAmelCase  :						Optional[Any]			     = PixaStructVisionConfig(**snake_case__							)

        _UpperCAmelCase  :						str			     = self.text_config.decoder_start_token_id
        _UpperCAmelCase  :						Optional[Any]			     = self.text_config.pad_token_id
        _UpperCAmelCase  :						str			     = self.text_config.eos_token_id

        _UpperCAmelCase  :						List[Any]			     = initializer_factor
        _UpperCAmelCase  :						List[str]			     = initializer_range

        _UpperCAmelCase  :						Optional[Any]			     = self.initializer_range
        _UpperCAmelCase  :						Optional[Any]			     = self.initializer_range

        _UpperCAmelCase  :						List[Any]			     = is_vqa







 @classmethod
 def _snake_case (							cls  ,a_  ,a_  ,**a_							)		->   List[Any]:
        return cls(text_config=text_config.to_dict()  ,vision_config=vision_config.to_dict()  ,**snake_case__							)







 def _snake_case (							self							)		->   Optional[int]:
        _UpperCAmelCase  :						Tuple			     = copy.deepcopy(self.__dict__							)
        _UpperCAmelCase  :						Tuple			     = self.text_config.to_dict()
        _UpperCAmelCase  :						Union[str, Any]			     = self.vision_config.to_dict()
        _UpperCAmelCase  :						str			     = self.__class__.model_type
        return output







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





