replay-rec 0.19.0__py3-none-any.whl → 0.20.0__py3-none-any.whl

replay/__init__.py

@@ -1,3 +1,7 @@
-""" RecSys library """
+"""RecSys library"""
 
-__version__ = "0.19.0"
+# NOTE: This ensures distutils monkey-patching is performed before any
+# functionality removed in Python 3.12 is used in downstream packages (like lightfm)
+import setuptools as _
+
+__version__ = "0.20.0"

replay/data/dataset.py

@@ -610,9 +610,9 @@ class Dataset:
             if self.is_pandas:
                 try:
                     data[column] = data[column].astype(int)
-                except Exception:
+                except Exception as exc:
                     msg = f"IDs in {source.name}.{column} are not encoded. They are not int."
-                    raise ValueError(msg)
+                    raise ValueError(msg) from exc
 
         if self.is_pandas:
             is_int = np.issubdtype(dict(data.dtypes)[column], int)
@@ -775,10 +775,10 @@ def check_dataframes_types_equal(dataframe: DataFrameLike, other: DataFrameLike)
 
     :returns: True if dataframes have same type.
     """
-    if isinstance(dataframe, PandasDataFrame) and isinstance(other, PandasDataFrame):
-        return True
-    if isinstance(dataframe, SparkDataFrame) and isinstance(other, SparkDataFrame):
-        return True
-    if isinstance(dataframe, PolarsDataFrame) and isinstance(other, PolarsDataFrame):
-        return True
-    return False
+    return any(
+        [
+            isinstance(dataframe, PandasDataFrame) and isinstance(other, PandasDataFrame),
+            isinstance(dataframe, SparkDataFrame) and isinstance(other, SparkDataFrame),
+            isinstance(dataframe, PolarsDataFrame) and isinstance(other, PolarsDataFrame),
+        ]
+    )

replay/data/nn/__init__.py

@@ -14,17 +14,17 @@ if TORCH_AVAILABLE:
     )
 
     __all__ = [
+        "DEFAULT_GROUND_TRUTH_PADDING_VALUE",
+        "DEFAULT_TRAIN_PADDING_VALUE",
         "MutableTensorMap",
+        "PandasSequentialDataset",
+        "PolarsSequentialDataset",
+        "SequenceTokenizer",
+        "SequentialDataset",
         "TensorFeatureInfo",
         "TensorFeatureSource",
         "TensorMap",
         "TensorSchema",
-        "SequenceTokenizer",
-        "PandasSequentialDataset",
-        "PolarsSequentialDataset",
-        "SequentialDataset",
-        "DEFAULT_GROUND_TRUTH_PADDING_VALUE",
-        "DEFAULT_TRAIN_PADDING_VALUE",
         "TorchSequentialBatch",
         "TorchSequentialDataset",
         "TorchSequentialValidationBatch",

replay/data/nn/sequence_tokenizer.py

@@ -3,7 +3,7 @@ import json
 import pickle
 import warnings
 from pathlib import Path
-from typing import Dict, Generic, List, Optional, Sequence, Set, Tuple, TypeVar, Union
+from typing import TYPE_CHECKING, Dict, Generic, List, Optional, Sequence, Set, Tuple, TypeVar, Union
 
 import numpy as np
 import polars as pl
@@ -14,11 +14,11 @@ from replay.data import Dataset, FeatureHint, FeatureSchema, FeatureSource, Feat
 from replay.data.dataset_utils import DatasetLabelEncoder
 from replay.preprocessing import LabelEncoder, LabelEncodingRule
 from replay.preprocessing.label_encoder import HandleUnknownStrategies
-from replay.utils.model_handler import deprecation_warning
+from replay.utils import deprecation_warning
 
-from .schema import TensorFeatureInfo, TensorFeatureSource, TensorSchema
-from .sequential_dataset import PandasSequentialDataset, PolarsSequentialDataset, SequentialDataset
-from .utils import ensure_pandas, groupby_sequences
+if TYPE_CHECKING:
+    from .schema import TensorFeatureInfo, TensorFeatureSource, TensorSchema
+    from .sequential_dataset import SequentialDataset
 
 SequenceDataFrameLike = Union[PandasDataFrame, PolarsDataFrame]
 _T = TypeVar("_T")
@@ -34,7 +34,7 @@ class SequenceTokenizer:
 
     def __init__(
         self,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
         handle_unknown_rule: HandleUnknownStrategies = "error",
         default_value_rule: Optional[Union[int, str]] = None,
         allow_collect_to_master: bool = False,
@@ -77,7 +77,7 @@ class SequenceTokenizer:
         self,
         dataset: Dataset,
         tensor_features_to_keep: Optional[Sequence[str]] = None,
-    ) -> SequentialDataset:
+    ) -> "SequentialDataset":
         """
         :param dataset: input dataset to transform
         :param tensor_features_to_keep: specified feature names to transform
