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

Author: Aethor

README.md

@@ -120,6 +120,7 @@ The following parameters can be set (taken from `./tibert/run_train.py` config f
 | `segment_size`               | `128`               |
 | `encoder`                    | `"bert-base-cased"` |
 | `out_model_dir`              | `"~/tibert/model"`  |
+| `checkpoint`                 | `None`              |
 
 
 One can monitor training metrics by adding run observers using command line flags - see `sacred` documentation for more details.

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tibert"
-version = "0.2.0"
+version = "0.2.2"
 description = "BERT for Coreference Resolution"
 authors = ["Arthur Amalvy <[email protected]>"]
 license = "GPL-3.0-only"

tests/test_train.py

@@ -0,0 +1,40 @@
+import tempfile
+import torch
+from torch.optim import optimizer
+from transformers import BertTokenizerFast
+from tibert.bertcoref import BertForCoreferenceResolutionConfig
+from tibert.train import _save_train_checkpoint, load_train_checkpoint
+from tibert import BertForCoreferenceResolution, predict_coref_simple
+
+
+def test_save_load_checkpoint():
+    model = BertForCoreferenceResolution(BertForCoreferenceResolutionConfig())
+    tokenizer = BertTokenizerFast.from_pretrained("bert-base-cased")
+
+    bert_lr = 1e-5
+    task_lr = 2e-4
+    optimizer = torch.optim.AdamW(
+        [
+            {"params": model.bert_parameters(), "lr": bert_lr},
+            {
+                "params": model.task_parameters(),
+                "lr": task_lr,
+            },
+        ],
+        lr=task_lr,
+    )
+
+    text = "Sli did not want the earpods. He didn't like them."
+    before_pred = predict_coref_simple(text, model, tokenizer)
+
+    with tempfile.TemporaryDirectory() as d:
+        checkpoint_f = f"{d}/checkpoint.pth"
+        _save_train_checkpoint(checkpoint_f, model, 1, optimizer, bert_lr, task_lr)
+        model, new_optimizer = load_train_checkpoint(
+            checkpoint_f, BertForCoreferenceResolution
+        )
+
+    assert new_optimizer.state_dict() == optimizer.state_dict()
+
+    after_pred = predict_coref_simple(text, model, tokenizer)
+    assert before_pred == after_pred

tibert/bertcoref.py

@@ -19,6 +19,7 @@
 from transformers.models.camembert.modeling_camembert import CamembertModel
 from transformers.models.camembert.configuration_camembert import CamembertConfig
 from transformers.tokenization_utils_base import BatchEncoding, PreTrainedTokenizerBase
+from transformers.utils import logging as transformers_logging
 from tqdm import tqdm
 from tibert.utils import spans_indexs, batch_index_select, spans
 
@@ -131,7 +132,14 @@ def prepared_document(
         """
         # (silly) exemple for the tokens ["I", "am", "PG"]
         # a BertTokenizer would produce ["[CLS]", "I", "am", "P", "##G", "[SEP]"]
-        batch = tokenizer(self.tokens, is_split_into_words=True, truncation=True)  # type: ignore
+        # NOTE: we disable tokenizer warning to avoid a length
+        # ----  warning. Usually, sequences should be truncated to a max
+        #       length (512 for BERT). However, in our case, the sequence is
+        #       later cut into segments of configurable size, so this does
+        #       not apply (see BertForCoreferenceResolutionConfig.segment_size)
+        transformers_logging.set_verbosity_error()
+        batch = tokenizer(self.tokens, is_split_into_words=True)
+        transformers_logging.set_verbosity_info()
         tokens = tokenizer.convert_ids_to_tokens(batch["input_ids"])  # type: ignore
 
         # words_ids is used to correspond post-tokenization word pieces

tibert/train.py

@@ -51,7 +51,7 @@ def load_train_checkpoint(
 
     model_config = config_class(**checkpoint["model_config"])
     model = model_class(model_config)
-    model.load_state_dict(checkpoint["model"])
+    model.load_state_dict(checkpoint["model"], strict=False)
 
     optimizer = torch.optim.AdamW(
         [
@@ -68,6 +68,17 @@ def load_train_checkpoint(
     return model, optimizer
 
 
+def _optimizer_to_(
+    optimizer: torch.optim.AdamW, device: torch.device
+) -> torch.optim.AdamW:
+    """From https://github.com/pytorch/pytorch/issues/2830"""
+    for state in optimizer.state.values():
+        for k, v in state.items():
+            if isinstance(v, torch.Tensor):
+                state[k] = v.to(device)
+    return optimizer
+
+
 def train_coref_model(
     model: Union[BertForCoreferenceResolution, CamembertForCoreferenceResolution],
     dataset: CoreferenceDataset,
@@ -161,8 +172,8 @@ def train_coref_model(
             ],
             lr=task_lr,
         )
+    optimizer = _optimizer_to_(optimizer, device)
 
-    # Best model saving
     # -----------------
     best_f1 = 0
     best_model = model