


"""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__       :							str        =  logging.getLogger(__name__)




@dataclass
class 				UpperCamelCase_			:

	"""simple docstring"""





	snake_case__       :						str					      =    field(
	    metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"})
	snake_case__       :						Optional[str]					      =    field(
	    default=UpperCamelCase    ,						metadata={"help": "Pretrained config name or path if not the same as model_name"})
	snake_case__       :						Optional[str]					      =    field(
	    default=UpperCamelCase    ,						metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"})
	snake_case__       :						Optional[str]					      =    field(
	    default=UpperCamelCase    ,						metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"}    ,						)
	snake_case__       :						bool					      =    field(default=UpperCamelCase    ,						metadata={"help": "Whether tp freeze the encoder."})
	snake_case__       :						bool					      =    field(default=UpperCamelCase    ,						metadata={"help": "Whether  to freeze the embeddings."})




@dataclass
class 				UpperCamelCase_			:

	"""simple docstring"""





	snake_case__       :						str					      =    field(
	    metadata={"help": "The input data dir. Should contain the .tsv files (or other data files) for the task."})
	snake_case__       :						Optional[str]					      =    field(
	    default="summarization"    ,						metadata={"help": "Task name, summarization (or summarization_{dataset} for pegasus) or translation"}    ,						)
	snake_case__       :						Optional[int]					      =    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."
	        )
	    }    ,						)
	snake_case__       :						Optional[int]					      =    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."
	        )
	    }    ,						)
	snake_case__       :						Optional[int]					      =    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``."
	        )
	    }    ,						)
	snake_case__       :						Optional[int]					      =    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."
	        )
	    }    ,						)
	snake_case__       :						Optional[int]					      =    field(default=-1    ,						metadata={"help": "# training examples. -1 means use all."})
	snake_case__       :						Optional[int]					      =    field(default=-1    ,						metadata={"help": "# validation examples. -1 means use all."})
	snake_case__       :						Optional[int]					      =    field(default=-1    ,						metadata={"help": "# test examples. -1 means use all."})
	snake_case__       :						Optional[str]					      =    field(default=UpperCamelCase    ,						metadata={"help": "Source language id for translation."})
	snake_case__       :						Optional[str]					      =    field(default=UpperCamelCase    ,						metadata={"help": "Target language id for translation."})
	snake_case__       :						Optional[int]					      =    field(default=UpperCamelCase    ,						metadata={"help": "# num_beams to use for evaluation."})
	snake_case__       :						bool					      =    field(
	    default=UpperCamelCase    ,						metadata={"help": "If only pad tokens should be ignored. This assumes that `config.pad_token_id` is defined."}    ,						)




def      UpperCAmelCase__ (lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_     ):




					'''simple docstring'''


					logger.info(f"""***** {split} metrics *****"""     )
					for key in sorted(metrics.keys()     ):
										logger.info(f"""  {key} = {metrics[key]}"""     )
					save_json(lowerCAmelCase_						,		os.path.join(lowerCAmelCase_						,		f"""{split}_results.json"""     )     )




def      UpperCAmelCase__ ():




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE		     =       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.
										__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()

					check_output_dir(lowerCAmelCase_     )

					# 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"						,		lowerCAmelCase_     )

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

					__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		     =       ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
					for p in extra_model_params:
										if getattr(lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_     ):
															assert hasattr(lowerCAmelCase_						,		lowerCAmelCase_     ), f"""({config.__class__.__name__}) doesn't have a `{p}` attribute"""
															setattr(lowerCAmelCase_						,		lowerCAmelCase_						,		getattr(lowerCAmelCase_						,		lowerCAmelCase_     )     )

					__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						,		)
					__SCREAMING_SNAKE_CASE		     =       AutoModelForSeqaSeqLM.from_pretrained(
					    model_args.model_name_or_path						,		from_tf=".ckpt" in model_args.model_name_or_path						,		config=lowerCAmelCase_						,		cache_dir=model_args.cache_dir						,		)

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

					# set num_beams for evaluation
					if data_args.eval_beams is None:
										__SCREAMING_SNAKE_CASE		     =       model.config.num_beams

					# set decoder_start_token_id for MBart
					if model.config.decoder_start_token_id is None and isinstance(lowerCAmelCase_						,		(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(lowerCAmelCase_						,		lowerCAmelCase_     ):
															__SCREAMING_SNAKE_CASE		     =       tokenizer.lang_code_to_id[data_args.tgt_lang]
										else:
															__SCREAMING_SNAKE_CASE		     =       tokenizer.convert_tokens_to_ids(data_args.tgt_lang     )

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

					__SCREAMING_SNAKE_CASE		     =       SeqaSeqDataset

					# Get datasets
					__SCREAMING_SNAKE_CASE		     =       (
					    dataset_class(
					        lowerCAmelCase_						,		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
					)
					__SCREAMING_SNAKE_CASE		     =       (
					    dataset_class(
					        lowerCAmelCase_						,		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
					)
					__SCREAMING_SNAKE_CASE		     =       (
					    dataset_class(
					        lowerCAmelCase_						,		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
					__SCREAMING_SNAKE_CASE		     =       (
					    build_compute_metrics_fn(data_args.task						,		lowerCAmelCase_     ) if training_args.predict_with_generate else None
					)
					__SCREAMING_SNAKE_CASE		     =       SeqaSeqTrainer(
					    model=lowerCAmelCase_						,		args=lowerCAmelCase_						,		data_args=lowerCAmelCase_						,		train_dataset=lowerCAmelCase_						,		eval_dataset=lowerCAmelCase_						,		data_collator=SeqaSeqDataCollator(
					        lowerCAmelCase_						,		lowerCAmelCase_						,		model.config.decoder_start_token_id						,		training_args.tpu_num_cores     )						,		compute_metrics=lowerCAmelCase_						,		tokenizer=lowerCAmelCase_						,		)

					__SCREAMING_SNAKE_CASE		     =       {}
					# Training
					if training_args.do_train:
										logger.info("*** Train ***"     )

										__SCREAMING_SNAKE_CASE		     =       trainer.train(
										    model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path     ) else None     )
										__SCREAMING_SNAKE_CASE		     =       train_result.metrics
										__SCREAMING_SNAKE_CASE		     =       data_args.n_train

										trainer.save_model()  # this also saves the tokenizer

										if trainer.is_world_process_zero():
															handle_metrics("train"						,		lowerCAmelCase_						,		training_args.output_dir     )
															all_metrics.update(lowerCAmelCase_     )

															# 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 ***"     )

										__SCREAMING_SNAKE_CASE		     =       trainer.evaluate(metric_key_prefix="val"     )
										__SCREAMING_SNAKE_CASE		     =       data_args.n_val
										__SCREAMING_SNAKE_CASE		     =       round(metrics["val_loss"]						,		4     )

										if trainer.is_world_process_zero():
															handle_metrics("val"						,		lowerCAmelCase_						,		training_args.output_dir     )
															all_metrics.update(lowerCAmelCase_     )

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

										__SCREAMING_SNAKE_CASE		     =       trainer.predict(test_dataset=lowerCAmelCase_						,		metric_key_prefix="test"     )
										__SCREAMING_SNAKE_CASE		     =       test_output.metrics
										__SCREAMING_SNAKE_CASE		     =       data_args.n_test

										if trainer.is_world_process_zero():
															__SCREAMING_SNAKE_CASE		     =       round(metrics["test_loss"]						,		4     )
															handle_metrics("test"						,		lowerCAmelCase_						,		training_args.output_dir     )
															all_metrics.update(lowerCAmelCase_     )

															if training_args.predict_with_generate:
																				__SCREAMING_SNAKE_CASE		     =       tokenizer.batch_decode(
																				    test_output.predictions						,		skip_special_tokens=lowerCAmelCase_						,		clean_up_tokenization_spaces=lowerCAmelCase_     )
																				__SCREAMING_SNAKE_CASE		     =       lmap(str.strip						,		lowerCAmelCase_     )
																				write_txt_file(lowerCAmelCase_						,		os.path.join(training_args.output_dir						,		"test_generations.txt"     )     )

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

					return all_metrics




def      UpperCAmelCase__ (lowerCAmelCase_     ):




					'''simple docstring'''


					main()


if __name__ == "__main__":
		main()




"""simple docstring"""







import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
		from sagemaker import Session, TrainingJobAnalytics
		from sagemaker.huggingface import HuggingFace




@pytest.mark.skipif(
    literal_eval(os.getenv("TEST_SAGEMAKER"    ,						"False")) is not True    ,						reason="Skipping test because should only be run when releasing minor transformers version"    ,						)
@pytest.mark.usefixtures("sm_env")
@parameterized_class(
    [
        {
            "framework": "pytorch",
            "script": "run_glue.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 650, "eval_accuracy": 0.7, "eval_loss": 0.6},
        },
        {
            "framework": "pytorch",
            "script": "run_ddp.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 600, "eval_accuracy": 0.7, "eval_loss": 0.6},
        },
        {
            "framework": "tensorflow",
            "script": "run_tf_dist.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 600, "eval_accuracy": 0.6, "eval_loss": 0.7},
        },
    ])
class 				UpperCamelCase_			(				unittest.TestCase):

	"""simple docstring"""





	def 						UpperCAmelCase_   (				self    :       Dict     ) ->			Any:
						if self.framework == "pytorch":
											subprocess.run(
											    F"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()    ,    encoding="utf-8"    ,    check=UpperCAmelCase__    ,    )
						assert hasattr(self    ,    "env"     )




	def 						UpperCAmelCase_   (				self    :       Tuple    ,    UpperCAmelCase__    :       int     ) ->			Any:
						__SCREAMING_SNAKE_CASE		     =       F"""{self.env.base_job_name}-{instance_count}-{'ddp' if 'ddp' in self.script else 'smd'}"""
						# distributed data settings
						__SCREAMING_SNAKE_CASE		     =       {"smdistributed": {"dataparallel": {"enabled": True}}} if self.script != "run_ddp.py" else None

						# creates estimator
						return HuggingFace(
						    entry_point=self.script    ,    source_dir=self.env.test_path    ,    role=self.env.role    ,    image_uri=self.env.image_uri    ,    base_job_name=UpperCAmelCase__    ,    instance_count=UpperCAmelCase__    ,    instance_type=self.instance_type    ,    debugger_hook_config=UpperCAmelCase__    ,    hyperparameters={**self.env.distributed_hyperparameters, "model_name_or_path": self.model_name_or_path}    ,    metric_definitions=self.env.metric_definitions    ,    distribution=UpperCAmelCase__    ,    py_version="py36"    ,    )




	def 						UpperCAmelCase_   (				self    :       Optional[int]    ,    UpperCAmelCase__    :       Optional[int]     ) ->			Optional[Any]:
						TrainingJobAnalytics(UpperCAmelCase__     ).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv"""     )




	@parameterized.expand([(2,)]     )
	def 						UpperCAmelCase_   (				self    :       Dict    ,    UpperCAmelCase__    :       Optional[Any]     ) ->			List[str]:
						# create estimator
						__SCREAMING_SNAKE_CASE		     =       self.create_estimator(UpperCAmelCase__     )

						# run training
						estimator.fit()

						# result dataframe
						__SCREAMING_SNAKE_CASE		     =       TrainingJobAnalytics(estimator.latest_training_job.name     ).dataframe()

						# extract kpis
						__SCREAMING_SNAKE_CASE		     =       list(result_metrics_df[result_metrics_df.metric_name == "eval_accuracy"]["value"]     )
						__SCREAMING_SNAKE_CASE		     =       list(result_metrics_df[result_metrics_df.metric_name == "eval_loss"]["value"]     )
						# get train time from SageMaker job, this includes starting, preprocessing, stopping
						__SCREAMING_SNAKE_CASE		     =       (
						    Session().describe_training_job(estimator.latest_training_job.name     ).get("TrainingTimeInSeconds"    ,    9_9_9_9_9_9     )
						)

						# assert kpis
						assert train_runtime <= self.results["train_runtime"]
						assert all(t >= self.results["eval_accuracy"] for t in eval_accuracy     )
						assert all(t <= self.results["eval_loss"] for t in eval_loss     )

						# dump tests result into json file to share in PR
						with open(F"""{estimator.latest_training_job.name}.json"""    ,    "w"     ) as outfile:
											json.dump({"train_time": train_runtime, "eval_accuracy": eval_accuracy, "eval_loss": eval_loss}    ,    UpperCAmelCase__     )
