
"""simple docstring"""




import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import transformers
from transformers import (
    AutoConfig,
    AutoModelForQuestionAnswering,
    AutoTokenizer,
    DataCollatorWithPadding,
    HfArgumentParser,
    SquadDataset,
    Trainer,
    TrainingArguments,
)
from transformers import SquadDataTrainingArguments as DataTrainingArguments
from transformers.trainer_utils import is_main_process


snake_case    =		logging.getLogger(__name__)




@dataclass
class      UpperCamelCase	:

				"""simple docstring"""



				UpperCAmelCase_				:      str						  =				field(
				    metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}      )
				UpperCAmelCase_				:      Optional[str]						  =				field(
				    default=__magic_name__  ,      metadata={"help": "Pretrained config name or path if not the same as model_name"}      )
				UpperCAmelCase_				:      Optional[str]						  =				field(
				    default=__magic_name__  ,      metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}      )
				UpperCAmelCase_				:      bool						  =				field(default=__magic_name__  ,      metadata={"help": "Set this flag to use fast tokenization."}      )
				# If you want to tweak more attributes on your tokenizer, you should do it in a distinct script,
				# or just modify its tokenizer_config.json.
				UpperCAmelCase_				:      Optional[str]						  =				field(
				    default=__magic_name__  ,      metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"}  ,      )






def   UpperCamelCase_  (		):
					# See all possible arguments in src/transformers/training_args.py
					# or by passing the --help flag to this script.
					# We now keep distinct sets of args, for a cleaner separation of concerns.

					SCREAMING_SNAKE_CASE				=	HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)			)

					if len(sys.argv			) == 2 and sys.argv[1].endswith('.json'			):
										# If we pass only one argument to the script and it's the path to a json file,
										# let's parse it to get our arguments.
										SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE				=	parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]			)			)
					else:
										SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE				=	parser.parse_args_into_dataclasses()

					if (
					    os.path.exists(training_args.output_dir			)
					    and os.listdir(training_args.output_dir			)
					    and training_args.do_train
					    and not training_args.overwrite_output_dir
					):
										raise ValueError(
										    f'''Output directory ({training_args.output_dir}) already exists and is not empty. Use'''
										    ' --overwrite_output_dir to overcome.'			)

					# Setup logging
					logging.basicConfig(
					    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',     datefmt='%m/%d/%Y %H:%M:%S',     level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,     )
					logger.warning(
					    'Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s',     training_args.local_rank,     training_args.device,     training_args.n_gpu,     bool(training_args.local_rank != -1			),     training_args.fpaa,     )
					# Set the verbosity to info of the Transformers logger (on main process only):
					if is_main_process(training_args.local_rank			):
										transformers.utils.logging.set_verbosity_info()
										transformers.utils.logging.enable_default_handler()
										transformers.utils.logging.enable_explicit_format()
					logger.info('Training/evaluation parameters %s',     SCREAMING_SNAKE_CASE_			)

					# Prepare Question-Answering task
					# Load pretrained model and tokenizer
					#
					# Distributed training:
					# The .from_pretrained methods guarantee that only one local process can concurrently
					# download model & vocab.

					SCREAMING_SNAKE_CASE				=	AutoConfig.from_pretrained(
					    model_args.config_name if model_args.config_name else model_args.model_name_or_path,     cache_dir=model_args.cache_dir,     )
					SCREAMING_SNAKE_CASE				=	AutoTokenizer.from_pretrained(
					    model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,     cache_dir=model_args.cache_dir,     use_fast=SCREAMING_SNAKE_CASE_,     )
					SCREAMING_SNAKE_CASE				=	AutoModelForQuestionAnswering.from_pretrained(
					    model_args.model_name_or_path,     from_tf=bool('.ckpt' in model_args.model_name_or_path			),     config=SCREAMING_SNAKE_CASE_,     cache_dir=model_args.cache_dir,     )

					# Get datasets
					SCREAMING_SNAKE_CASE				=	hasattr(model.config,     'lang2id'			)
					SCREAMING_SNAKE_CASE				=	(
					    SquadDataset(
					        SCREAMING_SNAKE_CASE_,     tokenizer=SCREAMING_SNAKE_CASE_,     is_language_sensitive=SCREAMING_SNAKE_CASE_,     cache_dir=model_args.cache_dir			)
					    if training_args.do_train
					    else None
					)
					SCREAMING_SNAKE_CASE				=	(
					    SquadDataset(
					        SCREAMING_SNAKE_CASE_,     tokenizer=SCREAMING_SNAKE_CASE_,     mode='dev',     is_language_sensitive=SCREAMING_SNAKE_CASE_,     cache_dir=model_args.cache_dir,     )
					    if training_args.do_eval
					    else None
					)

					# Data collator
					SCREAMING_SNAKE_CASE				=	DataCollatorWithPadding(SCREAMING_SNAKE_CASE_,     pad_to_multiple_of=8			) if training_args.fpaa else None

					# Initialize our Trainer
					SCREAMING_SNAKE_CASE				=	Trainer(
					    model=SCREAMING_SNAKE_CASE_,     args=SCREAMING_SNAKE_CASE_,     train_dataset=SCREAMING_SNAKE_CASE_,     eval_dataset=SCREAMING_SNAKE_CASE_,     data_collator=SCREAMING_SNAKE_CASE_,     )

					# Training
					if training_args.do_train:
										trainer.train(
										    model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path			) else None			)
										trainer.save_model()
										# For convenience, we also re-save the tokenizer to the same directory,
										# so that you can share your model easily on huggingface.co/models =)
										if trainer.is_world_master():
															tokenizer.save_pretrained(training_args.output_dir			)

