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

replay/experimental/models/lightfm_wrap.py

@@ -0,0 +1,303 @@
+import os
+from os.path import join
+from typing import Optional
+
+import numpy as np
+from scipy.sparse import csr_matrix, diags, hstack
+from sklearn.preprocessing import MinMaxScaler
+
+from replay.data import get_schema
+from replay.experimental.models.base_rec import HybridRecommender
+from replay.experimental.utils.session_handler import State
+from replay.preprocessing import CSRConverter
+from replay.utils import PYSPARK_AVAILABLE, PandasDataFrame, SparkDataFrame
+from replay.utils.spark_utils import check_numeric, load_pickled_from_parquet, save_picklable_to_parquet
+
+if PYSPARK_AVAILABLE:
+    import pyspark.sql.functions as sf
+
+
+class LightFMWrap(HybridRecommender):
+    """Wrapper for LightFM."""
+
+    epochs: int = 10
+    _search_space = {
+        "loss": {
+            "type": "categorical",
+            "args": ["logistic", "bpr", "warp", "warp-kos"],
+        },
+        "no_components": {"type": "loguniform_int", "args": [8, 512]},
+    }
+    user_feat_scaler: Optional[MinMaxScaler] = None
+    item_feat_scaler: Optional[MinMaxScaler] = None
+
+    def __init__(
+        self,
+        no_components: int = 128,
+        loss: str = "warp",
+        random_state: Optional[int] = None,
+    ):
+        np.random.seed(42)
+        self.no_components = no_components
+        self.loss = loss
+        self.random_state = random_state
+        cpu_count = os.cpu_count()
+        self.num_threads = cpu_count if cpu_count is not None else 1
+
+    @property
+    def _init_args(self):
+        return {
+            "no_components": self.no_components,
+            "loss": self.loss,
+            "random_state": self.random_state,
+        }
+
+    def _save_model(self, path: str):
+        save_picklable_to_parquet(self.model, join(path, "model"))
+        save_picklable_to_parquet(self.user_feat_scaler, join(path, "user_feat_scaler"))
+        save_picklable_to_parquet(self.item_feat_scaler, join(path, "item_feat_scaler"))
+
+    def _load_model(self, path: str):
+        self.model = load_pickled_from_parquet(join(path, "model"))
+        self.user_feat_scaler = load_pickled_from_parquet(join(path, "user_feat_scaler"))
+        self.item_feat_scaler = load_pickled_from_parquet(join(path, "item_feat_scaler"))
+
+    def _feature_table_to_csr(
+        self,
+        log_ids_list: SparkDataFrame,
+        feature_table: Optional[SparkDataFrame] = None,
+    ) -> Optional[csr_matrix]:
+        """
+        Transform features to sparse matrix
+        Matrix consists of two parts:
+        1) Left one is a ohe-hot encoding of user and item ids.
+        Matrix size is: number of users or items * number of user or items in fit.
+        Cold users and items are represented with empty strings
+        2) Right one is a numerical features, passed with feature_table.
+        MinMaxScaler is applied per column, and then value is divided by the row sum.
+
+        :param feature_table: dataframe with ``user_idx`` or ``item_idx``,
+            other columns are features.
+        :param log_ids_list: dataframe with ``user_idx`` or ``item_idx``,
+            containing unique ids from log.
+        :returns: feature matrix
+        """
+
+        if feature_table is None:
+            return None
+
+        check_numeric(feature_table)
+        log_ids_list = log_ids_list.distinct()
+        entity = "item" if "item_idx" in feature_table.columns else "user"
+        idx_col_name = f"{entity}_idx"
+
+        # filter features by log
+        feature_table = feature_table.join(log_ids_list, on=idx_col_name, how="inner")
+
+        fit_dim = getattr(self, f"_{entity}_dim")
+        matrix_height = max(
+            fit_dim,
+            log_ids_list.select(sf.max(idx_col_name)).first()[0] + 1,
+        )
+        if not feature_table.rdd.isEmpty():
+            matrix_height = max(
+                matrix_height,
+                feature_table.select(sf.max(idx_col_name)).first()[0] + 1,
+            )
+
+        features_np = (
+            feature_table.select(
+                idx_col_name,
+                # first column contains id, next contain features
+                *sorted(set(feature_table.columns).difference({idx_col_name})),
+            )
+            .toPandas()
+            .to_numpy()
+        )
+        entities_ids = features_np[:, 0]
+        features_np = features_np[:, 1:]
+        number_of_features = features_np.shape[1]
+
+        all_ids_list = log_ids_list.toPandas().to_numpy().ravel()
+        entities_seen_in_fit = all_ids_list[all_ids_list < fit_dim]
+
+        entity_id_features = csr_matrix(
+            (
+                [1.0] * entities_seen_in_fit.shape[0],
+                (entities_seen_in_fit, entities_seen_in_fit),
+            ),
+            shape=(matrix_height, fit_dim),
+        )
+
+        scaler_name = f"{entity}_feat_scaler"
+        if getattr(self, scaler_name) is None:
+            if not features_np.size:
+                msg = f"features for {entity}s from log are absent"
+                raise ValueError(msg)
+            setattr(self, scaler_name, MinMaxScaler().fit(features_np))
+
+        if features_np.size:
+            features_np = getattr(self, scaler_name).transform(features_np)
+            sparse_features = csr_matrix(
+                (
+                    features_np.ravel(),
+                    (
+                        np.repeat(entities_ids, number_of_features),
+                        np.tile(
+                            np.arange(number_of_features),
+                            entities_ids.shape[0],
+                        ),
+                    ),
+                ),
+                shape=(matrix_height, number_of_features),
+            )
+
+        else:
+            sparse_features = csr_matrix((matrix_height, number_of_features))
+
+        concat_features = hstack([entity_id_features, sparse_features])
+        concat_features_sum = concat_features.sum(axis=1).A.ravel()
+        mask = concat_features_sum != 0.0
+        concat_features_sum[mask] = 1.0 / concat_features_sum[mask]
+        return diags(concat_features_sum, format="csr") @ concat_features
+
+    def _fit(
+        self,
+        log: SparkDataFrame,
+        user_features: Optional[SparkDataFrame] = None,
+        item_features: Optional[SparkDataFrame] = None,
+    ) -> None:
+        from lightfm import LightFM
+
+        self.user_feat_scaler = None
+        self.item_feat_scaler = None
+
+        interactions_matrix = CSRConverter(
+            first_dim_column="user_idx",
+            second_dim_column="item_idx",
+            data_column="relevance",
+            row_count=self._user_dim,
+            column_count=self._item_dim,
+        ).transform(log)
+        csr_item_features = self._feature_table_to_csr(log.select("item_idx").distinct(), item_features)
+        csr_user_features = self._feature_table_to_csr(log.select("user_idx").distinct(), user_features)
+
+        if user_features is not None:
+            self.can_predict_cold_users = True
+        if item_features is not None:
+            self.can_predict_cold_items = True
+
+        self.model = LightFM(
+            loss=self.loss,
+            no_components=self.no_components,
+            random_state=self.random_state,
+        ).fit(
+            interactions=interactions_matrix,
+            epochs=self.epochs,
+            num_threads=self.num_threads,
+            item_features=csr_item_features,
+            user_features=csr_user_features,
+        )
+
+    def _predict_selected_pairs(
+        self,
+        pairs: SparkDataFrame,
+        user_features: Optional[SparkDataFrame] = None,
+        item_features: Optional[SparkDataFrame] = None,
+    ):
+        def predict_by_user(pandas_df: PandasDataFrame) -> PandasDataFrame:
+            pandas_df["relevance"] = model.predict(
+                user_ids=pandas_df["user_idx"].to_numpy(),
+                item_ids=pandas_df["item_idx"].to_numpy(),
+                item_features=csr_item_features,
+                user_features=csr_user_features,
+            )
+            return pandas_df
+
+        model = self.model
+
+        if self.can_predict_cold_users and user_features is None:
+            msg = "User features are missing for predict"
+            raise ValueError(msg)
+        if self.can_predict_cold_items and item_features is None:
+            msg = "Item features are missing for predict"
+            raise ValueError(msg)
+
+        csr_item_features = self._feature_table_to_csr(pairs.select("item_idx").distinct(), item_features)
+        csr_user_features = self._feature_table_to_csr(pairs.select("user_idx").distinct(), user_features)
+        rec_schema = get_schema(
+            query_column="user_idx",
+            item_column="item_idx",
+            rating_column="relevance",
+            has_timestamp=False,
+        )
+        return pairs.groupby("user_idx").applyInPandas(predict_by_user, rec_schema)
+
+    def _predict(
+        self,
+        log: SparkDataFrame,  # noqa: ARG002
+        k: int,  # noqa: ARG002
+        users: SparkDataFrame,
+        items: SparkDataFrame,
+        user_features: Optional[SparkDataFrame] = None,
+        item_features: Optional[SparkDataFrame] = None,
+        filter_seen_items: bool = True,  # noqa: ARG002
+    ) -> SparkDataFrame:
+        return self._predict_selected_pairs(users.crossJoin(items), user_features, item_features)
+
+    def _predict_pairs(
+        self,
+        pairs: SparkDataFrame,
+        log: Optional[SparkDataFrame] = None,  # noqa: ARG002
+        user_features: Optional[SparkDataFrame] = None,
+        item_features: Optional[SparkDataFrame] = None,
+    ) -> SparkDataFrame:
+        return self._predict_selected_pairs(pairs, user_features, item_features)
+
+    def _get_features(
+        self, ids: SparkDataFrame, features: Optional[SparkDataFrame]
+    ) -> tuple[Optional[SparkDataFrame], Optional[int]]:
+        """
+        Get features from LightFM.
+        LightFM has methods get_item_representations/get_user_representations,
+        which accept object matrix and return features.
+
+        :param ids: id item_idx/user_idx to get features for
+        :param features: features for item_idx/user_idx
+        :return: spark-dataframe with biases and vectors for users/items and vector size
+        """
+        entity = "item" if "item_idx" in ids.columns else "user"
+        ids_list = ids.toPandas()[f"{entity}_idx"]
+
+        # models without features use sparse matrix
+        if features is None:
+            matrix_width = getattr(self, f"fit_{entity}s").count()
+            warm_ids = ids_list[ids_list < matrix_width]
+            sparse_features = csr_matrix(
+                (
+                    [1] * warm_ids.shape[0],
+                    (warm_ids, warm_ids),
+                ),
+                shape=(ids_list.max() + 1, matrix_width),
+            )
+        else:
+            sparse_features = self._feature_table_to_csr(ids, features)
+
+        biases, vectors = getattr(self.model, f"get_{entity}_representations")(sparse_features)
+
+        embed_list = list(
+            zip(
+                ids_list,
+                biases[ids_list].tolist(),
+                vectors[ids_list].tolist(),
+            )
+        )
+        lightfm_factors = State().session.createDataFrame(
+            embed_list,
+            schema=[
+                f"{entity}_idx",
+                f"{entity}_bias",
+                f"{entity}_factors",
+            ],
+        )
+        return lightfm_factors, self.model.no_components

