





'''simple docstring'''






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

from seqaseq_trainer import SeqaSeqTrainer
from seqaseq_training_args import SeqaSeqTrainingArguments

import transformers
from transformers import (
    AutoConfig,
    AutoModelForSeqaSeqLM,
    AutoTokenizer,
    HfArgumentParser,
    MBartTokenizer,
    MBartTokenizerFast,
    set_seed,
)
from transformers.trainer_utils import EvaluationStrategy, is_main_process
from transformers.training_args import ParallelMode
from utils import (
    SeqaSeqDataCollator,
    SeqaSeqDataset,
    assert_all_frozen,
    build_compute_metrics_fn,
    check_output_dir,
    freeze_embeds,
    freeze_params,
    lmap,
    save_json,
    use_task_specific_params,
    write_txt_file,
)


A   : int        =			logging.getLogger(__name__)







@dataclass
class 			lowerCamelCase					:
					_SCREAMING_SNAKE_CASE							   =    field(
					    metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}  )
					_SCREAMING_SNAKE_CASE							   =    field(
					    default=UpperCAmelCase__     ,							metadata={"help": "Pretrained config name or path if not the same as model_name"}  )
					_SCREAMING_SNAKE_CASE							   =    field(
					    default=UpperCAmelCase__     ,							metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}  )
					_SCREAMING_SNAKE_CASE							   =    field(
					    default=UpperCAmelCase__     ,							metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"}     ,							)
					_SCREAMING_SNAKE_CASE							   =    field(default=UpperCAmelCase__     ,							metadata={"help": "Whether tp freeze the encoder."}  )
					_SCREAMING_SNAKE_CASE							   =    field(default=UpperCAmelCase__     ,							metadata={"help": "Whether  to freeze the embeddings."}  )







@dataclass
class 			lowerCamelCase					:
					_SCREAMING_SNAKE_CASE							   =    field(
					    metadata={"help": "The input data dir. Should contain the .tsv files (or other data files) for the task."}  )
					_SCREAMING_SNAKE_CASE							   =    field(
					    default="summarization"     ,							metadata={"help": "Task name, summarization (or summarization_{dataset} for pegasus) or translation"}     ,							)
					_SCREAMING_SNAKE_CASE							   =    field(
					    default=1024     ,							metadata={
					        "help": (
					            "The maximum total input sequence length after tokenization. Sequences longer "
					            "than this will be truncated, sequences shorter will be padded."
					        )
					    }     ,							)
					_SCREAMING_SNAKE_CASE							   =    field(
					    default=128     ,							metadata={
					        "help": (
					            "The maximum total sequence length for target text after tokenization. Sequences longer "
					            "than this will be truncated, sequences shorter will be padded."
					        )
					    }     ,							)
					_SCREAMING_SNAKE_CASE							   =    field(
					    default=142     ,							metadata={
					        "help": (
					            "The maximum total sequence length for validation target text after tokenization. Sequences longer "
					            "than this will be truncated, sequences shorter will be padded. "
					            "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
					            "during ``evaluate`` and ``predict``."
					        )
					    }     ,							)
					_SCREAMING_SNAKE_CASE							   =    field(
					    default=142     ,							metadata={
					        "help": (
					            "The maximum total sequence length for test target text after tokenization. Sequences longer "
					            "than this will be truncated, sequences shorter will be padded."
					        )
					    }     ,							)
					_SCREAMING_SNAKE_CASE							   =    field(default=-1     ,							metadata={"help": "# training examples. -1 means use all."}  )
					_SCREAMING_SNAKE_CASE							   =    field(default=-1     ,							metadata={"help": "# validation examples. -1 means use all."}  )
					_SCREAMING_SNAKE_CASE							   =    field(default=-1     ,							metadata={"help": "# test examples. -1 means use all."}  )
					_SCREAMING_SNAKE_CASE							   =    field(default=UpperCAmelCase__     ,							metadata={"help": "Source language id for translation."}  )
					_SCREAMING_SNAKE_CASE							   =    field(default=UpperCAmelCase__     ,							metadata={"help": "Target language id for translation."}  )
					_SCREAMING_SNAKE_CASE							   =    field(default=UpperCAmelCase__     ,							metadata={"help": "# num_beams to use for evaluation."}  )
					_SCREAMING_SNAKE_CASE							   =    field(
					    default=UpperCAmelCase__     ,							metadata={"help": "If only pad tokens should be ignored. This assumes that `config.pad_token_id` is defined."}     ,							)





def       lowercase_      (		lowercase__				,  lowercase__				,  lowercase__						)   ->Tuple:
						logger.info(F'''***** {split} metrics *****'''						)
						for key in sorted(metrics.keys()						):
												logger.info(F'''  {key} = {metrics[key]}'''						)
						save_json(UpperCAmelCase__				,  os.path.join(UpperCAmelCase__				,  F'''{split}_results.json'''						)						)





def       lowercase_      (		)   ->List[Any]:
						_snake_case:   Tuple  					=		HfArgumentParser((ModelArguments, DataTrainingArguments, SeqaSeqTrainingArguments)						)

						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.
												_snake_case					,						_snake_case					,						_snake_case:   List[Any]  					=		parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]						)						)
						else:
												_snake_case					,						_snake_case					,						_snake_case:   Union[str, Any]  					=		parser.parse_args_into_dataclasses()

						check_output_dir(UpperCAmelCase__						)

						# 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.parallel_mode == ParallelMode.DISTRIBUTED						)				,  training_args.fpaa				,  )
						transformers.utils.logging.enable_default_handler()
						transformers.utils.logging.enable_explicit_format()
						# 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()
						logger.info('Training/evaluation parameters %s'				,  UpperCAmelCase__						)

						# Set seed
						set_seed(training_args.seed						)

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

						_snake_case:   List[str]  					=		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				,  )

						_snake_case:   Dict  					=		('encoder_layerdrop', 'decoder_layerdrop', 'dropout', 'attention_dropout')
						for p in extra_model_params:
												if getattr(UpperCAmelCase__				,  UpperCAmelCase__				,  UpperCAmelCase__						):
																		assert hasattr(UpperCAmelCase__				,  UpperCAmelCase__						), F'''({config.__class__.__name__}) doesn\'t have a `{p}` attribute'''
																		setattr(UpperCAmelCase__				,  UpperCAmelCase__				,  getattr(UpperCAmelCase__				,  UpperCAmelCase__						)						)

						_snake_case:   Dict  					=		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				,  )
						_snake_case:   List[str]  					=		AutoModelForSeqaSeqLM.from_pretrained(
						    model_args.model_name_or_path				,  from_tf='.ckpt' in model_args.model_name_or_path				,  config=UpperCAmelCase__				,  cache_dir=model_args.cache_dir				,  )

						# use task specific params
						use_task_specific_params(UpperCAmelCase__				,  data_args.task						)

						# set num_beams for evaluation
						if data_args.eval_beams is None:
												_snake_case:   Union[str, Any]  					=		model.config.num_beams

						# set decoder_start_token_id for MBart
						if model.config.decoder_start_token_id is None and isinstance(UpperCAmelCase__				,  (MBartTokenizer, MBartTokenizerFast)						):
												assert (
												    data_args.tgt_lang is not None and data_args.src_lang is not None
												), "mBart requires --tgt_lang and --src_lang"
												if isinstance(UpperCAmelCase__				,  UpperCAmelCase__						):
																		_snake_case:   Optional[Any]  					=		tokenizer.lang_code_to_id[data_args.tgt_lang]
												else:
																		_snake_case:   Union[str, Any]  					=		tokenizer.convert_tokens_to_ids(data_args.tgt_lang						)

						if model_args.freeze_embeds:
												freeze_embeds(UpperCAmelCase__						)
						if model_args.freeze_encoder:
												freeze_params(model.get_encoder()						)
												assert_all_frozen(model.get_encoder()						)

						_snake_case:   Dict  					=		SeqaSeqDataset

						# Get datasets
						_snake_case:   List[Any]  					=		(
						    dataset_class(
						        UpperCAmelCase__				,  type_path='train'				,  data_dir=data_args.data_dir				,  n_obs=data_args.n_train				,  max_target_length=data_args.max_target_length				,  max_source_length=data_args.max_source_length				,  prefix=model.config.prefix or ''				,  )
						    if training_args.do_train
						    else None
						)
						_snake_case:   Optional[int]  					=		(
						    dataset_class(
						        UpperCAmelCase__				,  type_path='val'				,  data_dir=data_args.data_dir				,  n_obs=data_args.n_val				,  max_target_length=data_args.val_max_target_length				,  max_source_length=data_args.max_source_length				,  prefix=model.config.prefix or ''				,  )
						    if training_args.do_eval or training_args.evaluation_strategy != EvaluationStrategy.NO
						    else None
						)
						_snake_case:   List[str]  					=		(
						    dataset_class(
						        UpperCAmelCase__				,  type_path='test'				,  data_dir=data_args.data_dir				,  n_obs=data_args.n_test				,  max_target_length=data_args.test_max_target_length				,  max_source_length=data_args.max_source_length				,  prefix=model.config.prefix or ''				,  )
						    if training_args.do_predict
						    else None
						)

						# Initialize our Trainer
						_snake_case:   Optional[Any]  					=		(
						    build_compute_metrics_fn(data_args.task				,  UpperCAmelCase__						) if training_args.predict_with_generate else None
						)
						_snake_case:   Optional[Any]  					=		SeqaSeqTrainer(
						    model=UpperCAmelCase__				,  args=UpperCAmelCase__				,  data_args=UpperCAmelCase__				,  train_dataset=UpperCAmelCase__				,  eval_dataset=UpperCAmelCase__				,  data_collator=SeqaSeqDataCollator(
						        UpperCAmelCase__				,  UpperCAmelCase__				,  model.config.decoder_start_token_id				,  training_args.tpu_num_cores						)				,  compute_metrics=UpperCAmelCase__				,  tokenizer=UpperCAmelCase__				,  )

						_snake_case:   List[Any]  					=		{}
						# Training
						if training_args.do_train:
												logger.info('*** Train ***'						)

												_snake_case:   int  					=		trainer.train(
												    model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path						) else None						)
												_snake_case:   Any  					=		train_result.metrics
												_snake_case:   Optional[Any]  					=		data_args.n_train

												trainer.save_model()  # this also saves the tokenizer

												if trainer.is_world_process_zero():
																		handle_metrics('train'				,  UpperCAmelCase__				,  training_args.output_dir						)
																		all_metrics.update(UpperCAmelCase__						)

																		# Need to save the state, since Trainer.save_model saves only the tokenizer with the model
																		trainer.state.save_to_json(os.path.join(training_args.output_dir				,  'trainer_state.json'						)						)

																		# For convenience, we also re-save the tokenizer to the same directory,
																		# so that you can share your model easily on huggingface.co/models =)
																		tokenizer.save_pretrained(training_args.output_dir						)

																		 # Evaluation
						if training_args.do_eval:
												logger.info('*** Evaluate ***'						)

												_snake_case:   Optional[int]  					=		trainer.evaluate(metric_key_prefix='val'						)
												_snake_case:   int  					=		data_args.n_val
												_snake_case:   List[Any]  					=		round(metrics['val_loss']				,  4						)

												if trainer.is_world_process_zero():
																		handle_metrics('val'				,  UpperCAmelCase__				,  training_args.output_dir						)
																		all_metrics.update(UpperCAmelCase__						)

						if training_args.do_predict:
												logger.info('*** Predict ***'						)

												_snake_case:   Optional[Any]  					=		trainer.predict(test_dataset=UpperCAmelCase__				,  metric_key_prefix='test'						)
												_snake_case:   int  					=		test_output.metrics
												_snake_case:   Optional[int]  					=		data_args.n_test

												if trainer.is_world_process_zero():
																		_snake_case:   List[Any]  					=		round(metrics['test_loss']				,  4						)
																		handle_metrics('test'				,  UpperCAmelCase__				,  training_args.output_dir						)
																		all_metrics.update(UpperCAmelCase__						)

																		if training_args.predict_with_generate:
																								_snake_case:   int  					=		tokenizer.batch_decode(
																								    test_output.predictions				,  skip_special_tokens=UpperCAmelCase__				,  clean_up_tokenization_spaces=UpperCAmelCase__						)
																								_snake_case:   Optional[Any]  					=		lmap(str.strip				,  UpperCAmelCase__						)
																								write_txt_file(UpperCAmelCase__				,  os.path.join(training_args.output_dir				,  'test_generations.txt'						)						)

						if trainer.is_world_process_zero():
												save_json(UpperCAmelCase__				,  os.path.join(training_args.output_dir				,  'all_results.json'						)						)

						return all_metrics