def   UpperCamelCase_  (		SCREAMING_SNAKE_CASE_			):
					# For xla_spawn (TPUs)
					main()


if __name__ == "__main__":
	main()


"""simple docstring"""




import unittest

from transformers.testing_utils import is_torch_available, require_torch, tooslow

from ...generation.test_utils import torch_device
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    CpmAntConfig,
	    CpmAntForCausalLM,
	    CpmAntModel,
	    CpmAntTokenizer,
	)




@require_torch
class      UpperCamelCase	:

				"""simple docstring"""



				def __init__(				self		,	lowercase__		,	lowercase__=2		,	lowercase__=8		,	lowercase__=True		,	lowercase__=False		,	lowercase__=False		,	lowercase__=False		,	lowercase__=False		,	lowercase__=99		,	lowercase__=32		,	lowercase__=3		,	lowercase__=4		,	lowercase__=37		,	lowercase__=32		,	lowercase__=128		,	lowercase__=8		,	lowercase__=8		,	lowercase__=8		,	lowercase__=1.0		,	lowercase__=True		,	)    ->    str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	parent
									SCREAMING_SNAKE_CASE				=	batch_size
									SCREAMING_SNAKE_CASE				=	seq_length
									SCREAMING_SNAKE_CASE				=	is_training
									SCREAMING_SNAKE_CASE				=	use_token_type_ids
									SCREAMING_SNAKE_CASE				=	use_input_mask
									SCREAMING_SNAKE_CASE				=	use_labels
									SCREAMING_SNAKE_CASE				=	use_mc_token_ids
									SCREAMING_SNAKE_CASE				=	vocab_size
									SCREAMING_SNAKE_CASE				=	hidden_size
									SCREAMING_SNAKE_CASE				=	num_hidden_layers
									SCREAMING_SNAKE_CASE				=	num_attention_heads
									SCREAMING_SNAKE_CASE				=	intermediate_size
									SCREAMING_SNAKE_CASE				=	num_buckets
									SCREAMING_SNAKE_CASE				=	max_distance
									SCREAMING_SNAKE_CASE				=	prompt_length
									SCREAMING_SNAKE_CASE				=	prompt_types
									SCREAMING_SNAKE_CASE				=	segment_types
									SCREAMING_SNAKE_CASE				=	init_std
									SCREAMING_SNAKE_CASE				=	return_dict


				def 							A					(				self			)    ->    Tuple:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	{}
									SCREAMING_SNAKE_CASE				=	ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size			).type(torch.intaa			)
									SCREAMING_SNAKE_CASE				=	False

									SCREAMING_SNAKE_CASE				=	self.get_config()

									return (config, input_ids)


				def 							A					(				self			)    ->    List[Any]:
									"""simple docstring"""







									return CpmAntConfig(
									    vocab_size=self.vocab_size		,	hidden_size=self.hidden_size		,	num_hidden_layers=self.num_hidden_layers		,	num_attention_heads=self.num_attention_heads		,	dim_ff=self.intermediate_size		,	position_bias_num_buckets=self.num_buckets		,	position_bias_max_distance=self.max_distance		,	prompt_types=self.prompt_types		,	prompt_length=self.prompt_length		,	segment_types=self.segment_types		,	use_cache=lowercase__		,	init_std=self.init_std		,	return_dict=self.return_dict		,	)


				def 							A					(				self		,	lowercase__		,	lowercase__		,	*lowercase__			)    ->    Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	CpmAntModel(config=lowercase__			)
									model.to(lowercase__			)
									model.eval()

									SCREAMING_SNAKE_CASE				=	model(**lowercase__			).last_hidden_state

									self.parent.assertEqual(hidden_states.shape		,	(self.batch_size, self.seq_length, config.hidden_size)			)


				def 							A					(				self		,	lowercase__		,	lowercase__		,	*lowercase__			)    ->    int:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	CpmAntForCausalLM(lowercase__			)
									model.to(lowercase__			)
									SCREAMING_SNAKE_CASE				=	input_ids['input_ids'].to(lowercase__			)
									model.eval()

									SCREAMING_SNAKE_CASE				=	model(**lowercase__			)
									self.parent.assertEqual(
									    model_output.logits.shape		,	(self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length)		,	)




				def 							A					(				self			)    ->    Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE				=	self.prepare_config_and_inputs()
									return config, inputs_dict






