replay-rec 0.20.3__py3-none-any.whl → 0.20.3rc0__py3-none-any.whl

replay/experimental/models/neuromf.py

@@ -0,0 +1,406 @@
+"""
+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

@@ -0,0 +1,293 @@
+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")
+        )