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

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

@@ -11,22 +11,17 @@ if PYSPARK_AVAILABLE:
     from pyspark.sql.pandas.functions import pandas_udf
 
 
-# pylint: disable=too-few-public-methods
 class NmslibFilterIndexInferer(IndexInferer):
     """Nmslib index inferer with filter seen items. Infers nmslib hnsw index."""
 
-    def infer(
-        self, vectors: SparkDataFrame, features_col: str, k: int
-    ) -> SparkDataFrame:
+    def infer(self, vectors: SparkDataFrame, features_col: str, k: int) -> SparkDataFrame:  # noqa: ARG002
         _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(  # pylint: disable=too-many-locals
+        def infer_index_udf(
             user_idx: pd.Series,
             vector_items: pd.Series,
             vector_ratings: pd.Series,
@@ -36,12 +31,8 @@ class NmslibFilterIndexInferer(IndexInferer):
             index_store = index_store_broadcast.value
             index = index_store.load_index(
                 init_index=lambda: create_nmslib_index_instance(index_params),
-                load_index=lambda index, path: index.loadIndex(
-                    path, load_data=True
-                ),
-                configure_index=lambda index: index.setQueryTimeParams(
-                    {"efSearch": index_params.ef_s}
-                )
+                load_index=lambda index, path: index.loadIndex(path, load_data=True),
+                configure_index=lambda index: index.setQueryTimeParams({"efSearch": index_params.ef_s})
                 if index_params.ef_s
                 else None,
             )
@@ -49,9 +40,7 @@ class NmslibFilterIndexInferer(IndexInferer):
             # max number of items to retrieve per batch
             max_items_to_retrieve = num_items.max()
 
-            user_vectors = get_csr_matrix(
-                user_idx, vector_items, vector_ratings
-            )
+            user_vectors = get_csr_matrix(user_idx, vector_items, vector_ratings)
 
             neighbours = index.knnQueryBatch(
                 user_vectors[user_idx.values, :],
@@ -61,9 +50,7 @@ class NmslibFilterIndexInferer(IndexInferer):
 
             neighbours_filtered = []
             for i, (item_idxs, distances) in enumerate(neighbours):
-                non_seen_item_indexes = ~np.isin(
-                    item_idxs, seen_item_ids[i], assume_unique=True
-                )
+                non_seen_item_indexes = ~np.isin(item_idxs, seen_item_ids[i], assume_unique=True)
                 neighbours_filtered.append(
                     (
                         (item_idxs[non_seen_item_indexes])[:k],
@@ -71,14 +58,14 @@ class NmslibFilterIndexInferer(IndexInferer):
                     )
                 )
 
-            pd_res = PandasDataFrame(
-                neighbours_filtered, columns=["item_idx", "distance"]
-            )
+            pd_res = PandasDataFrame(neighbours_filtered, columns=["item_idx", "distance"])
 
-            # pd_res looks like
-            # item_idx       distances
-            # [1, 2, 3, ...] [-0.5, -0.3, -0.1, ...]
-            # [1, 3, 4, ...] [-0.1, -0.8, -0.2, ...]
+            """
+            pd_res looks like
+            item_idx       distances
+            [1, 2, 3, ...] [-0.5, -0.3, -0.1, ...]
+            [1, 3, 4, ...] [-0.1, -0.8, -0.2, ...]
+            """
 
             return pd_res
 
@@ -89,7 +76,6 @@ class NmslibFilterIndexInferer(IndexInferer):
             "num_items",
             "seen_item_idxs",
         ]
-        # cols = cols + ["num_items", "seen_item_idxs"]
 
         res = vectors.select(
             "user_idx",

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

@@ -1,28 +1,24 @@
 import pandas as pd
 
-from .base_inferer import IndexInferer
-from .utils import get_csr_matrix
 from replay.models.extensions.ann.utils import create_nmslib_index_instance
 from replay.utils import PYSPARK_AVAILABLE, PandasDataFrame, SparkDataFrame
 from replay.utils.session_handler import State
 
+from .base_inferer import IndexInferer
+from .utils import get_csr_matrix
+
 if PYSPARK_AVAILABLE:
     from pyspark.sql.pandas.functions import pandas_udf
 
 
-# pylint: disable=too-few-public-methods
 class NmslibIndexInferer(IndexInferer):
     """Nmslib index inferer without filter seen items. Infers nmslib hnsw index."""
 
-    def infer(
-        self, vectors: SparkDataFrame, features_col: str, k: int
-    ) -> SparkDataFrame:
+    def infer(self, vectors: SparkDataFrame, features_col: str, k: int) -> SparkDataFrame:  # noqa: ARG002
         _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(
@@ -33,29 +29,23 @@ class NmslibIndexInferer(IndexInferer):
             index_store = index_store_broadcast.value
             index = index_store.load_index(
                 init_index=lambda: create_nmslib_index_instance(index_params),
-                load_index=lambda index, path: index.loadIndex(
-                    path, load_data=True
-                ),
-                configure_index=lambda index: index.setQueryTimeParams(
-                    {"efSearch": index_params.ef_s}
-                )
+                load_index=lambda index, path: index.loadIndex(path, load_data=True),
+                configure_index=lambda index: index.setQueryTimeParams({"efSearch": index_params.ef_s})
                 if index_params.ef_s
                 else None,
             )
 
-            user_vectors = get_csr_matrix(
-                user_idx, vector_items, vector_ratings
-            )
-            neighbours = index.knnQueryBatch(
-                user_vectors[user_idx.values, :], k=k, num_threads=1
-            )
+            user_vectors = get_csr_matrix(user_idx, vector_items, vector_ratings)
+            neighbours = index.knnQueryBatch(user_vectors[user_idx.values, :], k=k, num_threads=1)
 
             pd_res = PandasDataFrame(neighbours, columns=["item_idx", "distance"])
 
-            # pd_res looks like
-            # item_idx       distances
-            # [1, 2, 3, ...] [-0.5, -0.3, -0.1, ...]
-            # [1, 3, 4, ...] [-0.1, -0.8, -0.2, ...]
+            """
+            pd_res looks like
+            item_idx       distances
+            [1, 2, 3, ...] [-0.5, -0.3, -0.1, ...]
+            [1, 3, 4, ...] [-0.1, -0.8, -0.2, ...]
+            """
 
             return pd_res
 

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

@@ -12,19 +12,12 @@ def get_csr_matrix(
         (
             vector_ratings.explode().values.astype(float),
             (
-                user_idx.repeat(
-                    vector_items.apply(
-                        lambda x: len(x)  # pylint: disable=unnecessary-lambda
-                    )
-                ).values,
+                user_idx.repeat(vector_items.apply(lambda x: len(x))).values,
                 vector_items.explode().values.astype(int),
             ),
         ),
         shape=(
             user_idx.max() + 1,
-            vector_items.apply(
-                lambda x: max(x)  # pylint: disable=unnecessary-lambda
-            ).max()
-            + 1,
+            vector_items.apply(lambda x: max(x)).max() + 1,
         ),
     )

replay/models/extensions/ann/index_stores/hdfs_index_store.py

@@ -19,12 +19,9 @@ class HdfsIndexStore(IndexStore):
         index_dir_path = os.path.join(warehouse_dir, index_dir)
         self._index_dir_info = get_filesystem(index_dir_path)
         if self._index_dir_info.filesystem != FileSystem.HDFS:
-            raise ValueError(
-                f"Can't recognize path {index_dir_path} as HDFS path!"
-            )
-        self._hadoop_fs = fs.HadoopFileSystem.from_uri(
-            self._index_dir_info.hdfs_uri
-        )
+            msg = f"Can't recognize path {index_dir_path} as HDFS path!"
+            raise ValueError(msg)
+        self._hadoop_fs = fs.HadoopFileSystem.from_uri(self._index_dir_info.hdfs_uri)
         super().__init__()
 
         if self.cleanup:
@@ -32,9 +29,7 @@ class HdfsIndexStore(IndexStore):
                 "Index directory %s is marked for deletion via weakref.finalize()",
                 self._index_dir_info.path,
             )
-            weakref.finalize(
-                self, self._hadoop_fs.delete_dir, self._index_dir_info.path
-            )
+            weakref.finalize(self, self._hadoop_fs.delete_dir, self._index_dir_info.path)
 
     def load_index(
         self,

replay/models/extensions/ann/index_stores/shared_disk_index_store.py

@@ -17,9 +17,7 @@ class SharedDiskIndexStore(IndexStore):
     It can also be used with a local disk when the driver and executors
     are running on the same machine."""
 
-    def __init__(
-        self, warehouse_dir: str, index_dir: str, cleanup: bool = True
-    ):
+    def __init__(self, warehouse_dir: str, index_dir: str, cleanup: bool = True):
         self.index_dir_path = os.path.join(warehouse_dir, index_dir)
         super().__init__(cleanup)
         if self.cleanup:
@@ -52,9 +50,7 @@ class SharedDiskIndexStore(IndexStore):
         save_index(temp_file_path)
 
     def dump_index(self, target_path: str):
-        destination_filesystem, target_path = fs.FileSystem.from_uri(
-            target_path
-        )
+        destination_filesystem, target_path = fs.FileSystem.from_uri(target_path)
         target_path = os.path.join(target_path, "index_files")
         destination_filesystem.create_dir(target_path)
         fs.copy_files(

replay/models/extensions/ann/index_stores/spark_files_index_store.py

@@ -7,11 +7,12 @@ from typing import Any, Callable
 
 from pyarrow import fs
 
-from .base_index_store import IndexStore
-from .utils import FileSystem, get_filesystem
 from replay.utils import PYSPARK_AVAILABLE
 from replay.utils.session_handler import State
 
+from .base_index_store import IndexStore
+from .utils import FileSystem, get_filesystem
+
 if PYSPARK_AVAILABLE:
     from pyspark import SparkFiles
 
@@ -20,6 +21,7 @@ logger = logging.getLogger("replay")
 
 
 if PYSPARK_AVAILABLE:
+
     class SparkFilesIndexStore(IndexStore):
         """Class that responsible for index store in spark files.
         Works through SparkContext.addFile()."""
@@ -62,14 +64,10 @@ if PYSPARK_AVAILABLE:
             for filename in os.listdir(self.index_dir_path):
                 index_file_path = os.path.join(self.index_dir_path, filename)
                 spark.sparkContext.addFile("file://" + index_file_path)
-                logger.info(
-                    "Index file %s transferred to executors", index_file_path
-                )
+                logger.info("Index file %s transferred to executors", index_file_path)
 
         def dump_index(self, target_path: str):
-            destination_filesystem, target_path = fs.FileSystem.from_uri(
-                target_path
-            )
+            destination_filesystem, target_path = fs.FileSystem.from_uri(target_path)
             target_path = os.path.join(target_path, "index_files")
             destination_filesystem.create_dir(target_path)
             fs.copy_files(
@@ -83,9 +81,7 @@ if PYSPARK_AVAILABLE:
             """Loads index from `path` directory to spark files."""
             path_info = get_filesystem(path)
             source_filesystem, path = fs.FileSystem.from_uri(
-                path_info.hdfs_uri + path_info.path
-                if path_info.filesystem == FileSystem.HDFS
-                else path_info.path
+                path_info.hdfs_uri + path_info.path if path_info.filesystem == FileSystem.HDFS else path_info.path
             )
             path = os.path.join(path, "index_files")
             self.index_dir_path: str = tempfile.mkdtemp()
@@ -100,6 +96,4 @@ if PYSPARK_AVAILABLE:
             for filename in os.listdir(self.index_dir_path):
                 index_file_path = os.path.join(self.index_dir_path, filename)
                 spark.sparkContext.addFile("file://" + index_file_path)
-                logger.info(
-                    "Index file %s transferred to executors", index_file_path
-                )
+                logger.info("Index file %s transferred to executors", index_file_path)

replay/models/extensions/ann/index_stores/utils.py

@@ -9,6 +9,7 @@ if PYSPARK_AVAILABLE:
 
 class FileSystem(Enum):
     """File system types"""
+
     HDFS = 1
     LOCAL = 2
 
@@ -24,6 +25,7 @@ def get_default_fs() -> str:
 @dataclass(frozen=True)
 class FileInfo:
     """File meta-information: filesystem, path and hdfs_uri (optional)"""
+
     path: str
     filesystem: FileSystem
     hdfs_uri: str = None
@@ -83,15 +85,16 @@ or set 'fs.defaultFS' in hadoop configuration.
             if default_fs.startswith("hdfs://"):
                 return FileInfo(path[prefix_len:], FileSystem.HDFS, default_fs)
             else:
-                raise ValueError(
+                msg = (
                     f"Can't get default hdfs uri for path = '{path}'. "
                     "Specify an explicit path, such as 'hdfs://host:port/dir/file', "
                     "or set 'fs.defaultFS' in hadoop configuration."
                 )
+                raise ValueError(msg)
         else:
             hostname = path[prefix_len:].split("/", 1)[0]
             hdfs_uri = "hdfs://" + hostname
-            return FileInfo(path[len(hdfs_uri):], FileSystem.HDFS, hdfs_uri)
+            return FileInfo(path[len(hdfs_uri) :], FileSystem.HDFS, hdfs_uri)
     elif path.startswith("file://"):
         return FileInfo(path[prefix_len:], FileSystem.LOCAL)
     else:

replay/models/extensions/ann/utils.py

@@ -15,9 +15,7 @@ def create_hnswlib_index_instance(params: HnswlibParam, init: bool = False):
         If `False` then the index will be used to load index data from a file.
     :return: `hnswlib` index instance
     """
-    index = hnswlib.Index(  # pylint: disable=c-extension-no-member
-        space=params.space, dim=params.dim
-    )
+    index = hnswlib.Index(space=params.space, dim=params.dim)
 
     if init:
         # Initializing index - the maximum number of elements should be known beforehand
@@ -37,10 +35,10 @@ def create_nmslib_index_instance(params: NmslibHnswParam):
     :param params: `NmslibHnswParam`
     :return: `nmslib` index
     """
-    index = nmslib.init(  # pylint: disable=c-extension-no-member
+    index = nmslib.init(
         method=params.method,
         space=params.space,
-        data_type=nmslib.DataType.SPARSE_VECTOR,  # pylint: disable=c-extension-no-member
+        data_type=nmslib.DataType.SPARSE_VECTOR,
     )
 
     return index

replay/models/kl_ucb.py

@@ -1,13 +1,15 @@
 import math
-
 from typing import Optional
-from .ucb import UCB
-from replay.utils import PYSPARK_AVAILABLE
+
 from scipy.optimize import root_scalar
 
+from replay.utils import PYSPARK_AVAILABLE
+
+from .ucb import UCB
+
 if PYSPARK_AVAILABLE:
-    from pyspark.sql.types import DoubleType
     from pyspark.sql.functions import udf
+    from pyspark.sql.types import DoubleType
 
 
 class KLUCB(UCB):
@@ -17,7 +19,7 @@ class KLUCB(UCB):
     computes item relevance as an upper confidence bound of true fraction of
     positive interactions.
 
-    In a nutshell, KL-UCB сonsiders the data as the history of interactions
+    In a nutshell, KL-UCB considers the data as the history of interactions
     with items. The interaction may be either positive or negative. For each
     item the model computes empirical frequency of positive interactions
     and estimates the true frequency with an upper confidence bound. The higher
@@ -137,14 +139,11 @@ class KLUCB(UCB):
         super().__init__(exploration_coef, sample, seed)
 
     def _calc_item_popularity(self):
-
-        right_hand_side = math.log(self.full_count) \
-            + self.coef * math.log(math.log(self.full_count))
+        right_hand_side = math.log(self.full_count) + self.coef * math.log(math.log(self.full_count))
         eps = 1e-12
 
         def bernoulli_kl(proba_p, proba_q):  # pragma: no cover
-            return proba_p * math.log(proba_p / proba_q) +\
-                (1 - proba_p) * math.log((1 - proba_p) / (1 - proba_q))
+            return proba_p * math.log(proba_p / proba_q) + (1 - proba_p) * math.log((1 - proba_p) / (1 - proba_q))
 
         @udf(returnType=DoubleType())
         def get_ucb(pos, total):  # pragma: no cover
@@ -152,27 +151,22 @@ class KLUCB(UCB):
 
             if proba == 0:
                 ucb = root_scalar(
-                    f=lambda qq: math.log(1 / (1 - qq)) - right_hand_side,
-                    bracket=[0, 1 - eps],
-                    method='brentq').root
+                    f=lambda qq: math.log(1 / (1 - qq)) - right_hand_side, bracket=[0, 1 - eps], method="brentq"
+                ).root
                 return ucb
 
             if proba == 1:
                 ucb = root_scalar(
-                    f=lambda qq: math.log(1 / qq) - right_hand_side,
-                    bracket=[0 + eps, 1],
-                    method='brentq').root
+                    f=lambda qq: math.log(1 / qq) - right_hand_side, bracket=[0 + eps, 1], method="brentq"
+                ).root
                 return ucb
 
             ucb = root_scalar(
-                f=lambda q: total * bernoulli_kl(proba, q) - right_hand_side,
-                bracket=[proba, 1 - eps],
-                method='brentq').root
+                f=lambda q: total * bernoulli_kl(proba, q) - right_hand_side, bracket=[proba, 1 - eps], method="brentq"
+            ).root
             return ucb
 
-        items_counts = self.items_counts_aggr.withColumn(
-            self.rating_column, get_ucb("pos", "total")
-        )
+        items_counts = self.items_counts_aggr.withColumn(self.rating_column, get_ucb("pos", "total"))
 
         self.item_popularity = items_counts.drop("pos", "total")
 

replay/models/knn.py

@@ -1,17 +1,17 @@
 from typing import Any, Dict, Optional
 
 from replay.data import Dataset
-from .base_neighbour_rec import NeighbourRec
-from .extensions.ann.index_builders.base_index_builder import IndexBuilder
 from replay.optimization.optuna_objective import ItemKNNObjective
 from replay.utils import PYSPARK_AVAILABLE, SparkDataFrame
 
+from .base_neighbour_rec import NeighbourRec
+from .extensions.ann.index_builders.base_index_builder import IndexBuilder
+
 if PYSPARK_AVAILABLE:
     from pyspark.sql import functions as sf
     from pyspark.sql.window import Window
 
 
-# pylint: disable=too-many-ancestors, too-many-instance-attributes
 class ItemKNN(NeighbourRec):
     """Item-based ItemKNN with modified cosine similarity measure."""
 
@@ -29,15 +29,15 @@ class ItemKNN(NeighbourRec):
     _search_space = {
         "num_neighbours": {"type": "int", "args": [1, 100]},
         "shrink": {"type": "int", "args": [0, 100]},
-        "weighting": {"type": "categorical", "args": [None, "tf_idf", "bm25"]}
+        "weighting": {"type": "categorical", "args": [None, "tf_idf", "bm25"]},
     }
 
-    def __init__(  # pylint: disable=too-many-arguments
+    def __init__(
         self,
         num_neighbours: int = 10,
         use_rating: bool = False,
         shrink: float = 0.0,
-        weighting: str = None,
+        weighting: Optional[str] = None,
         index_builder: Optional[IndexBuilder] = None,
     ):
         """
@@ -54,7 +54,8 @@ class ItemKNN(NeighbourRec):
 
         valid_weightings = self._search_space["weighting"]["args"]
         if weighting not in valid_weightings:
-            raise ValueError(f"weighting must be one of {valid_weightings}")
+            msg = f"weighting must be one of {valid_weightings}"
+            raise ValueError(msg)
         self.weighting = weighting
         if isinstance(index_builder, (IndexBuilder, type(None))):
             self.index_builder = index_builder
@@ -75,8 +76,7 @@ class ItemKNN(NeighbourRec):
     def _shrink(dot_products: SparkDataFrame, shrink: float) -> SparkDataFrame:
         return dot_products.withColumn(
             "similarity",
-            sf.col("dot_product")
-            / (sf.col("norm1") * sf.col("norm2") + shrink),
+            sf.col("dot_product") / (sf.col("norm1") * sf.col("norm2") + shrink),
         ).select("item_idx_one", "item_idx_two", "similarity")
 
     def _get_similarity(self, interactions: SparkDataFrame) -> SparkDataFrame:
@@ -116,25 +116,19 @@ class ItemKNN(NeighbourRec):
         :param interactions: SparkDataFrame with interactions, `[user_id, item_id, rating]`
         :return: interactions `[user_id, item_id, rating]`
         """
-        item_stats = interactions.groupBy(self.item_column).agg(
-            sf.count(self.query_column).alias("n_queries_per_item")
-        )
+        item_stats = interactions.groupBy(self.item_column).agg(sf.count(self.query_column).alias("n_queries_per_item"))
         avgdl = item_stats.select(sf.mean("n_queries_per_item")).take(1)[0][0]
         interactions = interactions.join(item_stats, how="inner", on=self.item_column)
 
-        interactions = (
-            interactions.withColumn(
-                self.rating_column,
-                sf.col(self.rating_column) * (self.bm25_k1 + 1) / (
-                    sf.col(self.rating_column) + self.bm25_k1 * (
-                        1 - self.bm25_b + self.bm25_b * (
-                            sf.col("n_queries_per_item") / avgdl
-                        )
-                    )
-                )
-            )
-            .drop("n_queries_per_item")
-        )
+        interactions = interactions.withColumn(
+            self.rating_column,
+            sf.col(self.rating_column)
+            * (self.bm25_k1 + 1)
+            / (
+                sf.col(self.rating_column)
+                + self.bm25_k1 * (1 - self.bm25_b + self.bm25_b * (sf.col("n_queries_per_item") / avgdl))
+            ),
+        ).drop("n_queries_per_item")
 
         return interactions
 
@@ -150,23 +144,15 @@ class ItemKNN(NeighbourRec):
         n_items = interactions.select(self.item_column).distinct().count()
 
         if self.weighting == "tf_idf":
-            idf = (
-                df.withColumn("idf", sf.log1p(sf.lit(n_items) / sf.col("DF")))
-                .drop("DF")
-            )
+            idf = df.withColumn("idf", sf.log1p(sf.lit(n_items) / sf.col("DF"))).drop("DF")
         elif self.weighting == "bm25":
-            idf = (
-                df.withColumn(
-                    "idf",
-                    sf.log1p(
-                        (sf.lit(n_items) - sf.col("DF") + 0.5)
-                        / (sf.col("DF") + 0.5)
-                    ),
-                )
-                .drop("DF")
-            )
+            idf = df.withColumn(
+                "idf",
+                sf.log1p((sf.lit(n_items) - sf.col("DF") + 0.5) / (sf.col("DF") + 0.5)),
+            ).drop("DF")
         else:
-            raise ValueError("weighting must be one of ['tf_idf', 'bm25']")
+            msg = "weighting must be one of ['tf_idf', 'bm25']"
+            raise ValueError(msg)
 
         return idf
 
@@ -180,12 +166,12 @@ class ItemKNN(NeighbourRec):
         if self.weighting:
             interactions = self._reweight_interactions(interactions)
 
-        left = interactions.withColumnRenamed(
-            self.item_column, "item_idx_one"
-        ).withColumnRenamed(self.rating_column, "rel_one")
-        right = interactions.withColumnRenamed(
-            self.item_column, "item_idx_two"
-        ).withColumnRenamed(self.rating_column, "rel_two")
+        left = interactions.withColumnRenamed(self.item_column, "item_idx_one").withColumnRenamed(
+            self.rating_column, "rel_one"
+        )
+        right = interactions.withColumnRenamed(self.item_column, "item_idx_two").withColumnRenamed(
+            self.rating_column, "rel_two"
+        )
 
         dot_products = (
             left.join(right, how="inner", on=self.query_column)
@@ -201,19 +187,11 @@ class ItemKNN(NeighbourRec):
             .agg(sf.sum(self.rating_column).alias("square_norm"))
             .select(sf.col(self.item_column), sf.sqrt("square_norm").alias("norm"))
         )
-        norm1 = item_norms.withColumnRenamed(
-            self.item_column, "item_id1"
-        ).withColumnRenamed("norm", "norm1")
-        norm2 = item_norms.withColumnRenamed(
-            self.item_column, "item_id2"
-        ).withColumnRenamed("norm", "norm2")
-
-        dot_products = dot_products.join(
-            norm1, how="inner", on=sf.col("item_id1") == sf.col("item_idx_one")
-        )
-        dot_products = dot_products.join(
-            norm2, how="inner", on=sf.col("item_id2") == sf.col("item_idx_two")
-        )
+        norm1 = item_norms.withColumnRenamed(self.item_column, "item_id1").withColumnRenamed("norm", "norm1")
+        norm2 = item_norms.withColumnRenamed(self.item_column, "item_id2").withColumnRenamed("norm", "norm2")
+
+        dot_products = dot_products.join(norm1, how="inner", on=sf.col("item_id1") == sf.col("item_idx_one"))
+        dot_products = dot_products.join(norm2, how="inner", on=sf.col("item_id2") == sf.col("item_idx_two"))
 
         return dot_products
 

replay/models/nn/optimizer_utils/__init__.py

@@ -1,9 +1,4 @@
 from replay.utils import TORCH_AVAILABLE
 
 if TORCH_AVAILABLE:
-    from .optimizer_factory import (
-        FatLRSchedulerFactory,
-        FatOptimizerFactory,
-        LRSchedulerFactory,
-        OptimizerFactory
-    )
+    from .optimizer_factory import FatLRSchedulerFactory, FatOptimizerFactory, LRSchedulerFactory, OptimizerFactory

replay/models/nn/optimizer_utils/optimizer_factory.py

@@ -4,7 +4,6 @@ from typing import Iterator, Tuple
 import torch
 
 
-# pylint: disable=too-few-public-methods
 class OptimizerFactory(abc.ABC):
     """
     Interface for optimizer factory
@@ -21,7 +20,6 @@ class OptimizerFactory(abc.ABC):
         """
 
 
-# pylint: disable=too-few-public-methods
 class LRSchedulerFactory(abc.ABC):
     """
     Interface for learning rate scheduler factory
@@ -38,13 +36,11 @@ class LRSchedulerFactory(abc.ABC):
         """
 
 
-# pylint: disable=too-few-public-methods
 class FatOptimizerFactory(OptimizerFactory):
     """
     Factory that creates optimizer depending on passed parameters
     """
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         optimizer: str = "adam",
@@ -74,10 +70,11 @@ class FatOptimizerFactory(OptimizerFactory):
             return torch.optim.SGD(
                 parameters, lr=self.learning_rate, weight_decay=self.weight_decay, momentum=self.sgd_momentum
             )
-        raise ValueError("Unexpected optimizer")
+
+        msg = "Unexpected optimizer"
+        raise ValueError(msg)
 
 
-# pylint: disable=too-few-public-methods
 class FatLRSchedulerFactory(LRSchedulerFactory):
     """
     Factory that creates learning rate schedule depending on passed parameters