
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


SCREAMING_SNAKE_CASE							:List[Any]										=    {
    "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 						UpperCAmelCase      (     a_		,       a_		,       a_		,       a_=None						)		->     str:







				"""simple docstring"""



				__A             =			XLNetConfig.from_json_file(UpperCAmelCase_						)

				__A             =			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}'''						)
								__A             =			finetuning_task
								__A             =			GLUE_TASKS_NUM_LABELS[finetuning_task]
								__A             =			XLNetForSequenceClassification(UpperCAmelCase_						)
				elif "squad" in finetuning_task:
								__A             =			finetuning_task
								__A             =			XLNetForQuestionAnswering(UpperCAmelCase_						)
				else:
								__A             =			XLNetLMHeadModel(UpperCAmelCase_						)

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

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


if __name__ == "__main__":
	SCREAMING_SNAKE_CASE							:Any										=    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',
	)
	SCREAMING_SNAKE_CASE							:Any										=    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
	)


# 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 torch

from ..models.auto import AutoModelForSequenceClassification, AutoTokenizer
from .base import PipelineTool





class 							UpperCAmelCase				(  __SCREAMING_SNAKE_CASE  ):




	'''simple docstring'''





	snake_case_    =					"facebook/bart-large-mnli"
	snake_case_    =					(
	    "This is a tool that classifies an English text using provided labels. It takes two inputs: `text`, which "
	    "should be the text to classify, and `labels`, which should be the list of labels to use for classification. "
	    "It returns the most likely label in the list of provided `labels` for the input text."
	)
	snake_case_    =					"text_classifier"
	snake_case_    =					AutoTokenizer
	snake_case_    =					AutoModelForSequenceClassification

	snake_case_    =					["text", ["text"]]
	snake_case_    =					["text"]






	def     UpperCamelCase_  (	self :			str			):
					super().setup()
					__A             =			self.model.config
					__A             =			-1
					for idx, label in config.idalabel.items():
									if label.lower().startswith("entail"			):
													__A             =			int(A			)
					if self.entailment_id == -1:
									raise ValueError("Could not determine the entailment ID from the model config, please pass it at init."			)






	def     UpperCamelCase_  (	self :			int							,A :			Optional[Any]							,A :			Dict			):
					__A             =			labels
					return self.pre_processor(
					    [text] * len(A			)							,[f'''This example is {label}''' for label in labels]							,return_tensors="pt"							,padding="max_length"							,)







	def     UpperCamelCase_  (	self :			Union[str, Any]							,A :			Tuple			):
					__A             =			outputs.logits
					__A             =			torch.argmax(logits[:, 2]			).item()
					return self._labels[label_id]
