replay-rec 0.20.1rc0__py3-none-any.whl → 0.20.2__py3-none-any.whl

replay/experimental/models/neuromf.py

@@ -1,406 +0,0 @@
-"""
-Generalized Matrix Factorization (GMF),
-Multi-Layer Perceptron (MLP),
-Neural Matrix Factorization (MLP + GMF).
-"""
-
-from typing import Optional
-
-import numpy as np
-import torch
-import torch.nn.functional as sf
-from sklearn.model_selection import train_test_split
-from torch import LongTensor, Tensor, nn
-from torch.optim import Adam
-from torch.optim.lr_scheduler import ReduceLROnPlateau
-from torch.utils.data import DataLoader, TensorDataset
-
-from replay.experimental.models.base_torch_rec import TorchRecommender
-from replay.utils import PandasDataFrame, SparkDataFrame
-
-EMBED_DIM = 128
-
-
-def xavier_init_(layer: nn.Module):
-    """
-    Xavier initialization
-
-    :param layer: net layer
-    """
-    if isinstance(layer, (nn.Embedding, nn.Linear)):
-        nn.init.xavier_normal_(layer.weight.data)
-
-    if isinstance(layer, nn.Linear):
-        layer.bias.data.normal_(0.0, 0.001)
-
-
-class GMF(nn.Module):
-    """Generalized Matrix Factorization"""
-
-    def __init__(self, user_count: int, item_count: int, embedding_dim: int):
-        """
-        :param user_count: number of users
-        :param item_count: number of items
-        :param embedding_dim: embedding size
-        """
-        super().__init__()
-        self.user_embedding = nn.Embedding(num_embeddings=user_count, embedding_dim=embedding_dim)
-        self.item_embedding = nn.Embedding(num_embeddings=item_count, embedding_dim=embedding_dim)
-        self.item_biases = nn.Embedding(num_embeddings=item_count, embedding_dim=1)
-        self.user_biases = nn.Embedding(num_embeddings=user_count, embedding_dim=1)
-
-        xavier_init_(self.user_embedding)
-        xavier_init_(self.item_embedding)
-        self.user_biases.weight.data.zero_()
-        self.item_biases.weight.data.zero_()
-
-    def forward(self, user: Tensor, item: Tensor) -> Tensor:
-        """
-        :param user: user id batch
-        :param item: item id batch
-        :return: model output
-        """
-        user_emb = self.user_embedding(user) + self.user_biases(user)
-        item_emb = self.item_embedding(item) + self.item_biases(item)
-        element_product = torch.mul(user_emb, item_emb)
-
-        return element_product
-
-
-class MLP(nn.Module):
-    """Multi-Layer Perceptron"""
-
-    def __init__(
-        self,
-        user_count: int,
-        item_count: int,
-        embedding_dim: int,
-        hidden_dims: Optional[list[int]] = None,
-    ):
-        """
-        :param user_count: number of users
-        :param item_count: number of items
-        :param embedding_dim: embedding size
-        :param hidden_dims: list of hidden dimension sizes
-        """
-        super().__init__()
-        self.user_embedding = nn.Embedding(num_embeddings=user_count, embedding_dim=embedding_dim)
-        self.item_embedding = nn.Embedding(num_embeddings=item_count, embedding_dim=embedding_dim)
-        self.item_biases = nn.Embedding(num_embeddings=item_count, embedding_dim=1)
-        self.user_biases = nn.Embedding(num_embeddings=user_count, embedding_dim=1)
-
-        if hidden_dims:
-            full_hidden_dims = [2 * embedding_dim, *hidden_dims]
-            self.hidden_layers = nn.ModuleList(
-                [nn.Linear(d_in, d_out) for d_in, d_out in zip(full_hidden_dims[:-1], full_hidden_dims[1:])]
-            )
-
-        else:
-            self.hidden_layers = nn.ModuleList()
-
-        self.activation = nn.ReLU()
-
-        xavier_init_(self.user_embedding)
-        xavier_init_(self.item_embedding)
-        self.user_biases.weight.data.zero_()
-        self.item_biases.weight.data.zero_()
-        for layer in self.hidden_layers:
-            xavier_init_(layer)
-
-    def forward(self, user: Tensor, item: Tensor) -> Tensor:
-        """
-        :param user: user id batch
-        :param item: item id batch
-        :return: output
-        """
-        user_emb = self.user_embedding(user) + self.user_biases(user)
-        item_emb = self.item_embedding(item) + self.item_biases(item)
-        hidden = torch.cat([user_emb, item_emb], dim=-1)
-        for layer in self.hidden_layers:
-            hidden = layer(hidden)
-            hidden = self.activation(hidden)
-        return hidden
-
-
-class NMF(nn.Module):
-    """NMF = MLP + GMF"""
-
-    def __init__(
-        self,
-        user_count: int,
-        item_count: int,
-        embedding_gmf_dim: Optional[int] = None,
-        embedding_mlp_dim: Optional[int] = None,
-        hidden_mlp_dims: Optional[list[int]] = None,
-    ):
-        """
-        :param user_count: number of users
-        :param item_count: number of items
-        :param embedding_gmf_dim: embedding size for gmf
-        :param embedding_mlp_dim: embedding size for mlp
-        :param hidden_mlp_dims: list of hidden dimension sizes for mlp
-        """
-        self.gmf: Optional[GMF] = None
-        self.mlp: Optional[MLP] = None
-
-        super().__init__()
-        merged_dim = 0
-        if embedding_gmf_dim:
-            self.gmf = GMF(user_count, item_count, embedding_gmf_dim)
-            merged_dim += embedding_gmf_dim
-
-        if embedding_mlp_dim:
-            self.mlp = MLP(user_count, item_count, embedding_mlp_dim, hidden_mlp_dims)
-            merged_dim += hidden_mlp_dims[-1] if hidden_mlp_dims else 2 * embedding_mlp_dim
-
-        self.last_layer = nn.Linear(merged_dim, 1)
-        xavier_init_(self.last_layer)
-
-    def forward(self, user: Tensor, item: Tensor) -> Tensor:
-        """
-        :param user: user id batch
-        :param item: item id batch
-        :return: output
-        """
-        batch_size = len(user)
-        gmf_vector = self.gmf(user, item) if self.gmf else torch.zeros(batch_size, 0).to(user.device)
-        mlp_vector = self.mlp(user, item) if self.mlp else torch.zeros(batch_size, 0).to(user.device)
-
-        merged_vector = torch.cat([gmf_vector, mlp_vector], dim=1)
-        merged_vector = self.last_layer(merged_vector).squeeze(dim=1)
-        merged_vector = torch.sigmoid(merged_vector)
-
-        return merged_vector
-
-
-class NeuroMF(TorchRecommender):
-    """
-    Neural Matrix Factorization model (NeuMF, NCF).
-
-    In this implementation MLP and GMF modules are optional.
-    """
-
-    num_workers: int = 0
-    batch_size_users: int = 100000
-    patience: int = 3
-    n_saved: int = 2
-    valid_split_size: float = 0.1
-    seed: int = 42
-    _search_space = {
-        "embedding_gmf_dim": {"type": "int", "args": [EMBED_DIM, EMBED_DIM]},
-        "embedding_mlp_dim": {"type": "int", "args": [EMBED_DIM, EMBED_DIM]},
-        "learning_rate": {"type": "loguniform", "args": [0.0001, 0.5]},
-        "l2_reg": {"type": "loguniform", "args": [1e-9, 5]},
-        "count_negative_sample": {"type": "int", "args": [1, 20]},
-        "factor": {"type": "uniform", "args": [0.2, 0.2]},
-        "patience": {"type": "int", "args": [3, 3]},
-    }
-
-    def __init__(
-        self,
-        learning_rate: float = 0.05,
-        epochs: int = 20,
-        embedding_gmf_dim: Optional[int] = None,
-        embedding_mlp_dim: Optional[int] = None,
-        hidden_mlp_dims: Optional[list[int]] = None,
-        l2_reg: float = 0,
-        count_negative_sample: int = 1,
-        factor: float = 0.2,
-        patience: int = 3,
-    ):
-        """
-        MLP or GMF model can be ignored if
-        its embedding size (embedding_mlp_dim or embedding_gmf_dim) is set to ``None``.
-        Default variant is MLP + GMF with embedding size 128.
-
-        :param learning_rate: learning rate
-        :param epochs: number of epochs to train model
-        :param embedding_gmf_dim: embedding size for gmf
-        :param embedding_mlp_dim: embedding size for mlp
-        :param hidden_mlp_dims: list of hidden dimension sized for mlp
-        :param l2_reg: l2 regularization term
-        :param count_negative_sample: number of negative samples to use
-        :param factor: ReduceLROnPlateau reducing factor. new_lr = lr * factor
-        :param patience: number of non-improved epochs before reducing lr
-        """
-        super().__init__()
-        if not embedding_gmf_dim and not embedding_mlp_dim:
-            embedding_gmf_dim, embedding_mlp_dim = EMBED_DIM, EMBED_DIM
-
-        if (embedding_gmf_dim is None or embedding_gmf_dim < 0) and (
-            embedding_mlp_dim is None or embedding_mlp_dim < 0
-        ):
-            msg = "embedding_gmf_dim and embedding_mlp_dim must be positive"
-            raise ValueError(msg)
-
-        self.learning_rate = learning_rate
-        self.epochs = epochs
-        self.embedding_gmf_dim = embedding_gmf_dim
-        self.embedding_mlp_dim = embedding_mlp_dim
-        self.hidden_mlp_dims = hidden_mlp_dims
-        self.l2_reg = l2_reg
-        self.count_negative_sample = count_negative_sample
-        self.factor = factor
-        self.patience = patience
-
-    @property
-    def _init_args(self):
-        return {
-            "learning_rate": self.learning_rate,
-            "epochs": self.epochs,
-            "embedding_gmf_dim": self.embedding_gmf_dim,
-            "embedding_mlp_dim": self.embedding_mlp_dim,
-            "hidden_mlp_dims": self.hidden_mlp_dims,
-            "l2_reg": self.l2_reg,
-            "count_negative_sample": self.count_negative_sample,
-            "factor": self.factor,
-            "patience": self.patience,
-        }
-
-    def _data_loader(self, data: PandasDataFrame, shuffle: bool = True) -> DataLoader:
-        user_batch = LongTensor(data["user_idx"].values)
-        item_batch = LongTensor(data["item_idx"].values)
-
-        dataset = TensorDataset(user_batch, item_batch)
-
-        loader = DataLoader(
-            dataset,
-            batch_size=self.batch_size_users,
-            shuffle=shuffle,
-            num_workers=self.num_workers,
-        )
-        return loader
-
-    def _get_neg_batch(self, batch: Tensor) -> Tensor:
-        return torch.from_numpy(np.random.choice(self._fit_items_np, batch.shape[0] * self.count_negative_sample))
-
-    def _fit(
-        self,
-        log: SparkDataFrame,
-        user_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-        item_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-    ) -> None:
-        self.logger.debug("Create DataLoaders")
-        tensor_data = log.select("user_idx", "item_idx").toPandas()
-        train_tensor_data, valid_tensor_data = train_test_split(
-            tensor_data,
-            test_size=self.valid_split_size,
-            random_state=self.seed,
-        )
-        train_data_loader = self._data_loader(train_tensor_data)
-        valid_data_loader = self._data_loader(valid_tensor_data)
-        self._fit_items_np = self.fit_items.toPandas().to_numpy().ravel()
-
-        self.logger.debug("Training NeuroMF")
-        self.model = NMF(
-            user_count=self._user_dim,
-            item_count=self._item_dim,
-            embedding_gmf_dim=self.embedding_gmf_dim,
-            embedding_mlp_dim=self.embedding_mlp_dim,
-            hidden_mlp_dims=self.hidden_mlp_dims,
-        ).to(self.device)
-        optimizer = Adam(
-            self.model.parameters(),
-            lr=self.learning_rate,
-            weight_decay=self.l2_reg / self.batch_size_users,
-        )
-        lr_scheduler = ReduceLROnPlateau(optimizer, factor=self.factor, patience=self.patience)
-
-        self.train(
-            train_data_loader,
-            valid_data_loader,
-            optimizer,
-            lr_scheduler,
-            self.epochs,
-            "neuromf",
-        )
-
-        del self._fit_items_np
-
-    @staticmethod
-    def _loss(y_pred, y_true):
-        return sf.binary_cross_entropy(y_pred, y_true).mean()
-
-    def _batch_pass(self, batch, model):
-        user_batch, pos_item_batch = batch
-        neg_item_batch = self._get_neg_batch(user_batch)
-        pos_relevance = model(user_batch.to(self.device), pos_item_batch.to(self.device))
-        neg_relevance = model(
-            user_batch.repeat([self.count_negative_sample]).to(self.device),
-            neg_item_batch.to(self.device),
-        )
-        y_pred = torch.cat((pos_relevance, neg_relevance), 0)
-        y_true_pos = torch.ones_like(pos_item_batch).to(self.device)
-        y_true_neg = torch.zeros_like(neg_item_batch).to(self.device)
-        y_true = torch.cat((y_true_pos, y_true_neg), 0).float()
-
-        return {"y_pred": y_pred, "y_true": y_true}
-
-    @staticmethod
-    def _predict_pairs_inner(
-        model: nn.Module,
-        user_idx: int,
-        items_np: np.ndarray,
-        cnt: Optional[int] = None,
-    ) -> SparkDataFrame:
-        model.eval()
-        with torch.no_grad():
-            user_batch = LongTensor([user_idx] * len(items_np))
-            item_batch = LongTensor(items_np)
-            user_recs = torch.reshape(
-                model(user_batch, item_batch).detach(),
-                [
-                    -1,
-                ],
-            )
-            if cnt is not None:
-                best_item_idx = (torch.argsort(user_recs, descending=True)[:cnt]).numpy()
-                user_recs = user_recs[best_item_idx]
-                items_np = items_np[best_item_idx]
-
-            return PandasDataFrame(
-                {
-                    "user_idx": user_recs.shape[0] * [user_idx],
-                    "item_idx": items_np,
-                    "relevance": user_recs,
-                }
-            )
-
-    @staticmethod
-    def _predict_by_user(
-        pandas_df: PandasDataFrame,
-        model: nn.Module,
-        items_np: np.ndarray,
-        k: int,
-        item_count: int,  # noqa: ARG004
-    ) -> PandasDataFrame:
-        return NeuroMF._predict_pairs_inner(
-            model=model,
-            user_idx=pandas_df["user_idx"][0],
-            items_np=items_np,
-            cnt=min(len(pandas_df) + k, len(items_np)),
-        )
-
-    @staticmethod
-    def _predict_by_user_pairs(
-        pandas_df: PandasDataFrame,
-        model: nn.Module,
-        item_count: int,  # noqa: ARG004
-    ) -> PandasDataFrame:
-        return NeuroMF._predict_pairs_inner(
-            model=model,
-            user_idx=pandas_df["user_idx"][0],
-            items_np=np.array(pandas_df["item_idx_to_pred"][0]),
-            cnt=None,
-        )
-
-    def _load_model(self, path: str):
-        self.model = NMF(
-            user_count=self._user_dim,
-            item_count=self._item_dim,
-            embedding_gmf_dim=self.embedding_gmf_dim,
-            embedding_mlp_dim=self.embedding_mlp_dim,
-            hidden_mlp_dims=self.hidden_mlp_dims,
-        ).to(self.device)
-        self.model.load_state_dict(torch.load(path))
-        self.model.eval()

