


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


A  :     Optional[int]    					=		logging.getLogger(__name__)
class 						_UpperCamelCase						(     _UpperCamelCase	):
							'''simple docstring'''







							__UpperCAmelCase		:					int          ="""token-classification"""

							def __init__(  self					,     __a			):
										if type(_UpperCAmelCase			) == dict:
													__lowerCAmelCase       					=    Namespace(**_UpperCAmelCase			)
										__lowerCAmelCase       					=    import_module("tasks"			)
										try:
													__lowerCAmelCase       					=    getattr(_UpperCAmelCase					,     hparams.task_type			)
													__lowerCAmelCase       					=    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       					=    self.token_classification_task.get_labels(hparams.labels			)
										__lowerCAmelCase       					=    CrossEntropyLoss().ignore_index
										super().__init__(_UpperCAmelCase					,     len(self.labels			)					,     self.mode			)

							def     snake_case       (  self					,     **__a			):
										return self.model(**_UpperCAmelCase			)

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

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

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

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

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

							def     snake_case       (  self					,     __a			):
										__lowerCAmelCase       					=    torch.stack([x["val_loss"] for x in outputs]			).mean()
										__lowerCAmelCase       					=    np.concatenate([x["pred"] for x in outputs]					,     axis=0			)
										__lowerCAmelCase       					=    np.argmax(_UpperCAmelCase					,     axis=2			)
										__lowerCAmelCase       					=    np.concatenate([x["target"] for x in outputs]					,     axis=0			)

										__lowerCAmelCase       					=    dict(enumerate(self.labels			)			)
										__lowerCAmelCase       					=    [[] for _ in range(out_label_ids.shape[0]			)]
										__lowerCAmelCase       					=    [[] 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       					=    {
										    '''val_loss''': val_loss_mean,
										    '''accuracy_score''': accuracy_score(_UpperCAmelCase					,     _UpperCAmelCase			),
										    '''precision''': precision_score(_UpperCAmelCase					,     _UpperCAmelCase			),
										    '''recall''': recall_score(_UpperCAmelCase					,     _UpperCAmelCase			),
										    '''f1''': fa_score(_UpperCAmelCase					,     _UpperCAmelCase			),
										}

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

							def     snake_case       (  self					,     __a			):
										# when stable
										__lowerCAmelCase       					=    self._eval_end(_UpperCAmelCase			)
										__lowerCAmelCase       					=    ret['''log''']
										return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}

							def     snake_case       (  self					,     __a			):
										# updating to test_epoch_end instead of deprecated test_end
										__lowerCAmelCase       					=    self._eval_end(_UpperCAmelCase			)

										# Converting to the dict required by pl
										# https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
										# pytorch_lightning/trainer/logging.py#L139
										__lowerCAmelCase       					=    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     snake_case       (  __a					,     __a			):
										# Add NER specific options
										BaseTransformer.add_model_specific_args(_UpperCAmelCase					,     _UpperCAmelCase			)
										parser.add_argument(
										    "--task_type"					,     default="NER"					,     type=_UpperCAmelCase					,     help="Task type to fine tune in training (e.g. NER, POS, etc)"			)
										parser.add_argument(
										    "--max_seq_length"					,     default=1_28					,     type=_UpperCAmelCase					,     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=_UpperCAmelCase					,     help="Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."					,     )
										parser.add_argument(
										    "--gpus"					,     default=0					,     type=_UpperCAmelCase					,     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__":
	A  :     Optional[Any]    					=		argparse.ArgumentParser()
	add_generic_args(parser, os.getcwd())
	A  :     Union[str, Any]    					=		NERTransformer.add_model_specific_args(parser, os.getcwd())
	A  :     Tuple    					=		parser.parse_args()
	A  :     List[str]    					=		NERTransformer(args)
	A  :     Tuple    					=		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
		A  :     List[str]    					=		sorted(glob.glob(os.path.join(args.output_dir, "checkpoint-epoch=*.ckpt"), recursive=True))
		A  :     List[str]    					=		model.load_from_checkpoint(checkpoints[-1])
		trainer.test(model)








"""simple docstring"""


import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import TimesformerConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch
	from torch import nn

	from transformers import (
	    MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
	    TimesformerForVideoClassification,
	    TimesformerModel,
	)
	from transformers.models.timesformer.modeling_timesformer import TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from transformers import VideoMAEImageProcessor
class 						_UpperCamelCase						:
							'''simple docstring'''







							def __init__(  self					,     __a					,     __a=13					,     __a=10					,     __a=3					,     __a=2					,     __a=2					,     __a=True					,     __a=True					,     __a=32					,     __a=5					,     __a=4					,     __a=37					,     __a="gelu"					,     __a=0.1					,     __a=0.1					,     __a=10					,     __a=0.0_2					,     __a="divided_space_time"					,     __a=None					,     ):
										__lowerCAmelCase       					=    parent
										__lowerCAmelCase       					=    batch_size
										__lowerCAmelCase       					=    image_size
										__lowerCAmelCase       					=    num_channels
										__lowerCAmelCase       					=    patch_size
										__lowerCAmelCase       					=    num_frames
										__lowerCAmelCase       					=    is_training
										__lowerCAmelCase       					=    use_labels
										__lowerCAmelCase       					=    hidden_size
										__lowerCAmelCase       					=    num_hidden_layers
										__lowerCAmelCase       					=    num_attention_heads
										__lowerCAmelCase       					=    intermediate_size
										__lowerCAmelCase       					=    hidden_act
										__lowerCAmelCase       					=    hidden_dropout_prob
										__lowerCAmelCase       					=    attention_probs_dropout_prob
										__lowerCAmelCase       					=    attention_type
										__lowerCAmelCase       					=    initializer_range
										__lowerCAmelCase       					=    scope
										__lowerCAmelCase       					=    num_labels

										# in TimeSformer, the number of spatial tokens equals num_frames * num_patches per frame + 1 CLS token
										__lowerCAmelCase       					=    (image_size // patch_size) ** 2
										__lowerCAmelCase       					=    (num_frames) * self.num_patches_per_frame + 1

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    floats_tensor(
										    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]			)

										__lowerCAmelCase       					=    None
										if self.use_labels:
													__lowerCAmelCase       					=    ids_tensor([self.batch_size]					,     self.num_labels			)

										__lowerCAmelCase       					=    self.get_config()

										return config, pixel_values, labels

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    TimesformerConfig(
										    image_size=self.image_size					,     patch_size=self.patch_size					,     num_channels=self.num_channels					,     num_frames=self.num_frames					,     hidden_size=self.hidden_size					,     num_hidden_layers=self.num_hidden_layers					,     num_attention_heads=self.num_attention_heads					,     intermediate_size=self.intermediate_size					,     hidden_act=self.hidden_act					,     hidden_dropout_prob=self.hidden_dropout_prob					,     attention_probs_dropout_prob=self.attention_probs_dropout_prob					,     initializer_range=self.initializer_range					,     attention_type=self.attention_type					,     )
										__lowerCAmelCase       					=    self.num_labels
										return config

							def     snake_case       (  self					,     __a					,     __a					,     __a			):
										__lowerCAmelCase       					=    TimesformerModel(config=__a			)
										model.to(__a			)
										model.eval()
										__lowerCAmelCase       					=    model(__a			)
										self.parent.assertEqual(result.last_hidden_state.shape					,     (self.batch_size, self.seq_length, self.hidden_size)			)

							def     snake_case       (  self					,     __a					,     __a					,     __a			):
										__lowerCAmelCase       					=    TimesformerForVideoClassification(__a			)
										model.to(__a			)
										model.eval()

										__lowerCAmelCase       					=    model(__a			)

										# verify the logits shape
										__lowerCAmelCase       					=    torch.Size((self.batch_size, self.num_labels)			)
										self.parent.assertEqual(result.logits.shape					,     __a			)



							def     snake_case       (  self			):
										__lowerCAmelCase       					=    self.prepare_config_and_inputs()
										__lowerCAmelCase					,			__lowerCAmelCase					,			__lowerCAmelCase       					=    config_and_inputs
										__lowerCAmelCase       					=    {"pixel_values": pixel_values}
										return config, inputs_dict
@require_torch
class 						_UpperCamelCase						(     lowerCAmelCase__ ,lowerCAmelCase__ ,unittest.TestCase	):
							'''simple docstring'''







							__UpperCAmelCase		:					Optional[int]          =(TimesformerModel, TimesformerForVideoClassification) if is_torch_available() else ()
							__UpperCAmelCase		:					Tuple          =(
							    {"""feature-extraction""": TimesformerModel, """video-classification""": TimesformerForVideoClassification}
							    if is_torch_available()
							    else {}
							)

							__UpperCAmelCase		:					Any          =False
							__UpperCAmelCase		:					Optional[int]          =False
							__UpperCAmelCase		:					Union[str, Any]          =False
							__UpperCAmelCase		:					Any          =False

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    TimesformerModelTester(self			)
										__lowerCAmelCase       					=    ConfigTester(
										    self					,     config_class=__a					,     has_text_modality=__a					,     hidden_size=37			)

							def     snake_case       (  self					,     __a					,     __a					,     __a=False			):
										__lowerCAmelCase       					=    copy.deepcopy(__a			)

										if return_labels:
													if model_class in get_values(__a			):
																__lowerCAmelCase       					=    torch.zeros(
																    self.model_tester.batch_size					,     dtype=torch.long					,     device=__a			)

										return inputs_dict

							def     snake_case       (  self			):
										self.config_tester.run_common_tests()

							@unittest.skip(reason="TimeSformer does not use inputs_embeds"			)
							def     snake_case       (  self			):
										pass

							def     snake_case       (  self			):
										__lowerCAmelCase					,			__lowerCAmelCase       					=    self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
													__lowerCAmelCase       					=    model_class(__a			)
													self.assertIsInstance(model.get_input_embeddings()					,     (nn.Module)			)
													__lowerCAmelCase       					=    model.get_output_embeddings()
													self.assertTrue(x is None or isinstance(__a					,     nn.Linear			)			)

							def     snake_case       (  self			):
										__lowerCAmelCase					,			__lowerCAmelCase       					=    self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
													__lowerCAmelCase       					=    model_class(__a			)
													__lowerCAmelCase       					=    inspect.signature(model.forward			)
													# signature.parameters is an OrderedDict => so arg_names order is deterministic
													__lowerCAmelCase       					=    [*signature.parameters.keys()]

													__lowerCAmelCase       					=    ["pixel_values"]
													self.assertListEqual(arg_names[:1]					,     __a			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*__a			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_video_classification(*__a			)

							@slow
							def     snake_case       (  self			):
										for model_name in TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													__lowerCAmelCase       					=    TimesformerModel.from_pretrained(__a			)
													self.assertIsNotNone(__a			)

							def     snake_case       (  self			):
										if not self.has_attentions:
													pass

										else:
													__lowerCAmelCase					,			__lowerCAmelCase       					=    self.model_tester.prepare_config_and_inputs_for_common()
													__lowerCAmelCase       					=    True

													for model_class in self.all_model_classes:
																__lowerCAmelCase       					=    self.model_tester.seq_length
																__lowerCAmelCase       					=    self.model_tester.num_frames

																__lowerCAmelCase       					=    True
																__lowerCAmelCase       					=    False
																__lowerCAmelCase       					=    True
																__lowerCAmelCase       					=    model_class(__a			)
																model.to(__a			)
																model.eval()
																with torch.no_grad():
																			__lowerCAmelCase       					=    model(**self._prepare_for_class(__a					,     __a			)			)
																__lowerCAmelCase       					=    outputs.attentions
																self.assertEqual(len(__a			)					,     self.model_tester.num_hidden_layers			)

																# check that output_attentions also work using config
																del inputs_dict["output_attentions"]
																__lowerCAmelCase       					=    True
																__lowerCAmelCase       					=    model_class(__a			)
																model.to(__a			)
																model.eval()
																with torch.no_grad():
																			__lowerCAmelCase       					=    model(**self._prepare_for_class(__a					,     __a			)			)
																__lowerCAmelCase       					=    outputs.attentions
																self.assertEqual(len(__a			)					,     self.model_tester.num_hidden_layers			)

																# attentions has shape (batch_size x num_frames) x num_heads x (num_patches per frame + 1) x (num_patches per frame + 1)
																self.assertListEqual(
																    list(attentions[0].shape[-3:]			)					,     [self.model_tester.num_attention_heads, seq_len // num_frames + 1, seq_len // num_frames + 1]					,     )
																__lowerCAmelCase       					=    len(__a			)

																# Check attention is always last and order is fine
																__lowerCAmelCase       					=    True
																__lowerCAmelCase       					=    True
																__lowerCAmelCase       					=    model_class(__a			)
																model.to(__a			)
																model.eval()
																with torch.no_grad():
																			__lowerCAmelCase       					=    model(**self._prepare_for_class(__a					,     __a			)			)

																self.assertEqual(out_len + 1					,     len(__a			)			)

																__lowerCAmelCase       					=    outputs.attentions

																self.assertEqual(len(__a			)					,     self.model_tester.num_hidden_layers			)

																# attentions has shape (batch_size x num_frames) x num_heads x (num_patches per frame + 1) x (num_patches per frame + 1)
																self.assertListEqual(
																    list(self_attentions[0].shape[-3:]			)					,     [self.model_tester.num_attention_heads, seq_len // num_frames + 1, seq_len // num_frames + 1]					,     )



							def     snake_case       (  self			):
										def check_hidden_states_output(__a					,     __a					,     __a			):
													__lowerCAmelCase       					=    model_class(__a			)
													model.to(__a			)
													model.eval()

													with torch.no_grad():
																__lowerCAmelCase       					=    model(**self._prepare_for_class(__a					,     __a			)			)

													__lowerCAmelCase       					=    outputs.hidden_states
													__lowerCAmelCase       					=    self.model_tester.num_hidden_layers + 1
													self.assertEqual(len(__a			)					,     __a			)

													__lowerCAmelCase       					=    self.model_tester.seq_length

													self.assertListEqual(
													    list(hidden_states[0].shape[-2:]			)					,     [seq_length, self.model_tester.hidden_size]					,     )

										__lowerCAmelCase					,			__lowerCAmelCase       					=    self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
													__lowerCAmelCase       					=    True
													check_hidden_states_output(__a					,     __a					,     __a			)

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													__lowerCAmelCase       					=    True

													check_hidden_states_output(__a					,     __a					,     __a			)


def _lowerCamelCase			(	):

			'''simple docstring'''




			__lowerCAmelCase       					=    hf_hub_download(
			    repo_id="hf-internal-testing/spaghetti-video"					,				filename="eating_spaghetti.npy"					,				repo_type="dataset"   )
			__lowerCAmelCase       					=    np.load(_UpperCamelCase   )
			return list(_UpperCamelCase   )


@require_torch
@require_vision
class 						_UpperCamelCase						(     unittest.TestCase	):
							'''simple docstring'''







							@cached_property
							def     snake_case       (  self			):
										# logits were tested with a different mean and std, so we use the same here
										return (
										    VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]					,     image_std=[0.5, 0.5, 0.5]			)
										    if is_vision_available()
										    else None
										)



							@slow
							def     snake_case       (  self			):
										__lowerCAmelCase       					=    TimesformerForVideoClassification.from_pretrained("facebook/timesformer-base-finetuned-k400"			).to(
										    __a			)

										__lowerCAmelCase       					=    self.default_image_processor
										__lowerCAmelCase       					=    prepare_video()
										__lowerCAmelCase       					=    image_processor(video[:8]					,     return_tensors="pt"			).to(__a			)

										# forward pass
										with torch.no_grad():
													__lowerCAmelCase       					=    model(**__a			)

										# verify the logits
										__lowerCAmelCase       					=    torch.Size((1, 4_00)			)
										self.assertEqual(outputs.logits.shape					,     __a			)

										__lowerCAmelCase       					=    torch.tensor([-0.3_0_1_6, -0.7_7_1_3, -0.4_2_0_5]			).to(__a			)

										self.assertTrue(torch.allclose(outputs.logits[0, :3]					,     __a					,     atol=1e-4			)			)




