



import os
import tempfile
import unittest

from transformers import NezhaConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    MODEL_FOR_PRETRAINING_MAPPING,
	    NezhaForMaskedLM,
	    NezhaForMultipleChoice,
	    NezhaForNextSentencePrediction,
	    NezhaForPreTraining,
	    NezhaForQuestionAnswering,
	    NezhaForSequenceClassification,
	    NezhaForTokenClassification,
	    NezhaModel,
	)
	from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST

class  __snake_case      :







							'''simple docstring'''






							def __init__(							self  ,  a_  ,  a_=13  ,  a_=7  ,  a_=True  ,  a_=True  ,  a_=True  ,  a_=True  ,  a_=99  ,  a_=32  ,  a_=5  ,  a_=4  ,  a_=37  ,  a_="gelu"  ,  a_=0.1  ,  a_=0.1  ,  a_=128  ,  a_=32  ,  a_=16  ,  a_=2  ,  a_=0.02  ,  a_=3  ,  a_=4  ,  a_=None  ,  ):
											a__				      =		parent
											a__				      =		batch_size
											a__				      =		seq_length
											a__				      =		is_training
											a__				      =		use_input_mask
											a__				      =		use_token_type_ids
											a__				      =		use_labels
											a__				      =		vocab_size
											a__				      =		hidden_size
											a__				      =		num_hidden_layers
											a__				      =		num_attention_heads
											a__				      =		intermediate_size
											a__				      =		hidden_act
											a__				      =		hidden_dropout_prob
											a__				      =		attention_probs_dropout_prob
											a__				      =		max_position_embeddings
											a__				      =		type_vocab_size
											a__				      =		type_sequence_label_size
											a__				      =		initializer_range
											a__				      =		num_labels
											a__				      =		num_choices
											a__				      =		scope


							def _a      (							self		):
											a__				      =		ids_tensor([self.batch_size, self.seq_length]  ,  self.vocab_size		)

											a__				      =		None
											if self.use_input_mask:
															a__				      =		random_attention_mask([self.batch_size, self.seq_length]		)

											a__				      =		None
											if self.use_token_type_ids:
															a__				      =		ids_tensor([self.batch_size, self.seq_length]  ,  self.type_vocab_size		)

											a__				      =		None
											a__				      =		None
											a__				      =		None
											if self.use_labels:
															a__				      =		ids_tensor([self.batch_size]  ,  self.type_sequence_label_size		)
															a__				      =		ids_tensor([self.batch_size, self.seq_length]  ,  self.num_labels		)
															a__				      =		ids_tensor([self.batch_size]  ,  self.num_choices		)

											a__				      =		self.get_config()

											return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


							def _a      (							self		):
											return NezhaConfig(
											    vocab_size=self.vocab_size  ,  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  ,  max_position_embeddings=self.max_position_embeddings  ,  type_vocab_size=self.type_vocab_size  ,  is_decoder=a_  ,  initializer_range=self.initializer_range  ,  )


							def _a      (							self		):
											(
											    (
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			
											)				      =		self.prepare_config_and_inputs()

											a__				      =		True
											a__				      =		floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
											a__				      =		ids_tensor([self.batch_size, self.seq_length]  ,  vocab_size=2		)

											return (
											    config,
											    input_ids,
											    token_type_ids,
											    input_mask,
											    sequence_labels,
											    token_labels,
											    choice_labels,
											    encoder_hidden_states,
											    encoder_attention_mask,
											)


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		NezhaModel(config=a_		)
											model.to(a_		)
											model.eval()
											a__				      =		model(a_  ,  attention_mask=a_  ,  token_type_ids=a_		)
											a__				      =		model(a_  ,  token_type_ids=a_		)
											a__				      =		model(a_		)
											self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)		)
											self.parent.assertEqual(result.pooler_output.shape  ,  (self.batch_size, self.hidden_size)		)


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  ):
											a__				      =		True
											a__				      =		NezhaModel(a_		)
											model.to(a_		)
											model.eval()
											a__				      =		model(
											    a_  ,  attention_mask=a_  ,  token_type_ids=a_  ,  encoder_hidden_states=a_  ,  encoder_attention_mask=a_  ,  )
											a__				      =		model(
											    a_  ,  attention_mask=a_  ,  token_type_ids=a_  ,  encoder_hidden_states=a_  ,  )
											a__				      =		model(a_  ,  attention_mask=a_  ,  token_type_ids=a_		)
											self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)		)
											self.parent.assertEqual(result.pooler_output.shape  ,  (self.batch_size, self.hidden_size)		)


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		NezhaForMaskedLM(config=a_		)
											model.to(a_		)
											model.eval()
											a__				      =		model(a_  ,  attention_mask=a_  ,  token_type_ids=a_  ,  labels=a_		)
											self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.vocab_size)		)


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		NezhaForNextSentencePrediction(config=a_		)
											model.to(a_		)
											model.eval()
											a__				      =		model(
											    a_  ,  attention_mask=a_  ,  token_type_ids=a_  ,  labels=a_  ,  )
											self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, 2)		)


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		NezhaForPreTraining(config=a_		)
											model.to(a_		)
											model.eval()
											a__				      =		model(
											    a_  ,  attention_mask=a_  ,  token_type_ids=a_  ,  labels=a_  ,  next_sentence_label=a_  ,  )
											self.parent.assertEqual(result.prediction_logits.shape  ,  (self.batch_size, self.seq_length, self.vocab_size)		)
											self.parent.assertEqual(result.seq_relationship_logits.shape  ,  (self.batch_size, 2)		)


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		NezhaForQuestionAnswering(config=a_		)
											model.to(a_		)
											model.eval()
											a__				      =		model(
											    a_  ,  attention_mask=a_  ,  token_type_ids=a_  ,  start_positions=a_  ,  end_positions=a_  ,  )
											self.parent.assertEqual(result.start_logits.shape  ,  (self.batch_size, self.seq_length)		)
											self.parent.assertEqual(result.end_logits.shape  ,  (self.batch_size, self.seq_length)		)


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		self.num_labels
											a__				      =		NezhaForSequenceClassification(a_		)
											model.to(a_		)
											model.eval()
											a__				      =		model(a_  ,  attention_mask=a_  ,  token_type_ids=a_  ,  labels=a_		)
											self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)		)


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		self.num_labels
											a__				      =		NezhaForTokenClassification(config=a_		)
											model.to(a_		)
											model.eval()
											a__				      =		model(a_  ,  attention_mask=a_  ,  token_type_ids=a_  ,  labels=a_		)
											self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.num_labels)		)


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		self.num_choices
											a__				      =		NezhaForMultipleChoice(config=a_		)
											model.to(a_		)
											model.eval()
											a__				      =		input_ids.unsqueeze(1		).expand(-1  ,  self.num_choices  ,  -1		).contiguous()
											a__				      =		token_type_ids.unsqueeze(1		).expand(-1  ,  self.num_choices  ,  -1		).contiguous()
											a__				      =		input_mask.unsqueeze(1		).expand(-1  ,  self.num_choices  ,  -1		).contiguous()
											a__				      =		model(
											    a_  ,  attention_mask=a_  ,  token_type_ids=a_  ,  labels=a_  ,  )
											self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_choices)		)







							def _a      (							self		):
											a__				      =		self.prepare_config_and_inputs()
											(
											    (
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			
											)				      =		config_and_inputs
											a__				      =		{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
											return config, inputs_dict




@require_torch
class  __snake_case      (   SCREAMING_SNAKE_CASE			,SCREAMING_SNAKE_CASE			,SCREAMING_SNAKE_CASE			,unittest.TestCase):







							'''simple docstring'''






							UpperCamelCase__				:      List[Any]         =					(
							    (
							        NezhaModel,
							        NezhaForMaskedLM,
							        NezhaForMultipleChoice,
							        NezhaForNextSentencePrediction,
							        NezhaForPreTraining,
							        NezhaForQuestionAnswering,
							        NezhaForSequenceClassification,
							        NezhaForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							UpperCamelCase__				:      Optional[int]         =					(
							    {
							        """feature-extraction""": NezhaModel,
							        """fill-mask""": NezhaForMaskedLM,
							        """question-answering""": NezhaForQuestionAnswering,
							        """text-classification""": NezhaForSequenceClassification,
							        """token-classification""": NezhaForTokenClassification,
							        """zero-shot""": NezhaForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							UpperCamelCase__				:      Optional[Any]         =					True


							def _a      (							self  ,  a_  ,  a_  ,  a_=False		):
											a__				      =		super()._prepare_for_class(a_  ,  a_  ,  return_labels=a_		)

											if return_labels:
															if model_class in get_values(a_		):
																			a__				      =		torch.zeros(
																			    (self.model_tester.batch_size, self.model_tester.seq_length)  ,  dtype=torch.long  ,  device=a_		)
																			a__				      =		torch.zeros(
																			    self.model_tester.batch_size  ,  dtype=torch.long  ,  device=a_		)
											return inputs_dict


							def _a      (							self		):
											a__				      =		NezhaModelTester(self		)
											a__				      =		ConfigTester(self  ,  config_class=a_  ,  hidden_size=37		)


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


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*a_		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*a_		)


							def _a      (							self		):
											# This regression test was failing with PyTorch < 1.3
											(
											    (
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			(
											    a__
											)      ,			
											)				      =		self.model_tester.prepare_config_and_inputs_for_decoder()

											a__				      =		None

											self.model_tester.create_and_check_model_as_decoder(
											    a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  )


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*a_		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_multiple_choice(*a_		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_next_sequence_prediction(*a_		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_pretraining(*a_		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_question_answering(*a_		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_sequence_classification(*a_		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*a_		)


							@slow
							def _a      (							self		):
											for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															a__				      =		NezhaModel.from_pretrained(a_		)
															self.assertIsNotNone(a_		)







							@slow
							@require_torch_gpu
							def _a      (							self		):
											a__      ,			a__				      =		self.model_tester.prepare_config_and_inputs_for_common()
											for model_class in self.all_model_classes:
															# NezhaForMultipleChoice behaves incorrectly in JIT environments.
															if model_class == NezhaForMultipleChoice:
																			return

															a__				      =		True
															a__				      =		model_class(config=a_		)

															a__				      =		self._prepare_for_class(a_  ,  a_		)
															a__				      =		torch.jit.trace(
															    a_  ,  (inputs_dict["""input_ids"""].to("""cpu"""		), inputs_dict["""attention_mask"""].to("""cpu"""		))		)

															with tempfile.TemporaryDirectory() as tmp:
																			torch.jit.save(a_  ,  os.path.join(a_  ,  """bert.pt"""		)		)
																			a__				      =		torch.jit.load(os.path.join(a_  ,  """bert.pt"""		)  ,  map_location=a_		)
																			loaded(inputs_dict["""input_ids"""].to(a_		)  ,  inputs_dict["""attention_mask"""].to(a_		)		)




@require_torch
class  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							@slow
							def _a      (							self		):
											a__				      =		NezhaModel.from_pretrained("""sijunhe/nezha-cn-base"""		)
											a__				      =		torch.tensor([[0, 1, 2, 3, 4, 5]]		)
											a__				      =		torch.tensor([[0, 1, 1, 1, 1, 1]]		)
											with torch.no_grad():
															a__				      =		model(a_  ,  attention_mask=a_		)[0]
											a__				      =		torch.Size((1, 6, 768)		)
											self.assertEqual(output.shape  ,  a_		)
											a__				      =		torch.tensor([[[0.0_685, 0.2_441, 0.1_102], [0.0_600, 0.1_906, 0.1_349], [0.0_221, 0.0_819, 0.0_586]]]		)

											self.assertTrue(torch.allclose(output[:, 1:4, 1:4]  ,  a_  ,  atol=1E-4		)		)







							@slow
							def _a      (							self		):
											a__				      =		NezhaForMaskedLM.from_pretrained("""sijunhe/nezha-cn-base"""		)
											a__				      =		torch.tensor([[0, 1, 2, 3, 4, 5]]		)
											a__				      =		torch.tensor([[1, 1, 1, 1, 1, 1]]		)
											with torch.no_grad():
															a__				      =		model(a_  ,  attention_mask=a_		)[0]
											a__				      =		torch.Size((1, 6, 21_128)		)
											self.assertEqual(output.shape  ,  a_		)
											a__				      =		torch.tensor(
											    [[-2.7_939, -1.7_902, -2.2_189], [-2.8_585, -1.8_908, -2.3_723], [-2.6_499, -1.7_750, -2.2_558]]		)

											self.assertTrue(torch.allclose(output[:, 1:4, 1:4]  ,  a_  ,  atol=1E-4		)		)






import os
from argparse import ArgumentParser, Namespace

from ..data import SingleSentenceClassificationProcessor as Processor
from ..pipelines import TextClassificationPipeline
from ..utils import is_tf_available, is_torch_available, logging
from . import BaseTransformersCLICommand


if not is_tf_available() and not is_torch_available():
	raise RuntimeError("""At least one of PyTorch or TensorFlow 2.0+ should be installed to use CLI training""")

# TF training parameters
UpperCAmelCase							 =					False
UpperCAmelCase							 =					False

def 			A_      (	__a     :				Namespace		):

				"""simple docstring"""
				return TrainCommand(__a		)




class  __snake_case      (   SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							@staticmethod
							def _a      (							a_		):
											a__				      =		parser.add_parser("""train"""  ,  help="""CLI tool to train a model on a task."""		)

											train_parser.add_argument(
											    """--train_data"""  ,  type=a_  ,  required=a_  ,  help="""path to train (and optionally evaluation) dataset as a csv with tab separated labels and sentences."""  ,  )
											train_parser.add_argument(
											    """--column_label"""  ,  type=a_  ,  default=0  ,  help="""Column of the dataset csv file with example labels."""		)
											train_parser.add_argument(
											    """--column_text"""  ,  type=a_  ,  default=1  ,  help="""Column of the dataset csv file with example texts."""		)
											train_parser.add_argument(
											    """--column_id"""  ,  type=a_  ,  default=2  ,  help="""Column of the dataset csv file with example ids."""		)
											train_parser.add_argument(
											    """--skip_first_row"""  ,  action="""store_true"""  ,  help="""Skip the first row of the csv file (headers)."""		)

											train_parser.add_argument("""--validation_data"""  ,  type=a_  ,  default=""""""  ,  help="""path to validation dataset."""		)
											train_parser.add_argument(
											    """--validation_split"""  ,  type=a_  ,  default=0.1  ,  help="""if validation dataset is not provided, fraction of train dataset to use as validation dataset."""  ,  )

											train_parser.add_argument("""--output"""  ,  type=a_  ,  default="""./"""  ,  help="""path to saved the trained model."""		)

											train_parser.add_argument(
											    """--task"""  ,  type=a_  ,  default="""text_classification"""  ,  help="""Task to train the model on."""		)
											train_parser.add_argument(
											    """--model"""  ,  type=a_  ,  default="""bert-base-uncased"""  ,  help="""Model's name or path to stored model."""		)
											train_parser.add_argument("""--train_batch_size"""  ,  type=a_  ,  default=32  ,  help="""Batch size for training."""		)
											train_parser.add_argument("""--valid_batch_size"""  ,  type=a_  ,  default=64  ,  help="""Batch size for validation."""		)
											train_parser.add_argument("""--learning_rate"""  ,  type=a_  ,  default=3E-5  ,  help="""Learning rate."""		)
											train_parser.add_argument("""--adam_epsilon"""  ,  type=a_  ,  default=1E-0_8  ,  help="""Epsilon for Adam optimizer."""		)
											train_parser.set_defaults(func=a_		)


							def __init__(							self  ,  a_		):
											a__				      =		logging.get_logger("""transformers-cli/training"""		)

											a__				      =		"""tf""" if is_tf_available() else """torch"""

											os.makedirs(args.output  ,  exist_ok=a_		)
											a__				      =		args.output

											a__				      =		args.column_label
											a__				      =		args.column_text
											a__				      =		args.column_id

											self.logger.info(F'''Loading {args.task} pipeline for {args.model}'''		)
											if args.task == "text_classification":
															a__				      =		TextClassificationPipeline.from_pretrained(args.model		)
											elif args.task == "token_classification":
															raise NotImplementedError
											elif args.task == "question_answering":
															raise NotImplementedError

											self.logger.info(F'''Loading dataset from {args.train_data}'''		)
											a__				      =		Processor.create_from_csv(
											    args.train_data  ,  column_label=args.column_label  ,  column_text=args.column_text  ,  column_id=args.column_id  ,  skip_first_row=args.skip_first_row  ,  )
											a__				      =		None
											if args.validation_data:
															self.logger.info(F'''Loading validation dataset from {args.validation_data}'''		)
															a__				      =		Processor.create_from_csv(
															    args.validation_data  ,  column_label=args.column_label  ,  column_text=args.column_text  ,  column_id=args.column_id  ,  skip_first_row=args.skip_first_row  ,  )

											a__				      =		args.validation_split
											a__				      =		args.train_batch_size
											a__				      =		args.valid_batch_size
											a__				      =		args.learning_rate
											a__				      =		args.adam_epsilon


							def _a      (							self		):
											if self.framework == "tf":
															return self.run_tf()
											return self.run_torch()


							def _a      (							self		):
											raise NotImplementedError







							def _a      (							self		):
											self.pipeline.fit(
											    self.train_dataset  ,  validation_data=self.valid_dataset  ,  validation_split=self.validation_split  ,  learning_rate=self.learning_rate  ,  adam_epsilon=self.adam_epsilon  ,  train_batch_size=self.train_batch_size  ,  valid_batch_size=self.valid_batch_size  ,  )

											# Save trained pipeline
											self.pipeline.save_pretrained(self.output		)