def       lowercase_      (		lowercase__						)   ->Optional[Any]:
						main()


if __name__ == "__main__":
		main()












'''simple docstring'''






import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from transformers import AutoConfig, TFAutoModel, is_tensorflow_text_available, is_tf_available
from transformers.models.bert.tokenization_bert import BertTokenizer
from transformers.testing_utils import require_tensorflow_text, require_tf, slow


if is_tf_available():
		import tensorflow as tf

if is_tensorflow_text_available():
		from transformers.models.bert import TFBertTokenizer


A   : Tuple             =			['bert-base-uncased', 'bert-base-cased']
A   : str             =			'hf-internal-testing/tiny-bert-tf-only'

if is_tf_available():







		class 			lowerCamelCase					(						tf.keras.Model  ):


							def __init__(					self	:  Union[str, Any]       ,     __snake_case	:  List[Any]       ):





													'''simple docstring'''





													super().__init__()
													_snake_case:   List[Any]  					=		tokenizer
													_snake_case:   str  					=		AutoConfig.from_pretrained(__snake_case       )
													_snake_case:   List[Any]  					=		TFAutoModel.from_config(__snake_case       )




							def    SCREAMING_SNAKE_CASE_		(					self	:  str       ,     __snake_case	:  List[str]       ):





													'''simple docstring'''





													_snake_case:   Optional[int]  					=		self.tokenizer(__snake_case       )
													_snake_case:   Tuple  					=		self.bert(**__snake_case       )
													return out["pooler_output"]