@@ -89,7 +89,7 @@ class SequenceTokenizer:
     def fit_transform(
         self,
         dataset: Dataset,
-    ) -> SequentialDataset:
+    ) -> "SequentialDataset":
         """
         :param dataset: input dataset to transform
         :returns: SequentialDataset
@@ -97,7 +97,7 @@ class SequenceTokenizer:
         return self.fit(dataset)._transform_unchecked(dataset)
 
     @property
-    def tensor_schema(self) -> TensorSchema:
+    def tensor_schema(self) -> "TensorSchema":
         """
         :returns: tensor schema
         """
@@ -149,7 +149,9 @@ class SequenceTokenizer:
         self,
         dataset: Dataset,
         tensor_features_to_keep: Optional[Sequence[str]] = None,
-    ) -> SequentialDataset:
+    ) -> "SequentialDataset":
+        from replay.data.nn.sequential_dataset import PandasSequentialDataset, PolarsSequentialDataset
+
         schema = self._tensor_schema
         if tensor_features_to_keep is not None:
             schema = schema.subset(tensor_features_to_keep)
@@ -186,6 +188,8 @@ class SequenceTokenizer:
         self,
         dataset: Dataset,
     ) -> Tuple[SequenceDataFrameLike, Optional[SequenceDataFrameLike], Optional[SequenceDataFrameLike]]:
+        from replay.data.nn.utils import ensure_pandas, groupby_sequences
+
         grouped_interactions = groupby_sequences(
             events=dataset.interactions,
             groupby_col=dataset.feature_schema.query_id_column,
@@ -218,7 +222,7 @@ class SequenceTokenizer:
 
     def _make_sequence_features(
         self,
-        schema: TensorSchema,
+        schema: "TensorSchema",
         feature_schema: FeatureSchema,
         grouped_interactions: SequenceDataFrameLike,
         query_features: Optional[SequenceDataFrameLike],
@@ -242,7 +246,7 @@ class SequenceTokenizer:
     def _match_features_with_tensor_schema(
         cls,
         dataset: Dataset,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
     ) -> Dataset:
         feature_subset_filter = cls._get_features_filter_from_schema(
             tensor_schema,
@@ -261,7 +265,7 @@ class SequenceTokenizer:
     @classmethod
     def _get_features_filter_from_schema(
         cls,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
         query_id_column: str,
         item_id_column: str,
     ) -> Set[str]:
@@ -291,7 +295,7 @@ class SequenceTokenizer:
         return set(features_subset)
 
     @classmethod
-    def _check_tensor_schema(cls, tensor_schema: TensorSchema) -> None:
+    def _check_tensor_schema(cls, tensor_schema: "TensorSchema") -> None:
         # Check consistency of sequential features
         for tensor_feature in tensor_schema.all_features:
             feature_sources = tensor_feature.feature_sources
@@ -319,11 +323,11 @@ class SequenceTokenizer:
     def _check_if_tensor_schema_matches_data(  # noqa: C901
         cls,
         dataset: Dataset,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
         tensor_features_to_keep: Optional[Sequence[str]] = None,
     ) -> None:
         # Check if all source columns specified in tensor schema exist in provided data frames
-        sources_for_tensors: List[TensorFeatureSource] = []
+        sources_for_tensors: List["TensorFeatureSource"] = []
         for tensor_feature_name, tensor_feature in tensor_schema.items():
             if tensor_features_to_keep is not None and tensor_feature_name not in tensor_features_to_keep:
                 continue
@@ -413,6 +417,8 @@ class SequenceTokenizer:
 
         :returns: Loaded tokenizer object.
         """
