Repository: aai-institute/nnbench

Issue:
Add example: Implement benchmarking on saved artifacts
- [ ] Performance diffing by materializing older versions of a model and then comparing accuracy on a different evaluation set (needs design of artifact sourcing, download, materialization).

Discussion:
maxmynter: > We could define a functools.partial alias for the parametrization decorator with all the class labels to refactor the raw injections a little bit.

Agree

> Could you share a screenshot on how the results look in a table?

There you go :) 

<img width="1120" alt="Screenshot 2024-03-06 at 12 37 01" src="https://github.com/aai-institute/nnbench/assets/32773644/b9dec649-8e68-4d7d-9143-6245232998c3">

maxmynter: Using `functools.partial` does not work straightforward. The collection then does not recognize the parametrized benchmarks anymore. Therefore, this is out of scope for this PR.

Patch:
diff --git a/docs/tutorials/artifact_benchmarking.md b/docs/tutorials/artifact_benchmarking.md
new file mode 100644
index 0000000..ae86097
--- /dev/null
+++ b/docs/tutorials/artifact_benchmarking.md
@@ -0,0 +1,115 @@
+# Benchmark on saved models
+There is a high likelihood that you, at some point, find yourself wanting to benchmark models that were trained previously.
+In this guide we will walk through how we can accomplish this with nnbench.
+
+## Example: Named Entity Recognition
+We will start with an aside that talks through the setup of the example we will use in this guide.
+If you are only interested in the application of nnbench, you can skip this section.
+
+There are lots of reasons why you could want to retrieve saved models for benchmarking. 
+Among them these are reviewing the work of colleagues, comparing experimental performances to an existing benchmark, or dealing with models that require significant compute such that in-place retraining is impractical.
+For this example, we deal with a named entity recognition (NER) model that is based on the pre-trained encoder-decoder transformer [BERT](https://arxiv.org/abs/1810.04805).
+The model is trained on the [CoNLLpp dataset](https://huggingface.co/datasets/conllpp) which consists of sentences from news stories where words were tagged with Person, Organization, Location, or Miscellaneous if they referred to entities. 
+Words are assigned an out-of-entity label if they do not represent an entity.
+
+### Model Training
+You find the code to train the model in the nnbench [repository](https://github.com/aai-institute/nnbench) in the directory `examples/artifact_benchmarking/src/training/training.py`.
+If you want to skip running the training script but still want to reproduce this example, you can take any BERT model fine tuned for NER with the CoNLL dataset family.
+You find many on the Huggingface model hub, for example [this one](https://huggingface.co/dslim/bert-base-NER). You need to download the `model.safetensors`, `config.json`, `tokenizer_config.json`, and `tokenizer.json` files.
+If you want to train your own model, continue below. 
+
+There is some necessary preprocessing and data wrangling to train the model. 
+We will not go into the details here. But if you are interested in a more thorough walkthrough, look into this [resource](https://huggingface.co/learn/nlp-course/chapter7/2?fw=pt) by Huggingface which served as the basis for this example. 
+It is not feasible to train the model on a CPU. If you do not have access to a GPU you can use free GPU instances on [Google Colab](https://colab.research.google.com/).
+When you open a new notebook there make sure to select a GPU instance in the upper right corner.
+The you can upload the `training.py`.
+You can ignore An eventual warning that the data is not persisted.
+Next, install the necessary dependencies: `!pip install datasets transformers[torch]`.
+Google Colab comes with some dependencies already installed in the environment.
+Hence, if you are working with a different GPU instance, make sure to install everything from the `pyproject.toml` in the `examples/artifact_benchmarking` folder. 
+Next you can execute the `training.py` with `!python training.py`.
+This will train two BERT models ("bert-base-uncased" and "distilbert-base-uncased") which we can compare using nnbench. 
+If you want, you can adapt the training script to train other models by editing the tuples in the `tokenizers_and_models` list at the bottom of the training script. 
+The training of the models takes around 10 minutes.
+Once it is done, download the respective files and save them to your disk.
+They should be the same mentioned above. 
+We will need the path to the files for benchmarking later.
+
+### The Benchmarks
+The benchmarking code is found in the `examples/artifact_benchmarking/benchmark.py`.
+We calculate precision, recall, accuracy, and f1 scores for the whole test set and specific labels.
+Additionally, we obtain information about the model such as its memory footprint and inference time.
+We are not walking through the whole file but instead point out certain design choices as an inspiration to you. 
+If you are interested in a more detailed walkthrough on how to set up benchmarks, you can find it [here](../guides/benchmarks.md).
+Notable design choices in this benchmark are that we factored out the evaluation loop as it is necessary for all evaluation metrics. We cache it using the `functools.lru_cache` decorator so the evaluation loop runs only once per benchmark run instead of once per metric which greatly reduces runtime.
+We use `nnbench.parametrize` to get the per-class metrics. 
+
+## Running Benchmarks with saved Artifacts
+Now that we have explained the example, let's jump into the benchmarking code.
+You find it in the nnbench repository in `examples/artifact_benchmarking/src/runner.py`.
+
+The benchmarking code is written to be executed as a script that consumes any number of file paths to models to benchmark as arguments, `python runner.py /path/to/model1 /path/to/model2`.
+The parsing of the arguments is handled by the last lines in the script which then calls the `main()` function:
+
+```python
+--8<-- "examples/artifact_benchmarking/runner.py:92:96"
+```
+
+### Artifact Classes
+The `main()` function first sets up the benchmark reporter and the runner. 
+Then we create a list of models. These models are instances of a `TokenClassificationModel`, a custom class we implemented which inherits from the `Artifact` base class.
+
+```python
+--8<-- "examples/artifact_benchmarking/runner.py:30:37"
+```
+
+The `Artifact` class is a typesafe wrapper around serialized data of any kind.
+It allows for lazy deserialization of artifacts from a  path attribute.
+This attribute is set by the `ArtifactLoader` (which we will cover in a moment) that supplied a path to the local artifact disk storage. 
+In our derived class, we have to override the `deserialize()` method to properly load the artifact value into memory.
+
+```python
+--8<-- "examples/artifact_benchmarking/runner.py:59"
+```
+
+The `deserialize()` method has to set the `self._value` attribute to the value we want to access later.
+In this case, we assign it a tuple containing the Huggingface Transformers model and tokenizer.
+
+We do similar with the CoNLLpp dataset.
+
+```python
+--8<-- "examples/artifact_benchmarking/runner.py:30:37"
+```
+
+In this case, we store the `datasets.Dataset` object as well ass a dictionary which maps the label id to a semantic string in the `_value` attribute. 
+The value that we store in the `_value` of an artifact can be of any kind and that we use tuples in both instances here is a circumstance.
+
+### Artifact Loaders
+Upon instantiation of an `Artifact` or derived classes we need to supply an `ArtifactLoader` (or a class derived from it). `ArtifactLoader`s are classes that implement a `load()` method that resolves to a path-like object or string which points to the local storage location of the artifact. This method is used by the Artifact class. 
+
+For our models, we use the provided `LocalArtifactLoader` which consumes a path and passes it on later.
+
+```python
+--8<-- "examples/artifact_benchmarking/runner.py:52:54"
+```
+
+We have a little more logic with respect to the dataset as we handle the train test split as well.
+
+```python
+--8<-- "examples/artifact_benchmarking/runner.py:19:27"
+```
+
+### Artifact Collections
+The `ArtifactCollection` object is a convenient wrapper around `Artifact`s to allow for iteration and abstract away the `deserialize()` call. 
+We can instantiate them by passing an iterable of `Artifact`s.
+
+```python
+--8<-- "examples/artifact_benchmarking/runner.py:56"
+```
+
+We can then either iterate over the `ArtifactCollection` directly, using the lazy deserialization, or we use the `.values()` method to return whatever we assigned to the `._values` attribute of the `Artifact`s.
+We then execute the benchmark in the loop for the different models we have supplied to the `ArtifactCollection`.
+
+```python
+--8<-- "examples/artifact_benchmarking/runner.py:61:89"
+```
diff --git a/examples/artifact_benchmarking/pyproject.toml b/examples/artifact_benchmarking/pyproject.toml
new file mode 100644
index 0000000..0ddb1b9
--- /dev/null
+++ b/examples/artifact_benchmarking/pyproject.toml
@@ -0,0 +1,19 @@
+[build-system]
+requires = ["setuptools>=45", "setuptools-scm[toml]>=7.1"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "nnbench-saved-artifacts"
+requires-python = ">= 3.8"
+version = "0.1.0"
+description = "Integration example of Prefect with nnbench."
+readme = "docs/guides/prefect.md"
+license = { text = "Apache-2.0" }
+dependencies = [
+    "transformers",
+    "datasets",
+    "torch",
+    "nnbench@git+https://github.com/aai-institute/nnbench.git"
+]
+maintainers = [{name="Max Mynter", email="m.mynter@appliedai-institute.de"}]
+authors = [{ name = "appliedAI Initiative", email = "info+oss@appliedai.de" }]
diff --git a/examples/artifact_benchmarking/requirements.txt b/examples/artifact_benchmarking/requirements.txt
new file mode 100644
index 0000000..c3ea68b
--- /dev/null
+++ b/examples/artifact_benchmarking/requirements.txt
@@ -0,0 +1,48 @@
+#
+# This file is autogenerated by pip-compile with Python 3.11
+# by the following command:
+#
+#    pip-compile --extra=dev --no-annotate
+#
+aiohttp==3.9.3
+aiosignal==1.3.1
+attrs==23.2.0
+certifi==2024.2.2
+charset-normalizer==3.3.2
+datasets==2.18.0
+dill==0.3.8
+filelock==3.13.1
+frozenlist==1.4.1
+fsspec[http]==2024.2.0
+huggingface-hub==0.21.3
+idna==3.6
+jinja2==3.1.3
+markupsafe==2.1.5
+mpmath==1.3.0
+multidict==6.0.5
+multiprocess==0.70.16
+networkx==3.2.1
+nnbench @ git+https://github.com/aai-institute/nnbench.git
+numpy==1.26.4
+packaging==23.2
+pandas==2.2.1
+pyarrow==15.0.0
+pyarrow-hotfix==0.6
+python-dateutil==2.9.0.post0
+pytz==2024.1
+pyyaml==6.0.1
+regex==2023.12.25
+requests==2.31.0
+safetensors==0.4.2
+six==1.16.0
+sympy==1.12
+tabulate==0.9.0
+tokenizers==0.15.2
+torch==2.2.1
+tqdm==4.66.2
+transformers==4.38.2
+typing-extensions==4.10.0
+tzdata==2024.1
+urllib3==2.2.1
+xxhash==3.4.1
+yarl==1.9.4
diff --git a/examples/artifact_benchmarking/src/__init__.py b/examples/artifact_benchmarking/src/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/examples/artifact_benchmarking/src/benchmark.py b/examples/artifact_benchmarking/src/benchmark.py
new file mode 100644
index 0000000..27904b8
--- /dev/null
+++ b/examples/artifact_benchmarking/src/benchmark.py
@@ -0,0 +1,226 @@
+import os
+import tempfile
+import time
+from functools import lru_cache
+
+import torch
+from torch.nn import Module
+from torch.utils.data import DataLoader
+
+import nnbench
+
+
+@lru_cache
+def run_eval_loop(model, dataloader, padding_id=-100):
+    true_positives = torch.zeros(model.config.num_labels)
+    false_positives = torch.zeros(model.config.num_labels)
+    true_negatives = torch.zeros(model.config.num_labels)
+    false_negatives = torch.zeros(model.config.num_labels)
+
+    model.eval()
+    with torch.no_grad():
+        for batch in dataloader:
+            outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            labels = batch["labels"]
+
+            valid_indices = labels.view(-1) != padding_id
+            predictions = predictions.view(-1)[valid_indices]
+            labels = labels.view(-1)[valid_indices]
+            for idx in range(model.config.num_labels):
+                tp = ((predictions == idx) & (labels == idx)).sum()
+                fp = ((predictions == idx) & (labels != idx)).sum()
+                fn = ((predictions != idx) & (labels == idx)).sum()
+                tn = ((predictions != idx) & (labels != idx)).sum()
+
+                true_positives[idx] += tp
+                false_positives[idx] += fp
+                false_negatives[idx] += fn
+                true_negatives[idx] += tn
+
+    return true_positives, false_positives, true_negatives, false_negatives
+
+
+@nnbench.benchmark(tags=("metric", "aggregate"))
+def precision(model: Module, test_dataloader: DataLoader, padding_id: int = -100) -> float:
+    tp, fp, tn, fn = run_eval_loop(model, test_dataloader, padding_id)
+    precision = tp / (tp + fp + 1e-6)
+    return torch.mean(precision).item()
+
+
+@nnbench.parametrize(
+    (
+        {"class_label": "O"},
+        {"class_label": "B-PER"},
+        {"class_label": "I-PER"},
+        {"class_label": "B-ORG"},
+        {"class_label": "I-ORG"},
+        {"class_label": "B-LOC"},
+        {"class_label": "I-LOC"},
+        {"class_label": "B-MISC"},
+        {"class_label": "I-MISC"},
+    ),
+    tags=("metric", "per-class"),
+)
+def precision_per_class(
+    class_label: str,
+    model: Module,
+    test_dataloader: DataLoader,
+    index_label_mapping: dict[int, str],
+    padding_id: int = -100,
+) -> float:
+    tp, fp, tn, fn = run_eval_loop(model, test_dataloader, padding_id)
+    precision_values = tp / (tp + fp + 1e-6)
+    for idx, lbl in index_label_mapping.items():
+        if lbl == class_label:
+            return precision_values[idx]
+    raise ValueError(f" Key {class_label} not in test labels")
+
+
+@nnbench.benchmark(tags=("metric", "aggregate"))
+def recall(model: Module, test_dataloader: DataLoader, padding_id: int = -100) -> float:
+    tp, fp, tn, fn = run_eval_loop(model, test_dataloader, padding_id)
+    recall = tp / (tp + fn + 1e-6)
+    return torch.mean(recall).item()
+
+
+@nnbench.parametrize(
+    (
+        {"class_label": "O"},
+        {"class_label": "B-PER"},
+        {"class_label": "I-PER"},
+        {"class_label": "B-ORG"},
+        {"class_label": "I-ORG"},
+        {"class_label": "B-LOC"},
+        {"class_label": "I-LOC"},
+        {"class_label": "B-MISC"},
+        {"class_label": "I-MISC"},
+    ),
+    tags=("metric", "per-class"),
+)
+def recall_per_class(
+    class_label: str,
+    model: Module,
+    test_dataloader: DataLoader,
+    index_label_mapping: dict[int, str],
+    padding_id: int = -100,
+) -> float:
+    tp, fp, tn, fn = run_eval_loop(model, test_dataloader, padding_id)
+    recall_values = tp / (tp + fn + 1e-6)
+    for idx, lbl in index_label_mapping.items():
+        if lbl == class_label:
+            return recall_values[idx]
+    raise ValueError(f" Key {class_label} not in test labels")
+
+
+@nnbench.benchmark(tags=("metric", "aggregate"))
+def f1(model: Module, test_dataloader: DataLoader, padding_id: int = -100) -> float:
+    tp, fp, tn, fn = run_eval_loop(model, test_dataloader, padding_id)
+    precision = tp / (tp + fp + 1e-6)
+    recall = tp / (tp + fn + 1e-6)
+    f1 = 2 * (precision * recall) / (precision + recall + 1e-6)
+    return torch.mean(f1).item()
+
+
+@nnbench.parametrize(
+    (
+        {"class_label": "O"},
+        {"class_label": "B-PER"},
+        {"class_label": "I-PER"},
+        {"class_label": "B-ORG"},
+        {"class_label": "I-ORG"},
+        {"class_label": "B-LOC"},
+        {"class_label": "I-LOC"},
+        {"class_label": "B-MISC"},
+        {"class_label": "I-MISC"},
+    ),
+    tags=("metric", "per-class"),
+)
+def f1_per_class(
+    class_label: str,
+    model: Module,
+    test_dataloader: DataLoader,
+    index_label_mapping: dict[int, str],
+    padding_id: int = -100,
+) -> float:
+    tp, fp, tn, fn = run_eval_loop(model, test_dataloader, padding_id)
+    precision = tp / (tp + fp + 1e-6)
+    recall = tp / (tp + fn + 1e-6)
+    f1_values = 2 * (precision * recall) / (precision + recall + 1e-6)
+    for idx, lbl in index_label_mapping.items():
+        if lbl == class_label:
+            return f1_values[idx]
+    raise ValueError(f" Key {class_label} not in test labels")
+
+
+@nnbench.benchmark(tags=("metric", "aggregate"))
+def accuracy(model: Module, test_dataloader: DataLoader, padding_id: int = -100) -> float:
+    tp, fp, tn, fn = run_eval_loop(model, test_dataloader, padding_id)
+    accuracy = (tp + tn) / (tp + tn + fp + fn + 1e-6)
+    return torch.mean(accuracy).item()
+
+
+@nnbench.parametrize(
+    (
+        {"class_label": "O"},
+        {"class_label": "B-PER"},
+        {"class_label": "I-PER"},
+        {"class_label": "B-ORG"},
+        {"class_label": "I-ORG"},
+        {"class_label": "B-LOC"},
+        {"class_label": "I-LOC"},
+        {"class_label": "B-MISC"},
+        {"class_label": "I-MISC"},
+    ),
+    tags=("metric", "per-class"),
+)
+def accuracy_per_class(
+    class_label: str,
+    model: Module,
+    test_dataloader: DataLoader,
+    index_label_mapping: dict[int, str],
+    padding_id: int = -100,
+) -> dict[str, float]:
+    tp, fp, tn, fn = run_eval_loop(model, test_dataloader, padding_id)
+    accuracy_values = (tp + tn) / (tp + tn + fp + fn + 1e-6)
+    for idx, lbl in index_label_mapping.items():
+        if lbl == class_label:
+            return accuracy_values[idx]
+    raise ValueError(f" Key {class_label} not in test labels")
+
+
+@nnbench.benchmark(tags=("config",))
+def model_configuration(model: Module) -> dict:
+    model.eval()
+    config = model.config.to_dict()
+    return config
+
+
+@nnbench.benchmark(tags=("model-meta", "inference-time"))
+def avg_inference_time_ns(model: Module, test_dataloader: DataLoader, avg_n: int = 100) -> float:
+    start_time = time.perf_counter()
+    model.eval()
+    num_datapoints = 0
+    with torch.no_grad():
+        for batch in test_dataloader:
+            if num_datapoints >= avg_n:
+                break
+            num_datapoints += len(batch)
+            _ = model(**batch)
+    end_time = time.perf_counter()
+
+    total_time = end_time - start_time
+    average_time = total_time / num_datapoints
+    return average_time
+
+
+@nnbench.benchmark(tags=("model-meta", "size-on-disk"))
+def model_size_mb(model: Module) -> float:
+    model.eval()
+    with tempfile.NamedTemporaryFile() as tmp:
+        torch.save(model.state_dict(), tmp.name)
+        tmp.flush()
+        tmp.seek(0, os.SEEK_END)
+        tmp_size = tmp.tell()
+    size_mb = tmp_size / (1024 * 1024)
+    return size_mb
diff --git a/examples/artifact_benchmarking/src/runner.py b/examples/artifact_benchmarking/src/runner.py
new file mode 100644
index 0000000..56fb7a4
--- /dev/null
+++ b/examples/artifact_benchmarking/src/runner.py
@@ -0,0 +1,96 @@
+import argparse
+import os
+from pathlib import Path
+
+from datasets import Dataset, load_dataset
+from torch.utils.data import DataLoader
+from training.training import tokenize_and_align_labels
+from transformers import (
+    AutoModelForTokenClassification,
+    AutoTokenizer,
+    DataCollatorForTokenClassification,
+)
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+
+import nnbench
+from nnbench import context, reporter, types
+
+
+class ConllppLoader(types.ArtifactLoader):
+    def __init__(self, path: str, split: str) -> None:
+        self.path = path
+        self.split = split
+
+    def load(self) -> os.PathLike[str]:
+        dataset = load_dataset(self.path)
+        path = dataset.cache_files[self.split][0]["filename"]
+        return path
+
+
+class TokenClassificationModel(types.Artifact):
+    def deserialize(self) -> None:
+        path = Path(self.path).resolve()
+        model: AutoModelForTokenClassification = AutoModelForTokenClassification.from_pretrained(
+            path, use_safetensors=True
+        )
+        tokenizer: PreTrainedTokenizerBase = AutoTokenizer.from_pretrained(path)
+        self._value = (model, tokenizer)
+
+
+class ConllValidationData(types.Artifact):
+    def deserialize(self) -> None:
+        dataset = Dataset.from_file(self.path)
+        label_names = dataset.features["ner_tags"].feature.names
+        id2label = {i: label for i, label in enumerate(label_names)}
+        self._value = dataset, id2label
+
+
+def main(model_paths: list[str]) -> None:
+    console_reporter = reporter.BenchmarkReporter()
+    runner = nnbench.BenchmarkRunner()
+
+    models = [
+        TokenClassificationModel(loader=types.LocalArtifactLoader(path)) for path in model_paths
+    ]
+
+    model_store = types.ArtifactCollection(*models)
+
+    conllpp = ConllValidationData(ConllppLoader("conllpp", split="validation"))
+    conllpp.deserialize()
+
+    for model, tokenizer in model_store.values():
+        dataset, id2label = conllpp.value
+        tokenized_dataset = dataset.map(
+            lambda examples: tokenize_and_align_labels(tokenizer, examples),
+            batched=True,
+            remove_columns=dataset.column_names,
+        )
+        dataloader = DataLoader(
+            tokenized_dataset,
+            shuffle=False,
+            collate_fn=DataCollatorForTokenClassification(tokenizer, padding=True),
+            batch_size=8,
+        )
+        result = runner.run(
+            "benchmark.py",
+            params={
+                "model": model,
+                "test_dataloader": dataloader,
+                "index_label_mapping": id2label,
+            },
+            context=context.Context.make(
+                {
+                    "model": model.name_or_path,
+                    "dataset": str(dataset),
+                }
+            ),
+            tags=("metric", "model-meta"),
+        )
+        console_reporter.display(result)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run named entity recognition model benchmark.")
+    parser.add_argument("model_paths", type=str, nargs="+", help="Filepaths to pretrained models.")
+    args = parser.parse_args()
+    main(args.model_paths)
diff --git a/examples/artifact_benchmarking/src/training/__init__.py b/examples/artifact_benchmarking/src/training/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/examples/artifact_benchmarking/src/training/training.py b/examples/artifact_benchmarking/src/training/training.py
new file mode 100644
index 0000000..c06c0c7
--- /dev/null
+++ b/examples/artifact_benchmarking/src/training/training.py
@@ -0,0 +1,116 @@
+from typing import Any
+
+import torch
+from datasets import load_dataset
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+from transformers import (
+    AutoModelForTokenClassification,
+    AutoTokenizer,
+    BatchEncoding,
+    DataCollatorForTokenClassification,
+    PreTrainedTokenizer,
+    get_scheduler,
+)
+
+OUTPUT_DIR = "artifacts"
+
+
+def align_labels_with_tokens(labels: list[int], word_ids: list[int]) -> list[int]:
+    new_labels = []
+    current_word = None
+    for word_id in word_ids:
+        if word_id != current_word:
+            # Start of a new word!
+            current_word = word_id
+            label = -100 if word_id is None else labels[word_id]
+            new_labels.append(label)
+        elif word_id is None:
+            # Special token
+            new_labels.append(-100)
+        else:
+            # Same word as previous token
+            label = labels[word_id]
+            # If the label is B-XXX we change it to I-XXX
+            if label % 2 == 1:
+                label += 1
+            new_labels.append(label)
+
+    return new_labels
+
+
+def tokenize_and_align_labels(
+    tokenizer: PreTrainedTokenizer, examples: dict[str, Any]
+) -> BatchEncoding:
+    tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
+
+    labels = []
+    for i, label in enumerate(examples["ner_tags"]):
+        word_ids = tokenized_inputs.word_ids(batch_index=i)
+        labels.append(align_labels_with_tokens(label, word_ids))
+
+    tokenized_inputs["labels"] = labels
+    return tokenized_inputs
+
+
+def train_model(tokenizer_name: str, model_name: str, output_dir: str) -> None:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    model = AutoModelForTokenClassification.from_pretrained(model_name, num_labels=9).to(device)
+
+    dataset = load_dataset("conllpp", split="train")
+    tokenized_datasets = dataset.map(
+        lambda examples: tokenize_and_align_labels(tokenizer, examples),
+        batched=True,
+        remove_columns=dataset.column_names,
+    )
+
+    optimizer = AdamW(model.parameters(), lr=2e-5)
+
+    data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer, padding=True)
+
+    train_dataloader = DataLoader(
+        tokenized_datasets,
+        shuffle=True,
+        collate_fn=data_collator,
+        batch_size=8,
+    )
+
+    num_train_epochs = 3
+    num_update_steps_per_epoch = len(train_dataloader)
+    num_training_steps = num_train_epochs * num_update_steps_per_epoch
+
+    lr_scheduler = get_scheduler(
+        "linear",
+        optimizer=optimizer,
+        num_warmup_steps=0,
+        num_training_steps=num_training_steps,
+    )
+
+    progress_bar = tqdm(range(num_training_steps))
+    for epoch in range(num_train_epochs):
+        model.train()
+        for batch in train_dataloader:
+            batch = {k: v.to(device) for k, v in batch.items()}
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss.backward()
+
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+            progress_bar.update(1)
+
+    model_save_path = f"{output_dir}/{model_name}"
+    model.save_pretrained(model_save_path)
+    tokenizer.save_pretrained(model_save_path)
+
+
+if __name__ == "__main__":
+    tokenizers_and_models = [
+        ("distilbert-base-uncased", "distilbert-base-uncased"),
+        ("bert-base-uncased", "bert-base-uncased"),
+    ]
+    for t, m in tokenizers_and_models:
+        train_model(t, m, OUTPUT_DIR)
diff --git a/mkdocs.yml b/mkdocs.yml
index 34b185d..4e5333f 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -21,9 +21,10 @@ nav:
     - guides/benchmarks.md
     - guides/customization.md
     - guides/organization.md
-    - guides/runners.md
+    - guides/runners.md 
   - Examples:
     - tutorials/index.md
+    - tutorials/artifact_benchmarking.md
     - tutorials/mnist.md
     - tutorials/prefect.md
     - tutorials/streamlit.md
diff --git a/src/nnbench/context.py b/src/nnbench/context.py
index 63578e5..b0ad493 100644
--- a/src/nnbench/context.py
+++ b/src/nnbench/context.py
@@ -277,21 +277,27 @@ class Context:
         """
         yield from self._ctx_items(d=self._data, prefix="", sep=sep)
 
-    def update(self, other: ContextProvider | dict[str, Any] | "Context") -> None:
+    def add(self, provider: ContextProvider) -> None:
+        """
+        Adds data from a provider to the context.
+
+        Parameters
+        ----------
+        provider : ContextProvider
+            The provider to inject into this context.
+        """
+        self._data.update(provider())
+
+    def update(self, other: "Context") -> None:
         """
         Updates the context.
-        If `sep` is None, perform shallow update, deep update otherwise.
 
         Parameters
         ----------
-        other : ContextProvider | dict[str, Any] | "Context"
-            The data to update the context with. This can be a dictionary, a Context instance, or a provider.
+        other : Context
+            The other context to update this context with.
         """
-        if callable(other):
-            other = other()
-        elif isinstance(other, Context):
-            other = other._data
-        self._data.update(other)
+        self._data.update(other._data)
 
     @staticmethod
     def _flatten_dict(d: dict[str, Any], prefix: str = "", sep: str = ".") -> dict[str, Any]:
@@ -300,11 +306,11 @@ class Context:
 
         Parameters
         ----------
-        d: dict[str, Any]
+        d : dict[str, Any]
             (Possibly) nested dictionary to flatten.
-        prefix: str
+        prefix : str
             Key prefix to apply at the top-level (nesting level 0).
-        sep: str
+        sep : str
             Separator on which to join keys, "." by default.
 
         Returns
diff --git a/src/nnbench/reporter/base.py b/src/nnbench/reporter/base.py
index d2f3567..d2204e7 100644
--- a/src/nnbench/reporter/base.py
+++ b/src/nnbench/reporter/base.py
@@ -116,7 +116,7 @@ class BenchmarkReporter:
             # only apply custom formatters after context merge
             #  to allow custom formatting of context values.
             filteredbm = {
-                k: custom_formatters.get(k, lambda x: x)(v) for k, v in filteredbm.items()
+                k: (custom_formatters or {}).get(k, lambda x: x)(v) for k, v in filteredbm.items()
             }
             filtered.append(filteredbm)
 
diff --git a/src/nnbench/runner.py b/src/nnbench/runner.py
index 5e5d325..1dacccb 100644
--- a/src/nnbench/runner.py
+++ b/src/nnbench/runner.py
@@ -188,7 +188,7 @@ class BenchmarkRunner:
         path_or_module: str | os.PathLike[str],
         params: dict[str, Any] | Parameters | None = None,
         tags: tuple[str, ...] = (),
-        context: Sequence[ContextProvider] = (),
+        context: Sequence[ContextProvider] | Context = (),
     ) -> BenchmarkRecord:
         """
         Run a previously collected benchmark workload.
@@ -204,7 +204,7 @@ class BenchmarkRunner:
         tags: tuple[str, ...]
             Tags to filter for when collecting benchmarks. Only benchmarks containing either of
             these tags are collected.
-        context: Sequence[ContextProvider]
+        context: Sequence[ContextProvider] | Context
             Additional context to log with the benchmark in the output JSON record. Useful for
             obtaining environment information and configuration, like CPU/GPU hardware info,
             ML model metadata, and more.
@@ -237,16 +237,20 @@ class BenchmarkRunner:
 
         self._check(dparams)
 
-        ctx = Context()
+        if isinstance(context, Context):
+            ctx = context
+        else:
+            ctx = Context()
+
+            for provider in context:
+                ctx_cand = Context.make(provider())
 
-        for provider in context:
-            ctxval = provider()
-            # we do not allow multiple values for a context key.
-            duplicates = set(ctx.keys()) & set(ctxval.keys())
-            if duplicates:
-                dupe, *_ = duplicates
-                raise ValueError(f"got multiple values for context key {dupe!r}")
-            ctx.update(ctxval)
+                # we do not allow multiple values for a context key.
+                duplicates = set(ctx.keys()) & set(ctx_cand.keys())
+                if duplicates:
+                    dupe, *_ = duplicates
+                    raise ValueError(f"got multiple values for context key {dupe!r}")
+                ctx.update(ctx_cand)
 
         results: list[dict[str, Any]] = []
         for benchmark in self.benchmarks:
diff --git a/src/nnbench/types.py b/src/nnbench/types.py
index ac9ce77..0d66885 100644
--- a/src/nnbench/types.py
+++ b/src/nnbench/types.py
@@ -2,11 +2,13 @@
 
 from __future__ import annotations
 
+import collections.abc
 import copy
 import inspect
 import os
+from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass, field
-from typing import Any, Callable, Generic, Literal, TypeVar
+from typing import Any, Callable, Generic, Iterable, Iterator, Literal, TypeVar
 
 from nnbench.context import Context
 
@@ -68,6 +70,7 @@ class BenchmarkRecord:
         Expand a list of deserialized JSON-like objects into a benchmark record.
         This is equivalent to extracting the context given by the method it was
         serialized with, and then returning the rest of the data as is.
+
         Parameters
         ----------
         bms: list[dict[str, Any]]
@@ -94,7 +97,38 @@ class BenchmarkRecord:
     #  context data.
 
 
-class Artifact(Generic[T]):
+class ArtifactLoader:
+    @abstractmethod
+    def load(self) -> os.PathLike[str]:
+        """Load the artifact"""
+
+
+class LocalArtifactLoader(ArtifactLoader):
+    """
+    ArtifactLoader for loading artifacts from a local file system.
+
+    Parameters
+    ----------
+    path : str | os.PathLike[str]
+        The file system pathto the artifact.
+    """
+
+    def __init__(self, path: str | os.PathLike[str]):
+        self._path = path
+
+    def load(self):
+        """
+        Returns the path to the artifact on the local file system.
+        """
+        return self._path
+
+
+class S3ArtifactLoader(ArtifactLoader):
+    # TODO: Implement this and other common ArtifactLoders here or in a util
+    pass
+
+
+class Artifact(Generic[T], metaclass=ABCMeta):
     """
     A base artifact class for loading (materializing) artifacts from disk or from remote storage.
 
@@ -102,34 +136,99 @@ class Artifact(Generic[T]):
     It is most useful when running models on already saved data or models, e.g. when
     comparing a newly trained model against a baseline in storage.
 
-    Subclasses need to implement the `Artifact.materialize()` API, telling nnbench how to
-    load the desired artifact from a path.
+    You need to supply an ArtifactLoader with a load() method to load the Artifact into the
+    local system storage.
+
+    Subclasses need to implement the `Artifact.deserialize()` API, telling nnbench to
+    load the desired artifact from their path.
 
     Parameters
     ----------
-    path: str | os.PathLike[str]
-        Path to the artifact files.
+    loader: ArtifactLoader
+        Loader to get the artifact.
     """
 
-    def __init__(self, path: str | os.PathLike[str]) -> None:
-        # Save the path for later just-in-time materialization.
-        self.path = path
+    def __init__(self, loader: ArtifactLoader) -> None:
+        # Save the path for later just-in-time deserialization.
+        self.path = loader.load()  # fetch the artifact from wherever it resides
         self._value: T | None = None
 
-    @classmethod
-    def materialize(cls) -> "Artifact":
-        """Load the artifact from storage."""
-        raise NotImplementedError
+    @abstractmethod
+    def deserialize(self) -> None:
+        """Deserialize the artifact."""
+
+    def is_deserialized(self) -> bool:
+        """Checks if the artifact is already deserialized."""
+        return self._value is not None
+
+    def __str__(self) -> str:
+        return f"Artifact(path={self.path!r}, is_deserialized={self.is_deserialized()})"
 
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(path={self.path!r}, is_deserialized={self.is_deserialized()})"
+
+    @property
     def value(self) -> T:
+        """
+        Returns the deserialized artifact value.
+
+        Returns
+        -------
+        T
+            The deserialized value of the artifact.
+        """
         if self._value is None:
-            raise ValueError(
-                f"artifact has not been instantiated yet, "
-                f"perhaps you forgot to call {self.__class__.__name__}.materialize()?"
-            )
+            self.deserialize()
         return self._value
 
 
+class ArtifactCollection(Generic[T], collections.abc.Iterable[Artifact[T]]):
+    """
+    A collection wrapper around multiple Artifact instances.
+    """
+
+    def __init__(self, *artifacts: Artifact[T]) -> None:
+        self._artifacts = [art for art in artifacts]
+
+    def add(self, artifacts: Artifact[T] | Iterable[Artifact[T]]) -> None:
+        """
+        Adds a single Artifact or an Iterable of Artifacts to the collection.
+
+        Parameters
+        ----------
+        artifacts : Artifact[T] | Iterable[Artifact[T]]
+            The Artifact instance(s) to add to the collection.
+        """
+        if isinstance(artifacts, Iterable):
+            self._artifacts.extend(artifacts)
+        else:
+            self._artifacts.append(artifacts)
+
+    def __iter__(self) -> Iterator[Artifact[T]]:
+        """
+        Iterator over deserialized artifacts in the collection.
+
+        Yields
+        -------
+        Iterator[Artifact[T]]
+            An iterator over deserialized artifacts in the collection.
+        """
+        for artifact in self._artifacts:
+            yield artifact
+
+    def values(self) -> Iterator[T]:
+        """
+        Iterator over the values of deserialized artifacts in the collection.
+
+        Yields
+        ------
+        Iterator[T]
+            An iterator over the values of the deserialized artifacts.
+        """
+        for artifact in self:
+            yield artifact.value
+
+
 @dataclass(init=False, frozen=True)
 class Parameters:
     """


Test patch:
diff --git a/tests/test_context.py b/tests/test_context.py
index 6b2848c..7c73d6d 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -71,9 +71,16 @@ def test_context_items():
     assert set(ctx.items()) == expected_items
 
 
-def test_update_with_unnested_dict():
+def test_update_with_context():
     ctx = Context({"a": 1, "b": 2})
-    ctx.update({"a": 3, "c": 4})
+    ctx.update(Context.make({"a": 3, "c": 4}))
+    expected_dict = {"a": 3, "b": 2, "c": 4}
+    assert ctx._data == expected_dict
+
+
+def test_add_with_provider():
+    ctx = Context({"a": 1, "b": 2})
+    ctx.add(lambda: {"a": 3, "c": 4})
     expected_dict = {"a": 3, "b": 2, "c": 4}
     assert ctx._data == expected_dict