replay/experimental/models/scala_als.py

@@ -1,293 +0,0 @@
-from typing import Any, Optional
-
-from replay.experimental.models.base_rec import ItemVectorModel, Recommender
-from replay.experimental.models.extensions.spark_custom_models.als_extension import ALS, ALSModel
-from replay.models.extensions.ann.ann_mixin import ANNMixin
-from replay.models.extensions.ann.index_builders.base_index_builder import IndexBuilder
-from replay.utils import OPTUNA_AVAILABLE, PYSPARK_AVAILABLE, SparkDataFrame
-from replay.utils.spark_utils import list_to_vector_udf
-
-if PYSPARK_AVAILABLE:
-    import pyspark.sql.functions as sf
-    from pyspark.sql.types import DoubleType
-
-
-class ALSWrap(Recommender, ItemVectorModel):
-    """Wrapper for `Spark ALS
-    <https://spark.apache.org/docs/latest/api/python/pyspark.mllib.html#pyspark.mllib.recommendation.ALS>`_.
-    """
-
-    _seed: Optional[int] = None
-    if OPTUNA_AVAILABLE:
-        _search_space = {
-            "rank": {"type": "loguniform_int", "args": [8, 256]},
-        }
-
-    def __init__(
-        self,
-        rank: int = 10,
-        implicit_prefs: bool = True,
-        seed: Optional[int] = None,
-        num_item_blocks: Optional[int] = None,
-        num_user_blocks: Optional[int] = None,
-    ):
-        """
-        :param rank: hidden dimension for the approximate matrix
-        :param implicit_prefs: flag to use implicit feedback
-        :param seed: random seed
-        :param num_item_blocks: number of blocks the items will be partitioned into in order
-            to parallelize computation.
-            if None then will be init with number of partitions of log.
-        :param num_user_blocks: number of blocks the users will be partitioned into in order
-            to parallelize computation.
-            if None then will be init with number of partitions of log.
-        """
-        self.rank = rank
-        self.implicit_prefs = implicit_prefs
-        self._seed = seed
-        self._num_item_blocks = num_item_blocks
-        self._num_user_blocks = num_user_blocks
-
-    @property
-    def _init_args(self):
-        return {
-            "rank": self.rank,
-            "implicit_prefs": self.implicit_prefs,
-            "seed": self._seed,
-        }
-
-    def _save_model(self, path: str):
-        self.model.write().overwrite().save(path)
-
-    def _load_model(self, path: str):
-        self.model = ALSModel.load(path)
-        self.model.itemFactors.cache()
-        self.model.userFactors.cache()
-
-    def _fit(
-        self,
-        log: SparkDataFrame,
-        user_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-        item_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-    ) -> None:
-        if self._num_item_blocks is None:
-            self._num_item_blocks = log.rdd.getNumPartitions()
-        if self._num_user_blocks is None:
-            self._num_user_blocks = log.rdd.getNumPartitions()
-
-        self.model = ALS(
-            rank=self.rank,
-            numItemBlocks=self._num_item_blocks,
-            numUserBlocks=self._num_user_blocks,
-            userCol="user_idx",
-            itemCol="item_idx",
-            ratingCol="relevance",
-            implicitPrefs=self.implicit_prefs,
-            seed=self._seed,
-            coldStartStrategy="drop",
-        ).fit(log)
-        self.model.itemFactors.cache()
-        self.model.userFactors.cache()
-        self.model.itemFactors.count()
-        self.model.userFactors.count()
-
-    def _clear_cache(self):
-        if hasattr(self, "model"):
-            self.model.itemFactors.unpersist()
-            self.model.userFactors.unpersist()
-
-    def _predict(
-        self,
-        log: Optional[SparkDataFrame],
-        k: int,
-        users: SparkDataFrame,
-        items: SparkDataFrame,
-        user_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-        item_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-        filter_seen_items: bool = True,
-    ) -> SparkDataFrame:
-        if (items.count() == self.fit_items.count()) and (
-            items.join(self.fit_items, on="item_idx", how="inner").count() == self.fit_items.count()
-        ):
-            max_seen = 0
-            if filter_seen_items and log is not None:
-                max_seen_in_log = (
-                    log.join(users, on="user_idx")
-                    .groupBy("user_idx")
-                    .agg(sf.count("user_idx").alias("num_seen"))
-                    .select(sf.max("num_seen"))
-                    .first()[0]
-                )
-                max_seen = max_seen_in_log if max_seen_in_log is not None else 0
-
-            recs_als = self.model.recommendForUserSubset(users, k + max_seen)
-            return (
-                recs_als.withColumn("recommendations", sf.explode("recommendations"))
-                .withColumn("item_idx", sf.col("recommendations.item_idx"))
-                .withColumn(
-                    "relevance",
-                    sf.col("recommendations.rating").cast(DoubleType()),
-                )
-                .select("user_idx", "item_idx", "relevance")
-            )
-
-        return self._predict_pairs(
-            pairs=users.crossJoin(items).withColumn("relevance", sf.lit(1)),
-            log=log,
-        )
-
-    def _predict_pairs(
-        self,
-        pairs: SparkDataFrame,
-        log: Optional[SparkDataFrame] = None,  # noqa: ARG002
-        user_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-        item_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-    ) -> SparkDataFrame:
-        return (
-            self.model.transform(pairs)
-            .withColumn("relevance", sf.col("prediction").cast(DoubleType()))
-            .drop("prediction")
-        )
-
-    def _get_features(
-        self, ids: SparkDataFrame, features: Optional[SparkDataFrame]  # noqa: ARG002
-    ) -> tuple[Optional[SparkDataFrame], Optional[int]]:
-        entity = "user" if "user_idx" in ids.columns else "item"
-        als_factors = getattr(self.model, f"{entity}Factors")
-        als_factors = als_factors.withColumnRenamed("id", f"{entity}_idx").withColumnRenamed(
-            "features", f"{entity}_factors"
-        )
-        return (
-            als_factors.join(ids, how="right", on=f"{entity}_idx"),
-            self.model.rank,
-        )
-
-    def _get_item_vectors(self):
-        return self.model.itemFactors.select(
-            sf.col("id").alias("item_idx"),
-            list_to_vector_udf(sf.col("features")).alias("item_vector"),
-        )
-
-
-class ScalaALSWrap(ALSWrap, ANNMixin):
-    """Wrapper for `Spark ALS
-    <https://spark.apache.org/docs/latest/api/python/pyspark.mllib.html#pyspark.mllib.recommendation.ALS>`_.
-    """
-
-    def _get_ann_infer_params(self) -> dict[str, Any]:
-        self.index_builder.index_params.dim = self.rank
-        return {
-            "features_col": "user_factors",
-        }
-
-    def _get_vectors_to_infer_ann_inner(
-        self, interactions: SparkDataFrame, queries: SparkDataFrame  # noqa: ARG002
-    ) -> SparkDataFrame:
-        user_vectors, _ = self.get_features(queries)
-        return user_vectors
-
-    def _configure_index_builder(self, interactions: SparkDataFrame):
-        item_vectors, _ = self.get_features(interactions.select("item_idx").distinct())
-
-        self.index_builder.index_params.dim = self.rank
-        self.index_builder.index_params.max_elements = interactions.select("item_idx").distinct().count()
-
-        return item_vectors, {
-            "features_col": "item_factors",
-            "ids_col": "item_idx",
-        }
-
-    def __init__(
-        self,
-        rank: int = 10,
-        implicit_prefs: bool = True,
-        seed: Optional[int] = None,
-        num_item_blocks: Optional[int] = None,
-        num_user_blocks: Optional[int] = None,
-        index_builder: Optional[IndexBuilder] = None,
-    ):
-        ALSWrap.__init__(self, rank, implicit_prefs, seed, num_item_blocks, num_user_blocks)
-        self.init_index_builder(index_builder)
-        self.num_elements = None
-
-    @property
-    def _init_args(self):
-        return {
-            "rank": self.rank,
-            "implicit_prefs": self.implicit_prefs,
-            "seed": self._seed,
-            "index_builder": self.index_builder.init_meta_as_dict() if self.index_builder else None,
-        }
-
-    def _fit(
-        self,
-        log: SparkDataFrame,
-        user_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-        item_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-    ) -> None:
-        if self._num_item_blocks is None:
-            self._num_item_blocks = log.rdd.getNumPartitions()
-        if self._num_user_blocks is None:
-            self._num_user_blocks = log.rdd.getNumPartitions()
-
-        self.model = ALS(
-            rank=self.rank,
-            numItemBlocks=self._num_item_blocks,
-            numUserBlocks=self._num_user_blocks,
-            userCol="user_idx",
-            itemCol="item_idx",
-            ratingCol="relevance",
-            implicitPrefs=self.implicit_prefs,
-            seed=self._seed,
-            coldStartStrategy="drop",
-        ).fit(log)
-        self.model.itemFactors.cache()
-        self.model.userFactors.cache()
-        self.model.itemFactors.count()
-        self.model.userFactors.count()
-
-    def _save_model(self, path: str):
-        self.model.write().overwrite().save(path)
-
-        if self._use_ann:
-            self._save_index(path)
-
-    def _load_model(self, path: str):
-        self.model = ALSModel.load(path)
-        self.model.itemFactors.cache()
-        self.model.userFactors.cache()
-
-        if self._use_ann:
-            self._load_index(path)
-
-    def _predict(
-        self,
-        log: Optional[SparkDataFrame],
-        k: int,
-        users: SparkDataFrame,
-        items: SparkDataFrame,
-        user_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-        item_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-        filter_seen_items: bool = True,
-    ) -> SparkDataFrame:
-        max_seen = 0
-        if filter_seen_items and log is not None:
-            max_seen_in_log = (
-                log.join(users, on="user_idx")
-                .groupBy("user_idx")
-                .agg(sf.count("user_idx").alias("num_seen"))
-                .select(sf.max("num_seen"))
-                .first()[0]
-            )
-            max_seen = max_seen_in_log if max_seen_in_log is not None else 0
-
-        recs_als = self.model.recommendItemsForUserItemSubset(users, items, k + max_seen)
-        return (
-            recs_als.withColumn("recommendations", sf.explode("recommendations"))
-            .withColumn("item_idx", sf.col("recommendations.item_idx"))
-            .withColumn(
-                "relevance",
-                sf.col("recommendations.rating").cast(DoubleType()),
-            )
-            .select("user_idx", "item_idx", "relevance")
-        )