Add the possibility to specifiy train and test dataset for training

Author: Aethor

tibert/bertcoref.py

@@ -21,7 +21,12 @@
 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
+from tibert.utils import (
+    spans_indexs,
+    batch_index_select,
+    spans,
+    split_coreference_document,
+)
 
 
 @dataclass
@@ -579,6 +584,21 @@ def merged_datasets(datasets: List[CoreferenceDataset]) -> CoreferenceDataset:
             datasets[0].max_span_size,
         )
 
+    def splitted(self, ratio: float) -> Tuple[CoreferenceDataset, CoreferenceDataset]:
+        first_docs = self.documents[: int(ratio * len(self))]
+        second_docs = self.documents[int(ratio * len(self)) :]
+        return (
+            CoreferenceDataset(first_docs, self.tokenizer, self.max_span_size),
+            CoreferenceDataset(second_docs, self.tokenizer, self.max_span_size),
+        )
+
+    def limit_doc_size_(self, sents_nb: int):
+        self.documents = list(
+            flatten(
+                [split_coreference_document(doc, sents_nb) for doc in self.documents]
+            )
+        )
+
     def __len__(self) -> int:
         return len(self.documents)
 

tibert/run_train.py

@@ -13,6 +13,7 @@
     train_coref_model,
     load_train_checkpoint,
 )
+from tibert.bertcoref import CoreferenceDataset
 
 ex = Experiment()
 
@@ -99,15 +100,20 @@ def main(
 
     tokenizer = config["tokenizer_class"].from_pretrained(encoder)
 
-    dataset = config["loading_function"](dataset_path, tokenizer, max_span_size)
+    dataset: CoreferenceDataset = config["loading_function"](
+        dataset_path, tokenizer, max_span_size
+    )
+    train_dataset, test_dataset = dataset.splitted(0.9)
+    train_dataset.limit_doc_size_(sents_per_documents_train)
+    test_dataset.limit_doc_size_(11)
 
     train_coref_model(
         model,
-        dataset,
+        train_dataset,
+        test_dataset,
         tokenizer,
         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,

tibert/train.py

@@ -1,7 +1,6 @@
 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
@@ -14,7 +13,7 @@
 )
 from tibert.score import score_coref_predictions, score_mention_detection
 from tibert.predict import predict_coref
-from tibert.utils import gpu_memory_usage, split_coreference_document
+from tibert.utils import gpu_memory_usage
 
 
 def _save_train_checkpoint(
@@ -81,11 +80,11 @@ def _optimizer_to_(
 
 def train_coref_model(
     model: Union[BertForCoreferenceResolution, CamembertForCoreferenceResolution],
-    dataset: CoreferenceDataset,
+    train_dataset: CoreferenceDataset,
+    test_dataset: CoreferenceDataset,
     tokenizer: Union[BertTokenizerFast, CamembertTokenizerFast],
     batch_size: int = 1,
     epochs_nb: int = 30,
-    sents_per_documents_train: int = 11,
     bert_lr: float = 1e-5,
     task_lr: float = 2e-4,
     model_save_dir: Optional[str] = None,
@@ -121,37 +120,6 @@ def train_coref_model(
     device = torch.device(device_str)
     model = model.to(device)
 
-    # Prepare datasets
-    # ----------------
-    train_dataset = CoreferenceDataset(
-        dataset.documents[: int(0.9 * len(dataset))],
-        dataset.tokenizer,
-        dataset.max_span_size,
-    )
-    train_dataset.documents = list(
-        flatten(
-            [
-                split_coreference_document(doc, sents_per_documents_train)
-                for doc in train_dataset.documents
-            ]
-        )
-    )
-
-    test_dataset = CoreferenceDataset(
-        dataset.documents[int(0.9 * len(dataset)) :],
-        dataset.tokenizer,
-        dataset.max_span_size,
-    )
-    test_dataset.documents = list(
-        flatten(
-            [
-                # HACK: test on full documents
-                split_coreference_document(doc, 11)
-                for doc in test_dataset.documents
-            ]
-        )
-    )
-
     data_collator = DataCollatorForSpanClassification(
         tokenizer, model.config.max_span_size
     )