





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 	UpperCamelCase_							(  UpperCamelCase__     ):
							lowerCamelCase_							     =       "token-classification"


							def __init__(				self     :Dict				,					__A     :int		) -> Union[str, Any]:



									"""simple docstring"""



									if type(__A		) == dict:
											SCREAMING_SNAKE_CASE__								=   Namespace(**__A		)
									SCREAMING_SNAKE_CASE__								=   import_module("""tasks"""		)
									try:
											SCREAMING_SNAKE_CASE__								=   getattr(__A				,					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__(__A				,					len(self.labels		)				,					self.mode		)


							def   _snake_case	(				self     :Tuple				,					**__A     :Optional[Any]		) -> List[str]:



									"""simple docstring"""



									return self.model(**__A		)


							def   _snake_case	(				self     :int				,					__A     :List[Any]				,					__A     :Any		) -> Optional[int]:



									"""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(**__A		)
									SCREAMING_SNAKE_CASE__								=   outputs[0]
									# tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
									return {"loss": loss}


							def   _snake_case	(				self     :List[Any]		) -> Dict:



									"""simple docstring"""



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


							def   _snake_case	(				self     :Any				,					__A     :int				,					__A     :int				,					__A     :bool = False		) -> DataLoader:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self._feature_file(__A		)
									logger.info("""Loading features from cached file %s"""				,					__A		)
									SCREAMING_SNAKE_CASE__								=   torch.load(__A		)
									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(__A				,					__A				,					__A				,					__A		)				,					batch_size=__A		)


							def   _snake_case	(				self     :List[Any]				,					__A     :Optional[Any]				,					__A     :Union[str, Any]		) -> Optional[Any]:



									"""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(**__A		)
									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   _snake_case	(				self     :Optional[int]				,					__A     :List[Any]		) -> Tuple:



									"""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(__A				,					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(__A				,					__A		),
									    """precision""": precision_score(__A				,					__A		),
									    """recall""": recall_score(__A				,					__A		),
									    """f1""": fa_score(__A				,					__A		),
									}

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


							def   _snake_case	(				self     :str				,					__A     :List[str]		) -> List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   self._eval_end(__A		)
									SCREAMING_SNAKE_CASE__								=   ret["""log"""]
									return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}


							def   _snake_case	(				self     :Any				,					__A     :List[Any]		) -> str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   self._eval_end(__A		)

									# 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   _snake_case	(				__A     :Optional[int]				,					__A     :int		) -> Tuple:



									"""simple docstring"""



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






import argparse

import torch
from datasets import load_dataset
from donut import DonutModel

from transformers import (
    DonutImageProcessor,
    DonutProcessor,
    DonutSwinConfig,
    DonutSwinModel,
    MBartConfig,
    MBartForCausalLM,
    VisionEncoderDecoderModel,
    XLMRobertaTokenizerFast,
)






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				str     ):
		SCREAMING_SNAKE_CASE__								=   model.config

		SCREAMING_SNAKE_CASE__								=   DonutSwinConfig(
		    image_size=original_config.input_size       , patch_size=4       , depths=original_config.encoder_layer       , num_heads=[4, 8, 16, 32]       , window_size=original_config.window_size       , embed_dim=128       , )
		SCREAMING_SNAKE_CASE__								=   MBartConfig(
		    is_decoder=UpperCamelCase__       , is_encoder_decoder=UpperCamelCase__       , add_cross_attention=UpperCamelCase__       , decoder_layers=original_config.decoder_layer       , max_position_embeddings=original_config.max_position_embeddings       , vocab_size=len(
		        model.decoder.tokenizer     )       , scale_embedding=UpperCamelCase__       , add_final_layer_norm=UpperCamelCase__       , )

		return encoder_config, decoder_config






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				List[str]     ):
		if "encoder.model" in name:
				SCREAMING_SNAKE_CASE__								=   name.replace("""encoder.model"""       , """encoder"""     )
		if "decoder.model" in name:
				SCREAMING_SNAKE_CASE__								=   name.replace("""decoder.model"""       , """decoder"""     )
		if "patch_embed.proj" in name:
				SCREAMING_SNAKE_CASE__								=   name.replace("""patch_embed.proj"""       , """embeddings.patch_embeddings.projection"""     )
		if "patch_embed.norm" in name:
				SCREAMING_SNAKE_CASE__								=   name.replace("""patch_embed.norm"""       , """embeddings.norm"""     )
		if name.startswith("""encoder"""     ):
				if "layers" in name:
						SCREAMING_SNAKE_CASE__								=   """encoder.""" + name
				if "attn.proj" in name:
						SCREAMING_SNAKE_CASE__								=   name.replace("""attn.proj"""       , """attention.output.dense"""     )
				if "attn" in name and "mask" not in name:
						SCREAMING_SNAKE_CASE__								=   name.replace("""attn"""       , """attention.self"""     )
				if "norm1" in name:
						SCREAMING_SNAKE_CASE__								=   name.replace("""norm1"""       , """layernorm_before"""     )
				if "norm2" in name:
						SCREAMING_SNAKE_CASE__								=   name.replace("""norm2"""       , """layernorm_after"""     )
				if "mlp.fc1" in name:
						SCREAMING_SNAKE_CASE__								=   name.replace("""mlp.fc1"""       , """intermediate.dense"""     )
				if "mlp.fc2" in name:
						SCREAMING_SNAKE_CASE__								=   name.replace("""mlp.fc2"""       , """output.dense"""     )

				if name == "encoder.norm.weight":
						SCREAMING_SNAKE_CASE__								=   """encoder.layernorm.weight"""
				if name == "encoder.norm.bias":
						SCREAMING_SNAKE_CASE__								=   """encoder.layernorm.bias"""

		return name






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				str       , UpperCamelCase__:				Optional[int]     ):
		for key in orig_state_dict.copy().keys():
				SCREAMING_SNAKE_CASE__								=   orig_state_dict.pop(UpperCamelCase__     )

				if "qkv" in key:
						SCREAMING_SNAKE_CASE__								=   key.split("""."""     )
						SCREAMING_SNAKE_CASE__								=   int(key_split[3]     )
						SCREAMING_SNAKE_CASE__								=   int(key_split[5]     )
						SCREAMING_SNAKE_CASE__								=   model.encoder.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

						if "weight" in key:
								SCREAMING_SNAKE_CASE__								=   val[:dim, :]
								SCREAMING_SNAKE_CASE__								=   val[dim : dim * 2, :]
								SCREAMING_SNAKE_CASE__								=   val[-dim:, :]
						else:
								SCREAMING_SNAKE_CASE__								=   val[:dim]
								SCREAMING_SNAKE_CASE__								=   val[dim : dim * 2]
								SCREAMING_SNAKE_CASE__								=   val[-dim:]
				elif "attn_mask" in key or key in ["encoder.model.norm.weight", "encoder.model.norm.bias"]:
						# HuggingFace implementation doesn't use attn_mask buffer
						# and model doesn't use final LayerNorms for the encoder
						pass
				else:
						SCREAMING_SNAKE_CASE__								=   val

		return orig_state_dict






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				Union[str, Any]       , UpperCamelCase__:				int=None       , UpperCamelCase__:				str=False     ):
		# load original model
		SCREAMING_SNAKE_CASE__								=   DonutModel.from_pretrained(UpperCamelCase__     ).eval()

		# load HuggingFace model
		SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   get_configs(UpperCamelCase__     )
		SCREAMING_SNAKE_CASE__								=   DonutSwinModel(UpperCamelCase__     )
		SCREAMING_SNAKE_CASE__								=   MBartForCausalLM(UpperCamelCase__     )
		SCREAMING_SNAKE_CASE__								=   VisionEncoderDecoderModel(encoder=UpperCamelCase__       , decoder=UpperCamelCase__     )
		model.eval()

		SCREAMING_SNAKE_CASE__								=   original_model.state_dict()
		SCREAMING_SNAKE_CASE__								=   convert_state_dict(UpperCamelCase__       , UpperCamelCase__     )
		model.load_state_dict(UpperCamelCase__     )

		# verify results on scanned document
		SCREAMING_SNAKE_CASE__								=   load_dataset("""hf-internal-testing/example-documents"""     )
		SCREAMING_SNAKE_CASE__								=   dataset["""test"""][0]["""image"""].convert("""RGB"""     )

		SCREAMING_SNAKE_CASE__								=   XLMRobertaTokenizerFast.from_pretrained(UpperCamelCase__       , from_slow=UpperCamelCase__     )
		SCREAMING_SNAKE_CASE__								=   DonutImageProcessor(
		    do_align_long_axis=original_model.config.align_long_axis       , size=original_model.config.input_size[::-1]     )
		SCREAMING_SNAKE_CASE__								=   DonutProcessor(UpperCamelCase__       , UpperCamelCase__     )
		SCREAMING_SNAKE_CASE__								=   processor(UpperCamelCase__       , return_tensors="""pt"""     ).pixel_values

		if model_name == "naver-clova-ix/donut-base-finetuned-docvqa":
				SCREAMING_SNAKE_CASE__								=   """<s_docvqa><s_question>{user_input}</s_question><s_answer>"""
				SCREAMING_SNAKE_CASE__								=   """When is the coffee break?"""
				SCREAMING_SNAKE_CASE__								=   task_prompt.replace("""{user_input}"""       , UpperCamelCase__     )
		elif model_name == "naver-clova-ix/donut-base-finetuned-rvlcdip":
				SCREAMING_SNAKE_CASE__								=   """<s_rvlcdip>"""
		elif model_name in [
		    "naver-clova-ix/donut-base-finetuned-cord-v1",
		    "naver-clova-ix/donut-base-finetuned-cord-v1-2560",
		]:
				SCREAMING_SNAKE_CASE__								=   """<s_cord>"""
		elif model_name == "naver-clova-ix/donut-base-finetuned-cord-v2":
				SCREAMING_SNAKE_CASE__								=   """s_cord-v2>"""
		elif model_name == "naver-clova-ix/donut-base-finetuned-zhtrainticket":
				SCREAMING_SNAKE_CASE__								=   """<s_zhtrainticket>"""
		elif model_name in ["naver-clova-ix/donut-proto", "naver-clova-ix/donut-base"]:
				# use a random prompt
				SCREAMING_SNAKE_CASE__								=   """hello world"""
		else:
				raise ValueError("""Model name not supported"""     )
		SCREAMING_SNAKE_CASE__								=   original_model.decoder.tokenizer(UpperCamelCase__       , add_special_tokens=UpperCamelCase__       , return_tensors="""pt"""     )[
		    """input_ids"""
		]

		SCREAMING_SNAKE_CASE__								=   original_model.encoder.model.patch_embed(UpperCamelCase__     )
		SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   model.encoder.embeddings(UpperCamelCase__     )
		assert torch.allclose(UpperCamelCase__       , UpperCamelCase__       , atol=1e-3     )

		# verify encoder hidden states
		SCREAMING_SNAKE_CASE__								=   original_model.encoder(UpperCamelCase__     )
		SCREAMING_SNAKE_CASE__								=   model.encoder(UpperCamelCase__     ).last_hidden_state
		assert torch.allclose(UpperCamelCase__       , UpperCamelCase__       , atol=1e-2     )

		# verify decoder hidden states
		SCREAMING_SNAKE_CASE__								=   original_model(UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__     ).logits
		SCREAMING_SNAKE_CASE__								=   model(UpperCamelCase__       , decoder_input_ids=UpperCamelCase__     ).logits
		assert torch.allclose(UpperCamelCase__       , UpperCamelCase__       , atol=1e-3     )
		print("""Looks ok!"""     )

		if pytorch_dump_folder_path is not None:
				print(f'''Saving model and processor to {pytorch_dump_folder_path}'''     )
				model.save_pretrained(UpperCamelCase__     )
				processor.save_pretrained(UpperCamelCase__     )

		if push_to_hub:
				model.push_to_hub("""nielsr/""" + model_name.split("""/"""     )[-1]       , commit_message="""Update model"""     )
				processor.push_to_hub("""nielsr/""" + model_name.split("""/"""     )[-1]       , commit_message="""Update model"""     )


if __name__ == "__main__":
	_lowerCamelCase 					=      argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '--model_name',
	    default='naver-clova-ix/donut-base-finetuned-docvqa',
	    required=False,
	    type=str,
	    help='Name of the original model you\'d like to convert.',
	)
	parser.add_argument(
	    '--pytorch_dump_folder_path',
	    default=None,
	    required=False,
	    type=str,
	    help='Path to the output PyTorch model directory.',
	)
	parser.add_argument(
	    '--push_to_hub',
	    action='store_true',
	    help='Whether or not to push the converted model and processor to the 🤗 hub.',
	)

	_lowerCamelCase 					=      parser.parse_args()
	convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)