Merge pull request #600 from mdekstrand/feature/config-cleanups

Migrate HPF to configuration object and fix BiasConfig serialization

lenskit-hpf/lenskit/hpf.py

@@ -8,14 +8,20 @@
 
 import hpfrec
 import numpy as np
-from typing_extensions import Any, override
+from pydantic import BaseModel, JsonValue
+from typing_extensions import override
 
 from lenskit.data import Dataset, ItemList, QueryInput, RecQuery, Vocabulary
 from lenskit.pipeline import Component, Trainable
 
 _logger = logging.getLogger(__name__)
 
 
+class HPFConfig(BaseModel, extra="allow"):
+    __pydantic_extra__: dict[str, JsonValue]
+    features: int = 50
+
+
 class HPFScorer(Component[ItemList], Trainable):
     """
     Hierarchical Poisson factorization, provided by
@@ -34,21 +40,13 @@ class HPFScorer(Component[ItemList], Trainable):
             additional arguments to pass to :class:`hpfrec.HPF`.
     """
 
-    features: int
-    _kwargs: dict[str, Any]
+    config: HPFConfig
 
     users_: Vocabulary
     user_features_: np.ndarray[tuple[int, int], np.dtype[np.float64]]
     items_: Vocabulary
     item_features_: np.ndarray[tuple[int, int], np.dtype[np.float64]]
 
-    def __init__(self, features: int = 50, **kwargs):
-        self.features = features
-        self._kwargs = kwargs
-
-    def get_config(self):
-        return {"features": self.features} | self._kwargs
-
     @property
     def is_trained(self) -> bool:
         return hasattr(self, "item_features_")
@@ -64,9 +62,9 @@ def train(self, data: Dataset):
             }
         )
 
-        hpf = hpfrec.HPF(self.features, reindex=False, **self._kwargs)
+        hpf = hpfrec.HPF(self.config.features, reindex=False, **self.config.__pydantic_extra__)  # type: ignore
 
-        _logger.info("fitting HPF model with %d features", self.features)
+        _logger.info("fitting HPF model with %d features", self.config.features)
         hpf.fit(log)
 
         self.users_ = data.users

lenskit-hpf/tests/test_hpf.py

@@ -28,7 +28,7 @@ class TestHPF(BasicComponentTests, ScorerTests):
 
 @mark.slow
 def test_hpf_train_large(tmp_path, ml_ratings):
-    algo = hpf.HPFScorer(20)
+    algo = hpf.HPFScorer(features=20)
     ratings = ml_ratings.assign(rating=ml_ratings.rating + 0.5)
     ds = from_interactions_df(ratings)
     algo.train(ds)
@@ -57,7 +57,7 @@ def test_hpf_train_large(tmp_path, ml_ratings):
 
 @mark.slow
 def test_hpf_train_binary(tmp_path, ml_ratings):
-    algo = hpf.HPFScorer(20)
+    algo = hpf.HPFScorer(features=20)
     ratings = ml_ratings.drop(columns=["timestamp", "rating"])
     ds = from_interactions_df(ratings)
     algo.train(ds)

lenskit/lenskit/basic/bias.py

@@ -12,11 +12,11 @@
 import logging
 from collections.abc import Container
 from dataclasses import dataclass
-from typing import Literal
+from typing import Annotated, Literal
 
 import numpy as np
 import torch
-from pydantic import BaseModel, NonNegativeFloat
+from pydantic import BaseModel, NonNegativeFloat, PlainSerializer
 from typing_extensions import Self, TypeAlias, overload
 
 from lenskit.data import ID, Dataset, ItemList, QueryInput, RecQuery, Vocabulary
@@ -256,7 +256,9 @@ class BiasConfig(BaseModel, extra="forbid"):
     Configuration for :class:`BiasScorer`.
     """
 
-    entities: set[Literal["user", "item"]] = {"user", "item"}
+    entities: Annotated[
+        set[Literal["user", "item"]], PlainSerializer(lambda s: sorted(s), return_type=list[str])
+    ] = {"user", "item"}
     """
     The entities to compute biases for, in addition to global bais.  Defaults to
     users and items.