😮🪑 Add OGB Evaluator (pykeen#948)

Add an evaluator to call OGB evaluation from within the pipeline.

Currently blocked by `Evaluator.evaluate` receiving multiple `kwargs`
(including `additional_filter_triples`) to be passed to `evaluate`,
which are not supported by the OGB evaluate method.

#### Dependencies

* [x] pykeen#1088

---------

Co-authored-by: Charles Tapley Hoyt <[email protected]>

Author: mberr

README.md

@@ -287,12 +287,13 @@ The following 2 stoppers are implemented in PyKEEN.
 
 ### Evaluators
 
-The following 4 evaluators are implemented in PyKEEN.
+The following 5 evaluators are implemented in PyKEEN.
 
 | Name             | Reference                                                                                                                                     | Description                                                              |
 |------------------|-----------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------|
 | classification   | [`pykeen.evaluation.ClassificationEvaluator`](https://pykeen.readthedocs.io/en/latest/api/pykeen.evaluation.ClassificationEvaluator.html)     | An evaluator that uses a classification metrics.                         |
 | macrorankbased   | [`pykeen.evaluation.MacroRankBasedEvaluator`](https://pykeen.readthedocs.io/en/latest/api/pykeen.evaluation.MacroRankBasedEvaluator.html)     | Macro-average rank-based evaluation.                                     |
+| ogb              | [`pykeen.evaluation.OGBEvaluator`](https://pykeen.readthedocs.io/en/latest/api/pykeen.evaluation.OGBEvaluator.html)                           | A sampled, rank-based evaluator that applies a custom OGB evaluation.    |
 | rankbased        | [`pykeen.evaluation.RankBasedEvaluator`](https://pykeen.readthedocs.io/en/latest/api/pykeen.evaluation.RankBasedEvaluator.html)               | A rank-based evaluator for KGE models.                                   |
 | sampledrankbased | [`pykeen.evaluation.SampledRankBasedEvaluator`](https://pykeen.readthedocs.io/en/latest/api/pykeen.evaluation.SampledRankBasedEvaluator.html) | A rank-based evaluator using sampled negatives instead of all negatives. |
 

src/pykeen/evaluation/__init__.py

@@ -7,6 +7,7 @@
 from .classification_evaluator import ClassificationEvaluator, ClassificationMetricResults
 from .evaluation_loop import LCWAEvaluationLoop
 from .evaluator import Evaluator, MetricResults, evaluate
+from .ogb_evaluator import OGBEvaluator
 from .rank_based_evaluator import (
     MacroRankBasedEvaluator,
     RankBasedEvaluator,
@@ -23,6 +24,7 @@
     "MacroRankBasedEvaluator",
     "LCWAEvaluationLoop",
     "SampledRankBasedEvaluator",
+    "OGBEvaluator",
     "ClassificationEvaluator",
     "ClassificationMetricResults",
     "evaluator_resolver",

src/pykeen/evaluation/ogb_evaluator.py

@@ -14,12 +14,44 @@
 from ..typing import RANK_REALISTIC, SIDE_BOTH, ExtendedTarget, MappedTriples, RankType, Target
 
 __all__ = [
+    "OGBEvaluator",
     "evaluate_ogb",
 ]
 
 logger = logging.getLogger(__name__)
 
 
+class OGBEvaluator(SampledRankBasedEvaluator):
+    """A sampled, rank-based evaluator that applies a custom OGB evaluation."""
+
+    # docstr-coverage: inherited
+    def __init__(self, filtered: bool = False, **kwargs):
+        if filtered:
+            raise ValueError(
+                "OGB evaluator is already filtered, but not dynamically like other evaluators because "
+                "it requires pre-calculated filtered negative triples. Therefore, it is not allowed to "
+                "accept filtered=True"
+            )
+        super().__init__(**kwargs, filtered=filtered)
+
+    def evaluate(
+        self,
+        model: Model,
+        mapped_triples: MappedTriples,
+        batch_size: Optional[int] = None,
+        slice_size: Optional[int] = None,
+        **kwargs,
+    ) -> MetricResults:
+        """Run :func:`evaluate_ogb` with this evaluator."""
+        return evaluate_ogb(
+            evaluator=self,
+            model=model,
+            mapped_triples=mapped_triples,
+            batch_size=batch_size,
+            **kwargs,
+        )
+
+
 def evaluate_ogb(
     evaluator: SampledRankBasedEvaluator,
     model: Model,
@@ -52,6 +84,8 @@ def evaluate_ogb(
         if ogb is not installed
     :raises NotImplementedError:
         if `batch_size` is None, i.e., automatic batch size selection is selected
+    :raises ValueError:
+        if illegal ``additional_filter_triples`` argument is given in the kwargs
     """
     try:
         import ogb.linkproppred
@@ -61,6 +95,13 @@ def evaluate_ogb(
     if batch_size is None:
         raise NotImplementedError("Automatic batch size selection not available for OGB evaluation.")
 
+    additional_filter_triples = kwargs.pop("additional_filter_triples", None)
+    if additional_filter_triples is not None:
+        raise ValueError(
+            f"evaluate_ogb received additional_filter_triples={additional_filter_triples}. However, it uses "
+            f"explicitly given filtered negative triples, and therefore shouldn't be passed any additional ones"
+        )
+
     class _OGBEvaluatorBridge(ogb.linkproppred.Evaluator):
         """A wrapper around OGB's evaluator to support evaluation on non-OGB datasets."""
 

src/pykeen/evaluation/rank_based_evaluator.py

@@ -587,6 +587,8 @@ def __init__(
                 LABEL_HEAD: head_negatives,
                 LABEL_TAIL: tail_negatives,
             }
+            if additional_filter_triples is not None:
+                logger.warning(f"Ignoring parameter additional_filter_triples={additional_filter_triples}")
 
         # verify input
         for side, side_negatives in negatives.items():
@@ -630,42 +632,6 @@ def process_scores_(
         # TODO: should we give num_entities in the constructor instead of inferring it every time ranks are processed?
         self.num_entities = num_entities
 
-    def evaluate_ogb(
-        self,
-        model,
-        mapped_triples: MappedTriples,
-        batch_size: Optional[int] = None,
-        **kwargs,
-    ) -> MetricResults:
-        """
-        Evaluate a model using OGB's evaluator.
-
-        :param model:
-            the model; will be set to evaluation mode.
-        :param mapped_triples:
-            the evaluation triples
-
-            .. note ::
-                the evaluation triples have to match with the stored explicit negatives
-
-        :param batch_size:
-            the batch size
-        :param kwargs:
-            additional keyword-based parameters passed to :meth:`pykeen.nn.Model.predict`
-
-        :return:
-            the evaluation results
-        """
-        from .ogb_evaluator import evaluate_ogb
-
-        return evaluate_ogb(
-            evaluator=self,
-            model=model,
-            mapped_triples=mapped_triples,
-            batch_size=batch_size,
-            **kwargs,
-        )
-
 
 class MacroRankBasedEvaluator(RankBasedEvaluator):
     """Macro-average rank-based evaluation."""

tests/test_evaluation/test_evaluators.py

@@ -18,7 +18,7 @@
 
 from pykeen.constants import COLUMN_LABELS
 from pykeen.datasets import Nations
-from pykeen.evaluation import Evaluator, MetricResults, RankBasedEvaluator, RankBasedMetricResults
+from pykeen.evaluation import Evaluator, MetricResults, OGBEvaluator, RankBasedEvaluator, RankBasedMetricResults
 from pykeen.evaluation.classification_evaluator import (
     CLASSIFICATION_METRICS,
     ClassificationEvaluator,
@@ -139,12 +139,25 @@ def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMa
         kwargs["additional_filter_triples"] = self.dataset.training.mapped_triples
         return kwargs
 
-    @needs_packages("ogb")
-    def test_ogb_evaluate(self):
+
+@needs_packages("ogb")
+class OGBEvaluatorTests(RankBasedEvaluatorTests):
+    """Unit test for OGB evaluator."""
+
+    cls = OGBEvaluator
+    kwargs = dict(num_negatives=3)
+
+    def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]:  # noqa: D102
+        kwargs = super()._pre_instantiation_hook(kwargs=kwargs)
+        kwargs["evaluation_factory"] = self.factory
+        kwargs["batch_size"] = 1
+        return kwargs
+
+    def test_ogb_evaluate_alternate(self):
         """Test OGB evaluation."""
         self.instance: SampledRankBasedEvaluator
         model = FixedModel(triples_factory=self.factory)
-        result = self.instance.evaluate_ogb(model=model, mapped_triples=self.factory.mapped_triples, batch_size=1)
+        result = self.instance.evaluate(model=model, mapped_triples=self.factory.mapped_triples, batch_size=1)
         assert isinstance(result, MetricResults)