replay/experimental/models/mult_vae.py

@@ -0,0 +1,332 @@
+"""
+MultVAE implementation
+(Variational Autoencoders for Collaborative Filtering)
+"""
+
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as sf
+from scipy.sparse import csr_matrix
+from sklearn.model_selection import GroupShuffleSplit
+from torch import 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
+
+
+class VAE(nn.Module):
+    """Base variational autoencoder"""
+
+    def __init__(
+        self,
+        item_count: int,
+        latent_dim: int,
+        hidden_dim: int = 600,
+        dropout: float = 0.3,
+    ):
+        """
+        :param item_count: number of items
+        :param latent_dim: latent dimension size
+        :param hidden_dim: hidden dimension size for encoder and decoder
+        :param dropout: dropout coefficient
+        """
+        super().__init__()
+
+        self.latent_dim = latent_dim
+        self.encoder_dims = [item_count, hidden_dim, latent_dim * 2]
+        self.decoder_dims = [latent_dim, hidden_dim, item_count]
+
+        self.encoder = nn.ModuleList(
+            [nn.Linear(d_in, d_out) for d_in, d_out in zip(self.encoder_dims[:-1], self.encoder_dims[1:])]
+        )
+        self.decoder = nn.ModuleList(
+            [nn.Linear(d_in, d_out) for d_in, d_out in zip(self.decoder_dims[:-1], self.decoder_dims[1:])]
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.activation = torch.nn.ReLU()
+
+        for layer in self.encoder:
+            self.weight_init(layer)
+
+        for layer in self.decoder:
+            self.weight_init(layer)
+
+    def encode(self, batch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Encode"""
+        hidden = sf.normalize(batch, p=2, dim=1)
+        hidden = self.dropout(hidden)
+
+        for layer in self.encoder[:-1]:
+            hidden = layer(hidden)
+            hidden = self.activation(hidden)
+
+        hidden = self.encoder[-1](hidden)
+        mu_latent = hidden[:, : self.latent_dim]
+        logvar_latent = hidden[:, self.latent_dim :]
+        return mu_latent, logvar_latent
+
+    def reparameterize(self, mu_latent: torch.Tensor, logvar_latent: torch.Tensor) -> torch.Tensor:
+        """Reparametrization trick"""
+
+        if self.training:
+            std = torch.exp(0.5 * logvar_latent)
+            eps = torch.randn_like(std)
+            return eps * std + mu_latent
+        return mu_latent
+
+    def decode(self, z_latent: torch.Tensor) -> torch.Tensor:
+        """Decode"""
+        hidden = z_latent
+        for layer in self.decoder[:-1]:
+            hidden = layer(hidden)
+            hidden = self.activation(hidden)
+        return self.decoder[-1](hidden)
+
+    def forward(self, batch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        :param batch: user batch
+        :return: output, expectation and logarithm of variation
+        """
+        mu_latent, logvar_latent = self.encode(batch)
+        z_latent = self.reparameterize(mu_latent, logvar_latent)
+        return self.decode(z_latent), mu_latent, logvar_latent
+
+    @staticmethod
+    def weight_init(layer: nn.Module):
+        """
+        Xavier initialization
+
+        :param layer: layer of a model
+        """
+        if isinstance(layer, nn.Linear):
+            nn.init.xavier_normal_(layer.weight.data)
+            layer.bias.data.normal_(0.0, 0.001)
+
+
+class MultVAE(TorchRecommender):
+    """`Variational Autoencoders for Collaborative Filtering
+    <https://arxiv.org/pdf/1802.05814.pdf>`_"""
+
+    num_workers: int = 0
+    batch_size_users: int = 5000
+    patience: int = 10
+    n_saved: int = 2
+    valid_split_size: float = 0.1
+    seed: int = 42
+    can_predict_cold_users = True
+    train_user_batch: csr_matrix
+    valid_user_batch: csr_matrix
+    _search_space = {
+        "learning_rate": {"type": "loguniform", "args": [0.0001, 0.5]},
+        "epochs": {"type": "int", "args": [100, 100]},
+        "latent_dim": {"type": "int", "args": [200, 200]},
+        "hidden_dim": {"type": "int", "args": [600, 600]},
+        "dropout": {"type": "uniform", "args": [0, 0.5]},
+        "anneal": {"type": "uniform", "args": [0.2, 1]},
+        "l2_reg": {"type": "loguniform", "args": [1e-9, 5]},
+        "factor": {"type": "uniform", "args": [0.2, 0.2]},
+        "patience": {"type": "int", "args": [3, 3]},
+    }
+
+    def __init__(
+        self,
+        learning_rate: float = 0.01,
+        epochs: int = 100,
+        latent_dim: int = 200,
+        hidden_dim: int = 600,
+        dropout: float = 0.3,
+        anneal: float = 0.1,
+        l2_reg: float = 0,
+        factor: float = 0.2,
+        patience: int = 3,
+    ):
+        """
+        :param learning_rate: learning rate
+        :param epochs: number of epochs to train model
+        :param latent_dim: latent dimension size for user vectors
+        :param hidden_dim: hidden dimension size for encoder and decoder
+        :param dropout: dropout coefficient
+        :param anneal: anneal coefficient [0,1]
+        :param l2_reg: l2 regularization term
+        :param factor: ReduceLROnPlateau reducing factor. new_lr = lr * factor
+        :param patience: number of non-improved epochs before reducing lr
+        """
+        super().__init__()
+        self.learning_rate = learning_rate
+        self.epochs = epochs
+        self.latent_dim = latent_dim
+        self.hidden_dim = hidden_dim
+        self.dropout = dropout
+        self.anneal = anneal
+        self.l2_reg = l2_reg
+        self.factor = factor
+        self.patience = patience
+
+    @property
+    def _init_args(self):
+        return {
+            "learning_rate": self.learning_rate,
+            "epochs": self.epochs,
+            "latent_dim": self.latent_dim,
+            "hidden_dim": self.hidden_dim,
+            "dropout": self.dropout,
+            "anneal": self.anneal,
+            "l2_reg": self.l2_reg,
+            "factor": self.factor,
+            "patience": self.patience,
+        }
+
+    def _get_data_loader(
+        self, data: PandasDataFrame, shuffle: bool = True
+    ) -> tuple[csr_matrix, DataLoader, np.ndarray]:
+        """get data loader and matrix with data"""
+        users_count = data["user_idx"].value_counts().count()
+        user_idx = data["user_idx"].astype("category").cat
+        user_batch = csr_matrix(
+            (
+                np.ones(len(data["user_idx"])),
+                ([user_idx.codes.values, data["item_idx"].values]),
+            ),
+            shape=(users_count, self._item_dim),
+        )
+        data_loader = DataLoader(
+            TensorDataset(torch.arange(users_count).long()),
+            batch_size=self.batch_size_users,
+            shuffle=shuffle,
+            num_workers=self.num_workers,
+        )
+
+        return user_batch, data_loader, user_idx.categories.values
+
+    def _fit(
+        self,
+        log: SparkDataFrame,
+        user_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
+        item_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
+    ) -> None:
+        self.logger.debug("Creating batch")
+        data = log.select("user_idx", "item_idx").toPandas()
+        splitter = GroupShuffleSplit(n_splits=1, test_size=self.valid_split_size, random_state=self.seed)
+        train_idx, valid_idx = next(splitter.split(data, groups=data["user_idx"]))
+        train_data, valid_data = data.iloc[train_idx], data.iloc[valid_idx]
+
+        self.train_user_batch, train_data_loader, _ = self._get_data_loader(train_data)
+        self.valid_user_batch, valid_data_loader, _ = self._get_data_loader(valid_data, False)
+
+        self.logger.debug("Training VAE")
+        self.model = VAE(
+            item_count=self._item_dim,
+            latent_dim=self.latent_dim,
+            hidden_dim=self.hidden_dim,
+            dropout=self.dropout,
+        ).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,
+            "multvae",
+        )
+
+    def _loss(self, y_pred, y_true, mu_latent, logvar_latent):
+        log_softmax_var = sf.log_softmax(y_pred, dim=1)
+        bce = -(log_softmax_var * y_true).sum(dim=1).mean()
+        kld = (
+            -0.5
+            * torch.sum(
+                1 + logvar_latent - mu_latent.pow(2) - logvar_latent.exp(),
+                dim=1,
+            ).mean()
+        )
+        return bce + self.anneal * kld
+
+    def _batch_pass(self, batch, model):
+        full_batch = self.train_user_batch if model.training else self.valid_user_batch
+        user_batch = torch.FloatTensor(full_batch[batch[0]].toarray()).to(self.device)
+        pred_user_batch, latent_mu, latent_logvar = self.model.forward(user_batch)
+        return {
+            "y_pred": pred_user_batch,
+            "y_true": user_batch,
+            "mu_latent": latent_mu,
+            "logvar_latent": latent_logvar,
+        }
+
+    @staticmethod
+    def _predict_pairs_inner(
+        model: nn.Module,
+        user_idx: int,
+        items_np_history: np.ndarray,
+        items_np_to_pred: np.ndarray,
+        item_count: int,
+        cnt: Optional[int] = None,
+    ) -> SparkDataFrame:
+        model.eval()
+        with torch.no_grad():
+            user_batch = torch.zeros((1, item_count))
+            user_batch[0, items_np_history] = 1
+            user_recs = sf.softmax(model(user_batch)[0][0].detach(), dim=0)
+            if cnt is not None:
+                best_item_idx = (torch.argsort(user_recs[items_np_to_pred], descending=True)[:cnt]).numpy()
+                items_np_to_pred = items_np_to_pred[best_item_idx]
+            return PandasDataFrame(
+                {
+                    "user_idx": np.array(items_np_to_pred.shape[0] * [user_idx]),
+                    "item_idx": items_np_to_pred,
+                    "relevance": user_recs[items_np_to_pred],
+                }
+            )
+
+    @staticmethod
+    def _predict_by_user(
+        pandas_df: PandasDataFrame,
+        model: nn.Module,
+        items_np: np.ndarray,
+        k: int,
+        item_count: int,
+    ) -> PandasDataFrame:
+        return MultVAE._predict_pairs_inner(
+            model=model,
+            user_idx=pandas_df["user_idx"][0],
+            items_np_history=pandas_df["item_idx"].values,
+            items_np_to_pred=items_np,
+            item_count=item_count,
+            cnt=min(len(pandas_df) + k, len(items_np)),
+        )
+
+    @staticmethod
+    def _predict_by_user_pairs(
+        pandas_df: PandasDataFrame,
+        model: nn.Module,
+        item_count: int,
+    ) -> PandasDataFrame:
+        return MultVAE._predict_pairs_inner(
+            model=model,
+            user_idx=pandas_df["user_idx"][0],
+            items_np_history=np.array(pandas_df["item_idx_history"][0]),
+            items_np_to_pred=np.array(pandas_df["item_idx_to_pred"][0]),
+            item_count=item_count,
+            cnt=None,
+        )
+
+    def _load_model(self, path: str):
+        self.model = VAE(
+            item_count=self._item_dim,
+            latent_dim=self.latent_dim,
+            hidden_dim=self.hidden_dim,
+            dropout=self.dropout,
+        ).to(self.device)
+        self.model.load_state_dict(torch.load(path))
+        self.model.eval()