replay-rec 0.16.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

replay/models/cat_pop_rec.py

@@ -2,9 +2,10 @@ from os.path import join
 from typing import Iterable, Optional, Union
 
 from replay.data import Dataset
-from .base_rec import IsSavable, RecommenderCommons
 from replay.utils import PYSPARK_AVAILABLE, SparkDataFrame
 
+from .base_rec import IsSavable, RecommenderCommons
+
 if PYSPARK_AVAILABLE:
     from pyspark.sql import functions as sf
 
@@ -18,7 +19,6 @@ if PYSPARK_AVAILABLE:
     )
 
 
-# pylint: disable=too-many-instance-attributes
 class CatPopRec(IsSavable, RecommenderCommons):
     """
     CatPopRec generate recommendation for item categories.
@@ -35,9 +35,7 @@ class CatPopRec(IsSavable, RecommenderCommons):
     can_predict_cold_items: bool = False
     fit_items: SparkDataFrame
 
-    def _generate_mapping(
-        self, cat_tree: SparkDataFrame, max_iter: int = 20
-    ) -> SparkDataFrame:
+    def _generate_mapping(self, cat_tree: SparkDataFrame, max_iter: int = 20) -> SparkDataFrame:
         """
         Create SparkDataFrame with mapping [`category`, `leaf_cat`]
         where `leaf_cat` is the lowest level categories of category tree,
@@ -49,9 +47,7 @@ class CatPopRec(IsSavable, RecommenderCommons):
         :param max_iter: maximal number of iteration of descend through the category tree
         :return: SparkDataFrame with mapping [`category`, `leaf_cat`]
         """
-        current_res = cat_tree.select(
-            sf.col("category"), sf.col("category").alias("leaf_cat")
-        )
+        current_res = cat_tree.select(sf.col("category"), sf.col("category").alias("leaf_cat"))
 
         i = 0
         res_size_growth = current_res.count()
@@ -108,9 +104,7 @@ class CatPopRec(IsSavable, RecommenderCommons):
         """
         self.max_iter = max_iter
         if cat_tree is not None:
-            self.leaf_cat_mapping = self._generate_mapping(
-                cat_tree, max_iter=max_iter
-            )
+            self.leaf_cat_mapping = self._generate_mapping(cat_tree, max_iter=max_iter)
 
     @property
     def _init_args(self):
@@ -165,7 +159,6 @@ class CatPopRec(IsSavable, RecommenderCommons):
         if hasattr(self, "leaf_cat_mapping"):
             self.leaf_cat_mapping.unpersist()
 
-    # pylint: disable=arguments-differ
     def predict(
         self,
         categories: Union[SparkDataFrame, Iterable],
@@ -219,9 +212,7 @@ class CatPopRec(IsSavable, RecommenderCommons):
         item_data = items or self.fit_items
         items = get_unique_entities(item_data, self.item_column)
 
-        num_new, items = filter_cold(
-            items, self.fit_items, col_name=self.item_column
-        )
+        num_new, items = filter_cold(items, self.fit_items, col_name=self.item_column)
         if num_new > 0:
             self.logger.info(
                 "%s model can't predict cold items, they will be ignored",
@@ -267,9 +258,7 @@ class CatPopRec(IsSavable, RecommenderCommons):
         # find number of interactions in all leaf categories after filtering
         num_interactions_in_cat = (
             res.join(
-                unique_leaf_cat_items.groupBy("leaf_cat").agg(
-                    sf.sum(self.rating_column).alias("sum_rating")
-                ),
+                unique_leaf_cat_items.groupBy("leaf_cat").agg(sf.sum(self.rating_column).alias("sum_rating")),
                 on="leaf_cat",
             )
             .groupBy("category")
@@ -284,9 +273,7 @@ class CatPopRec(IsSavable, RecommenderCommons):
             .groupBy("category", self.item_column)
             .agg(sf.sum(self.rating_column).alias(self.rating_column))
             .join(num_interactions_in_cat, on="category")
-            .withColumn(
-                self.rating_column, sf.col(self.rating_column) / sf.col("sum_rating")
-            )
+            .withColumn(self.rating_column, sf.col(self.rating_column) / sf.col("sum_rating"))
         )
 
     def _save_model(self, path: str):
@@ -296,7 +283,7 @@ class CatPopRec(IsSavable, RecommenderCommons):
                 "item_column": self.item_column,
                 "rating_column": self.rating_column,
             },
-            join(path, "params.dump")
+            join(path, "params.dump"),
         )
 
     def _load_model(self, path: str):

