


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


_UpperCAmelCase			:  Optional[int]     =logging.getLogger(__name__)







class 				snake_case__(							UpperCAmelCase__							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  str					=							"""token-classification"""




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




						def 						lowercase_ (				self   ,						**__lowercase							)					->    Any:
										return self.model(**__lowercase							)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase							)					->    int:
										lowerCAmelCase_ :      List[Any]    				=				{'''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_ :      Dict    				=				self(**__lowercase							)
										lowerCAmelCase_ :      int    				=				outputs[0]
										# tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
										return {"loss": loss}




						def 						lowercase_ (				self							)					->    str:
										lowerCAmelCase_ :      List[str]    				=				self.hparams
										for mode in ["train", "dev", "test"]:
														lowerCAmelCase_ :      List[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_ :      Union[str, Any]    				=				torch.load(__lowercase							)
														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   ,						__lowercase							)
																		lowerCAmelCase_ :      str    				=				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 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase = False							)					->    DataLoader:
										lowerCAmelCase_ :      List[Any]    				=				self._feature_file(__lowercase							)
										logger.info('''Loading features from cached file %s'''   ,						__lowercase							)
										lowerCAmelCase_ :      List[str]    				=				torch.load(__lowercase							)
										lowerCAmelCase_ :      List[str]    				=				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_ :      Optional[Any]    				=				torch.tensor([f.token_type_ids for f in features]   ,						dtype=torch.long							)
										else:
														lowerCAmelCase_ :      Optional[int]    				=				torch.tensor([0 for f in features]   ,						dtype=torch.long							)
														# HACK(we will not use this anymore soon)
										lowerCAmelCase_ :      Tuple    				=				torch.tensor([f.label_ids for f in features]   ,						dtype=torch.long							)
										return DataLoader(
										    TensorDataset(__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase							)   ,						batch_size=__lowercase							)




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




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

										lowerCAmelCase_ :      List[str]    				=				dict(enumerate(self.labels							)							)
										lowerCAmelCase_ :      List[str]    				=				[[] 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_ :      str    				=				{
										    '''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_ :      List[str]    				=				dict(results.items()							)
										lowerCAmelCase_ :      Union[str, Any]    				=				results
										return ret, preds_list, out_label_list




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




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



from manim import *







class 				snake_case__(							UpperCAmelCase__							):







						'''simple docstring'''




						def 						lowercase_ (				self							)					->    Tuple:
										lowerCAmelCase_ :      Dict    				=				Rectangle(height=0.5   ,						width=0.5							)
										lowerCAmelCase_ :      Tuple    				=				Rectangle(height=0.25   ,						width=0.25							)
										lowerCAmelCase_ :      Tuple    				=				Rectangle(height=0.46   ,						width=0.46							).set_stroke(width=0							)

										lowerCAmelCase_ :      Optional[int]    				=				[mem.copy() for i in range(6							)]
										lowerCAmelCase_ :      int    				=				[mem.copy() for i in range(6							)]
										lowerCAmelCase_ :      Optional[int]    				=				VGroup(*__lowercase							).arrange(__lowercase   ,						buff=0							)
										lowerCAmelCase_ :      List[str]    				=				VGroup(*__lowercase							).arrange(__lowercase   ,						buff=0							)
										lowerCAmelCase_ :      int    				=				VGroup(__lowercase   ,						__lowercase							).arrange(__lowercase   ,						buff=0							)
										lowerCAmelCase_ :      Tuple    				=				Text('''CPU'''   ,						font_size=2_4							)
										lowerCAmelCase_ :      Union[str, Any]    				=				Group(__lowercase   ,						__lowercase							).arrange(__lowercase   ,						buff=0.5   ,						aligned_edge=__lowercase							)
										cpu.move_to([-2.5, -0.5, 0]							)
										self.add(__lowercase							)

										lowerCAmelCase_ :      List[str]    				=				[mem.copy() for i in range(4							)]
										lowerCAmelCase_ :      Any    				=				VGroup(*__lowercase							).arrange(__lowercase   ,						buff=0							)
										lowerCAmelCase_ :      List[Any]    				=				Text('''GPU'''   ,						font_size=2_4							)
										lowerCAmelCase_ :      int    				=				Group(__lowercase   ,						__lowercase							).arrange(__lowercase   ,						buff=0.5   ,						aligned_edge=__lowercase							)
										gpu.move_to([-1, -1, 0]							)
										self.add(__lowercase							)

										lowerCAmelCase_ :      str    				=				[mem.copy() for i in range(6							)]
										lowerCAmelCase_ :      Dict    				=				VGroup(*__lowercase							).arrange(__lowercase   ,						buff=0							)

										lowerCAmelCase_ :      Dict    				=				Text('''Model'''   ,						font_size=2_4							)
										lowerCAmelCase_ :      str    				=				Group(__lowercase   ,						__lowercase							).arrange(__lowercase   ,						buff=0.5   ,						aligned_edge=__lowercase							)
										model.move_to([3, -1.0, 0]							)
										self.add(__lowercase							)

										lowerCAmelCase_ :      int    				=				[]
										lowerCAmelCase_ :      int    				=				[]
										lowerCAmelCase_ :      Dict    				=				[]

										for i, rect in enumerate(__lowercase							):
														rect.set_stroke(__lowercase							)

														lowerCAmelCase_ :      Any    				=				Rectangle(height=0.46 / 4   ,						width=0.46 / 3							).set_stroke(width=0.0							).set_fill(__lowercase   ,						opacity=0.7							)

														if i == 0:
																		cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT							)   ,						buff=0.02   ,						direction=__lowercase							)
																		cpu_target.set_x(cpu_target.get_x() + 0.1							)
														elif i == 3:
																		cpu_target.next_to(model_cpu_arr[0]   ,						direction=__lowercase   ,						buff=0.0							)
														else:
																		cpu_target.next_to(model_cpu_arr[i - 1]   ,						direction=__lowercase   ,						buff=0.0							)
														self.add(__lowercase							)
														model_cpu_arr.append(__lowercase							)

										self.add(*__lowercase   ,						*__lowercase   ,						*__lowercase							)

										lowerCAmelCase_ :      Union[str, Any]    				=				[mem.copy() for i in range(6							)]
										lowerCAmelCase_ :      List[str]    				=				VGroup(*__lowercase							).arrange(__lowercase   ,						buff=0							)

										lowerCAmelCase_ :      Union[str, Any]    				=				Text('''Loaded Checkpoint'''   ,						font_size=2_4							)
										lowerCAmelCase_ :      int    				=				Group(__lowercase   ,						__lowercase							).arrange(__lowercase   ,						buff=0.5   ,						aligned_edge=__lowercase							)
										checkpoint.move_to([3, 0.5, 0]							)

										self.add(__lowercase							)

										lowerCAmelCase_ :      Optional[Any]    				=				[]
										lowerCAmelCase_ :      Dict    				=				[]

										for i, rect in enumerate(__lowercase							):
														lowerCAmelCase_ :      str    				=				fill.copy().set_fill(__lowercase   ,						opacity=0.7							)
														target.move_to(__lowercase							)
														ckpt_arr.append(__lowercase							)

														lowerCAmelCase_ :      Union[str, Any]    				=				target.copy()
														if i < 5:
																		cpu_target.move_to(cpu_left_col_base[i + 1]							)
														else:
																		cpu_target.move_to(cpu_right_col_base[i - 5]							)
														ckpt_cpu_arr.append(__lowercase							)
										self.add(*__lowercase   ,						*__lowercase							)

										lowerCAmelCase_ :      Union[str, Any]    				=				Square(side_length=2.2							)
										key.move_to([-5, 2, 0]							)

										lowerCAmelCase_ :      str    				=				MarkupText(
										    f"""<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model"""   ,						font_size=1_8   ,						)

										key_text.move_to([-5, 2.4, 0]							)

										self.add(__lowercase   ,						__lowercase							)

										lowerCAmelCase_ :      str    				=				MarkupText(
										    f"""<span fgcolor='{BLUE}'>●</span> Checkpoint"""   ,						font_size=1_8   ,						)

										blue_text.next_to(__lowercase   ,						DOWN * 2.4   ,						aligned_edge=key_text.get_left()							)
										self.add(__lowercase							)

										lowerCAmelCase_ :      str    				=				MarkupText(
										    f"""Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device."""   ,						font_size=2_4   ,						)
										step_a.move_to([2, 2, 0]							)

										lowerCAmelCase_ :      List[Any]    				=				[meta_mem.copy() for i in range(6							)]
										lowerCAmelCase_ :      Any    				=				[meta_mem.copy() for i in range(6							)]
										lowerCAmelCase_ :      Any    				=				VGroup(*__lowercase							).arrange(__lowercase   ,						buff=0							)
										lowerCAmelCase_ :      Union[str, Any]    				=				VGroup(*__lowercase							).arrange(__lowercase   ,						buff=0							)
										lowerCAmelCase_ :      int    				=				VGroup(__lowercase   ,						__lowercase							).arrange(__lowercase   ,						buff=0							)
										lowerCAmelCase_ :      List[str]    				=				Text('''Disk'''   ,						font_size=2_4							)
										lowerCAmelCase_ :      Optional[int]    				=				Group(__lowercase   ,						__lowercase							).arrange(__lowercase   ,						buff=0.5   ,						aligned_edge=__lowercase							)
										disk.move_to([-4.0, -1.25, 0]							)
										self.play(Write(__lowercase   ,						run_time=3							)   ,						Write(__lowercase   ,						run_time=1							)   ,						Create(__lowercase   ,						run_time=1							)							)

										lowerCAmelCase_ :      int    				=				[]
										for i, rect in enumerate(__lowercase							):
														lowerCAmelCase_ :      int    				=				rect.copy()
														target.generate_target()
														target.target.move_to(disk_left_col_base[i]							).scale(0.5							)
														animations.append(MoveToTarget(__lowercase   ,						run_time=1.5							)							)
										self.play(*__lowercase							)

										self.play(FadeOut(__lowercase							)							)

										lowerCAmelCase_ :      Union[str, Any]    				=				MarkupText(f"""Then, the checkpoint is removed from memory\nthrough garbage collection."""   ,						font_size=2_4							)
										step_a.move_to([2, 2, 0]							)

										self.play(Write(__lowercase   ,						run_time=3							)							)

										self.play(
										    FadeOut(__lowercase   ,						__lowercase   ,						*__lowercase   ,						*__lowercase							)   ,						)

										self.wait()