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

replay/experimental/models/neural_ts.py

@@ -1,986 +0,0 @@
-import os
-from typing import Optional, Union
-
-import joblib
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-import torch
-from IPython.display import clear_output
-from pyspark.sql import DataFrame
-from sklearn.preprocessing import MinMaxScaler, OneHotEncoder
-from torch import Tensor, nn
-from torch.utils.data import DataLoader, Dataset, SequentialSampler
-from tqdm import tqdm
-
-from replay.experimental.models.base_rec import HybridRecommender
-from replay.splitters import TimeSplitter
-from replay.utils.spark_utils import convert2spark
-
-pd.options.mode.chained_assignment = None
-
-
-def cartesian_product(left, right):
-    """
-    This function computes cartesian product.
-    """
-    return left.assign(key=1).merge(right.assign(key=1), on="key").drop(columns=["key"])
-
-
-def num_tries_gt_zero(scores, batch_size, max_trials, max_num, device):
-    """
-    scores: [batch_size x N] float scores
-    returns: [batch_size x 1] the lowest indice per row where scores were first greater than 0. plus 1
-    """
-    tmp = scores.gt(0).nonzero().t()
-    # We offset these values by 1 to look for unset values (zeros) later
-    values = tmp[1] + 1
-    # Sparse tensors can't be moved with .to() or .cuda() if you want to send in cuda variables first
-    if device.type == "cuda":
-        tau = torch.cuda.sparse.LongTensor(tmp, values, torch.Size((batch_size, max_trials + 1))).to_dense()
-    else:
-        tau = torch.sparse.LongTensor(tmp, values, torch.Size((batch_size, max_trials + 1))).to_dense()
-    tau[(tau == 0)] += max_num  # set all unused indices to be max possible number so its not picked by min() call
-
-    tries = torch.min(tau, dim=1)[0]
-    return tries
-
-
-def w_log_loss(output, target, device):
-    """
-    This function computes weighted logistic loss.
-    """
-    output = torch.nn.functional.sigmoid(output)
-    output = torch.clamp(output, min=1e-7, max=1 - 1e-7)
-    count_1 = target.sum().item()
-    count_0 = target.shape[0] - count_1
-    class_count = np.array([count_0, count_1])
-    if count_1 == 0 or count_0 == 0:  # noqa: SIM108
-        weight = np.array([1.0, 1.0])
-    else:
-        weight = np.max(class_count) / class_count
-    weight = Tensor(weight).to(device)
-    loss = weight[1] * target * torch.log(output) + weight[0] * (1 - target) * torch.log(1 - output)
-    return -loss.mean()
-
-
-def warp_loss(positive_predictions, negative_predictions, num_labels, device):
-    """
-    positive_predictions: [batch_size x 1] floats between -1 to 1
-    negative_predictions: [batch_size x N] floats between -1 to 1
-    num_labels: int total number of labels in dataset (not just the subset you're using for the batch)
-    device: pytorch.device
-    """
-    batch_size, max_trials = negative_predictions.size(0), negative_predictions.size(1)
-
-    offsets, ones, max_num = (
-        torch.arange(0, batch_size, 1).long().to(device) * (max_trials + 1),
-        torch.ones(batch_size, 1).float().to(device),
-        batch_size * (max_trials + 1),
-    )
-
-    sample_scores = 1 + negative_predictions - positive_predictions
-    # Add column of ones so we know when we used all our attempts.
-    # This is used for indexing and computing should_count_loss if no real value is above 0
-    sample_scores, negative_predictions = (
-        torch.cat([sample_scores, ones], dim=1),
-        torch.cat([negative_predictions, ones], dim=1),
-    )
-
-    tries = num_tries_gt_zero(sample_scores, batch_size, max_trials, max_num, device)
-    attempts, trial_offset = tries.float(), (tries - 1) + offsets
-    # Don't count loss if we used max number of attempts
-    loss_weights = torch.log(torch.floor((num_labels - 1) / attempts))
-    should_count_loss = (attempts <= max_trials).float()
-    losses = (
-        loss_weights
-        * ((1 - positive_predictions.view(-1)) + negative_predictions.view(-1)[trial_offset])
-        * should_count_loss
-    )
-
-    return losses.sum()
-
-
-class SamplerWithReset(SequentialSampler):
-    """
-    Sampler class for train dataloader.
-    """
-
-    def __iter__(self):
-        self.data_source.reset()
-        return super().__iter__()
-
-
-class UserDatasetWithReset(Dataset):
-    """
-    Dataset class that takes data for a single user and
-    column names for continuous data, categorical data and data for
-    Wide model as well as the name of the target column.
-    The class also supports sampling of negative examples.
-    """
-
-    def __init__(
-        self,
-        idx,
-        log_train,
-        user_features,
-        item_features,
-        list_items,
-        union_cols,
-        cnt_neg_samples,
-        device,
-        target: Optional[str] = None,
-    ):
-        if cnt_neg_samples is not None:
-            self.cnt_neg_samples = cnt_neg_samples
-            self.user_features = user_features
-            self.item_features = item_features
-            item_idx_user = log_train["item_idx"].values.tolist()
-            self.item_idx_not_user = list(set(list_items).difference(set(item_idx_user)))
-        else:
-            self.cnt_neg_samples = cnt_neg_samples
-            self.user_features = None
-            self.item_features = None
-            self.item_idx_not_user = None
-        self.device = device
-        self.union_cols = union_cols
-        dataframe = log_train.merge(user_features, on="user_idx", how="inner")
-        self.dataframe = dataframe.merge(item_features, on="item_idx", how="inner")
-        self.user_idx = idx
-        self.data_sample = None
-        self.wide_part = Tensor(self.dataframe[self.union_cols["wide_cols"]].to_numpy().astype("float32")).to(
-            self.device
-        )
-        self.continuous_part = Tensor(
-            self.dataframe[self.union_cols["continuous_cols"]].to_numpy().astype("float32")
-        ).to(self.device)
-        self.cat_part = Tensor(self.dataframe[self.union_cols["cat_embed_cols"]].to_numpy().astype("float32")).to(
-            self.device
-        )
-        self.users = Tensor(self.dataframe[["user_idx"]].to_numpy().astype("int")).to(torch.long).to(self.device)
-        self.items = Tensor(self.dataframe[["item_idx"]].to_numpy().astype("int")).to(torch.long).to(self.device)
-        if target is not None:
-            self.target = Tensor(dataframe[target].to_numpy().astype("int")).to(self.device)
-        else:
-            self.target = target
-        self.target_column = target
-
-    def get_parts(self, data_sample):
-        """
-        Dataset method that selects user index, item indexes, categorical data,
-        continuous data, data for wide model, and target value.
-        """
-        self.wide_part = Tensor(data_sample[self.union_cols["wide_cols"]].to_numpy().astype("float32")).to(self.device)
-        self.continuous_part = Tensor(data_sample[self.union_cols["continuous_cols"]].to_numpy().astype("float32")).to(
-            self.device
-        )
-        self.cat_part = Tensor(data_sample[self.union_cols["cat_embed_cols"]].to_numpy().astype("float32")).to(
-            self.device
-        )
-        self.users = Tensor(data_sample[["user_idx"]].to_numpy().astype("int")).to(torch.long).to(self.device)
-        self.items = Tensor(data_sample[["item_idx"]].to_numpy().astype("int")).to(torch.long).to(self.device)
-        if self.target_column is not None:
-            self.target = Tensor(data_sample[self.target_column].to_numpy().astype("int")).to(self.device)
-        else:
-            self.target = self.target_column
-
-    def __getitem__(self, idx):
-        target = -1
-        if self.target is not None:
-            target = self.target[idx]
-        return (
-            self.wide_part[idx],
-            self.continuous_part[idx],
-            self.cat_part[idx],
-            self.users[idx],
-            self.items[idx],
-            target,
-        )
-
-    def __len__(self):
-        if self.data_sample is not None:
-            return self.data_sample.shape[0]
-        else:
-            return self.dataframe.shape[0]
-
-    def get_size_features(self):
-        """
-        Dataset method that gets the size of features after encoding/scaling.
-        """
-        return self.wide_part.shape[1], self.continuous_part.shape[1], self.cat_part.shape[1]
-
-    def reset(self):
-        """
-        Dataset methos that samples new negative examples..
-        """
-        n_samples = min(len(self.item_idx_not_user), self.cnt_neg_samples)
-        if n_samples > 0:
-            sample_item = np.random.choice(self.item_idx_not_user, n_samples, replace=False)
-            sample_item_feat = self.item_features.loc[self.item_features["item_idx"].isin(sample_item)]
-            sample_item_feat = sample_item_feat.set_axis(range(sample_item_feat.shape[0]), axis="index")
-            df_sample = cartesian_product(
-                self.user_features.loc[self.user_features["user_idx"] == self.user_idx], sample_item_feat
-            )
-            df_sample[self.target_column] = 0
-            self.data_sample = pd.concat([self.dataframe, df_sample], axis=0, ignore_index=True)
-            self.get_parts(self.data_sample)
-
-
-class Wide(nn.Module):
-    """
-    Wide model based on https://arxiv.org/abs/1606.07792
-    """
-
-    def __init__(self, input_dim: int, out_dim: int = 1):
-        super().__init__()
-
-        self.linear = nn.Sequential(nn.Linear(input_dim, out_dim), nn.ReLU(), nn.BatchNorm1d(out_dim))
-        self.out_dim = out_dim
-
-    def forward(self, input_data):
-        """
-        :param input_data: wide features
-        :return: torch tensor with shape batch_size*out_dim
-        """
-        output = self.linear(input_data)
-        return output
-
-
-class Deep(nn.Module):
-    """
-    Deep model based on https://arxiv.org/abs/1606.07792
-    """
-
-    def __init__(self, input_dim: int, out_dim: int, hidden_layers: list[int], deep_dropout: float):
-        super().__init__()
-        model = []
-        last_size = input_dim
-        for cur_size in hidden_layers:
-            model += [nn.Linear(last_size, cur_size)]
-            model += [nn.ReLU()]
-            model += [nn.BatchNorm1d(cur_size)]
-            model += [nn.Dropout(deep_dropout)]
-            last_size = cur_size
-        model += [nn.Linear(last_size, out_dim)]
-        model += [nn.ReLU()]
-        model += [nn.BatchNorm1d(out_dim)]
-        model += [nn.Dropout(deep_dropout)]
-        self.deep_model = nn.Sequential(*model)
-
-    def forward(self, input_data):
-        """
-        :param input_data: deep features
-        :return: torch tensor with shape batch_size*out_dim
-        """
-        output = self.deep_model(input_data)
-        return output
-
-
-class EmbedModel(nn.Module):
-    """
-    Model for getting embeddings for user and item indexes.
-    """
-
-    def __init__(self, cnt_users: int, cnt_items: int, user_embed: int, item_embed: int, crossed_embed: int):
-        super().__init__()
-        self.user_embed = nn.Embedding(num_embeddings=cnt_users, embedding_dim=user_embed)
-        self.item_embed = nn.Embedding(num_embeddings=cnt_items, embedding_dim=item_embed)
-        self.user_crossed_embed = nn.Embedding(num_embeddings=cnt_users, embedding_dim=crossed_embed)
-        self.item_crossed_embed = nn.Embedding(num_embeddings=cnt_items, embedding_dim=crossed_embed)
-
-    def forward(self, users, items):
-        """
-        :param users: user indexes
-        :param items: item indexes
-        :return: torch tensors: embedings for users, embedings for items,
-        embedings for users for wide model,
-        embedings for items for wide model,
-        embedings for pairs (users, items) for wide model
-        """
-        users_to_embed = self.user_embed(users).squeeze()
-        items_to_embed = self.item_embed(items).squeeze()
-        cross_users = self.user_crossed_embed(users).squeeze()
-        cross_items = self.item_crossed_embed(items).squeeze()
-        cross = (cross_users * cross_items).sum(dim=-1).unsqueeze(-1)
-        return users_to_embed, items_to_embed, cross_users, cross_items, cross
-
-
-class WideDeep(nn.Module):
-    """
-    Wide&Deep model based on https://arxiv.org/abs/1606.07792
-    """
-
-    def __init__(
-        self,
-        dim_head: int,
-        deep_out_dim: int,
-        hidden_layers: list[int],
-        size_wide_features: int,
-        size_continuous_features: int,
-        size_cat_features: int,
-        wide_out_dim: int,
-        head_dropout: float,
-        deep_dropout: float,
-        cnt_users: int,
-        cnt_items: int,
-        user_embed: int,
-        item_embed: int,
-        crossed_embed: int,
-    ):
-        super().__init__()
-        self.embed_model = EmbedModel(cnt_users, cnt_items, user_embed, item_embed, crossed_embed)
-        self.wide = Wide(size_wide_features + crossed_embed * 2 + 1, wide_out_dim)
-        self.deep = Deep(
-            size_cat_features + size_continuous_features + user_embed + item_embed,
-            deep_out_dim,
-            hidden_layers,
-            deep_dropout,
-        )
-        self.head_model = nn.Sequential(nn.Linear(wide_out_dim + deep_out_dim, dim_head), nn.ReLU())
-        self.last_layer = nn.Sequential(nn.Linear(dim_head, 1))
-        self.head_dropout = head_dropout
-
-    def forward_for_predict(self, wide_part, continuous_part, cat_part, users, items):
-        """
-        Forward method without last layer and dropout that is used for prediction.
-        """
-        users_to_embed, items_to_embed, cross_users, cross_items, cross = self.embed_model(users, items)
-        input_deep = torch.cat((cat_part, continuous_part, users_to_embed, items_to_embed), dim=-1).squeeze()
-        out_deep = self.deep(input_deep)
-        wide_part = torch.cat((wide_part, cross_users, cross_items, cross), dim=-1)
-        out_wide = self.wide(wide_part)
-        input_data = torch.cat((out_wide, out_deep), dim=-1)
-        output = self.head_model(input_data)
-        return output
-
-    def forward_dropout(self, input_data):
-        """
-        Forward method for multiple prediction with active dropout
-        :param input_data: output of forward_for_predict
-        :return: torch tensor after dropout and last linear layer
-        """
-        output = nn.functional.dropout(input_data, p=self.head_dropout, training=True)
-        output = self.last_layer(output)
-        return output
-
-    def forward_for_embeddings(
-        self, wide_part, continuous_part, cat_part, users_to_embed, items_to_embed, cross_users, cross_items, cross
-    ):
-        """
-        Forward method after getting emdeddings for users and items.
-        """
-        input_deep = torch.cat((cat_part, continuous_part, users_to_embed, items_to_embed), dim=-1).squeeze()
-        out_deep = self.deep(input_deep)
-        wide_part = torch.cat((wide_part, cross_users, cross_items, cross), dim=-1)
-        out_wide = self.wide(wide_part)
-        input_data = torch.cat((out_wide, out_deep), dim=-1)
-        output = self.head_model(input_data)
-        output = nn.functional.dropout(output, p=self.head_dropout, training=True)
-        output = self.last_layer(output)
-        return output
-
-    def forward(self, wide_part, continuous_part, cat_part, users, items):
-        """
-        :param wide_part: features for wide model
-        :param continuous_part: continuous features
-        :param cat_part: torch categorical features
-        :param users: user indexes
-        :param items: item indexes
-        :return: relevances for pair (users, items)
-
-        """
-        users_to_embed, items_to_embed, cross_users, cross_items, cross = self.embed_model(users, items)
-        output = self.forward_for_embeddings(
-            wide_part, continuous_part, cat_part, users_to_embed, items_to_embed, cross_users, cross_items, cross
-        )
-        return output
-
-
-class NeuralTS(HybridRecommender):
-    """
-    'Neural Thompson sampling recommender
-    <https://dl.acm.org/doi/pdf/10.1145/3383313.3412214>`_  based on `Wide&Deep model
-    <https://arxiv.org/abs/1606.07792>`_.
-
-    :param user_cols: user_cols = {'continuous_cols':List[str], 'cat_embed_cols':List[str], 'wide_cols': List[str]},
-        where List[str] -- some column names from user_features dataframe, which is input to the fit method,
-        or empty List
-    :param item_cols: item_cols = {'continuous_cols':List[str], 'cat_embed_cols':List[str], 'wide_cols': List[str]},
-        where List[str] -- some column names from item_features dataframe, which is input to the fit method,
-        or empty List
-    :param embedding_sizes: list of length three in which
-        embedding_sizes[0] = embedding size for users,
-        embedding_sizes[1] = embedding size for items,
-        embedding_sizes[2] = embedding size for pair (users, items)
-    :param hidden_layers: list of hidden layer sizes for Deep model
-    :param wide_out_dim: output size for the Wide model
-    :param deep_out_dim: output size for the Deep model
-    :param head_dropout: probability of an element to be zeroed for WideDeep model head
-    :param deep_dropout: probability of an element to be zeroed for Deep model
-    :param dim_head: output size for WideDeep model head
-    :param n_epochs: number of epochs for model training
-    :param opt_lr: learning rate for the AdamW optimizer
-    :param lr_min: minimum learning rate value for the CosineAnnealingLR learning rate scheduler
-    :param use_gpu: if true, the model will be trained on the GPU
-    :param use_warp_loss: if true, then warp loss will be used otherwise weighted logistic loss.
-    :param cnt_neg_samples: number of additional negative examples for each user
-    :param cnt_samples_for_predict: number of sampled predictions for one user,
-        which are used to estimate the mean and variance of relevance
-    :param exploration_coef:  exploration coefficient
-    :param plot_dir: file name where the training graphs will be saved, if None, the graphs will not be saved
-    :param cnt_users: number of users, used in Wide&Deep model initialization
-    :param cnt_items: number of items, used in Wide&Deep model initialization
-
-    """
-
-    def __init__(
-        self,
-        user_cols: dict[str, list[str]] = {"continuous_cols": [], "cat_embed_cols": [], "wide_cols": []},
-        item_cols: dict[str, list[str]] = {"continuous_cols": [], "cat_embed_cols": [], "wide_cols": []},
-        embedding_sizes: list[int] = [32, 32, 64],
-        hidden_layers: list[int] = [32, 20],
-        wide_out_dim: int = 1,
-        deep_out_dim: int = 20,
-        head_dropout: float = 0.8,
-        deep_dropout: float = 0.4,
-        dim_head: int = 20,
-        n_epochs: int = 2,
-        opt_lr: float = 3e-4,
-        lr_min: float = 1e-5,
-        use_gpu: bool = False,
-        use_warp_loss: bool = True,
-        cnt_neg_samples: int = 100,
-        cnt_samples_for_predict: int = 10,
-        exploration_coef: float = 1.0,
-        cnt_users: Optional[int] = None,
-        cnt_items: Optional[int] = None,
-        plot_dir: Optional[str] = None,
-    ):
-        self.user_cols = user_cols
-        self.item_cols = item_cols
-        self.embedding_sizes = embedding_sizes
-        self.hidden_layers = hidden_layers
-        self.wide_out_dim = wide_out_dim
-        self.deep_out_dim = deep_out_dim
-        self.head_dropout = head_dropout
-        self.deep_dropout = deep_dropout
-        self.dim_head = dim_head
-        self.n_epochs = n_epochs
-        self.opt_lr = opt_lr
-        self.lr_min = lr_min
-        self.device = torch.device("cpu")
-        if use_gpu:
-            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.use_warp_loss = use_warp_loss
-        self.cnt_neg_samples = cnt_neg_samples
-        self.cnt_samples_for_predict = cnt_samples_for_predict
-        self.exploration_coef = exploration_coef
-        self.cnt_users = cnt_users
-        self.cnt_items = cnt_items
-        self.plot_dir = plot_dir
-
-        self.size_wide_features = None
-        self.size_continuous_features = None
-        self.size_cat_features = None
-        self.scaler_user = None
-        self.encoder_intersept_user = None
-        self.encoder_diff_user = None
-        self.scaler_item = None
-        self.encoder_intersept_item = None
-        self.encoder_diff_item = None
-        self.union_cols = None
-        self.num_of_train_labels = None
-        self.dict_true_items_val = None
-        self.lr_scheduler = None
-        self.model = None
-        self.criterion = None
-        self.optimizer = None
-
-    def preprocess_features_fit(self, train, item_features, user_features):
-        """
-        This function initializes all ecoders for the features.
-        """
-        train_users = user_features.loc[user_features["user_idx"].isin(train["user_idx"].values.tolist())]
-        wide_cols_cat = list(set(self.user_cols["cat_embed_cols"]) & set(self.user_cols["wide_cols"]))
-        cat_embed_cols_not_wide = list(set(self.user_cols["cat_embed_cols"]).difference(set(wide_cols_cat)))
-        if len(self.user_cols["continuous_cols"]) != 0:
-            self.scaler_user = MinMaxScaler()
-            self.scaler_user.fit(train_users[self.user_cols["continuous_cols"]])
-        else:
-            self.scaler_user = None
-        if len(wide_cols_cat) != 0:
-            self.encoder_intersept_user = OneHotEncoder(sparse_output=False, handle_unknown="ignore")
-            self.encoder_intersept_user.fit(train_users[wide_cols_cat])
-        else:
-            self.encoder_intersept_user = None
-        if len(cat_embed_cols_not_wide) != 0:
-            self.encoder_diff_user = OneHotEncoder(sparse_output=False, handle_unknown="ignore")
-            self.encoder_diff_user.fit(train_users[cat_embed_cols_not_wide])
-        else:
-            self.encoder_diff_user = None
-        train_items = item_features.loc[item_features["item_idx"].isin(train["item_idx"].values.tolist())]
-        wide_cols_cat = list(set(self.item_cols["cat_embed_cols"]) & set(self.item_cols["wide_cols"]))
-        cat_embed_cols_not_wide = list(set(self.item_cols["cat_embed_cols"]).difference(set(wide_cols_cat)))
-        if len(self.item_cols["continuous_cols"]) != 0:
-            self.scaler_item = MinMaxScaler()
-            self.scaler_item.fit(train_items[self.item_cols["continuous_cols"]])
-        else:
-            self.scaler_item = None
-        if len(wide_cols_cat) != 0:
-            self.encoder_intersept_item = OneHotEncoder(sparse_output=False, handle_unknown="ignore")
-            self.encoder_intersept_item.fit(train_items[wide_cols_cat])
-        else:
-            self.encoder_intersept_item = None
-        if len(cat_embed_cols_not_wide) != 0:
-            self.encoder_diff_item = OneHotEncoder(sparse_output=False, handle_unknown="ignore")
-            self.encoder_diff_item.fit(train_items[cat_embed_cols_not_wide])
-        else:
-            self.encoder_diff_item = None
-
-    def preprocess_features_transform(self, item_features, user_features):
-        """
-        This function performs the transformation for all features.
-        """
-        self.union_cols = {"continuous_cols": [], "cat_embed_cols": [], "wide_cols": []}
-        wide_cols_cat = list(set(self.user_cols["cat_embed_cols"]) & set(self.user_cols["wide_cols"]))
-        cat_embed_cols_not_wide = list(set(self.user_cols["cat_embed_cols"]).difference(set(wide_cols_cat)))
-        if len(self.user_cols["continuous_cols"]) != 0:
-            users_continuous = pd.DataFrame(
-                self.scaler_user.transform(user_features[self.user_cols["continuous_cols"]]),
-                columns=self.user_cols["continuous_cols"],
-            )
-            self.union_cols["continuous_cols"] += self.user_cols["continuous_cols"]
-        else:
-            users_continuous = user_features[[]]
-        if len(wide_cols_cat) != 0:
-            users_wide_cat = pd.DataFrame(
-                self.encoder_intersept_user.transform(user_features[wide_cols_cat]),
-                columns=list(self.encoder_intersept_user.get_feature_names_out(wide_cols_cat)),
-            )
-            self.union_cols["cat_embed_cols"] += list(self.encoder_intersept_user.get_feature_names_out(wide_cols_cat))
-            self.union_cols["wide_cols"] += list(
-                set(self.user_cols["wide_cols"]).difference(set(self.user_cols["cat_embed_cols"]))
-            ) + list(self.encoder_intersept_user.get_feature_names_out(wide_cols_cat))
-        else:
-            users_wide_cat = user_features[[]]
-        if len(cat_embed_cols_not_wide) != 0:
-            users_cat = pd.DataFrame(
-                self.encoder_diff_user.transform(user_features[cat_embed_cols_not_wide]),
-                columns=list(self.encoder_diff_user.get_feature_names_out(cat_embed_cols_not_wide)),
-            )
-            self.union_cols["cat_embed_cols"] += list(
-                self.encoder_diff_user.get_feature_names_out(cat_embed_cols_not_wide)
-            )
-        else:
-            users_cat = user_features[[]]
-
-        transform_user_features = pd.concat(
-            [user_features[["user_idx"]], users_continuous, users_wide_cat, users_cat], axis=1
-        )
-
-        wide_cols_cat = list(set(self.item_cols["cat_embed_cols"]) & set(self.item_cols["wide_cols"]))
-        cat_embed_cols_not_wide = list(set(self.item_cols["cat_embed_cols"]).difference(set(wide_cols_cat)))
-        if len(self.item_cols["continuous_cols"]) != 0:
-            items_continuous = pd.DataFrame(
-                self.scaler_item.transform(item_features[self.item_cols["continuous_cols"]]),
-                columns=self.item_cols["continuous_cols"],
-            )
-            self.union_cols["continuous_cols"] += self.item_cols["continuous_cols"]
-        else:
-            items_continuous = item_features[[]]
-        if len(wide_cols_cat) != 0:
-            items_wide_cat = pd.DataFrame(
-                self.encoder_intersept_item.transform(item_features[wide_cols_cat]),
-                columns=list(self.encoder_intersept_item.get_feature_names_out(wide_cols_cat)),
-            )
-            self.union_cols["cat_embed_cols"] += list(self.encoder_intersept_item.get_feature_names_out(wide_cols_cat))
-            self.union_cols["wide_cols"] += list(
-                set(self.item_cols["wide_cols"]).difference(set(self.item_cols["cat_embed_cols"]))
-            ) + list(self.encoder_intersept_item.get_feature_names_out(wide_cols_cat))
-        else:
-            items_wide_cat = item_features[[]]
-        if len(cat_embed_cols_not_wide) != 0:
-            items_cat = pd.DataFrame(
-                self.encoder_diff_item.transform(item_features[cat_embed_cols_not_wide]),
-                columns=list(self.encoder_diff_item.get_feature_names_out(cat_embed_cols_not_wide)),
-            )
-            self.union_cols["cat_embed_cols"] += list(
-                self.encoder_diff_item.get_feature_names_out(cat_embed_cols_not_wide)
-            )
-        else:
-            items_cat = item_features[[]]
-
-        transform_item_features = pd.concat(
-            [item_features[["item_idx"]], items_continuous, items_wide_cat, items_cat], axis=1
-        )
-        return transform_user_features, transform_item_features
-
-    def _data_loader(
-        self, idx, log_train, transform_user_features, transform_item_features, list_items, train=False
-    ) -> Union[tuple[UserDatasetWithReset, DataLoader], DataLoader]:
-        if train:
-            train_dataset = UserDatasetWithReset(
-                idx=idx,
-                log_train=log_train,
-                user_features=transform_user_features,
-                item_features=transform_item_features,
-                list_items=list_items,
-                union_cols=self.union_cols,
-                cnt_neg_samples=self.cnt_neg_samples,
-                device=self.device,
-                target="relevance",
-            )
-            sampler = SamplerWithReset(train_dataset)
-            train_dataloader = DataLoader(
-                train_dataset, batch_size=log_train.shape[0] + self.cnt_neg_samples, sampler=sampler
-            )
-            return train_dataset, train_dataloader
-        else:
-            dataset = UserDatasetWithReset(
-                idx=idx,
-                log_train=log_train,
-                user_features=transform_user_features,
-                item_features=transform_item_features,
-                list_items=list_items,
-                union_cols=self.union_cols,
-                cnt_neg_samples=None,
-                device=self.device,
-                target=None,
-            )
-            dataloader = DataLoader(dataset, batch_size=log_train.shape[0], shuffle=False)
-            return dataloader
-
-    def _fit(
-        self,
-        log: DataFrame,
-        user_features: Optional[DataFrame] = None,
-        item_features: Optional[DataFrame] = None,
-    ) -> None:
-        if user_features is None:
-            msg = "User features are missing for fitting"
-            raise ValueError(msg)
-        if item_features is None:
-            msg = "Item features are missing for fitting"
-            raise ValueError(msg)
-
-        train_spl = TimeSplitter(
-            time_threshold=0.2,
-            drop_cold_items=True,
-            drop_cold_users=True,
-            query_column="user_idx",
-            item_column="item_idx",
-        )
-        train, val = train_spl.split(log)
-        train = train.drop("timestamp")
-        val = val.drop("timestamp")
-
-        train = train.toPandas()
-        val = val.toPandas()
-        pd_item_features = item_features.toPandas()
-        pd_user_features = user_features.toPandas()
-        if self.cnt_users is None:
-            self.cnt_users = pd_user_features.shape[0]
-        if self.cnt_items is None:
-            self.cnt_items = pd_item_features.shape[0]
-        self.num_of_train_labels = self.cnt_items
-
-        self.preprocess_features_fit(train, pd_item_features, pd_user_features)
-        transform_user_features, transform_item_features = self.preprocess_features_transform(
-            pd_item_features, pd_user_features
-        )
-
-        list_items = pd_item_features["item_idx"].values.tolist()
-
-        dataloader_train_users = []
-        train = train.set_axis(range(train.shape[0]), axis="index")
-        train_group_by_users = train.groupby("user_idx")
-        for idx, df_train_idx in tqdm(train_group_by_users):
-            df_train_idx = df_train_idx.loc[df_train_idx["relevance"] == 1]
-            if df_train_idx.shape[0] == 0:
-                continue
-            df_train_idx = df_train_idx.set_axis(range(df_train_idx.shape[0]), axis="index")
-            train_dataset, train_dataloader = self._data_loader(
-                idx, df_train_idx, transform_user_features, transform_item_features, list_items, train=True
-            )
-            dataloader_train_users.append(train_dataloader)
-
-        dataloader_val_users = []
-        self.dict_true_items_val = {}
-        transform_item_features.sort_values(by=["item_idx"], inplace=True, ignore_index=True)
-        val = val.set_axis(range(val.shape[0]), axis="index")
-        val_group_by_users = val.groupby("user_idx")
-        for idx, df_val_idx in tqdm(val_group_by_users):
-            self.dict_true_items_val[idx] = df_val_idx.loc[(df_val_idx["relevance"] == 1)]["item_idx"].values.tolist()
-            df_val = cartesian_product(pd.DataFrame({"user_idx": [idx]}), transform_item_features[["item_idx"]])
-            df_val = df_val.set_axis(range(df_val.shape[0]), axis="index")
-            dataloader_val_users.append(
-                self._data_loader(
-                    idx, df_val, transform_user_features, transform_item_features, list_items, train=False
-                )
-            )
-
-        self.size_wide_features, self.size_continuous_features, self.size_cat_features = (
-            train_dataset.get_size_features()
-        )
-        self.model = WideDeep(
-            dim_head=self.dim_head,
-            deep_out_dim=self.deep_out_dim,
-            hidden_layers=self.hidden_layers,
-            size_wide_features=self.size_wide_features,
-            size_continuous_features=self.size_continuous_features,
-            size_cat_features=self.size_cat_features,
-            wide_out_dim=self.wide_out_dim,
-            head_dropout=self.head_dropout,
-            deep_dropout=self.deep_dropout,
-            cnt_users=self.cnt_users,
-            cnt_items=self.cnt_items,
-            user_embed=self.embedding_sizes[0],
-            item_embed=self.embedding_sizes[1],
-            crossed_embed=self.embedding_sizes[2],
-        )
-        if self.use_warp_loss:
-            self.criterion = warp_loss
-        else:
-            self.criterion = w_log_loss
-        self.optimizer = torch.optim.AdamW(self.model.parameters(), self.opt_lr)
-        self.lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, self.n_epochs, self.lr_min)
-
-        self.train(self.model, dataloader_train_users, dataloader_val_users)
-
-    def train(self, model, train_dataloader, val_dataloader):
-        """
-        Run training loop.
-        """
-        train_losses = []
-        val_ndcg = []
-        model = model.to(self.device)
-        for epoch in range(self.n_epochs):
-            train_loss = self._batch_pass(model, train_dataloader)
-            ndcg = self.predict_val_with_ndcg(model, val_dataloader, k=10)
-            train_losses.append(train_loss)
-            val_ndcg.append(ndcg)
-
-            if self.plot_dir is not None and epoch > 0:
-                clear_output(wait=True)
-                _, (ax1, ax2) = plt.subplots(1, 2, figsize=(30, 15))
-                ax1.plot(train_losses, label="train", color="b")
-                ax2.plot(val_ndcg, label="val_ndcg", color="g")
-                size = max(1, round(epoch / 10))
-                plt.xticks(range(epoch - 1)[::size])
-                ax1.set_ylabel("loss")
-                ax1.set_xlabel("epoch")
-                ax2.set_ylabel("ndcg")
-                ax2.set_xlabel("epoch")
-                plt.legend()
-                plt.savefig(self.plot_dir)
-                plt.show()
-                self.logger.info("ndcg val =%.4f", ndcg)
-
-    def _loss(self, preds, labels):
-        if self.use_warp_loss:
-            ind_pos = torch.where(labels == 1)[0]
-            ind_neg = torch.where(labels == 0)[0]
-            min_batch = ind_pos.shape[0]
-            if ind_pos.shape[0] == 0 or ind_neg.shape[0] == 0:
-                return
-            indexes_pos = ind_pos
-            pos = preds.squeeze()[indexes_pos].unsqueeze(-1)
-            list_neg = []
-            for _ in range(min_batch):
-                indexes_neg = ind_neg[torch.randperm(ind_neg.shape[0])]
-                list_neg.append(preds.squeeze()[indexes_neg].unsqueeze(-1))
-            neg = torch.cat(list_neg, dim=-1)
-            neg = neg.transpose(0, 1)
-            loss = self.criterion(pos, neg, self.num_of_train_labels, self.device)
-        else:
-            loss = self.criterion(preds.squeeze(), labels, self.device)
-        return loss
-
-    def _batch_pass(self, model, train_dataloader):
-        """
-        Run training one epoch loop.
-        """
-        model.train()
-        idx = 0
-        cumulative_loss = 0
-        preds = None
-        for user_dataloader in tqdm(train_dataloader):
-            for batch in user_dataloader:
-                wide_part, continuous_part, cat_part, users, items, labels = batch
-                self.optimizer.zero_grad()
-                preds = model(wide_part, continuous_part, cat_part, users, items)
-                loss = self._loss(preds, labels)
-                if loss is not None:
-                    loss.backward()
-                    self.optimizer.step()
-                    cumulative_loss += loss.item()
-                    idx += 1
-
-        self.lr_scheduler.step()
-        return cumulative_loss / idx
-
-    def predict_val_with_ndcg(self, model, val_dataloader, k):
-        """
-        This function returns the NDCG metric for the validation data.
-        """
-        if len(val_dataloader) == 0:
-            return 0
-
-        ndcg = 0
-        idx = 0
-        model = model.to(self.device)
-        for user_dataloader in tqdm(val_dataloader):
-            _, _, _, users, _, _ = next(iter(user_dataloader))
-            user = int(users[0])
-            sample_pred = np.array(self.predict_val(model, user_dataloader))
-            top_k_predicts = (-sample_pred).argsort()[:k]
-            ndcg += (np.isin(top_k_predicts, self.dict_true_items_val[user]).sum()) / k
-            idx += 1
-
-        metric = ndcg / idx
-        return metric
-
-    def predict_val(self, model, val_dataloader):
-        """
-        This function returns the relevances for the validation data.
-        """
-        probs = []
-        model = model.to(self.device)
-        model.eval()
-        with torch.no_grad():
-            for wide_part, continuous_part, cat_part, users, items, _ in val_dataloader:
-                preds = model(wide_part, continuous_part, cat_part, users, items)
-                probs += (preds.squeeze()).tolist()
-        return probs
-
-    def predict_test(self, model, test_dataloader, cnt_samples_for_predict):
-        """
-        This function returns a list of cnt_samples_for_predict relevancies for each pair (users, items)
-        in val_dataloader
-        """
-        probs = []
-        model = model.to(self.device)
-        model.eval()
-        with torch.no_grad():
-            for wide_part, continuous_part, cat_part, users, items, _ in test_dataloader:
-                preds = model.forward_for_predict(wide_part, continuous_part, cat_part, users, items)
-                probs.extend(model.forward_dropout(preds).squeeze().tolist() for __ in range(cnt_samples_for_predict))
-        return probs
-
-    def _predict(
-        self,
-        log: DataFrame,  # noqa: ARG002
-        k: int,  # noqa: ARG002
-        users: DataFrame,
-        items: DataFrame,
-        user_features: Optional[DataFrame] = None,
-        item_features: Optional[DataFrame] = None,
-        filter_seen_items: bool = True,  # noqa: ARG002
-    ) -> DataFrame:
-        if user_features is None:
-            msg = "User features are missing for predict"
-            raise ValueError(msg)
-        if item_features is None:
-            msg = "Item features are missing for predict"
-            raise ValueError(msg)
-
-        pd_users = users.toPandas()
-        pd_items = items.toPandas()
-        pd_user_features = user_features.toPandas()
-        pd_item_features = item_features.toPandas()
-
-        list_items = pd_item_features["item_idx"].values.tolist()
-
-        transform_user_features, transform_item_features = self.preprocess_features_transform(
-            pd_item_features, pd_user_features
-        )
-
-        preds = []
-        users_ans = []
-        items_ans = []
-        for idx in tqdm(pd_users["user_idx"].unique()):
-            df_test_idx = cartesian_product(pd.DataFrame({"user_idx": [idx]}), pd_items)
-            df_test_idx = df_test_idx.set_axis(range(df_test_idx.shape[0]), axis="index")
-            test_dataloader = self._data_loader(
-                idx, df_test_idx, transform_user_features, transform_item_features, list_items, train=False
-            )
-
-            samples = np.array(self.predict_test(self.model, test_dataloader, self.cnt_samples_for_predict))
-            sample_pred = np.mean(samples, axis=0) + self.exploration_coef * np.sqrt(np.var(samples, axis=0))
-
-            preds += sample_pred.tolist()
-            users_ans += [idx] * df_test_idx.shape[0]
-            items_ans += df_test_idx["item_idx"].values.tolist()
-
-        res_df = pd.DataFrame({"user_idx": users_ans, "item_idx": items_ans, "relevance": preds})
-        pred = convert2spark(res_df)
-        return pred
-
-    @property
-    def _init_args(self):
-        return {
-            "n_epochs": self.n_epochs,
-            "union_cols": self.union_cols,
-            "cnt_users": self.cnt_users,
-            "cnt_items": self.cnt_items,
-            "size_wide_features": self.size_wide_features,
-            "size_continuous_features": self.size_continuous_features,
-            "size_cat_features": self.size_cat_features,
-        }
-
-    def model_save(self, dir_name):
-        """
-        This function saves the model.
-        """
-        os.makedirs(dir_name, exist_ok=True)
-
-        joblib.dump(self.scaler_user, os.path.join(dir_name, "scaler_user.joblib"))
-        joblib.dump(self.encoder_intersept_user, os.path.join(dir_name, "encoder_intersept_user.joblib"))
-        joblib.dump(self.encoder_diff_user, os.path.join(dir_name, "encoder_diff_user.joblib"))
-
-        joblib.dump(self.scaler_item, os.path.join(dir_name, "scaler_item.joblib"))
-        joblib.dump(self.encoder_intersept_item, os.path.join(dir_name, "encoder_intersept_item.joblib"))
-        joblib.dump(self.encoder_diff_item, os.path.join(dir_name, "encoder_diff_item.joblib"))
-
-        torch.save(self.model.state_dict(), os.path.join(dir_name, "model_weights.pth"))
-        torch.save(
-            {
-                "fit_users": self.fit_users.toPandas(),
-                "fit_items": self.fit_items.toPandas(),
-            },
-            os.path.join(dir_name, "fit_info.pth"),
-        )
-
-    def model_load(self, dir_name):
-        """
-        This function loads the model.
-        """
-        self.scaler_user = joblib.load(os.path.join(dir_name, "scaler_user.joblib"))
-        self.encoder_intersept_user = joblib.load(os.path.join(dir_name, "encoder_intersept_user.joblib"))
-        self.encoder_diff_user = joblib.load(os.path.join(dir_name, "encoder_diff_user.joblib"))
-
-        self.scaler_item = joblib.load(os.path.join(dir_name, "scaler_item.joblib"))
-        self.encoder_intersept_item = joblib.load(os.path.join(dir_name, "encoder_intersept_item.joblib"))
-        self.encoder_diff_item = joblib.load(os.path.join(dir_name, "encoder_diff_item.joblib"))
-
-        self.model = WideDeep(
-            dim_head=self.dim_head,
-            deep_out_dim=self.deep_out_dim,
-            hidden_layers=self.hidden_layers,
-            size_wide_features=self.size_wide_features,
-            size_continuous_features=self.size_continuous_features,
-            size_cat_features=self.size_cat_features,
-            wide_out_dim=self.wide_out_dim,
-            head_dropout=self.head_dropout,
-            deep_dropout=self.deep_dropout,
-            cnt_users=self.cnt_users,
-            cnt_items=self.cnt_items,
-            user_embed=self.embedding_sizes[0],
-            item_embed=self.embedding_sizes[1],
-            crossed_embed=self.embedding_sizes[2],
-        )
-        self.model.load_state_dict(torch.load(os.path.join(dir_name, "model_weights.pth")))
-
-        checkpoint = torch.load(os.path.join(dir_name, "fit_info.pth"), weights_only=False)
-        self.fit_users = convert2spark(checkpoint["fit_users"])
-        self.fit_items = convert2spark(checkpoint["fit_items"])