@require_tf
@require_tensorflow_text
class 			lowerCamelCase					(						unittest.TestCase  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											super().setUp()

											_snake_case:   List[Any]  					=		[
											    BertTokenizer.from_pretrained(__snake_case       ) for checkpoint in (TOKENIZER_CHECKPOINTS * 2)
											]  # repeat for when fast_bert_tokenizer=false
											_snake_case:   List[Any]  					=		[TFBertTokenizer.from_pretrained(__snake_case       ) for checkpoint in TOKENIZER_CHECKPOINTS] + [
											    TFBertTokenizer.from_pretrained(__snake_case       ,     use_fast_bert_tokenizer=__snake_case       )
											    for checkpoint in TOKENIZER_CHECKPOINTS
											]
											assert len(self.tokenizers       ) == len(self.tf_tokenizers       )

											_snake_case:   int  					=		[
											    'This is a straightforward English test sentence.',
											    'This one has some weird characters\rto\nsee\r\nif  those\u00E9break things.',
											    'Now we\'re going to add some Chinese: 一 二 三 一二三',
											    'And some much more rare Chinese: 齉 堃 齉堃',
											    'Je vais aussi écrire en français pour tester les accents',
											    'Classical Irish also has some unusual characters, so in they go: Gaelaċ, ꝼ',
											]
											_snake_case:   str  					=		list(zip(self.test_sentences       ,     self.test_sentences[::-1]       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





											for tokenizer, tf_tokenizer in zip(self.tokenizers       ,     self.tf_tokenizers       ):
																	for test_inputs in (self.test_sentences, self.paired_sentences):
																							_snake_case:   List[Any]  					=		tokenizer(__snake_case       ,     return_tensors='tf'       ,     padding='longest'       )
																							_snake_case:   Optional[Any]  					=		tf_tokenizer(__snake_case       )

																							for key in python_outputs.keys():
																													self.assertTrue(tf.reduce_all(python_outputs[key].shape == tf_outputs[key].shape       )       )
																													self.assertTrue(tf.reduce_all(tf.cast(python_outputs[key]       ,     tf.intaa       ) == tf_outputs[key]       )       )


					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											for tf_tokenizer in self.tf_tokenizers:
																	_snake_case:   str  					=		tf_tokenizer(self.paired_sentences       )
																	_snake_case:   Any  					=		tf_tokenizer(
																	    text=[sentence[0] for sentence in self.paired_sentences]       ,     text_pair=[sentence[1] for sentence in self.paired_sentences]       ,     )
																	for key in merged_outputs.keys():
																							self.assertTrue(tf.reduce_all(tf.cast(merged_outputs[key]       ,     tf.intaa       ) == separated_outputs[key]       )       )


					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											for tf_tokenizer in self.tf_tokenizers:
																	_snake_case:   List[str]  					=		tf.function(__snake_case       )
																	for test_inputs in (self.test_sentences, self.paired_sentences):
																							_snake_case:   Union[str, Any]  					=		tf.constant(__snake_case       )
																							_snake_case:   Any  					=		compiled_tokenizer(__snake_case       )
																							_snake_case:   Union[str, Any]  					=		tf_tokenizer(__snake_case       )

																							for key in eager_outputs.keys():
																													self.assertTrue(tf.reduce_all(eager_outputs[key] == compiled_outputs[key]       )       )




					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											for tf_tokenizer in self.tf_tokenizers:
																	_snake_case:   Optional[Any]  					=		ModelToSave(tokenizer=__snake_case       )
																	_snake_case:   Tuple  					=		tf.convert_to_tensor(self.test_sentences       )
																	_snake_case:   List[str]  					=		model(__snake_case       )  # Build model with some sample inputs
																	with TemporaryDirectory() as tempdir:
																							_snake_case:   List[str]  					=		Path(__snake_case       ) / 'saved.model'
																							model.save(__snake_case       )
																							_snake_case:   int  					=		tf.keras.models.load_model(__snake_case       )
																	_snake_case:   int  					=		loaded_model(__snake_case       )
																	# We may see small differences because the loaded model is compiled, so we need an epsilon for the test
																	self.assertLessEqual(tf.reduce_max(tf.abs(out - loaded_output       )       )       ,     1e-5       )





