run_seq2seq.py (31228B)
1 #!/usr/bin/env python 2 # coding=utf-8 3 # Copyright The HuggingFace Team and The HuggingFace Inc. team. All rights reserved. 4 # 5 # Licensed under the Apache License, Version 2.0 (the "License"); 6 # you may not use this file except in compliance with the License. 7 # You may obtain a copy of the License at 8 # 9 # http://www.apache.org/licenses/LICENSE-2.0 10 # 11 # Unless required by applicable law or agreed to in writing, software 12 # distributed under the License is distributed on an "AS IS" BASIS, 13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 # See the License for the specific language governing permissions and 15 # limitations under the License. 16 """ 17 Fine-tuning the library models for sequence to sequence. 18 """ 19 # You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments. 20 21 import logging 22 import os 23 import re 24 import sys 25 from dataclasses import dataclass, field 26 from typing import Optional 27 28 import nltk # Here to have a nice missing dependency error message early on 29 import numpy as np 30 from datasets import load_dataset, load_metric 31 32 import transformers 33 from filelock import FileLock 34 from transformers import ( 35 AutoConfig, 36 AutoModelForSeq2SeqLM, 37 AutoTokenizer, 38 DataCollatorForSeq2Seq, 39 HfArgumentParser, 40 MBartTokenizer, 41 default_data_collator, 42 set_seed 43 ) 44 from transformers.trainer_utils import get_last_checkpoint, is_main_process 45 46 from extraction.event_schema import EventSchema 47 from extraction.extraction_metrics import decoding_format_dict, get_extract_metrics 48 from seq2seq.constrained_seq2seq import ConstraintSeq2SeqTrainingArguments, ConstraintSeq2SeqTrainer 49 50 # if not os.path.exists("~/nltk_data/tokenizers/punkt"): 51 # print('Start download punk') 52 # with FileLock(".lock") as lock: 53 # nltk.download("punkt", quiet=True) 54 55 logger = logging.getLogger(__name__) 56 57 58 @dataclass 59 class ModelArguments: 60 """ 61 Arguments pertaining to which model/config/tokenizer we are going to fine-tune from. 62 """ 63 64 model_name_or_path: str = field( 65 metadata={ 66 "help": "Path to pretrained model or model identifier from huggingface.co/models"} 67 ) 68 config_name: Optional[str] = field( 69 default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"} 70 ) 71 tokenizer_name: Optional[str] = field( 72 default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"} 73 ) 74 cache_dir: Optional[str] = field( 75 default=None, 76 metadata={ 77 "help": "Where to store the pretrained models downloaded from huggingface.co"}, 78 ) 79 use_fast_tokenizer: bool = field( 80 default=False, 81 metadata={ 82 "help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."}, 83 ) 84 # !!! must use non-fast version 85 # fast: "<extra_id_0> <extra_id_1>" -> [32099, 3, 32098, 1] 86 # non-fast: "<extra_id_0> <extra_id_1>" -> [32099, 32098, 1] 87 model_revision: str = field( 88 default="main", 89 metadata={ 90 "help": "The specific model version to use (can be a branch name, tag name or commit id)."}, 91 ) 92 use_auth_token: bool = field( 93 default=False, 94 metadata={ 95 "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script " 96 "with private models)." 97 }, 98 ) 99 100 101 @dataclass 102 class DataTrainingArguments: 103 """ 104 Arguments pertaining to what data we are going to input our model for training and eval. 105 """ 106 107 task: str = field( 108 default="summarization", 109 metadata={ 110 "help": "The name of the task, should be summarization (or summarization_{dataset} for evaluating " 111 "pegasus) or translation (or translation_{xx}_to_{yy})." 112 }, 113 ) 114 dataset_name: Optional[str] = field( 115 default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."} 116 ) 117 dataset_config_name: Optional[str] = field( 118 default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."} 119 ) 120 text_column: Optional[str] = field( 121 default=None, 122 metadata={ 123 "help": "The name of the column in the datasets containing the full texts (for summarization)."}, 124 ) 125 summary_column: Optional[str] = field( 126 default=None, 127 metadata={ 128 "help": "The name of the column in the datasets containing the summaries (for summarization)."}, 129 ) 130 train_file: Optional[str] = field( 131 default=None, metadata={"help": "The input training data file (a jsonlines or csv file)."} 132 ) 133 validation_file: Optional[str] = field( 134 default=None, 135 metadata={ 136 "help": "An optional input evaluation data file to evaluate the metrics (rouge/sacreblue) on " 137 "(a jsonlines or csv file)." 138 }, 139 ) 140 test_file: Optional[str] = field( 141 default=None, 142 metadata={ 143 "help": "An optional input test data file to evaluate the metrics (rouge/sacreblue) on " 144 "(a jsonlines or csv file)." 145 }, 146 ) 147 overwrite_cache: bool = field( 148 default=False, metadata={"help": "Overwrite the cached training and evaluation sets"} 149 ) 150 preprocessing_num_workers: Optional[int] = field( 151 default=None, 152 metadata={"help": "The number of processes to use for the preprocessing."}, 153 ) 154 max_source_length: Optional[int] = field( 155 default=1024, 156 metadata={ 157 "help": "The maximum total input sequence length after tokenization. Sequences longer " 158 "than this will be truncated, sequences shorter will be padded." 159 }, 160 ) 161 max_target_length: Optional[int] = field( 162 default=128, 163 metadata={ 164 "help": "The maximum total sequence length for target text after tokenization. Sequences longer " 165 "than this will be truncated, sequences shorter will be padded." 166 }, 167 ) 168 val_max_target_length: Optional[int] = field( 169 default=None, 170 metadata={ 171 "help": "The maximum total sequence length for validation target text after tokenization. Sequences longer " 172 "than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`." 173 "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used " 174 "during ``evaluate`` and ``predict``." 175 }, 176 ) 177 pad_to_max_length: bool = field( 178 default=False, 179 metadata={ 180 "help": "Whether to pad all samples to model maximum sentence length. " 181 "If False, will pad the samples dynamically when batching to the maximum length in the batch. More " 182 "efficient on GPU but very bad for TPU." 183 }, 184 ) 185 max_train_samples: Optional[int] = field( 186 default=None, 187 metadata={ 188 "help": "For debugging purposes or quicker training, truncate the number of training examples to this " 189 "value if set." 190 }, 191 ) 192 max_val_samples: Optional[int] = field( 193 default=None, 194 metadata={ 195 "help": "For debugging purposes or quicker training, truncate the number of validation examples to this " 196 "value if set." 197 }, 198 ) 199 max_test_samples: Optional[int] = field( 200 default=None, 201 metadata={ 202 "help": "For debugging purposes or quicker training, truncate the number of test examples to this " 203 "value if set." 204 }, 205 ) 206 source_lang: Optional[str] = field( 207 default=None, metadata={"help": "Source language id for translation."}) 208 target_lang: Optional[str] = field( 209 default=None, metadata={"help": "Target language id for translation."}) 210 num_beams: Optional[int] = field( 211 default=None, 212 metadata={ 213 "help": "Number of beams to use for evaluation. This argument will be passed to ``model.generate``, " 214 "which is used during ``evaluate`` and ``predict``." 215 }, 216 ) 217 ignore_pad_token_for_loss: bool = field( 218 default=True, 219 metadata={ 220 "help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not." 221 }, 222 ) 223 source_prefix: Optional[str] = field( 224 default=None, metadata={"help": "A prefix to add before every source text (useful for T5 models)."} 225 ) 226 227 def __post_init__(self): 228 if self.dataset_name is None and self.train_file is None and self.validation_file is None: 229 raise ValueError( 230 "Need either a dataset name or a training/validation file.") 231 else: 232 if self.train_file is not None: 233 extension = self.train_file.split(".")[-1] 234 assert extension in [ 235 "csv", "json"], "`train_file` should be a csv or a json file." 236 if self.validation_file is not None: 237 extension = self.validation_file.split(".")[-1] 238 assert extension in [ 239 "csv", "json"], "`validation_file` should be a csv or a json file." 240 if not self.task.startswith("summarization") and not self.task.startswith( 241 "translation") and not self.task.startswith('event'): 242 raise ValueError( 243 "`task` should be summarization, summarization_{dataset}, translation or translation_{xx}_to_{yy}." 244 ) 245 if self.val_max_target_length is None: 246 self.val_max_target_length = self.max_target_length 247 248 # Start Code for Event Extraction 249 decoding_format: str = field( 250 default='tree', 251 metadata={"help": "Decoding Format, valid in %s" % 252 decoding_format_dict.keys()} 253 ) 254 event_schema: str = field( 255 default=None, metadata={"help": "The input event schema file."} 256 ) 257 # End Code for Event Extraction 258 259 260 summarization_name_mapping = { 261 "amazon_reviews_multi": ("review_body", "review_title"), 262 "big_patent": ("description", "abstract"), 263 "cnn_dailymail": ("article", "highlights"), 264 "orange_sum": ("text", "summary"), 265 "pn_summary": ("article", "summary"), 266 "psc": ("extract_text", "summary_text"), 267 "samsum": ("dialogue", "summary"), 268 "thaisum": ("body", "summary"), 269 "xglue": ("news_body", "news_title"), 270 "xsum": ("document", "summary"), 271 "wiki_summary": ("article", "highlights"), 272 } 273 274 event_extraction_name_mapping = { 275 "event": ("text", "event") 276 } 277 278 279 def main(): 280 # See all possible arguments in src/transformers/training_args.py 281 # or by passing the --help flag to this script. 282 # We now keep distinct sets of args, for a cleaner separation of concerns. 283 284 parser = HfArgumentParser( 285 (ModelArguments, DataTrainingArguments, ConstraintSeq2SeqTrainingArguments)) 286 if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): 287 # If we pass only one argument to the script and it's the path to a json file, 288 # let's parse it to get our arguments. 289 model_args, data_args, training_args = parser.parse_json_file( 290 json_file=os.path.abspath(sys.argv[1])) 291 else: 292 model_args, data_args, training_args = parser.parse_args_into_dataclasses() 293 294 print(model_args) 295 print(data_args) 296 print(training_args) 297 298 # Detecting last checkpoint. 299 last_checkpoint = None 300 if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir: 301 last_checkpoint = get_last_checkpoint(training_args.output_dir) 302 if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0: 303 raise ValueError( 304 f"Output directory ({training_args.output_dir}) already exists and is not empty. " 305 "Use --overwrite_output_dir to overcome." 306 ) 307 elif last_checkpoint is not None: 308 logger.info( 309 f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change " 310 "the `--output_dir` or add `--overwrite_output_dir` to train from scratch." 311 ) 312 313 # Setup logging 314 logging.basicConfig( 315 format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", 316 datefmt="%m/%d/%Y %H:%M:%S", 317 handlers=[logging.StreamHandler(sys.stdout)], 318 ) 319 logger.setLevel(logging.INFO if is_main_process( 320 training_args.local_rank) else logging.WARN) 321 322 # Log on each process the small summary: 323 logger.warning( 324 f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}" 325 + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}" 326 ) 327 # Set the verbosity to info of the Transformers logger (on main process only): 328 if is_main_process(training_args.local_rank): 329 transformers.utils.logging.set_verbosity_info() 330 logger.info("Training/evaluation parameters %s", training_args) 331 332 # Set seed before initializing model. 333 set_seed(training_args.seed) 334 335 # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below) 336 # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/ 337 # (the dataset will be downloaded automatically from the datasets Hub). 338 # 339 # For CSV/JSON files in the summarization task, this script will use the first column for the full texts and the 340 # second column for the summaries (unless you specify column names for this with the `text_column` and 341 # `summary_column` arguments). 342 # For translation, only JSON files are supported, with one field named "translation" containing two keys for the 343 # source and target languages (unless you adapt what follows). 344 # 345 # In distributed training, the load_dataset function guarantee that only one local process can concurrently 346 # download the dataset. 347 if data_args.dataset_name is not None: 348 # Downloading and loading a dataset from the hub. 349 datasets = load_dataset(data_args.dataset_name, 350 data_args.dataset_config_name) 351 else: 352 data_files = {} 353 if data_args.train_file is not None: 354 data_files["train"] = data_args.train_file 355 extension = data_args.train_file.split(".")[-1] 356 if data_args.validation_file is not None: 357 data_files["validation"] = data_args.validation_file 358 extension = data_args.validation_file.split(".")[-1] 359 if data_args.test_file is not None: 360 data_files["test"] = data_args.test_file 361 extension = data_args.test_file.split(".")[-1] 362 datasets = load_dataset(extension, data_files=data_files) 363 # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at 364 # https://huggingface.co/docs/datasets/loading_datasets.html. 365 366 # Load pretrained model and tokenizer 367 # 368 # Distributed training: 369 # The .from_pretrained methods guarantee that only one local process can concurrently 370 # download model & vocab. 371 config = AutoConfig.from_pretrained( 372 model_args.config_name if model_args.config_name else model_args.model_name_or_path, 373 cache_dir=model_args.cache_dir, 374 revision=model_args.model_revision, 375 use_auth_token=True if model_args.use_auth_token else None, 376 ) 377 378 # !!! Sometimes default max_length is setting to 20. 379 config.max_length = data_args.max_target_length 380 381 tokenizer = AutoTokenizer.from_pretrained( 382 model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path, 383 cache_dir=model_args.cache_dir, 384 use_fast=model_args.use_fast_tokenizer, 385 revision=model_args.model_revision, 386 use_auth_token=True if model_args.use_auth_token else None, 387 ) 388 389 to_remove_token_list = list() 390 if tokenizer.bos_token: 391 to_remove_token_list += [tokenizer.bos_token] 392 if tokenizer.eos_token: 393 to_remove_token_list += [tokenizer.eos_token] 394 if tokenizer.pad_token: 395 to_remove_token_list += [tokenizer.pad_token] 396 397 model = AutoModelForSeq2SeqLM.from_pretrained( 398 model_args.model_name_or_path, 399 from_tf=bool(".ckpt" in model_args.model_name_or_path), 400 config=config, 401 cache_dir=model_args.cache_dir, 402 revision=model_args.model_revision, 403 use_auth_token=True if model_args.use_auth_token else None, 404 ) 405 406 if tokenizer.encode("<extra_id_0> <extra_id_1>") != [32099, 32098, 1]: 407 # For non-t5 tokenizer 408 tokenizer.add_special_tokens( 409 {"additional_special_tokens": ["<extra_id_0>", "<extra_id_1>"]}) 410 model.resize_token_embeddings(len(tokenizer)) 411 412 # Set decoder_start_token_id 413 if model.config.decoder_start_token_id is None and isinstance(tokenizer, MBartTokenizer): 414 model.config.decoder_start_token_id = tokenizer.lang_code_to_id[data_args.target_lang] 415 if model.config.decoder_start_token_id is None: 416 raise ValueError( 417 "Make sure that `config.decoder_start_token_id` is correctly defined") 418 419 prefix = data_args.source_prefix if data_args.source_prefix is not None else "" 420 421 # Preprocessing the datasets. 422 # We need to tokenize inputs and targets. 423 if training_args.do_train: 424 column_names = datasets["train"].column_names 425 elif training_args.do_eval: 426 column_names = datasets["validation"].column_names 427 elif training_args.do_predict: 428 column_names = datasets["test"].column_names 429 else: 430 logger.info( 431 "There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.") 432 return 433 434 # For translation we set the codes of our source and target languages (only useful for mBART, the others will 435 # ignore those attributes). 436 if data_args.task.startswith("translation"): 437 if data_args.source_lang is not None: 438 tokenizer.src_lang = data_args.source_lang 439 if data_args.target_lang is not None: 440 tokenizer.tgt_lang = data_args.target_lang 441 442 # Start Code for Event Extraction 443 if data_args.task.startswith("event"): 444 decoding_type_schema = EventSchema.read_from_file( 445 data_args.event_schema) 446 else: 447 decoding_type_schema = None 448 # End Code for Event Extraction 449 450 # To serialize preprocess_function below, each of those four variables needs to be defined (even if we won't use 451 # them all). 452 source_lang, target_lang, text_column, summary_column = None, None, None, None 453 454 if data_args.task.startswith("summarization"): 455 # Get the column names for input/target. 456 dataset_columns = summarization_name_mapping.get( 457 data_args.dataset_name, None) 458 if data_args.text_column is None: 459 text_column = dataset_columns[0] if dataset_columns is not None else column_names[0] 460 else: 461 text_column = data_args.text_column 462 if data_args.summary_column is None: 463 summary_column = dataset_columns[1] if dataset_columns is not None else column_names[1] 464 else: 465 summary_column = data_args.summary_column 466 # Start Code for Event Extraction 467 elif data_args.task.startswith("event"): 468 dataset_columns = event_extraction_name_mapping.get( 469 data_args.dataset_name, None) 470 if data_args.text_column is None: 471 text_column = dataset_columns[0] if dataset_columns is not None else column_names[0] 472 else: 473 text_column = data_args.text_column 474 if data_args.summary_column is None: 475 summary_column = dataset_columns[1] if dataset_columns is not None else column_names[1] 476 else: 477 summary_column = data_args.summary_column 478 # End Code for Event Extraction 479 else: 480 # Get the language codes for input/target. 481 lang_search = re.match( 482 "translation_([a-z]+)_to_([a-z]+)", data_args.task) 483 if data_args.source_lang is not None: 484 source_lang = data_args.source_lang.split("_")[0] 485 else: 486 assert ( 487 lang_search is not None 488 ), "Provide a source language via --source_lang or rename your task 'translation_xx_to_yy'." 489 source_lang = lang_search.groups()[0] 490 491 if data_args.target_lang is not None: 492 target_lang = data_args.target_lang.split("_")[0] 493 else: 494 assert ( 495 lang_search is not None 496 ), "Provide a target language via --target_lang or rename your task 'translation_xx_to_yy'." 497 target_lang = lang_search.groups()[1] 498 499 # Temporarily set max_target_length for training. 500 max_target_length = data_args.max_target_length 501 padding = "max_length" if data_args.pad_to_max_length else False 502 503 if training_args.label_smoothing_factor > 0 and not hasattr(model, "prepare_decoder_input_ids_from_labels"): 504 logger.error( 505 "label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for" 506 f"`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory" 507 ) 508 509 def preprocess_function(examples): 510 if data_args.task.startswith("translation"): 511 inputs = [ex[source_lang] for ex in examples["translation"]] 512 targets = [ex[target_lang] for ex in examples["translation"]] 513 else: 514 inputs = examples[text_column] 515 targets = examples[summary_column] 516 inputs = [prefix + inp for inp in inputs] 517 model_inputs = tokenizer( 518 inputs, max_length=data_args.max_source_length, padding=padding, truncation=True) 519 520 # Setup the tokenizer for targets 521 with tokenizer.as_target_tokenizer(): 522 labels = tokenizer( 523 targets, max_length=max_target_length, padding=padding, truncation=True) 524 525 # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore 526 # padding in the loss. 527 if padding == "max_length" and data_args.ignore_pad_token_for_loss: 528 labels["input_ids"] = [ 529 [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"] 530 ] 531 532 model_inputs["labels"] = labels["input_ids"] 533 return model_inputs 534 535 if training_args.do_train: 536 train_dataset = datasets["train"] 537 if data_args.max_train_samples is not None: 538 train_dataset = train_dataset.select( 539 range(data_args.max_train_samples)) 540 train_dataset = train_dataset.map( 541 preprocess_function, 542 batched=True, 543 num_proc=data_args.preprocessing_num_workers, 544 remove_columns=column_names, 545 load_from_cache_file=not data_args.overwrite_cache, 546 ) 547 548 if training_args.do_eval: 549 max_target_length = data_args.val_max_target_length 550 eval_dataset = datasets["validation"] 551 if data_args.max_val_samples is not None: 552 eval_dataset = eval_dataset.select( 553 range(data_args.max_val_samples)) 554 eval_dataset = eval_dataset.map( 555 preprocess_function, 556 batched=True, 557 num_proc=data_args.preprocessing_num_workers, 558 remove_columns=column_names, 559 load_from_cache_file=not data_args.overwrite_cache, 560 ) 561 562 if training_args.do_predict: 563 max_target_length = data_args.val_max_target_length 564 test_dataset = datasets["test"] 565 if data_args.max_test_samples is not None: 566 test_dataset = test_dataset.select( 567 range(data_args.max_test_samples)) 568 test_dataset = test_dataset.map( 569 preprocess_function, 570 batched=True, 571 num_proc=data_args.preprocessing_num_workers, 572 remove_columns=column_names, 573 load_from_cache_file=not data_args.overwrite_cache, 574 ) 575 576 # Data collator 577 label_pad_token_id = - \ 578 100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id 579 if data_args.pad_to_max_length: 580 data_collator = default_data_collator 581 else: 582 data_collator = DataCollatorForSeq2Seq( 583 tokenizer, 584 model=model, 585 label_pad_token_id=label_pad_token_id, 586 pad_to_multiple_of=8 if training_args.fp16 else None, 587 ) 588 589 def compute_metrics(eval_preds): 590 preds, labels = eval_preds 591 if isinstance(preds, tuple): 592 preds = preds[0] 593 decoded_preds = tokenizer.batch_decode( 594 preds, skip_special_tokens=False) 595 if data_args.ignore_pad_token_for_loss: 596 # Replace -100 in the labels as we can't decode them. 597 labels = np.where(labels != -100, labels, tokenizer.pad_token_id) 598 decoded_labels = tokenizer.batch_decode( 599 labels, skip_special_tokens=False) 600 601 def clean_str(x_str): 602 for to_remove_token in to_remove_token_list: 603 x_str = x_str.replace(to_remove_token, '') 604 return x_str.strip() 605 606 decoded_preds = [clean_str(x) for x in decoded_preds] 607 decoded_labels = [clean_str(x) for x in decoded_labels] 608 609 result = get_extract_metrics( 610 pred_lns=decoded_preds, 611 tgt_lns=decoded_labels, 612 label_constraint=decoding_type_schema, 613 decoding_format=data_args.decoding_format, 614 ) 615 616 prediction_lens = [np.count_nonzero( 617 pred != tokenizer.pad_token_id) for pred in preds] 618 result["gen_len"] = np.mean(prediction_lens) 619 result = {k: round(v, 4) for k, v in result.items()} 620 return result 621 622 # Initialize our Trainer 623 trainer = ConstraintSeq2SeqTrainer( 624 model=model, 625 args=training_args, 626 train_dataset=train_dataset if training_args.do_train else None, 627 eval_dataset=eval_dataset if training_args.do_eval else None, 628 tokenizer=tokenizer, 629 data_collator=data_collator, 630 compute_metrics=compute_metrics if training_args.predict_with_generate else None, 631 decoding_type_schema=decoding_type_schema, 632 decoding_format=data_args.decoding_format, 633 source_prefix=prefix, 634 ) 635 636 # Training 637 if training_args.do_train: 638 # if last_checkpoint is not None: 639 # checkpoint = last_checkpoint 640 # elif os.path.isdir(model_args.model_name_or_path): 641 # checkpoint = model_args.model_name_or_path 642 # else: 643 # checkpoint = None 644 # TODO fix better about max_length 645 checkpoint = None 646 647 train_result = trainer.train(resume_from_checkpoint=checkpoint) 648 trainer.save_model() # Saves the tokenizer too for easy upload 649 decoding_type_schema.write_to_file( 650 os.path.join( 651 training_args.output_dir, 652 "event.schema", 653 ) 654 ) 655 656 output_train_file = os.path.join( 657 training_args.output_dir, "train_results.txt") 658 if trainer.is_world_process_zero(): 659 with open(output_train_file, "w") as writer: 660 logger.info("***** Train results *****") 661 for key, value in sorted(train_result.metrics.items()): 662 logger.info(f" {key} = {value}") 663 writer.write(f"{key} = {value}\n") 664 665 # Need to save the state, since Trainer.save_model saves only the tokenizer with the model 666 trainer.state.save_to_json(os.path.join( 667 training_args.output_dir, "trainer_state.json")) 668 669 670 # Evaluation 671 results = {} 672 if training_args.do_eval: 673 logger.info("*** Evaluate ***") 674 675 results = trainer.evaluate( 676 max_length=data_args.val_max_target_length, num_beams=data_args.num_beams) 677 results = {k: round(v, 4) for k, v in results.items()} 678 679 eval_results = trainer.predict( 680 eval_dataset, 681 metric_key_prefix="eval", 682 max_length=data_args.val_max_target_length, 683 num_beams=data_args.num_beams, 684 ) 685 686 output_eval_file = os.path.join( 687 training_args.output_dir, "eval_results_seq2seq.txt") 688 if trainer.is_world_process_zero(): 689 with open(output_eval_file, "w") as writer: 690 logger.info("***** Eval results *****") 691 for key, value in sorted(results.items()): 692 logger.info(f" {key} = {value}") 693 writer.write(f"{key} = {value}\n") 694 695 if training_args.predict_with_generate: 696 eval_preds = tokenizer.batch_decode( 697 eval_results.predictions, skip_special_tokens=False, clean_up_tokenization_spaces=True 698 ) 699 eval_preds = [pred.replace('<pad>', '').replace('<s>', '').replace('</s>', '').strip() 700 for pred in eval_preds] 701 output_test_preds_file = os.path.join( 702 training_args.output_dir, "eval_preds_seq2seq.txt") 703 with open(output_test_preds_file, "w") as writer: 704 writer.write("\n".join(eval_preds)) 705 706 if training_args.do_predict: 707 logger.info("*** Test ***") 708 709 test_results = trainer.predict( 710 test_dataset, 711 metric_key_prefix="test", 712 max_length=data_args.val_max_target_length, 713 num_beams=data_args.num_beams, 714 ) 715 test_metrics = test_results.metrics 716 test_metrics["test_loss"] = round(test_metrics["test_loss"], 4) 717 718 output_test_result_file = os.path.join( 719 training_args.output_dir, "test_results_seq2seq.txt") 720 if trainer.is_world_process_zero(): 721 with open(output_test_result_file, "w") as writer: 722 logger.info("***** Test results *****") 723 for key, value in sorted(test_metrics.items()): 724 logger.info(f" {key} = {value}") 725 writer.write(f"{key} = {value}\n") 726 727 if training_args.predict_with_generate: 728 test_preds = tokenizer.batch_decode( 729 test_results.predictions, skip_special_tokens=False, clean_up_tokenization_spaces=True 730 ) 731 test_preds = [pred.replace('<pad>', '').replace('<s>', '').replace('</s>', '').strip() 732 for pred in test_preds] 733 output_test_preds_file = os.path.join( 734 training_args.output_dir, "test_preds_seq2seq.txt") 735 with open(output_test_preds_file, "w") as writer: 736 writer.write("\n".join(test_preds)) 737 738 return results 739 740 741 def _mp_fn(index): 742 # For xla_spawn (TPUs) 743 main() 744 745 746 if __name__ == "__main__": 747 main()