



import argparse
from copy import deepcopy

import numpy as np
from datasets import ClassLabel, DatasetDict, load_dataset
from evaluate import load

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    set_seed,
)






def 					__lowerCAmelCase		(  )  ->						Optional[Any]:
					__lowerCAmelCase	   =argparse.ArgumentParser()
					parser.add_argument("""--model_ckpt"""				,   type=__lowerCamelCase				,   default="""microsoft/unixcoder-base-nine"""	)
					parser.add_argument("""--num_epochs"""				,   type=__lowerCamelCase				,   default=5	)
					parser.add_argument("""--batch_size"""				,   type=__lowerCamelCase				,   default=6	)
					parser.add_argument("""--gradient_accumulation_steps"""				,   type=__lowerCamelCase				,   default=1	)
					parser.add_argument("""--freeze"""				,   type=__lowerCamelCase				,   default=__lowerCamelCase	)
					parser.add_argument("""--learning_rate"""				,   type=__lowerCamelCase				,   default=5E-4	)
					parser.add_argument("""--seed"""				,   type=__lowerCamelCase				,   default=0	)
					parser.add_argument("""--lr_scheduler_type"""				,   type=__lowerCamelCase				,   default="""cosine"""	)
					parser.add_argument("""--num_warmup_steps"""				,   type=__lowerCamelCase				,   default=10	)
					parser.add_argument("""--weight_decay"""				,   type=__lowerCamelCase				,   default=0.01	)
					parser.add_argument("""--output_dir"""				,   type=__lowerCamelCase				,   default="""./results"""	)
					return parser.parse_args()


lowercase_	       =			load('''accuracy''')






def 					__lowerCAmelCase		(  __lowerCamelCase	:    Tuple	)  ->						List[Any]:
					__lowerCAmelCase   ,	__lowerCAmelCase	   =eval_pred
					__lowerCAmelCase	   =np.argmax(__lowerCamelCase				,   axis=1	)
					return metric.compute(predictions=__lowerCamelCase				,   references=__lowerCamelCase	)






class 						__a (     SCREAMING_SNAKE_CASE				):



	def __init__(						self							:    Any				,	snake_case_							:    Union[str, Any])->							None:
						super().__init__()
						__lowerCAmelCase	   =trainer



	def 		UpperCamelCase							(						self							:    Union[str, Any]				,	snake_case_							:    Union[str, Any]				,	snake_case_							:    Union[str, Any]				,	snake_case_							:    int				,	**snake_case_							:    Tuple)->							Optional[int]:
						if control.should_evaluate:
											__lowerCAmelCase	   =deepcopy(snake_case_)
											self._trainer.evaluate(eval_dataset=self._trainer.train_dataset				,	metric_key_prefix="""train""")
											return control_copy






def 					__lowerCAmelCase		(  )  ->						Tuple:
					__lowerCAmelCase	   =get_args()
					set_seed(args.seed	)

					__lowerCAmelCase	   =load_dataset("""codeparrot/codecomplex"""				,   split="""train"""	)
					__lowerCAmelCase	   =dataset.train_test_split(test_size=0.2	)
					__lowerCAmelCase	   =train_test["""test"""].train_test_split(test_size=0.5	)
					__lowerCAmelCase	   =DatasetDict(
					    {
					        """train""": train_test["""train"""],
					        """test""": test_validation["""train"""],
					        """valid""": test_validation["""test"""],
					    }	)

					print("""Loading tokenizer and model"""	)
					__lowerCAmelCase	   =AutoTokenizer.from_pretrained(args.model_ckpt	)
					__lowerCAmelCase	   =tokenizer.eos_token
					__lowerCAmelCase	   =AutoModelForSequenceClassification.from_pretrained(args.model_ckpt				,   num_labels=7	)
					__lowerCAmelCase	   =model.config.eos_token_id

					if args.freeze:
										for param in model.roberta.parameters():
															__lowerCAmelCase	   =False

					__lowerCAmelCase	   =ClassLabel(num_classes=7				,   names=list(set(train_test_validation["""train"""]["""complexity"""]	)	)	)

					def tokenize(__lowerCamelCase	:    Any	):
										__lowerCAmelCase	   =tokenizer(example["""src"""]				,   truncation=__lowerCamelCase				,   max_length=1024	)
										__lowerCAmelCase	   =labels.straint(example["""complexity"""]	)
										return {
										    "input_ids": inputs["input_ids"],
										    "attention_mask": inputs["attention_mask"],
										    "label": label,
										}

					__lowerCAmelCase	   =train_test_validation.map(
					    __lowerCamelCase				,   batched=__lowerCamelCase				,   remove_columns=train_test_validation["""train"""].column_names				,   )
					__lowerCAmelCase	   =DataCollatorWithPadding(tokenizer=__lowerCamelCase	)

					__lowerCAmelCase	   =TrainingArguments(
					    output_dir=args.output_dir				,   learning_rate=args.learning_rate				,   lr_scheduler_type=args.lr_scheduler_type				,   evaluation_strategy="""epoch"""				,   save_strategy="""epoch"""				,   logging_strategy="""epoch"""				,   per_device_train_batch_size=args.batch_size				,   per_device_eval_batch_size=args.batch_size				,   num_train_epochs=args.num_epochs				,   gradient_accumulation_steps=args.gradient_accumulation_steps				,   weight_decay=0.01				,   metric_for_best_model="""accuracy"""				,   run_name="""complexity-java"""				,   report_to="""wandb"""				,   )

					__lowerCAmelCase	   =Trainer(
					    model=__lowerCamelCase				,   args=__lowerCamelCase				,   train_dataset=tokenized_datasets["""train"""]				,   eval_dataset=tokenized_datasets["""valid"""]				,   tokenizer=__lowerCamelCase				,   data_collator=__lowerCamelCase				,   compute_metrics=__lowerCamelCase				,   )

					print("""Training..."""	)
					trainer.add_callback(CustomCallback(__lowerCamelCase	)	)
					trainer.train()


if __name__ == "__main__":
				main()









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 						__a (     unittest.TestCase				):



	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=snake_case_				,	)
						assert hasattr(self				,	"""env""")



	def 		UpperCamelCase							(						self							:    Dict				,	snake_case_							:    str)->							Dict:
						__lowerCAmelCase	   =F"""{self.env.base_job_name}-{instance_count}-{'ddp' if 'ddp' in self.script else 'smd'}"""
						# distributed data settings
						__lowerCAmelCase	   ={"""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=snake_case_				,	instance_count=snake_case_				,	instance_type=self.instance_type				,	debugger_hook_config=snake_case_				,	hyperparameters={**self.env.distributed_hyperparameters, """model_name_or_path""": self.model_name_or_path}				,	metric_definitions=self.env.metric_definitions				,	distribution=snake_case_				,	py_version="""py36"""				,	)



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



	@parameterized.expand([(2,)])
	def 		UpperCamelCase							(						self							:    Any				,	snake_case_							:    Dict)->							Optional[Any]:
						# create estimator
						__lowerCAmelCase	   =self.create_estimator(snake_case_)

						# run training
						estimator.fit()

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

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

						# 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}				,	snake_case_)




