
# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re

from ..models.auto import AutoProcessor
from ..models.vision_encoder_decoder import VisionEncoderDecoderModel
from ..utils import is_vision_available
from .base import PipelineTool


if is_vision_available():
			from PIL import Image




class        _SCREAMING_SNAKE_CASE	( snake_case_		):
	lowerCAmelCase__   =					'naver-clova-ix/donut-base-finetuned-docvqa'
	lowerCAmelCase__   =					(
	    'This is a tool that answers a question about an document (pdf). It takes an input named `document` which '
	    'should be the document containing the information, as well as a `question` that is the question about the '
	    'document. It returns a text that contains the answer to the question.'
	)
	lowerCAmelCase__   =					'document_qa'
	lowerCAmelCase__   =					AutoProcessor
	lowerCAmelCase__   =					VisionEncoderDecoderModel

	lowerCAmelCase__   =					['image', 'text']
	lowerCAmelCase__   =					['text']



	def __init__(       self		,    *lowercase		,    **lowercase					)							->			int:
					if not is_vision_available():
									raise ValueError("Pillow must be installed to use the DocumentQuestionAnsweringTool."					)

					super().__init__(*lowercase		,    **lowercase					)



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase		,    lowercase					)							->			Optional[Any]:
					lowerCamelCase_         =     "<s_docvqa><s_question>{user_input}</s_question><s_answer>"
					lowerCamelCase_         =     task_prompt.replace("{user_input}"		,    lowercase					)
					lowerCamelCase_         =     self.pre_processor.tokenizer(
					    lowercase		,    add_special_tokens=lowercase		,    return_tensors="pt"					).input_ids
					lowerCamelCase_         =     self.pre_processor(lowercase		,    return_tensors="pt"					).pixel_values

					return {"decoder_input_ids": decoder_input_ids, "pixel_values": pixel_values}



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase					)							->			Union[str, Any]:
					return self.model.generate(
					    inputs["pixel_values"].to(self.device					)		,    decoder_input_ids=inputs["decoder_input_ids"].to(self.device					)		,    max_length=self.model.decoder.config.max_position_embeddings		,    early_stopping=lowercase		,    pad_token_id=self.pre_processor.tokenizer.pad_token_id		,    eos_token_id=self.pre_processor.tokenizer.eos_token_id		,    use_cache=lowercase		,    num_beams=1		,    bad_words_ids=[[self.pre_processor.tokenizer.unk_token_id]]		,    return_dict_in_generate=lowercase		,    ).sequences



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase					)							->			List[str]:
					lowerCamelCase_         =     self.pre_processor.batch_decode(lowercase					)[0]
					lowerCamelCase_         =     sequence.replace(self.pre_processor.tokenizer.eos_token		,    ""					)
					lowerCamelCase_         =     sequence.replace(self.pre_processor.tokenizer.pad_token		,    ""					)
					lowerCamelCase_         =     re.sub(R"<.*?>"		,    ""		,    lowercase		,    count=1					).strip()  # remove first task start token
					lowerCamelCase_         =     self.pre_processor.tokenajson(lowercase					)

					return sequence["answer"]




import argparse
import os

import torch

from transformers import (
    XLNetConfig,
    XLNetForQuestionAnswering,
    XLNetForSequenceClassification,
    XLNetLMHeadModel,
    load_tf_weights_in_xlnet,
)
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


__A							={
    '''cola''': 2,
    '''mnli''': 3,
    '''mrpc''': 2,
    '''sst-2''': 2,
    '''sts-b''': 1,
    '''qqp''': 2,
    '''qnli''': 2,
    '''rte''': 2,
    '''wnli''': 2,
}


logging.set_verbosity_info()





def      lowerCamelCase_					(							lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__=None       ):
				# Initialise PyTorch model
				lowerCamelCase_         =     XLNetConfig.from_json_file(lowerCamelCase__       )

				lowerCamelCase_         =     finetuning_task.lower() if finetuning_task is not None else ""
				if finetuning_task in GLUE_TASKS_NUM_LABELS:
								print(F'Building PyTorch XLNetForSequenceClassification model from configuration: {config}'       )
								lowerCamelCase_         =     finetuning_task
								lowerCamelCase_         =     GLUE_TASKS_NUM_LABELS[finetuning_task]
								lowerCamelCase_         =     XLNetForSequenceClassification(lowerCamelCase__       )
				elif "squad" in finetuning_task:
								lowerCamelCase_         =     finetuning_task
								lowerCamelCase_         =     XLNetForQuestionAnswering(lowerCamelCase__       )
				else:
								lowerCamelCase_         =     XLNetLMHeadModel(lowerCamelCase__       )

				# Load weights from tf checkpoint
				load_tf_weights_in_xlnet(lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__       )

				# Save pytorch-model
				lowerCamelCase_         =     os.path.join(lowerCamelCase__      ,					lowerCamelCase__       )
				lowerCamelCase_         =     os.path.join(lowerCamelCase__      ,					lowerCamelCase__       )
				print(F'Save PyTorch model to {os.path.abspath(lowerCamelCase__       )}'       )
				torch.save(model.state_dict()      ,					lowerCamelCase__       )
				print(F'Save configuration file to {os.path.abspath(lowerCamelCase__       )}'       )
				with open(lowerCamelCase__      ,					"w"      ,					encoding="utf-8"       ) as f:
								f.write(config.to_json_string()       )


if __name__ == "__main__":
			__A							=argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    '''--tf_checkpoint_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
			)
			parser.add_argument(
			    '''--xlnet_config_file''',
			    default=None,
			    type=str,
			    required=True,
			    help=(
			        '''The config json file corresponding to the pre-trained XLNet model. \n'''
			        '''This specifies the model architecture.'''
			    ),
			)
			parser.add_argument(
			    '''--pytorch_dump_folder_path''',
			    default=None,
			    type=str,
			    required=True,
			    help='''Path to the folder to store the PyTorch model or dataset/vocab.''',
			)
			parser.add_argument(
			    '''--finetuning_task''',
			    default=None,
			    type=str,
			    help='''Name of a task on which the XLNet TensorFlow model was fine-tuned''',
			)
			__A							=parser.parse_args()
			print(args)

			convert_xlnet_checkpoint_to_pytorch(
			    args.tf_checkpoint_path, args.xlnet_config_file, args.pytorch_dump_folder_path, args.finetuning_task
			)


