train & eval scripts

.gitignore

@@ -167,3 +167,6 @@ local_notebooks
 
 # test caches
 tests/testdata/*/*/cache*
+
+# tensorboard logs
+logs/

config/topic_classifier_training.yml

@@ -0,0 +1,18 @@
+data_root: /path/to/data
+
+training_args:
+  output_dir: my_topic_classifier
+  learning_rate: 0.00002
+  per_device_train_batch_size: 32
+  per_device_eval_batch_size: 32
+  num_train_epochs: 2
+  weight_decay: 0.01
+  eval_strategy: "steps"
+  save_strategy: "steps"
+  eval_steps: 500
+  save_steps: 500
+  logging_steps: 50
+  load_best_model_at_end: True
+  eval_on_start: True
+  report_to: "tensorboard"
+  seed: 42

pyproject.toml

@@ -37,14 +37,18 @@ dependencies = [
   "torch",
   "torcheval",
   "pillow",
-  "unbabel-comet"
+  "unbabel-comet",
+  "accelerate",
+  "tensorboard",
+  "scikit-learn"
 ]
 
 [project.optional-dependencies]
 dev = [
   "pytest >=6",
   "pytest-cov >=3",
   "pre-commit",
+  "mypy",
 ]
 
 [project.urls]

scripts/README.md

@@ -0,0 +1,31 @@
+# Scripts
+
+## create_test_ds.py
+
+Create dataset used in testing. This has been run and output stored in GitHub.
+
+## train_topic_classifier.py
+
+Train a topic classifier over MultiEURLEX data.
+
+Sample config for training stored at [../config/topic_classifier_training.yaml](../config/topic_classifier_training.yml). Edit this first.
+
+run with (from project root):
+```bash
+python scripts/train_topic_classifier.py --config config/topic_classifier_training.yml
+```
+
+## eval_topic_classifier.py
+
+Evaluate a trained topic classifier over MultiEURLEX data.
+
+e.g. from project root:
+```bash
+python scripts/eval_topic_classifier.py \
+    --ckpt_path ./path/to/output/checkpoint-3286 \
+    --data_root ./data \
+    --batch_size 32 \
+    --eval_output_dir ./path/to/output \
+    --report_to tensorboard \
+    --dataset_name validation # change to "test" for test set evaluation
+```

scripts/eval_topic_classifier.py

@@ -0,0 +1,69 @@
+from jsonargparse import cli
+from transformers import (
+    DataCollatorWithPadding,
+    DistilBertForSequenceClassification,
+    DistilBertTokenizer,
+    Trainer,
+    TrainingArguments,
+)
+
+from arc_spice.data.multieurlex_utils import load_multieurlex
+from arc_spice.topic.utils import compute_metrics, preprocess_function
+
+
+def main(
+    ckpt_path: str,
+    data_root: str,
+    batch_size: int,
+    eval_output_dir: str,
+    report_to: str = "tensorboard",
+    dataset_name: str = "validation",
+) -> None:
+    model = DistilBertForSequenceClassification.from_pretrained(ckpt_path)
+    tokenizer = DistilBertTokenizer.from_pretrained(ckpt_path)
+
+    if dataset_name not in ["validation", "test"]:
+        msg = (
+            f"dataset_name should be one of 'validation' or 'test', got {dataset_name}"
+        )
+        raise ValueError(msg)
+
+    dataset_dict, metadata = load_multieurlex(
+        data_dir=data_root, level=1, languages=["en"], split=dataset_name
+    )
+    id2label = dict(enumerate(metadata["class_labels"]))
+
+    dataset = dataset_dict[dataset_name]
+
+    tokenized_dataset = dataset.map(
+        preprocess_function,
+        batched=True,
+        fn_kwargs={"tokenizer": tokenizer, "id2label": id2label},
+        remove_columns=["text", "labels"],
+    )
+
+    model = model.eval()
+
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+
+    training_args = TrainingArguments(
+        output_dir=eval_output_dir,
+        per_device_eval_batch_size=batch_size,
+        report_to=report_to,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=None,
+        eval_dataset=tokenized_dataset,
+        processing_class=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics,
+    )
+
+    trainer.evaluate()
+
+
+if __name__ == "__main__":
+    cli.CLI(main, as_positional=False)