+        from replay.data.nn import TensorSchema
+
         if not use_pickle:
             base_path = Path(path).with_suffix(".replay").resolve()
             with open(base_path / "init_args.json", "r") as file:
@@ -500,7 +506,7 @@ class _BaseSequenceProcessor(Generic[_T]):
 
     def __init__(
         self,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
         query_id_column: str,
         item_id_column: str,
         grouped_interactions: _T,
@@ -535,7 +541,7 @@ class _BaseSequenceProcessor(Generic[_T]):
             return self._process_num_feature(tensor_feature)
         assert False, "Unknown tensor feature type"
 
-    def _process_num_feature(self, tensor_feature: TensorFeatureInfo) -> _T:
+    def _process_num_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:
         """
         Process numerical tensor feature depends on it source.
         """
@@ -548,7 +554,7 @@ class _BaseSequenceProcessor(Generic[_T]):
             return self._process_num_item_feature(tensor_feature)
         assert False, "Unknown tensor feature source table"
 
-    def _process_cat_feature(self, tensor_feature: TensorFeatureInfo) -> _T:
+    def _process_cat_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:
         """
         Process categorical tensor feature depends on it source.
         """
@@ -562,27 +568,27 @@ class _BaseSequenceProcessor(Generic[_T]):
         assert False, "Unknown tensor feature source table"
 
     @abc.abstractmethod
-    def _process_cat_interaction_feature(self, tensor_feature: TensorFeatureInfo) -> _T:  # pragma: no cover
+    def _process_cat_interaction_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:  # pragma: no cover
         pass
 
     @abc.abstractmethod
-    def _process_cat_query_feature(self, tensor_feature: TensorFeatureInfo) -> _T:  # pragma: no cover
+    def _process_cat_query_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:  # pragma: no cover
         pass
 
     @abc.abstractmethod
-    def _process_cat_item_feature(self, tensor_feature: TensorFeatureInfo) -> _T:  # pragma: no cover
+    def _process_cat_item_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:  # pragma: no cover
         pass
 
     @abc.abstractmethod
-    def _process_num_interaction_feature(self, tensor_feature: TensorFeatureInfo) -> _T:  # pragma: no cover
+    def _process_num_interaction_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:  # pragma: no cover
         pass
 
     @abc.abstractmethod
-    def _process_num_query_feature(self, tensor_feature: TensorFeatureInfo) -> _T:  # pragma: no cover
+    def _process_num_query_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:  # pragma: no cover
         pass
 
     @abc.abstractmethod
-    def _process_num_item_feature(self, tensor_feature: TensorFeatureInfo) -> _T:  # pragma: no cover
+    def _process_num_item_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:  # pragma: no cover
         pass
 
 
@@ -597,7 +603,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
 
     def __init__(
         self,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
         query_id_column: str,
         item_id_column: str,
         grouped_interactions: PandasDataFrame,
@@ -628,7 +634,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
         return PandasDataFrame(all_features)
 
     def _process_num_interaction_feature(
-        self, tensor_feature: TensorFeatureInfo
+        self, tensor_feature: "TensorFeatureInfo"
     ) -> Union[List[np.ndarray], List[List]]:
         """
         Process numerical interaction feature.
@@ -650,7 +656,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
                 values.append(np.array(sequence))
         return values
 
-    def _process_num_item_feature(self, tensor_feature: TensorFeatureInfo) -> Union[List[np.ndarray], List[List]]:
+    def _process_num_item_feature(self, tensor_feature: "TensorFeatureInfo") -> Union[List[np.ndarray], List[List]]:
         """
         Process numerical feature from item features dataset.
 
@@ -676,7 +682,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
 
         return values
 
-    def _process_num_query_feature(self, tensor_feature: TensorFeatureInfo) -> List[np.ndarray]:
+    def _process_num_query_feature(self, tensor_feature: "TensorFeatureInfo") -> List[np.ndarray]:
         """
         Process numerical feature from query features dataset.
 
@@ -687,7 +693,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
         return self._process_cat_query_feature(tensor_feature)
 
     def _process_cat_interaction_feature(
-        self, tensor_feature: TensorFeatureInfo
+        self, tensor_feature: "TensorFeatureInfo"
     ) -> Union[List[np.ndarray], List[List]]:
         """
         Process categorical interaction feature.
@@ -709,7 +715,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
                 values.append(np.array(sequence))
         return values
 
-    def _process_cat_query_feature(self, tensor_feature: TensorFeatureInfo) -> List[np.ndarray]:
+    def _process_cat_query_feature(self, tensor_feature: "TensorFeatureInfo") -> List[np.ndarray]:
         """
         Process categorical feature from query features dataset.
 
@@ -738,7 +744,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
                 ]
         return [np.array([query_feature[i]]).reshape(-1) for i in range(len(self._grouped_interactions))]
 
-    def _process_cat_item_feature(self, tensor_feature: TensorFeatureInfo) -> Union[List[np.ndarray], List[List]]:
+    def _process_cat_item_feature(self, tensor_feature: "TensorFeatureInfo") -> Union[List[np.ndarray], List[List]]:
         """
         Process categorical feature from item features dataset.
 
@@ -784,7 +790,7 @@ class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
             data = data.join(self._process_feature(tensor_feature_name), on=self._query_id_column, how="left")
         return data
 
-    def _process_num_interaction_feature(self, tensor_feature: TensorFeatureInfo) -> PolarsDataFrame:
+    def _process_num_interaction_feature(self, tensor_feature: "TensorFeatureInfo") -> PolarsDataFrame:
         """
         Process numerical interaction feature.
 
@@ -794,7 +800,7 @@ class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
         """
         return self._process_cat_interaction_feature(tensor_feature)
 
-    def _process_num_query_feature(self, tensor_feature: TensorFeatureInfo) -> PolarsDataFrame:
+    def _process_num_query_feature(self, tensor_feature: "TensorFeatureInfo") -> PolarsDataFrame:
         """
         Process numerical feature from query features dataset.
 
@@ -805,7 +811,7 @@ class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
         """
         return self._process_cat_query_feature(tensor_feature)
 
-    def _process_num_item_feature(self, tensor_feature: TensorFeatureInfo) -> PolarsDataFrame:
+    def _process_num_item_feature(self, tensor_feature: "TensorFeatureInfo") -> PolarsDataFrame:
         """
         Process numerical feature from item features dataset.
 
@@ -816,7 +822,7 @@ class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
         """
         return self._process_cat_item_feature(tensor_feature)
 
-    def _process_cat_interaction_feature(self, tensor_feature: TensorFeatureInfo) -> PolarsDataFrame:
+    def _process_cat_interaction_feature(self, tensor_feature: "TensorFeatureInfo") -> PolarsDataFrame:
         """
         Process categorical interaction feature.
 
@@ -833,7 +839,7 @@ class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
             {source.column: tensor_feature.name}
         )
 
-    def _process_cat_query_feature(self, tensor_feature: TensorFeatureInfo) -> PolarsDataFrame:
+    def _process_cat_query_feature(self, tensor_feature: "TensorFeatureInfo") -> PolarsDataFrame:
         """
         Process categorical feature from query features dataset.
 
@@ -877,7 +883,7 @@ class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
             {source.column: tensor_feature.name}
         )
 
-    def _process_cat_item_feature(self, tensor_feature: TensorFeatureInfo) -> PolarsDataFrame:
+    def _process_cat_item_feature(self, tensor_feature: "TensorFeatureInfo") -> PolarsDataFrame:
         """
         Process categorical feature from item features dataset.
 

replay/data/nn/sequential_dataset.py

@@ -1,7 +1,7 @@
 import abc
 import json
 from pathlib import Path
-from typing import Tuple, Union
+from typing import TYPE_CHECKING, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -9,7 +9,8 @@ import polars as pl
 from pandas import DataFrame as PandasDataFrame
 from polars import DataFrame as PolarsDataFrame
 
-from .schema import TensorSchema
+if TYPE_CHECKING:
+    from .schema import TensorSchema
 
 
 class SequentialDataset(abc.ABC):
@@ -81,7 +82,7 @@ class SequentialDataset(abc.ABC):
 
     @property
     @abc.abstractmethod
-    def schema(self) -> TensorSchema:  # pragma: no cover
+    def schema(self) -> "TensorSchema":  # pragma: no cover
         """
         :returns: List of tensor features.
         """
@@ -128,7 +129,7 @@ class PandasSequentialDataset(SequentialDataset):
 
     def __init__(
         self,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
         query_id_column: str,
         item_id_column: str,
         sequences: PandasDataFrame,
@@ -184,11 +185,11 @@ class PandasSequentialDataset(SequentialDataset):
         )
 
     @property
-    def schema(self) -> TensorSchema:
+    def schema(self) -> "TensorSchema":
         return self._tensor_schema
 
     @classmethod
-    def _check_if_schema_matches_data(cls, tensor_schema: TensorSchema, data: PandasDataFrame) -> None:
+    def _check_if_schema_matches_data(cls, tensor_schema: "TensorSchema", data: PandasDataFrame) -> None:
         for tensor_feature_name in tensor_schema:
             if tensor_feature_name not in data:
                 msg = "Tensor schema does not match with provided data frame"
@@ -199,6 +200,8 @@ class PandasSequentialDataset(SequentialDataset):
         """
         Method for loading PandasSequentialDataset object from `.replay` directory.
         """
+        from replay.data.nn import TensorSchema
+
         base_path = Path(path).with_suffix(".replay").resolve()
         with open(base_path / "init_args.json", "r") as file:
             sequential_dict = json.loads(file.read())
@@ -221,7 +224,7 @@ class PolarsSequentialDataset(PandasSequentialDataset):
 
     def __init__(
         self,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
         query_id_column: str,
         item_id_column: str,
         sequences: PolarsDataFrame,
@@ -270,7 +273,7 @@ class PolarsSequentialDataset(PandasSequentialDataset):
         return pl.from_dict(df.to_dict("list"))
 
     @classmethod
-    def _check_if_schema_matches_data(cls, tensor_schema: TensorSchema, data: PolarsDataFrame) -> None:
+    def _check_if_schema_matches_data(cls, tensor_schema: "TensorSchema", data: PolarsDataFrame) -> None:
         for tensor_feature_name in tensor_schema:
             if tensor_feature_name not in data:
                 msg = "Tensor schema does not match with provided data frame"
@@ -281,6 +284,8 @@ class PolarsSequentialDataset(PandasSequentialDataset):
         """
         Method for loading PandasSequentialDataset object from `.replay` directory.
         """
+        from replay.data.nn import TensorSchema
+
         base_path = Path(path).with_suffix(".replay").resolve()
         with open(base_path / "init_args.json", "r") as file:
             sequential_dict = json.loads(file.read())

replay/data/nn/torch_sequential_dataset.py

@@ -1,13 +1,14 @@
-from typing import Generator, NamedTuple, Optional, Sequence, Tuple, Union, cast
+from typing import TYPE_CHECKING, Generator, NamedTuple, Optional, Sequence, Tuple, Union, cast
 
 import numpy as np
 import torch
 from torch.utils.data import Dataset as TorchDataset
 
-from replay.utils.model_handler import deprecation_warning
+from replay.utils import deprecation_warning
 
-from .schema import TensorFeatureInfo, TensorMap, TensorSchema
-from .sequential_dataset import SequentialDataset
+if TYPE_CHECKING:
+    from .schema import TensorFeatureInfo, TensorMap, TensorSchema
+    from .sequential_dataset import SequentialDataset
 
 
 # We do not use dataclasses as PyTorch default collate
@@ -19,7 +20,7 @@ class TorchSequentialBatch(NamedTuple):
 
     query_id: torch.LongTensor
     padding_mask: torch.BoolTensor
-    features: TensorMap
+    features: "TensorMap"
 
 
 class TorchSequentialDataset(TorchDataset):
@@ -33,7 +34,7 @@ class TorchSequentialDataset(TorchDataset):
     )
     def __init__(
         self,
-        sequential: SequentialDataset,
+        sequential: "SequentialDataset",
         max_sequence_length: int,
         sliding_window_step: Optional[int] = None,
         padding_value: int = 0,
@@ -89,7 +90,7 @@ class TorchSequentialDataset(TorchDataset):
 
     def _generate_tensor_feature(
         self,
-        feature: TensorFeatureInfo,
+        feature: "TensorFeatureInfo",
         sequence_index: int,
         sequence_offset: int,
     ) -> torch.Tensor:
@@ -161,7 +162,7 @@ class TorchSequentialValidationBatch(NamedTuple):
 
     query_id: torch.LongTensor
     padding_mask: torch.BoolTensor
-    features: TensorMap
+    features: "TensorMap"
     ground_truth: torch.LongTensor
     train: torch.LongTensor
 
@@ -181,9 +182,9 @@ class TorchSequentialValidationDataset(TorchDataset):
     )
     def __init__(
         self,
-        sequential: SequentialDataset,
-        ground_truth: SequentialDataset,
-        train: SequentialDataset,
+        sequential: "SequentialDataset",
+        ground_truth: "SequentialDataset",
+        train: "SequentialDataset",
         max_sequence_length: int,
         padding_value: int = 0,
         sliding_window_step: Optional[int] = None,
@@ -280,8 +281,8 @@ class TorchSequentialValidationDataset(TorchDataset):
     @classmethod
     def _check_if_schema_match(
         cls,
-        sequential_schema: TensorSchema,
-        ground_truth_schema: TensorSchema,
+        sequential_schema: "TensorSchema",
+        ground_truth_schema: "TensorSchema",
     ) -> None:
         sequential_item_feature = sequential_schema.item_id_features.item()
         ground_truth_item_feature = ground_truth_schema.item_id_features.item()

replay/data/nn/utils.py

@@ -30,7 +30,7 @@ def groupby_sequences(events: DataFrameLike, groupby_col: str, sort_col: Optiona
             events = events.sort_values(event_cols_without_iterable)
 
         grouped_sequences = (
-            events.groupby(groupby_col).agg({col: list for col in event_cols_without_groupby}).reset_index()
+            events.groupby(groupby_col).agg(dict.fromkeys(event_cols_without_groupby, list)).reset_index()
         )
     elif isinstance(events, PolarsDataFrame):
         event_cols_without_groupby = events.columns

replay/metrics/base_metric.py

@@ -145,7 +145,7 @@ class Metric(ABC):
 
     def _convert_pandas_to_dict_with_score(self, data: PandasDataFrame) -> Dict:
         return (
-            data.sort_values(by=self.rating_column, ascending=False)
+            data.sort_values(by=[self.rating_column, self.item_column], ascending=False, kind="stable")
             .groupby(self.query_column)[self.item_column]
             .apply(list)
             .to_dict()

replay/metrics/coverage.py

@@ -102,20 +102,16 @@ class Coverage(Metric):
         return grouped_recs
 
     def _get_enriched_recommendations_polars(self, recommendations: PolarsDataFrame) -> PolarsDataFrame:
-        sorted_by_score_recommendations = recommendations.select(
-            pl.all().sort_by(self.rating_column, descending=True).over(self.query_column)
-        )
-        sorted_by_score_recommendations = sorted_by_score_recommendations.with_columns(
-            sorted_by_score_recommendations.select(
-                pl.col(self.query_column).cum_count().over(self.query_column).alias("rank")
+        return (
+            recommendations.with_columns(
+                pl.col(self.rating_column)
+                .rank(method="ordinal", descending=True)
+                .over(self.query_column)
+                .alias("__rank")
             )
-        )
-        grouped_recs = (
-            sorted_by_score_recommendations.select(self.item_column, "rank")
             .group_by(self.item_column)
-            .agg(pl.col("rank").min().alias("best_position"))
+            .agg(pl.col("__rank").min().alias("best_position"))
         )
-        return grouped_recs
 
     def _spark_compute(self, recs: SparkDataFrame, train: SparkDataFrame) -> MetricsMeanReturnType:
         """

replay/metrics/experiment.py

@@ -84,12 +84,12 @@ class Experiment:
     >>> ex.add_result("model", recommendations)
     >>> ex.results
                 NDCG@2    NDCG@3  Surprisal@3
-    baseline  0.204382  0.234639     0.608476
-    model     0.333333  0.489760     0.719587
+    baseline  0.333333  0.25512     0.608476
+    model     0.333333  0.48976     0.719587
     >>> ex.compare("baseline")
               NDCG@2   NDCG@3 Surprisal@3
     baseline       –        –           –
-    model     63.09%  108.73%      18.26%
+    model     0.0%  91.97%      18.26%
     >>> ex = Experiment([Precision(3, mode=Median()), Precision(3, mode=ConfidenceInterval(0.95))], groundtruth)
     >>> ex.add_result("baseline", base_rec)
     >>> ex.add_result("model", recommendations)

replay/metrics/offline_metrics.py

@@ -121,11 +121,11 @@ class OfflineMetrics:
     ...     base_recommendations={"ALS": base_rec, "KNN": recommendations}
     ... )
     {'Precision@2': 0.3333333333333333,
-     'Unexpectedness_ALS@1': 0.3333333333333333,
+     'Unexpectedness_ALS@1': 0.6666666666666666,
      'Unexpectedness_ALS@2': 0.16666666666666666,
      'Unexpectedness_KNN@1': 0.0,
      'Unexpectedness_KNN@2': 0.0,
-     'Unexpectedness-PerUser_ALS@1': {1: 1.0, 2: 0.0, 3: 0.0},
+     'Unexpectedness-PerUser_ALS@1': {1: 1.0, 2: 1.0, 3: 0.0},
      'Unexpectedness-PerUser_ALS@2': {1: 0.5, 2: 0.0, 3: 0.0},
      'Unexpectedness-PerUser_KNN@1': {1: 0.0, 2: 0.0, 3: 0.0},
      'Unexpectedness-PerUser_KNN@2': {1: 0.0, 2: 0.0, 3: 0.0}}

