trainer.py (5033B)
1 """Define class for model training.""" 2 3 from __future__ import annotations 4 5 import datetime 6 import itertools 7 import json 8 import logging 9 from typing import TYPE_CHECKING, Any, Iterator 10 11 import transformers 12 13 import t5.collator 14 import t5.dataset 15 import t5.experiment 16 import t5.metric 17 import t5.model 18 19 20 if TYPE_CHECKING: 21 import pathlib 22 23 import numpy.typing 24 25 26 logger: logging.Logger = logging.getLogger(__name__) 27 28 29 class Trainer(t5.experiment.BaseExperiment): 30 """ 31 Train a model. 32 33 Config 34 ------ 35 dataset: Config 36 configuration for the t5.Dataset object 37 38 """ 39 40 def init(self: Trainer) -> None: 41 """Create objects needed for training.""" 42 self.config.declare_option("dataset", t5.config.Config, "Configuration of the dataset (see `t5.Dataset`).") 43 self.dataset = t5.dataset.Dataset(self.config.dataset) 44 self.config.declare_option("transformer", str, "Name of the huggingface model to use.") 45 self.config.declare_option("model", t5.config.Config, "Configuration of the model (see `t5.Model`).") 46 self.model = t5.model.Model(self.config.model, self.dataset) 47 self.data_collator = t5.collator.DataCollator(self.dataset, self.model) 48 self.metrics = t5.metric.Metrics(self.config, self.dataset, self.model, self.modeldir) 49 self.config.declare_option("num_train_epochs", float, "Number of training epochs.", 18.0) 50 self.config.declare_option("learning_rate", float, "The initial learning rate.", 3e-5) 51 self.modeldir.mkdir(parents=True, exist_ok=False) 52 training_arguments = transformers.Seq2SeqTrainingArguments( 53 self.modeldir, 54 report_to=["tensorboard"], 55 logging_dir=self.modeldir / "runs", 56 per_device_train_batch_size=8, 57 per_device_eval_batch_size=8, 58 bf16=True, 59 eval_strategy=transformers.IntervalStrategy.STEPS, 60 eval_steps=0.1, 61 save_steps=0.1, 62 load_best_model_at_end=True, 63 save_total_limit=3, 64 metric_for_best_model=f"{self.dataset.config.name}/aggregate/mean string accuracy", 65 predict_with_generate=False, 66 eval_do_concat_batches=False, 67 greater_is_better=True, 68 num_train_epochs=self.config.num_train_epochs, 69 learning_rate=self.config.learning_rate, 70 remove_unused_columns=False, 71 include_inputs_for_metrics=True, 72 batch_eval_metrics=True, 73 log_level="info") 74 self.trainer = transformers.Seq2SeqTrainer( 75 args=training_arguments, 76 train_dataset=self.dataset.train, 77 eval_dataset=self.dataset.validation, 78 model=self.model, 79 data_collator=self.data_collator, 80 compute_metrics=self.metrics 81 ) 82 self.config.declare_option("eval", bool, "Whether to run only evaluation.", default=False) 83 84 def _write_predictions(self: Trainer, split: str, predictions: Iterator[numpy.typing.NDArray[numpy.int64]]) -> None: 85 output_path: pathlib.Path = self.logdir / f"{split}_output.jsonl" 86 with output_path.open("w") as file: 87 for sample, generated_text in zip(self.dataset.datasets[split], predictions): 88 prediction_json: dict[str, Any] = self.dataset.recover_data_from_output(sample, generated_text) 89 prediction_json["generated_text"] = self.dataset.tokenizer.decode(generated_text) 90 json_sample: dict[str, Any] = { # write dates in files in string format 91 key: (str(value) if isinstance(value, (datetime.datetime, datetime.date)) else value) 92 for key, value in sample.items() 93 } 94 json.dump({"sample": json_sample, "prediction": prediction_json}, file) 95 print("\n", end="", file=file) 96 97 def evaluate(self: Trainer, split: str) -> dict[str, float]: 98 """Evaluate the model on a given split.""" 99 self.trainer.args.predict_with_generate = True 100 output: transformers.PredictionOutput = self.trainer.predict(self.dataset.datasets[split], metric_key_prefix=split) 101 unbatched_predictions: Iterator[numpy.typing.NDArray[numpy.int64]] = itertools.chain.from_iterable(output.predictions) 102 self._write_predictions(split, unbatched_predictions) 103 metrics: dict[str, float] = output.metrics 104 for key, value in metrics.items(): 105 logger.info("metric\t%s\t%s", key, value) 106 print(key, value) 107 return metrics 108 109 def run(self: Trainer) -> None: 110 """Train.""" 111 if not self.config.eval: 112 self.trainer.train() 113 metrics: dict[str, float] = self.evaluate("validation") 114 metrics.update(self.evaluate("test")) 115 with (self.logdir / "results.json").open("w") as results_file: 116 json.dump(metrics, results_file)