Merge branch 'main' of github.com:CompNet/Tibert

Aethor · Aethor · commit d6fa8ce73833 · 2023-09-02T13:03:50.000+02:00
diff --git a/README.md b/README.md
@@ -97,7 +97,7 @@ Aside from the `tibert.train.train_coref_model` function, it is possible to trai
 ```sh
 python -m tibert.run_train with\
        dataset_path=/path/to/litbank/repository\
-       out_model_path=/path/to/output/model/directory
+       out_model_dir=/path/to/output/model/directory
 ```
 
 The following parameters can be set (taken from `./tibert/run_train.py` config function):
@@ -119,7 +119,7 @@ The following parameters can be set (taken from `./tibert/run_train.py` config f
 | `dropout`                    | `0.3`               |
 | `segment_size`               | `128`               |
 | `encoder`                    | `"bert-base-cased"` |
-| `out_model_path`             | `"~/tibert/model"`  |
+| `out_model_dir`              | `"~/tibert/model"`  |
 
 
 One can monitor training metrics by adding run observers using command line flags - see `sacred` documentation for more details.
diff --git a/tibert/bertcoref.py b/tibert/bertcoref.py
@@ -218,7 +218,14 @@ def from_wpieced_to_tokenized(
             for mention in chain:
 
                 new_start_idx = wp_to_token[mention.start_idx]
-                new_end_idx = wp_to_token[mention.end_idx - 1] + 1
+                new_end_idx = wp_to_token[mention.end_idx - 1]
+                # NOTE: this happens in case the model has predicted
+                # an erroneous mention such as '[CLS]' or '[SEP]'. In
+                # that case, we simply ignore the mention.
+                if new_start_idx is None or new_end_idx is None:
+                    continue
+                new_end_idx += 1
+
                 new_mention = Mention(
                     tokens[new_start_idx:new_end_idx],
                     new_start_idx,
diff --git a/tibert/run_train.py b/tibert/run_train.py
@@ -1,5 +1,6 @@
-from typing import Literal, cast
+from typing import Literal, Optional, cast
 import os
+from torch.optim import optimizer
 from transformers import BertTokenizerFast, CamembertTokenizerFast  # type: ignore
 from sacred.experiment import Experiment
 from sacred.run import Run
@@ -10,6 +11,7 @@
     BertForCoreferenceResolution,
     CamembertForCoreferenceResolution,
     train_coref_model,
+    load_train_checkpoint,
 )
 
 ex = Experiment()
@@ -33,7 +35,8 @@ def config():
     dropout: float = 0.3
     segment_size: int = 128
     encoder: str = "bert-base-cased"
-    out_model_path: str = os.path.expanduser("~/tibert/model")
+    out_model_dir: str = os.path.expanduser("~/tibert/model")
+    checkpoint: Optional[str] = None
 
 
 @ex.main
@@ -54,7 +57,8 @@ def main(
     dropout: float,
     segment_size: int,
     encoder: str,
-    out_model_path: str,
+    out_model_dir: str,
+    checkpoint: Optional[str],
 ):
     print_config(_run)
 
@@ -76,39 +80,42 @@ def main(
 
     config = dataset_configs[dataset_name]
 
-    model = config["model_class"].from_pretrained(
-        encoder,
-        mentions_per_tokens=mentions_per_tokens,
-        antecedents_nb=antecedents_nb,
-        max_span_size=max_span_size,
-        segment_size=segment_size,
-        mention_scorer_hidden_size=mention_scorer_hidden_size,
-        mention_scorer_dropout=dropout,
-        hidden_dropout_prob=dropout,
-        attention_probs_dropout_prob=dropout,
-        mention_loss_coeff=mention_loss_coeff,
-    )
+    if not checkpoint is None:
+        model, optimizer = load_train_checkpoint(checkpoint, config["model_class"])
+    else:
+        model = config["model_class"].from_pretrained(
+            encoder,
+            mentions_per_tokens=mentions_per_tokens,
+            antecedents_nb=antecedents_nb,
+            max_span_size=max_span_size,
+            segment_size=segment_size,
+            mention_scorer_hidden_size=mention_scorer_hidden_size,
+            mention_scorer_dropout=dropout,
+            hidden_dropout_prob=dropout,
+            attention_probs_dropout_prob=dropout,
+            mention_loss_coeff=mention_loss_coeff,
+        )
+        optimizer = None
 
     tokenizer = config["tokenizer_class"].from_pretrained(encoder)
 
     dataset = config["loading_function"](dataset_path, tokenizer, max_span_size)
 
-    model = train_coref_model(
+    train_coref_model(
         model,
         dataset,
         tokenizer,
-        batch_size,
-        epochs_nb,
-        sents_per_documents_train,
-        bert_lr,
-        task_lr,
-        out_model_path,
-        "auto",
-        _run,
+        batch_size=batch_size,
+        epochs_nb=epochs_nb,
+        sents_per_documents_train=sents_per_documents_train,
+        bert_lr=bert_lr,
+        task_lr=task_lr,
+        model_save_dir=out_model_dir,
+        device_str="auto",
+        _run=_run,
+        optimizer=optimizer,
     )
 
-    model.save_pretrained(out_model_path)
-
 
 if __name__ == "__main__":
     ex.run_commandline()
diff --git a/tibert/train.py b/tibert/train.py
@@ -1,22 +1,71 @@
-from typing import Optional, Union, Literal
-import traceback, copy
+from typing import Optional, Tuple, Type, Union, Literal
+import traceback, copy, os
 from statistics import mean
 from more_itertools.recipes import flatten
 import torch
 from torch.utils.data.dataloader import DataLoader
 from transformers import BertTokenizerFast, CamembertTokenizerFast  # type: ignore
 from tqdm import tqdm
-from tibert import (
+from tibert.bertcoref import (
     BertForCoreferenceResolution,
     CamembertForCoreferenceResolution,
     CoreferenceDataset,
-    split_coreference_document,
     DataCollatorForSpanClassification,
-    score_coref_predictions,
-    score_mention_detection,
 )
+from tibert.score import score_coref_predictions, score_mention_detection
 from tibert.predict import predict_coref
-from tibert.utils import gpu_memory_usage
+from tibert.utils import gpu_memory_usage, split_coreference_document
+
+
+def _save_train_checkpoint(
+    path: str,
+    model: Union[BertForCoreferenceResolution, CamembertForCoreferenceResolution],
+    epoch: int,
+    optimizer: torch.optim.AdamW,
+    bert_lr: float,
+    task_lr: float,
+):
+    checkpoint = {
+        "model": model.state_dict(),
+        "model_config": vars(model.config),
+        "epoch": epoch,
+        "optimizer": optimizer.state_dict(),
+        "bert_lr": bert_lr,
+        "task_lr": task_lr,
+    }
+    torch.save(checkpoint, path)
+
+
+def load_train_checkpoint(
+    checkpoint_path: str,
+    model_class: Union[
+        Type[BertForCoreferenceResolution], Type[CamembertForCoreferenceResolution]
+    ],
+) -> Tuple[
+    Union[BertForCoreferenceResolution, CamembertForCoreferenceResolution],
+    torch.optim.AdamW,
+]:
+    config_class = model_class.config_class
+
+    checkpoint = torch.load(checkpoint_path)
+
+    model_config = config_class(**checkpoint["model_config"])
+    model = model_class(model_config)
+    model.load_state_dict(checkpoint["model"])
+
+    optimizer = torch.optim.AdamW(
+        [
+            {"params": model.bert_parameters(), "lr": checkpoint["bert_lr"]},
+            {
+                "params": model.task_parameters(),
+                "lr": checkpoint["task_lr"],
+            },
+        ],
+        lr=checkpoint["task_lr"],
+    )
+    optimizer.load_state_dict(checkpoint["optimizer"])
+
+    return model, optimizer
 
 
 def train_coref_model(
@@ -28,14 +77,41 @@ def train_coref_model(
     sents_per_documents_train: int = 11,
     bert_lr: float = 1e-5,
     task_lr: float = 2e-4,
-    model_save_path: Optional[str] = None,
+    model_save_dir: Optional[str] = None,
     device_str: Literal["cpu", "cuda", "auto"] = "auto",
     _run: Optional["sacred.run.Run"] = None,
+    optimizer: Optional[torch.optim.AdamW] = None,
 ) -> BertForCoreferenceResolution:
+    """
+    :param model: model to train
+    :param dataset: dataset to train on.  90% of that dataset will be
+        used for training, 10% for testing
+    :param tokenizer: tokenizer associated with ``model``
+    :param batch_size: batch_size during training and testing
+    :param epochs_nb: number of epochs to train for
+    :param sents_per_documents_train: max number of sentences in each
+        train document
+    :param bert_lr: learning rate of the BERT encoder
+    :param task_lr: learning rate for other parts of the network
+    :param model_save_dir: directory in which to save the final model
+        (under 'model') and checkpoints ('checkpoint.pth')
+    :param device_str:
+    :param _run: sacred run, used to log metrics
+    :param optimizer: a torch optimizer to use.  Can be useful to
+        resume training.
+
+    :return: the best trained model, according to CoNLL-F1 on the test
+             set
+    """
+    # Get torch device and send model to it
+    # -------------------------------------
     if device_str == "auto":
         device_str = "cuda" if torch.cuda.is_available() else "cpu"
     device = torch.device(device_str)
+    model = model.to(device)
 
+    # Prepare datasets
+    # ----------------
     train_dataset = CoreferenceDataset(
         dataset.documents[: int(0.9 * len(dataset))],
         dataset.tokenizer,
@@ -72,23 +148,28 @@ def train_coref_model(
         train_dataset, batch_size=batch_size, shuffle=True, collate_fn=data_collator
     )
 
-    optimizer = torch.optim.AdamW(
-        [
-            {"params": model.bert_parameters(), "lr": bert_lr},
-            {
-                "params": model.task_parameters(),
-                "lr": task_lr,
-            },
-        ],
-        lr=task_lr,
-    )
+    # Optimizer initialization
+    # ------------------------
+    if optimizer is None:
+        optimizer = torch.optim.AdamW(
+            [
+                {"params": model.bert_parameters(), "lr": bert_lr},
+                {
+                    "params": model.task_parameters(),
+                    "lr": task_lr,
+                },
+            ],
+            lr=task_lr,
+        )
 
+    # Best model saving
+    # -----------------
     best_f1 = 0
     best_model = model
 
-    model = model.to(device)
-
-    for _ in range(epochs_nb):
+    # Training loop
+    # -------------
+    for epoch_i in range(epochs_nb):
         model = model.train()
 
         epoch_losses = []
@@ -158,10 +239,22 @@ def train_coref_model(
             f"mention detection metrics: (precision: {m_precision}, recall: {m_recall}, f1: {m_f1})"
         )
 
+        # Model saving
+        # ------------
+        if not model_save_dir is None:
+            os.makedirs(model_save_dir, exist_ok=True)
+            _save_train_checkpoint(
+                os.path.join(model_save_dir, "checkpoint.pth"),
+                model,
+                epoch_i,
+                optimizer,
+                bert_lr,
+                task_lr,
+            )
         if conll_f1 > best_f1 or best_f1 == 0:
             best_model = copy.deepcopy(model).to("cpu")
-            if not model_save_path is None:
-                best_model.save_pretrained(model_save_path)
             best_f1 = conll_f1
+            if not model_save_dir is None:
+                model.save_pretrained(os.path.join(model_save_dir, "model"))
 
     return best_model
