



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_						=       logging.getLogger(__name__)
class 			_snake_case     (   lowercase__):
				UpperCamelCase__			:   List[str]       						="""token-classification"""




				def __init__(  self   :  Tuple,   __lowercase   :  Optional[int]  ):
							if type(__lowercase  ) == dict:
										lowercase__   					=			Namespace(**__lowercase  )
							lowercase__   					=			import_module("tasks"  )
							try:
										lowercase__   					=			getattr(__lowercase,   hparams.task_type  )
										lowercase__   					=			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__()}'''  )
							lowercase__   					=			self.token_classification_task.get_labels(hparams.labels  )
							lowercase__   					=			CrossEntropyLoss().ignore_index
							super().__init__(__lowercase,   len(self.labels  ),   self.mode  )




				def 							A__ (  self   :  List[str],   **__lowercase   :  Any  ):
							return self.model(**__lowercase  )




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

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




				def 							A__ (  self   :  Tuple  ):
							lowercase__   					=			self.hparams
							for mode in ["train", "dev", "test"]:
										lowercase__   					=			self._feature_file(__lowercase  )
										if os.path.exists(__lowercase  ) and not args.overwrite_cache:
													logger.info("Loading features from cached file %s",   __lowercase  )
													lowercase__   					=			torch.load(__lowercase  )
										else:
													logger.info("Creating features from dataset file at %s",   args.data_dir  )
													lowercase__   					=			self.token_classification_task.read_examples_from_file(args.data_dir,   __lowercase  )
													lowercase__   					=			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 							A__ (  self   :  List[str],   __lowercase   :  int,   __lowercase   :  int,   __lowercase   :  bool = False  ):
							lowercase__   					=			self._feature_file(__lowercase  )
							logger.info("Loading features from cached file %s",   __lowercase  )
							lowercase__   					=			torch.load(__lowercase  )
							lowercase__   					=			torch.tensor([f.input_ids for f in features],   dtype=torch.long  )
							lowercase__   					=			torch.tensor([f.attention_mask for f in features],   dtype=torch.long  )
							if features[0].token_type_ids is not None:
										lowercase__   					=			torch.tensor([f.token_type_ids for f in features],   dtype=torch.long  )
							else:
										lowercase__   					=			torch.tensor([0 for f in features],   dtype=torch.long  )
										# HACK(we will not use this anymore soon)
							lowercase__   					=			torch.tensor([f.label_ids for f in features],   dtype=torch.long  )
							return DataLoader(
							    TensorDataset(__lowercase,   __lowercase,   __lowercase,   __lowercase  ),   batch_size=__lowercase  )




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




				def 							A__ (  self   :  str,   __lowercase   :  Union[str, Any]  ):
							lowercase__   					=			torch.stack([x["val_loss"] for x in outputs]  ).mean()
							lowercase__   					=			np.concatenate([x["pred"] for x in outputs],   axis=0  )
							lowercase__   					=			np.argmax(__lowercase,   axis=2  )
							lowercase__   					=			np.concatenate([x["target"] for x in outputs],   axis=0  )

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

							lowercase__   					=			{
							    "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  ),
							}

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




				def 							A__ (  self   :  Union[str, Any],   __lowercase   :  Any  ):
							# when stable
							lowercase__					,    lowercase__					,    lowercase__   					=			self._eval_end(__lowercase  )
							lowercase__   					=			ret["log"]
							return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}




				def 							A__ (  self   :  Optional[Any],   __lowercase   :  Dict  ):
							# updating to test_epoch_end instead of deprecated test_end
							lowercase__					,    lowercase__					,    lowercase__   					=			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
							lowercase__   					=			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 							A__ (  __lowercase   :  Tuple,   __lowercase   :  Union[str, Any]  ):
							# Add NER specific options
							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__":
							lowercase_						=       argparse.ArgumentParser()
							add_generic_args(parser, os.getcwd())
							lowercase_						=       NERTransformer.add_model_specific_args(parser, os.getcwd())
							lowercase_						=       parser.parse_args()
							lowercase_						=       NERTransformer(args)
							lowercase_						=       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_						=       sorted(glob.glob(os.path.join(args.output_dir, """checkpoint-epoch=*.ckpt"""), recursive=True))
														lowercase_						=       model.load_from_checkpoint(checkpoints[-1])
														trainer.test(model)









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


lowercase_						=       logging.get_logger(__name__)

lowercase_						=       {
    """facebook/nllb-moe-54B""": """https://huggingface.co/facebook/nllb-moe-54b/resolve/main/config.json""",
}
class 			_snake_case     (   lowercase__):
				UpperCamelCase__			:   Tuple       						="""nllb-moe"""
				UpperCamelCase__			:   Dict       						=["""past_key_values"""]
				UpperCamelCase__			:   List[Any]       						={"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}




				def __init__(  self   :  List[str],   __lowercase   :  str=12_8112,   __lowercase   :  List[Any]=1024,   __lowercase   :  int=12,   __lowercase   :  Union[str, Any]=4096,   __lowercase   :  Union[str, Any]=16,   __lowercase   :  Tuple=12,   __lowercase   :  int=4096,   __lowercase   :  Any=16,   __lowercase   :  List[str]=0.05,   __lowercase   :  Optional[int]=0.05,   __lowercase   :  int=True,   __lowercase   :  Optional[Any]=True,   __lowercase   :  Optional[Any]="relu",   __lowercase   :  str=1024,   __lowercase   :  Optional[Any]=0.1,   __lowercase   :  List[Any]=0.1,   __lowercase   :  Union[str, Any]=0.0,   __lowercase   :  Tuple=0.02,   __lowercase   :  Optional[Any]=2,   __lowercase   :  Optional[int]=True,   __lowercase   :  List[Any]=False,   __lowercase   :  Any="float32",   __lowercase   :  Union[str, Any]=False,   __lowercase   :  Optional[int]=128,   __lowercase   :  Union[str, Any]=64,   __lowercase   :  Optional[int]=4,   __lowercase   :  Dict=4,   __lowercase   :  Any=0.001,   __lowercase   :  Optional[int]=0.001,   __lowercase   :  int="all",   __lowercase   :  Any=False,   __lowercase   :  Optional[int]=False,   __lowercase   :  Dict=1.0,   __lowercase   :  int=0.2,   __lowercase   :  str=1,   __lowercase   :  Optional[Any]=0,   __lowercase   :  Optional[int]=2,   __lowercase   :  Any=False,   **__lowercase   :  List[str],   ):
							lowercase__   					=			vocab_size
							lowercase__   					=			max_position_embeddings
							lowercase__   					=			d_model
							lowercase__   					=			encoder_ffn_dim
							lowercase__   					=			encoder_layers
							lowercase__   					=			encoder_attention_heads
							lowercase__   					=			decoder_ffn_dim
							lowercase__   					=			decoder_layers
							lowercase__   					=			decoder_attention_heads
							lowercase__   					=			dropout
							lowercase__   					=			attention_dropout
							lowercase__   					=			activation_dropout
							lowercase__   					=			activation_function
							lowercase__   					=			init_std
							lowercase__   					=			encoder_layerdrop
							lowercase__   					=			decoder_layerdrop
							lowercase__   					=			use_cache
							lowercase__   					=			encoder_layers
							lowercase__   					=			scale_embedding  # scale factor will be sqrt(d_model) if True
							lowercase__   					=			router_z_loss_coef
							lowercase__   					=			router_aux_loss_coef
							lowercase__   					=			decoder_sparse_step
							lowercase__   					=			encoder_sparse_step
							lowercase__   					=			num_experts
							lowercase__   					=			expert_capacity
							lowercase__   					=			router_bias
							if router_dtype not in ["float32", "float16", "bfloat16"]:
										raise ValueError(F'''`router_dtype` must be one of \'float32\', \'float16\' or \'bfloat16\', got {router_dtype}'''  )
							lowercase__   					=			router_dtype

							lowercase__   					=			router_ignore_padding_tokens
							lowercase__   					=			batch_prioritized_routing
							lowercase__   					=			second_expert_policy
							lowercase__   					=			normalize_router_prob_before_dropping
							lowercase__   					=			moe_eval_capacity_token_fraction
							lowercase__   					=			moe_token_dropout
							lowercase__   					=			output_router_logits
							super().__init__(
							    pad_token_id=__lowercase,   bos_token_id=__lowercase,   eos_token_id=__lowercase,   is_encoder_decoder=__lowercase,   decoder_start_token_id=__lowercase,   **__lowercase,   )