@require_torch
class      UpperCamelCase	(       __magic_name__  ,      __magic_name__  ,      unittest.TestCase      ):

				"""simple docstring"""



				UpperCAmelCase_				:      List[str]						  =				(CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
				UpperCAmelCase_				:      List[Any]						  =				(
				    {"feature-extraction": CpmAntModel, "text-generation": CpmAntForCausalLM} if is_torch_available() else {}
				)

				UpperCAmelCase_				:      List[Any]						  =				False
				UpperCAmelCase_				:      Optional[int]						  =				False
				UpperCAmelCase_				:      Optional[Any]						  =				False
				UpperCAmelCase_				:      List[str]						  =				False
				UpperCAmelCase_				:      Dict						  =				False


				def 							A					(				self			)    ->    Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	CpmAntModelTester(self			)
									SCREAMING_SNAKE_CASE				=	ConfigTester(self		,	config_class=lowercase__			)


				def 							A					(				self			)    ->    Tuple:
									"""simple docstring"""







									self.config_tester.create_and_test_config_common_properties()
									self.config_tester.create_and_test_config_to_json_string()
									self.config_tester.create_and_test_config_to_json_file()
									self.config_tester.create_and_test_config_from_and_save_pretrained()
									self.config_tester.check_config_can_be_init_without_params()
									self.config_tester.check_config_arguments_init()


				def 							A					(				self			)    ->    List[str]:
									"""simple docstring"""







									unittest.skip('CPMAnt doesn\'t support input_embeds.'			)(self.test_inputs_embeds			)


				def 							A					(				self			)    ->    Optional[int]:
									"""simple docstring"""







									unittest.skip(
									    'CPMAnt doesn\'t support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
                 So is attentions. We strongly recommand you use loss to tune model.'			)(self.test_retain_grad_hidden_states_attentions			)


				def 							A					(				self			)    ->    List[str]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE				=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_cpmant_model(lowercase__		,	lowercase__			)




				def 							A					(				self			)    ->    str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE				=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_lm_head_model(lowercase__		,	lowercase__			)






@require_torch
class      UpperCamelCase	(       unittest.TestCase      ):

				"""simple docstring"""



				@tooslow
				def 							A					(				self			)    ->    Any:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	'今天天气真好！'
									SCREAMING_SNAKE_CASE				=	'openbmb/cpm-ant-10b'
									SCREAMING_SNAKE_CASE				=	CpmAntModel.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	CpmAntTokenizer.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	tokenizer(lowercase__		,	return_tensors='pt'			)
									SCREAMING_SNAKE_CASE				=	model(**lowercase__			).last_hidden_state

									SCREAMING_SNAKE_CASE				=	torch.tensor(
									    [[[6.1_708, 5.9_244, 1.0_835], [6.5_207, 6.2_893, -11.3_324], [-1.0_107, -0.0_576, -5.9_577]]]		,	)
									self.assertTrue(torch.allclose(hidden_states[:, :3, :3]		,	lowercase__		,	atol=1E-2			)			)






@require_torch
class      UpperCamelCase	(       unittest.TestCase      ):

				"""simple docstring"""



				@tooslow
				def 							A					(				self			)    ->    str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	'今天天气真好！'
									SCREAMING_SNAKE_CASE				=	'openbmb/cpm-ant-10b'
									SCREAMING_SNAKE_CASE				=	CpmAntForCausalLM.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	CpmAntTokenizer.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	tokenizer(lowercase__		,	return_tensors='pt'			)
									SCREAMING_SNAKE_CASE				=	model(**lowercase__			).logits

									SCREAMING_SNAKE_CASE				=	torch.tensor(
									    [[[-6.4_267, -6.4_083, -6.3_958], [-5.8_802, -5.9_447, -5.7_811], [-5.3_896, -5.4_820, -5.4_295]]]		,	)
									self.assertTrue(torch.allclose(hidden_states[:, :3, :3]		,	lowercase__		,	atol=1E-2			)			)


				@tooslow
				def 							A					(				self			)    ->    Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	'openbmb/cpm-ant-10b'
									SCREAMING_SNAKE_CASE				=	CpmAntForCausalLM.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	CpmAntTokenizer.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	'今天天气不错，'
									SCREAMING_SNAKE_CASE				=	'今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的'
									SCREAMING_SNAKE_CASE				=	tokenizer(lowercase__		,	return_tensors='pt'			)
									SCREAMING_SNAKE_CASE				=	model.generate(**lowercase__			)
									SCREAMING_SNAKE_CASE				=	tokenizer.batch_decode(lowercase__			)
									self.assertEqual(lowercase__		,	lowercase__			)




				@tooslow
				def 							A					(				self			)    ->    str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	'openbmb/cpm-ant-10b'
									SCREAMING_SNAKE_CASE				=	CpmAntForCausalLM.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	CpmAntTokenizer.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	['今天天气不错，', '新年快乐，万事如意！']
									SCREAMING_SNAKE_CASE				=	[
									    '今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的',
									    '新年快乐，万事如意！在这辞旧迎新的美好时刻，我谨代表《农村新技术》杂志社全体同仁，向一直以来关心、支持《农村新技术》杂志发展的各级领导、各界朋友和广大读者致以最诚挚的',
									]
									SCREAMING_SNAKE_CASE				=	tokenizer(lowercase__		,	return_tensors='pt'		,	padding=lowercase__			)
									SCREAMING_SNAKE_CASE				=	model.generate(**lowercase__			)
									SCREAMING_SNAKE_CASE				=	tokenizer.batch_decode(lowercase__			)
									self.assertEqual(lowercase__		,	lowercase__			)