replay/models/cluster.py

@@ -2,9 +2,10 @@ from os.path import join
 from typing import Optional
 
 from replay.data.dataset import Dataset
-from .base_rec import QueryRecommender
 from replay.utils import PYSPARK_AVAILABLE, SparkDataFrame
 
+from .base_rec import QueryRecommender
+
 if PYSPARK_AVAILABLE:
     from pyspark.ml.clustering import KMeans, KMeansModel
     from pyspark.ml.feature import VectorAssembler
@@ -58,12 +59,10 @@ class ClusterRec(QueryRecommender):
             sf.count(self.item_column).alias("item_count")
         )
 
-        max_count_per_cluster = self.item_rel_in_cluster.groupby(
-            "cluster"
-        ).agg(sf.max("item_count").alias("max_count_in_cluster"))
-        self.item_rel_in_cluster = self.item_rel_in_cluster.join(
-            max_count_per_cluster, on="cluster"
+        max_count_per_cluster = self.item_rel_in_cluster.groupby("cluster").agg(
+            sf.max("item_count").alias("max_count_in_cluster")
         )
+        self.item_rel_in_cluster = self.item_rel_in_cluster.join(max_count_per_cluster, on="cluster")
         self.item_rel_in_cluster = self.item_rel_in_cluster.withColumn(
             self.rating_column, sf.col("item_count") / sf.col("max_count_in_cluster")
         ).drop("item_count", "max_count_in_cluster")
@@ -83,47 +82,38 @@ class ClusterRec(QueryRecommender):
         return vec.transform(query_features).select(self.query_column, "features")
 
     def _make_query_clusters(self, queries, query_features):
-
         query_cnt_in_fv = (
-            query_features
-            .select(self.query_column)
-            .distinct()
-            .join(queries.distinct(), on=self.query_column)
-            .count()
+            query_features.select(self.query_column).distinct().join(queries.distinct(), on=self.query_column).count()
         )
 
         query_cnt = queries.distinct().count()
 
         if query_cnt_in_fv < query_cnt:
-            self.logger.info("%s query(s) don't "
-                             "have a feature vector. "
-                             "The results will not be calculated for them.",
-                             query_cnt - query_cnt_in_fv)
+            self.logger.info(
+                "%s query(s) don't have a feature vector. The results will not be calculated for them.",
+                query_cnt - query_cnt_in_fv,
+            )
 
-        query_features_vector = self._transform_features(
-            query_features.join(queries, on=self.query_column)
-        )
+        query_features_vector = self._transform_features(query_features.join(queries, on=self.query_column))
         return (
             self.model.transform(query_features_vector)
             .select(self.query_column, "prediction")
             .withColumnRenamed("prediction", "cluster")
         )
 
-    # pylint: disable=too-many-arguments
     def _predict(
         self,
         dataset: Dataset,
-        k: int,
+        k: int,  # noqa: ARG002
         queries: SparkDataFrame,
         items: SparkDataFrame,
-        filter_seen_items: bool = True,
+        filter_seen_items: bool = True,  # noqa: ARG002
     ) -> SparkDataFrame:
         query_clusters = self._make_query_clusters(queries, dataset.query_features)
         filtered_items = self.item_rel_in_cluster.join(items, on=self.item_column)
         pred = query_clusters.join(filtered_items, on="cluster").drop("cluster")
         return pred
 
-    # pylint: disable=signature-differs
     def _predict_pairs(
         self,
         pairs: SparkDataFrame,
@@ -131,9 +121,8 @@ class ClusterRec(QueryRecommender):
     ) -> SparkDataFrame:
         query_clusters = self._make_query_clusters(pairs.select(self.query_column).distinct(), dataset.query_features)
         pairs_with_clusters = pairs.join(query_clusters, on=self.query_column)
-        filtered_items = (self.item_rel_in_cluster
-                          .join(pairs.select(self.item_column).distinct(), on=self.item_column))
-        pred = (pairs_with_clusters
-                .join(filtered_items, on=["cluster", self.item_column])
-                .select(self.query_column,self.item_column,self.rating_column))
+        filtered_items = self.item_rel_in_cluster.join(pairs.select(self.item_column).distinct(), on=self.item_column)
+        pred = pairs_with_clusters.join(filtered_items, on=["cluster", self.item_column]).select(
+            self.query_column, self.item_column, self.rating_column
+        )
         return pred

replay/models/extensions/ann/ann_mixin.py

@@ -5,15 +5,17 @@ from typing import Any, Dict, Iterable, Optional, Union
 
 from replay.data import Dataset
 from replay.models.base_rec import BaseRecommender
-from .index_builders.base_index_builder import IndexBuilder
 from replay.utils import PYSPARK_AVAILABLE, SparkDataFrame
 
+from .index_builders.base_index_builder import IndexBuilder
+
 if PYSPARK_AVAILABLE:
     from pyspark.sql import functions as sf
 
-    from .index_stores.spark_files_index_store import SparkFilesIndexStore
     from replay.utils.spark_utils import get_top_k_recs, return_recs
 
+    from .index_stores.spark_files_index_store import SparkFilesIndexStore
+
 
 logger = logging.getLogger("replay")
 
@@ -82,9 +84,7 @@ class ANNMixin(BaseRecommender):
             self.index_builder.build_index(vectors, **ann_params)
 
     @abstractmethod
-    def _get_vectors_to_infer_ann_inner(
-        self, interactions: SparkDataFrame, queries: SparkDataFrame
-    ) -> SparkDataFrame:
+    def _get_vectors_to_infer_ann_inner(self, interactions: SparkDataFrame, queries: SparkDataFrame) -> SparkDataFrame:
         """Implementations of this method must return a dataframe with user vectors.
         User vectors from this method are used to infer the index.
 
@@ -134,7 +134,6 @@ class ANNMixin(BaseRecommender):
 
         """
 
-    # pylint: disable=too-many-arguments, too-many-locals
     def _predict_wrap(
         self,
         dataset: Optional[Dataset],
@@ -144,14 +143,10 @@ class ANNMixin(BaseRecommender):
         filter_seen_items: bool = True,
         recs_file_path: Optional[str] = None,
     ) -> Optional[SparkDataFrame]:
-        dataset, queries, items = self._filter_interactions_queries_items_dataframes(
-            dataset, k, queries, items
-        )
+        dataset, queries, items = self._filter_interactions_queries_items_dataframes(dataset, k, queries, items)
 
         if self._use_ann:
-            vectors = self._get_vectors_to_infer_ann(
-                dataset.interactions, queries, filter_seen_items
-            )
+            vectors = self._get_vectors_to_infer_ann(dataset.interactions, queries, filter_seen_items)
             ann_params = self._get_ann_infer_params()
             inferer = self.index_builder.produce_inferer(filter_seen_items)
             recs = inferer.infer(vectors, ann_params["features_col"], k)

replay/models/extensions/ann/entities/base_hnsw_param.py

@@ -9,7 +9,7 @@ class BaseHnswParam:
     """
 
     space: str
-    m: int = 200  # pylint: disable=invalid-name
+    m: int = 200
     ef_c: int = 20000
     post: int = 0
     ef_s: Optional[int] = None

replay/models/extensions/ann/entities/hnswlib_param.py

@@ -59,9 +59,3 @@ class HnswlibParam(BaseHnswParam):
     dim: int = field(default=None, init=False)
     # Max number of elements that will be stored in the index
     max_elements: int = field(default=None, init=False)
-
-    # def init_args_as_dict(self):
-    #     # union dicts
-    #     return dict(
-    #         super().init_args_as_dict()["init_args"], **{"space": self.space}
-    #     )

replay/models/extensions/ann/entities/nmslib_hnsw_param.py

@@ -65,9 +65,3 @@ class NmslibHnswParam(BaseHnswParam):
     items_count: Optional[int] = field(default=None, init=False)
 
     method: ClassVar[str] = "hnsw"
-
-    # def init_args_as_dict(self):
-    #     # union dicts
-    #     return dict(
-    #         super().init_args_as_dict()["init_args"], **{"space": self.space}
-    #     )

replay/models/extensions/ann/index_builders/driver_hnswlib_index_builder.py

@@ -3,7 +3,6 @@ from typing import Optional
 
 import numpy as np
 
-from .base_index_builder import IndexBuilder
 from replay.models.extensions.ann.index_inferers.base_inferer import IndexInferer
 from replay.models.extensions.ann.index_inferers.hnswlib_filter_index_inferer import HnswlibFilterIndexInferer
 from replay.models.extensions.ann.index_inferers.hnswlib_index_inferer import HnswlibIndexInferer
@@ -11,6 +10,8 @@ from replay.models.extensions.ann.utils import create_hnswlib_index_instance
 from replay.utils import SparkDataFrame
 from replay.utils.spark_utils import spark_to_pandas
 
+from .base_index_builder import IndexBuilder
+
 logger = logging.getLogger("replay")
 
 
@@ -21,13 +22,10 @@ class DriverHnswlibIndexBuilder(IndexBuilder):
 
     def produce_inferer(self, filter_seen_items: bool) -> IndexInferer:
         if filter_seen_items:
-            return HnswlibFilterIndexInferer(
-                self.index_params, self.index_store
-            )
+            return HnswlibFilterIndexInferer(self.index_params, self.index_store)
         else:
             return HnswlibIndexInferer(self.index_params, self.index_store)
 
-    # pylint: disable=no-member
     def build_index(
         self,
         vectors: SparkDataFrame,
@@ -43,8 +41,4 @@ class DriverHnswlibIndexBuilder(IndexBuilder):
         else:
             index.add_items(np.stack(vectors_np))
 
-        self.index_store.save_to_store(
-            lambda path: index.save_index(  # pylint: disable=unnecessary-lambda)
-                path
-            )
-        )
+        self.index_store.save_to_store(lambda path: index.save_index(path))

replay/models/extensions/ann/index_builders/driver_nmslib_index_builder.py

@@ -1,14 +1,15 @@
 import logging
 from typing import Optional
 
-from .base_index_builder import IndexBuilder
-from .nmslib_index_builder_mixin import NmslibIndexBuilderMixin
 from replay.models.extensions.ann.index_inferers.base_inferer import IndexInferer
 from replay.models.extensions.ann.index_inferers.nmslib_filter_index_inferer import NmslibFilterIndexInferer
 from replay.models.extensions.ann.index_inferers.nmslib_index_inferer import NmslibIndexInferer
 from replay.utils import SparkDataFrame
 from replay.utils.spark_utils import spark_to_pandas
 
+from .base_index_builder import IndexBuilder
+from .nmslib_index_builder_mixin import NmslibIndexBuilderMixin
+
 logger = logging.getLogger("replay")
 
 
@@ -19,20 +20,15 @@ class DriverNmslibIndexBuilder(IndexBuilder):
 
     def produce_inferer(self, filter_seen_items: bool) -> IndexInferer:
         if filter_seen_items:
-            return NmslibFilterIndexInferer(
-                self.index_params, self.index_store
-            )
+            return NmslibFilterIndexInferer(self.index_params, self.index_store)
         else:
             return NmslibIndexInferer(self.index_params, self.index_store)
 
-    # pylint: disable=no-member
     def build_index(
         self,
         vectors: SparkDataFrame,
-        features_col: str,
-        ids_col: Optional[str] = None,
+        features_col: str,  # noqa: ARG002
+        ids_col: Optional[str] = None,  # noqa: ARG002
     ):
         vectors = spark_to_pandas(vectors, self.allow_collect_to_master)
-        NmslibIndexBuilderMixin.build_and_save_index(
-            vectors, self.index_params, self.index_store
-        )
+        NmslibIndexBuilderMixin.build_and_save_index(vectors, self.index_params, self.index_store)

replay/models/extensions/ann/index_builders/executor_hnswlib_index_builder.py

@@ -3,13 +3,14 @@ from typing import Iterator, Optional
 
 import numpy as np
 
-from .base_index_builder import IndexBuilder
 from replay.models.extensions.ann.index_inferers.base_inferer import IndexInferer
 from replay.models.extensions.ann.index_inferers.hnswlib_filter_index_inferer import HnswlibFilterIndexInferer
 from replay.models.extensions.ann.index_inferers.hnswlib_index_inferer import HnswlibIndexInferer
 from replay.models.extensions.ann.utils import create_hnswlib_index_instance
 from replay.utils import PandasDataFrame, SparkDataFrame
 
+from .base_index_builder import IndexBuilder
+
 logger = logging.getLogger("replay")
 
 
@@ -20,9 +21,7 @@ class ExecutorHnswlibIndexBuilder(IndexBuilder):
 
     def produce_inferer(self, filter_seen_items: bool) -> IndexInferer:
         if filter_seen_items:
-            return HnswlibFilterIndexInferer(
-                self.index_params, self.index_store
-            )
+            return HnswlibFilterIndexInferer(self.index_params, self.index_store)
         else:
             return HnswlibIndexInferer(self.index_params, self.index_store)
 
@@ -56,17 +55,11 @@ class ExecutorHnswlibIndexBuilder(IndexBuilder):
                     # ids will be from [0, ..., len(vectors_np)]
                     index.add_items(np.stack(vectors_np))
 
-            _index_store.save_to_store(
-                lambda path: index.save_index(  # pylint: disable=unnecessary-lambda)
-                    path
-                )
-            )
+            _index_store.save_to_store(lambda path: index.save_index(path))
 
             yield PandasDataFrame(data={"_success": 1}, index=[0])
 
         # Here we perform materialization (`.collect()`) to build the hnsw index.
         cols = [ids_col, features_col] if ids_col else [features_col]
 
-        vectors.select(*cols).mapInPandas(
-            build_index_udf, "_success int"
-        ).collect()
+        vectors.select(*cols).mapInPandas(build_index_udf, "_success int").collect()

replay/models/extensions/ann/index_builders/executor_nmslib_index_builder.py

@@ -3,13 +3,14 @@ from typing import Iterator, Optional
 
 import pandas as pd
 
-from .base_index_builder import IndexBuilder
-from .nmslib_index_builder_mixin import NmslibIndexBuilderMixin
 from replay.models.extensions.ann.index_inferers.base_inferer import IndexInferer
 from replay.models.extensions.ann.index_inferers.nmslib_filter_index_inferer import NmslibFilterIndexInferer
 from replay.models.extensions.ann.index_inferers.nmslib_index_inferer import NmslibIndexInferer
 from replay.utils import PandasDataFrame, SparkDataFrame
 
+from .base_index_builder import IndexBuilder
+from .nmslib_index_builder_mixin import NmslibIndexBuilderMixin
+
 logger = logging.getLogger("replay")
 
 
@@ -20,9 +21,7 @@ class ExecutorNmslibIndexBuilder(IndexBuilder):
 
     def produce_inferer(self, filter_seen_items: bool) -> IndexInferer:
         if filter_seen_items:
-            return NmslibFilterIndexInferer(
-                self.index_params, self.index_store
-            )
+            return NmslibFilterIndexInferer(self.index_params, self.index_store)
         else:
             return NmslibIndexInferer(self.index_params, self.index_store)
 
@@ -47,15 +46,9 @@ class ExecutorNmslibIndexBuilder(IndexBuilder):
             # with the same `item_idx_two`.
             # And therefore we cannot call the `addDataPointBatch` iteratively
             # (in build_and_save_index).
-            pdfs = []
-            for pdf in iterator:
-                pdfs.append(pdf)
-
-            pdf = pd.concat(pdfs)
+            pdf = pd.concat(list(iterator))
 
-            NmslibIndexBuilderMixin.build_and_save_index(
-                pdf, index_params, index_store
-            )
+            NmslibIndexBuilderMixin.build_and_save_index(pdf, index_params, index_store)
 
             yield PandasDataFrame(data={"_success": 1}, index=[0])
 
@@ -64,8 +57,8 @@ class ExecutorNmslibIndexBuilder(IndexBuilder):
     def build_index(
         self,
         vectors: SparkDataFrame,
-        features_col: str,
-        ids_col: Optional[str] = None,
+        features_col: str,  # noqa: ARG002
+        ids_col: Optional[str] = None,  # noqa: ARG002
     ):
         # to execution in one executor
         vectors = vectors.repartition(1)
@@ -74,6 +67,6 @@ class ExecutorNmslibIndexBuilder(IndexBuilder):
         build_index_udf = self.make_build_index_udf()
 
         # Here we perform materialization (`.collect()`) to build the hnsw index.
-        vectors.select(
-            "similarity", "item_idx_one", "item_idx_two"
-        ).mapInPandas(build_index_udf, "_success int").collect()
+        vectors.select("similarity", "item_idx_one", "item_idx_two").mapInPandas(
+            build_index_udf, "_success int"
+        ).collect()

replay/models/extensions/ann/index_builders/nmslib_index_builder_mixin.py

@@ -6,7 +6,6 @@ from replay.models.extensions.ann.utils import create_nmslib_index_instance
 from replay.utils import PandasDataFrame
 
 
-# pylint: disable=too-few-public-methods
 class NmslibIndexBuilderMixin:
     """Provides nmslib index building method for different nmslib index builders"""
 
@@ -49,6 +48,4 @@ class NmslibIndexBuilderMixin:
         index.addDataPointBatch(data=sim_matrix)
         index.createIndex(creation_index_params)
 
-        index_store.save_to_store(
-            lambda path: index.saveIndex(path, save_data=True)
-        )  # pylint: disable=unnecessary-lambda)
+        index_store.save_to_store(lambda path: index.saveIndex(path, save_data=True))

replay/models/extensions/ann/index_inferers/base_inferer.py

@@ -8,7 +8,6 @@ if PYSPARK_AVAILABLE:
     from pyspark.sql import functions as sf
 
 
-# pylint: disable=too-few-public-methods
 class IndexInferer(ABC):
     """Abstract base class that describes a common interface for index inferers
     and provides common methods for them."""
@@ -21,9 +20,7 @@ class IndexInferer(ABC):
         self.index_store = index_store
 
     @abstractmethod
-    def infer(
-        self, vectors: SparkDataFrame, features_col: str, k: int
-    ) -> SparkDataFrame:
+    def infer(self, vectors: SparkDataFrame, features_col: str, k: int) -> SparkDataFrame:
         """Infers index"""
 
     @staticmethod
@@ -51,9 +48,7 @@ class IndexInferer(ABC):
         """
         res = inference_result.select(
             "user_idx",
-            sf.explode(
-                sf.arrays_zip("neighbours.item_idx", "neighbours.distance")
-            ).alias("zip_exp"),
+            sf.explode(sf.arrays_zip("neighbours.item_idx", "neighbours.distance")).alias("zip_exp"),
         )
 
         # Fix arrays_zip random behavior.
@@ -65,8 +60,6 @@ class IndexInferer(ABC):
         res = res.select(
             "user_idx",
             sf.col(f"zip_exp.{item_idx_field_name}").alias("item_idx"),
-            (sf.lit(-1.0) * sf.col(f"zip_exp.{distance_field_name}")).alias(
-                "relevance"
-            ),
+            (sf.lit(-1.0) * sf.col(f"zip_exp.{distance_field_name}")).alias("relevance"),
         )
         return res

replay/models/extensions/ann/index_inferers/hnswlib_filter_index_inferer.py

@@ -1,28 +1,24 @@
 import numpy as np
 import pandas as pd
 
-from .base_inferer import IndexInferer
 from replay.models.extensions.ann.utils import create_hnswlib_index_instance
 from replay.utils import PYSPARK_AVAILABLE, PandasDataFrame, SparkDataFrame
 from replay.utils.session_handler import State
 
+from .base_inferer import IndexInferer
+
 if PYSPARK_AVAILABLE:
     from pyspark.sql.pandas.functions import pandas_udf
 
 
-# pylint: disable=too-few-public-methods
 class HnswlibFilterIndexInferer(IndexInferer):
     """Hnswlib index inferer with filter seen items. Infers hnswlib index."""
 
-    def infer(
-        self, vectors: SparkDataFrame, features_col: str, k: int
-    ) -> SparkDataFrame:
+    def infer(self, vectors: SparkDataFrame, features_col: str, k: int) -> SparkDataFrame:
         _index_store = self.index_store
         index_params = self.index_params
 
-        index_store_broadcast = State().session.sparkContext.broadcast(
-            _index_store
-        )
+        index_store_broadcast = State().session.sparkContext.broadcast(_index_store)
 
         @pandas_udf(self.udf_return_type)
         def infer_index_udf(
@@ -34,9 +30,7 @@ class HnswlibFilterIndexInferer(IndexInferer):
             index = index_store.load_index(
                 init_index=lambda: create_hnswlib_index_instance(index_params),
                 load_index=lambda index, path: index.load_index(path),
-                configure_index=lambda index: index.set_ef(index_params.ef_s)
-                if index_params.ef_s
-                else None,
+                configure_index=lambda index: index.set_ef(index_params.ef_s) if index_params.ef_s else None,
             )
 
             # max number of items to retrieve per batch
@@ -51,13 +45,9 @@ class HnswlibFilterIndexInferer(IndexInferer):
             filtered_labels = []
             filtered_distances = []
             for i, item_ids in enumerate(labels):
-                non_seen_item_indexes = ~np.isin(
-                    item_ids, seen_item_ids[i], assume_unique=True
-                )
+                non_seen_item_indexes = ~np.isin(item_ids, seen_item_ids[i], assume_unique=True)
                 filtered_labels.append((item_ids[non_seen_item_indexes])[:k])
-                filtered_distances.append(
-                    (distances[i][non_seen_item_indexes])[:k]
-                )
+                filtered_distances.append((distances[i][non_seen_item_indexes])[:k])
 
             pd_res = pd.DataFrame(
                 {

replay/models/extensions/ann/index_inferers/hnswlib_index_inferer.py

@@ -1,28 +1,24 @@
 import numpy as np
 import pandas as pd
 
-from .base_inferer import IndexInferer
 from replay.models.extensions.ann.utils import create_hnswlib_index_instance
 from replay.utils import PYSPARK_AVAILABLE, PandasDataFrame, SparkDataFrame
 from replay.utils.session_handler import State
 
+from .base_inferer import IndexInferer
+
 if PYSPARK_AVAILABLE:
     from pyspark.sql.pandas.functions import pandas_udf
 
 
-# pylint: disable=too-few-public-methods
 class HnswlibIndexInferer(IndexInferer):
     """Hnswlib index inferer without filter seen items. Infers hnswlib index."""
 
-    def infer(
-        self, vectors: SparkDataFrame, features_col: str, k: int
-    ) -> SparkDataFrame:
+    def infer(self, vectors: SparkDataFrame, features_col: str, k: int) -> SparkDataFrame:
         _index_store = self.index_store
         index_params = self.index_params
 
-        index_store_broadcast = State().session.sparkContext.broadcast(
-            _index_store
-        )
+        index_store_broadcast = State().session.sparkContext.broadcast(_index_store)
 
         @pandas_udf(self.udf_return_type)
         def infer_index_udf(vectors: pd.Series) -> PandasDataFrame:  # pragma: no cover
@@ -30,9 +26,7 @@ class HnswlibIndexInferer(IndexInferer):
             index = index_store.load_index(
                 init_index=lambda: create_hnswlib_index_instance(index_params),
                 load_index=lambda index, path: index.load_index(path),
-                configure_index=lambda index: index.set_ef(index_params.ef_s)
-                if index_params.ef_s
-                else None,
+                configure_index=lambda index: index.set_ef(index_params.ef_s) if index_params.ef_s else None,
             )
 
             labels, distances = index.knn_query(
@@ -41,9 +35,7 @@ class HnswlibIndexInferer(IndexInferer):
                 num_threads=1,
             )
 
-            pd_res = pd.DataFrame(
-                {"item_idx": list(labels), "distance": list(distances)}
-            )
+            pd_res = pd.DataFrame({"item_idx": list(labels), "distance": list(distances)})
 
             return pd_res