replay/models/__init__.py

@@ -23,3 +23,22 @@ from .thompson_sampling import ThompsonSampling
 from .ucb import UCB
 from .wilson import Wilson
 from .word2vec import Word2VecRec
+
+__all__ = [
+    "KLUCB",
+    "SLIM",
+    "UCB",
+    "ALSWrap",
+    "AssociationRulesItemRec",
+    "CatPopRec",
+    "ClusterRec",
+    "ItemKNN",
+    "LinUCB",
+    "PopRec",
+    "QueryPopRec",
+    "RandomRec",
+    "Recommender",
+    "ThompsonSampling",
+    "Wilson",
+    "Word2VecRec",
+]

replay/models/association_rules.py

@@ -142,6 +142,7 @@ class AssociationRulesItemRec(NeighbourRec):
         :param index_builder: `IndexBuilder` instance that adds ANN functionality.
             If not set, then ann will not be used.
         """
+        self.init_index_builder(index_builder)
 
         self.session_column = session_column
         self.min_item_count = min_item_count
@@ -149,10 +150,6 @@ class AssociationRulesItemRec(NeighbourRec):
         self.num_neighbours = num_neighbours
         self.use_rating = use_rating
         self.similarity_metric = similarity_metric
-        if isinstance(index_builder, (IndexBuilder, type(None))):
-            self.index_builder = index_builder
-        elif isinstance(index_builder, dict):
-            self.init_builder_from_dict(index_builder)
 
     @property
     def _init_args(self):