scripts/train_topic_classifier.py

@@ -0,0 +1,59 @@
+# finetune DistilBERT for topic classification over MultiEURLEX dataset
+
+
+from jsonargparse import cli
+from transformers import (
+    DataCollatorWithPadding,
+    DistilBertForSequenceClassification,
+    DistilBertTokenizer,
+    Trainer,
+    TrainingArguments,
+)
+
+from arc_spice.data.multieurlex_utils import load_multieurlex
+from arc_spice.topic.utils import compute_metrics, preprocess_function
+
+
+def main(data_root: str, training_args: dict):
+    dataset, metadata = load_multieurlex(data_dir=data_root, level=1, languages=["en"])
+
+    id2label = dict(enumerate(metadata["class_labels"]))
+    label2id = {v: k for k, v in id2label.items()}
+
+    model = DistilBertForSequenceClassification.from_pretrained(
+        "distilbert/distilbert-base-uncased",
+        num_labels=len(id2label),
+        id2label=id2label,
+        label2id=label2id,
+        problem_type="multi_label_classification",
+    )
+    tokenizer = DistilBertTokenizer.from_pretrained(
+        "distilbert/distilbert-base-uncased"
+    )
+
+    tokenized_dataset = dataset.map(
+        preprocess_function,
+        batched=True,
+        fn_kwargs={"tokenizer": tokenizer, "id2label": id2label},
+        remove_columns=["text", "labels"],
+    )
+
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+
+    training_args = TrainingArguments(**training_args)
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset["train"],
+        eval_dataset=tokenized_dataset["validation"],
+        processing_class=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics,
+    )
+
+    trainer.train()
+
+
+if __name__ == "__main__":
+    cli.CLI(main)

src/arc_spice/data/multieurlex_utils.py

@@ -2,7 +2,7 @@
 from typing import Any
 
 import torch
-from datasets import DatasetDict, load_dataset
+from datasets import Dataset, DatasetDict, load_dataset
 from datasets.formatting.formatting import LazyRow
 from torch.nn.functional import one_hot
 
@@ -120,6 +120,7 @@ def load_multieurlex(
     level: int,
     languages: list[str],
     drop_empty: bool = True,
+    split: str | None = None,
 ) -> tuple[DatasetDict, dict[str, Any]]:
     """
     load the multieurlex dataset
@@ -130,7 +131,7 @@ def load_multieurlex(
         languages: a list of iso codes for languages to be used
 
     Returns:
-        List of datasets and a dictionary with some metadata information
+        dataset dict and metdata.
     """
     metadata = load_mutlieurlex_metadata(data_dir=data_dir, level=level)
 
@@ -147,7 +148,19 @@ def load_multieurlex(
         load_langs,
         label_level=f"level_{level}",
         trust_remote_code=True,
+        split=split,
     )
+    # ensure we always return dataset dict even if only single split
+    if split is not None:
+        if not isinstance(dataset_dict, Dataset):
+            msg = (
+                "Error. load_dataset should return a Dataset object if split specified"
+            )
+            raise ValueError(msg)
+
+        tmp = DatasetDict()
+        tmp[split] = dataset_dict
+        dataset_dict = tmp
 
     dataset_dict = dataset_dict.map(
         extract_articles,

src/arc_spice/topic/utils.py

@@ -0,0 +1,22 @@
+import torch
+from sklearn.metrics import hamming_loss, zero_one_loss
+from transformers import DistilBertTokenizer
+
+
+def preprocess_function(
+    examples, tokenizer: DistilBertTokenizer, id2label: dict[str, int]
+):
+    output = tokenizer(examples["text"], truncation=True, padding=True)
+    labels = torch.zeros((len(examples["labels"]), len(id2label)), dtype=torch.float32)
+    for i, label in enumerate(examples["labels"]):
+        labels[i][label] = 1.0
+    output["label"] = labels
+    return output
+
+
+def compute_metrics(eval_preds):
+    logits, labels = eval_preds
+    preds = (torch.sigmoid(torch.from_numpy(logits)) > 0.5).numpy().astype(int)
+    zo_acc = 1 - zero_one_loss(labels, preds)
+    ham_acc = 1 - hamming_loss(labels, preds)
+    return {"zero_one_acc": zo_acc, "hamming_acc": ham_acc}