



'''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,
)


lowercase							       =			logging.getLogger(__name__)







@dataclass
class 	UpperCAmelCase	:


	'''simple docstring'''




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







@dataclass
class 	UpperCAmelCase	:


	'''simple docstring'''




	__magic_name__					: str  =      field(
	    metadata={"help": "The input data dir. Should contain the .tsv files (or other data files) for the task."})
	__magic_name__					: Optional[str]  =      field(
	    default="summarization"					,	metadata={"help": "Task name, summarization (or summarization_{dataset} for pegasus) or translation"}					,	)
	__magic_name__					: Optional[int]  =      field(
	    default=1_024					,	metadata={
	        "help": (
	            "The maximum total input sequence length after tokenization. Sequences longer "
	            "than this will be truncated, sequences shorter will be padded."
	        )
	    }					,	)
	__magic_name__					: 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."
	        )
	    }					,	)
	__magic_name__					: 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``."
	        )
	    }					,	)
	__magic_name__					: 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."
	        )
	    }					,	)
	__magic_name__					: Optional[int]  =      field(default=-1					,	metadata={"help": "# training examples. -1 means use all."})
	__magic_name__					: Optional[int]  =      field(default=-1					,	metadata={"help": "# validation examples. -1 means use all."})
	__magic_name__					: Optional[int]  =      field(default=-1					,	metadata={"help": "# test examples. -1 means use all."})
	__magic_name__					: Optional[str]  =      field(default=__a					,	metadata={"help": "Source language id for translation."})
	__magic_name__					: Optional[str]  =      field(default=__a					,	metadata={"help": "Target language id for translation."})
	__magic_name__					: Optional[int]  =      field(default=__a					,	metadata={"help": "# num_beams to use for evaluation."})
	__magic_name__					: bool  =      field(
	    default=__a					,	metadata={"help": "If only pad tokens should be ignored. This assumes that `config.pad_token_id` is defined."}					,	)






def  UpperCAmelCase_       (		lowercase__       ,	lowercase__       ,	lowercase__       ):


					'''simple docstring'''







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






def  UpperCAmelCase_       (		):


					'''simple docstring'''







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

					check_output_dir(lowercase__       )

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

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

					a_       		=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       ,	)

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

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

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

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

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

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

					a_       		=SeqaSeqDataset

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

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

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

										trainer.save_model()  # this also saves the tokenizer

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

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

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

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

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

										a_       		=trainer.predict(test_dataset=lowercase__       ,	metric_key_prefix="test"       )
										a_       		=test_output.metrics
										a_       		=data_args.n_test

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

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

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

					return all_metrics






def  UpperCAmelCase_       (		lowercase__       ):


					'''simple docstring'''







					main()


if __name__ == "__main__":
		main()







'''simple docstring'''


import json
import logging
import os
import sys
from time import time
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, require_torch_tpu


logging.basicConfig(level=logging.DEBUG)

lowercase							       =			logging.getLogger()






def  UpperCAmelCase_       (		lowercase__       ):


     '''simple docstring'''







     a_       		={}
     a_       		=os.path.join(lowercase__       ,	"all_results.json"       )
     if os.path.exists(lowercase__       ):
          with open(lowercase__       ,	"r"       ) as f:
               a_       		=json.load(lowercase__       )
     else:
          raise ValueError(F"""can't find {path}"""       )
     return results


lowercase							       =			logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







@require_torch_tpu
class 	UpperCAmelCase	(		__a):


 '''simple docstring'''




 def       lowercase_  (   self) ->      List[Any]:



      """simple docstring"""







      import xla_spawn

      a_       		=self.get_auto_remove_tmp_dir()
      a_       		=f"""
            ./examples/pytorch/text-classification/run_glue.py
            --num_cores=8
            ./examples/pytorch/text-classification/run_glue.py
            --model_name_or_path distilbert-base-uncased
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --do_train
            --do_eval
            --debug tpu_metrics_debug
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --max_steps=10
            --warmup_steps=2
            --seed=42
            --max_seq_length=128
            """.split()

      with patch.object(lowerCAmelCase_							,  "argv"							,  lowerCAmelCase_):
           a_       		=time()
           xla_spawn.main()
           a_       		=time()

           a_       		=get_results(lowerCAmelCase_)
           self.assertGreaterEqual(result["eval_accuracy"]							,  0.7_5)

           # Assert that the script takes less than 500 seconds to make sure it doesn't hang.
           self.assertLess(end - start							,  5_0_0)




 def       lowercase_  (   self) ->      Tuple:



      """simple docstring"""







      import xla_spawn

      a_       		="\n            ./tests/test_trainer_tpu.py\n            --num_cores=8\n            ./tests/test_trainer_tpu.py\n            ".split()
      with patch.object(lowerCAmelCase_							,  "argv"							,  lowerCAmelCase_):
           xla_spawn.main()


