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

replay/experimental/models/dt4rec/dt4rec.py

@@ -1,189 +0,0 @@
-from typing import Optional
-
-import pandas as pd
-from tqdm import tqdm
-
-from replay.experimental.models.base_rec import Recommender
-from replay.utils import PYSPARK_AVAILABLE, TORCH_AVAILABLE, SparkDataFrame
-
-if PYSPARK_AVAILABLE:
-    from replay.utils.spark_utils import convert2spark
-
-if TORCH_AVAILABLE:
-    import torch
-    from torch.utils.data.dataloader import DataLoader
-
-    from .gpt1 import GPT, GPTConfig
-    from .trainer import Trainer, TrainerConfig
-    from .utils import (
-        Collator,
-        StateActionReturnDataset,
-        ValidateDataset,
-        WarmUpScheduler,
-        create_dataset,
-        matrix2df,
-        set_seed,
-    )
-
-
-class DT4Rec(Recommender):
-    """
-    Decision Transformer for Recommendations
-
-    General Idea:
-        `Decision Transformer: Reinforcement Learning
-        via Sequence Modeling <https://arxiv.org/pdf/2106.01345.pdf>`_.
-
-    Ideas for improvements:
-        `User Retention-oriented Recommendation with Decision
-        Transformer <https://arxiv.org/pdf/2303.06347.pdf>`_.
-
-    Also, some sources are listed in their respective classes
-    """
-
-    optimizer = None
-    train_batch_size = 128
-    val_batch_size = 128
-    lr_scheduler = None
-
-    def __init__(
-        self,
-        item_num,
-        user_num,
-        seed=123,
-        trajectory_len=30,
-        epochs=1,
-        batch_size=64,
-        use_cuda=True,
-    ):
-        self.item_num = item_num
-        self.user_num = user_num
-        self.seed = seed
-        self.trajectory_len = trajectory_len
-        self.epochs = epochs
-        self.batch_size = batch_size
-        self.tconf: TrainerConfig = TrainerConfig(epochs=epochs)
-        self.mconf: GPTConfig = GPTConfig(
-            user_num=user_num,
-            item_num=item_num,
-            vocab_size=self.item_num + 1,
-            block_size=self.trajectory_len * 3,
-            max_timestep=self.item_num,
-        )
-        self.model: GPT
-        self.user_trajectory: list
-        self.trainer: Trainer
-        self.use_cuda = use_cuda
-        set_seed(self.seed)
-
-    def _init_args(self):
-        pass
-
-    def _update_mconf(self, **kwargs):
-        self.mconf.update(**kwargs)
-
-    def _update_tconf(self, **kwargs):
-        self.tconf.update(**kwargs)
-
-    def _make_prediction_dataloader(self, users, items, max_context_len=30):
-        val_dataset = ValidateDataset(
-            self.user_trajectory,
-            max_context_len=max_context_len - 1,
-            val_items=items,
-            val_users=users,
-        )
-
-        val_dataloader = DataLoader(
-            val_dataset,
-            pin_memory=True,
-            batch_size=self.val_batch_size,
-            collate_fn=Collator(self.item_num),
-        )
-
-        return val_dataloader
-
-    def train(
-        self,
-        log,
-        val_users=None,
-        val_items=None,
-        experiment=None,
-    ):
-        """
-        Run training loop
-        """
-        assert (val_users is None) == (val_items is None) == (experiment is None)
-        with_validate = experiment is not None
-        df = log.toPandas()[["user_idx", "item_idx", "relevance", "timestamp"]]
-        self.user_trajectory = create_dataset(df, user_num=self.user_num, item_pad=self.item_num)
-
-        train_dataset = StateActionReturnDataset(self.user_trajectory, self.trajectory_len)
-
-        train_dataloader = DataLoader(
-            train_dataset,
-            shuffle=True,
-            pin_memory=True,
-            batch_size=self.train_batch_size,
-            collate_fn=Collator(self.item_num),
-        )
-
-        if with_validate:
-            val_dataloader = self._make_prediction_dataloader(val_users, val_items, max_context_len=self.trajectory_len)
-        else:
-            val_dataloader = None
-
-        self.model = GPT(self.mconf)
-
-        optimizer = torch.optim.AdamW(
-            self.model.configure_optimizers(),
-            lr=3e-4,
-            betas=(0.9, 0.95),
-        )
-        lr_scheduler = WarmUpScheduler(optimizer, dim_embed=768, warmup_steps=4000)
-
-        self.tconf.update(optimizer=optimizer, lr_scheduler=lr_scheduler)
-        self.trainer = Trainer(
-            self.model,
-            train_dataloader,
-            self.tconf,
-            val_dataloader,
-            experiment,
-            self.use_cuda,
-        )
-        self.trainer.train()
-
-    def _fit(
-        self,
-        log: SparkDataFrame,
-        user_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-        item_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-    ) -> None:
-        self.train(log)
-
-    def _predict(
-        self,
-        log: SparkDataFrame,  # noqa: ARG002
-        k: int,  # noqa: ARG002
-        users: SparkDataFrame,
-        items: SparkDataFrame,
-        user_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-        item_features: Optional[SparkDataFrame] = None,  # noqa: ARG002
-        filter_seen_items: bool = True,  # noqa: ARG002
-    ) -> SparkDataFrame:
-        items_consider_in_pred = items.toPandas()["item_idx"].values
-        users_consider_in_pred = users.toPandas()["user_idx"].values
-        ans = self._predict_helper(users_consider_in_pred, items_consider_in_pred)
-        return convert2spark(ans)
-
-    def _predict_helper(self, users, items, max_context_len=30):
-        predict_dataloader = self._make_prediction_dataloader(users, items, max_context_len)
-        self.model.eval()
-        ans_df = pd.DataFrame(columns=["user_idx", "item_idx", "relevance"])
-        with torch.no_grad():
-            for batch in tqdm(predict_dataloader):
-                states, actions, rtgs, timesteps, users = self.trainer._move_batch(batch)
-                logits = self.model(states, actions, rtgs, timesteps, users)
-                items_relevances = logits[:, -1, :][:, items]
-                ans_df = pd.concat([ans_df, matrix2df(items_relevances, users.squeeze(), items)])
-
-        return ans_df

replay/experimental/models/dt4rec/gpt1.py

@@ -1,401 +0,0 @@
-import logging
-import math
-
-from replay.utils import TORCH_AVAILABLE
-
-if TORCH_AVAILABLE:
-    import torch
-    from torch import nn
-    from torch.nn import functional as func
-
-logger = logging.getLogger(__name__)
-
-
-class GELU(nn.Module):
-    """
-    GELU callable class
-    """
-
-    def forward(self, x):
-        """
-        Apply GELU
-        """
-        return func.gelu(x)
-
-
-class GPTConfig:
-    """base GPT config, params common to all GPT versions"""
-
-    embd_pdrop = 0.1
-    resid_pdrop = 0.1
-    attn_pdrop = 0.1
-    n_layer = 6
-    n_head = 8
-    n_embd = 128
-    memory_size = 3
-
-    def __init__(self, vocab_size, block_size, **kwargs):
-        self.vocab_size = vocab_size
-        self.block_size = block_size
-        for k, v in kwargs.items():
-            setattr(self, k, v)
-
-    def update(self, **kwargs):
-        """
-        Arguments setter
-        """
-        for k, v in kwargs.items():
-            setattr(self, k, v)
-
-
-class GPT1Config(GPTConfig):
-    """GPT-1 like network roughly 125M params"""
-
-    n_layer = 12
-    n_head = 12
-    n_embd = 768
-
-
-class CausalSelfAttention(nn.Module):
-    """
-    A vanilla multi-head masked self-attention layer with a projection at the end.
-    It is possible to use torch.nn.MultiheadAttention here but I am including an
-    explicit implementation here to show that there is nothing too scary here.
-    """
-
-    def __init__(self, config):
-        super().__init__()
-        assert config.n_embd % config.n_head == 0
-        # key, query, value projections for all heads
-        self.key = nn.Linear(config.n_embd, config.n_embd)
-        self.query = nn.Linear(config.n_embd, config.n_embd)
-        self.value = nn.Linear(config.n_embd, config.n_embd)
-        # regularization
-        self.attn_drop = nn.Dropout(config.attn_pdrop)
-        self.resid_drop = nn.Dropout(config.resid_pdrop)
-        # output projection
-        self.proj = nn.Linear(config.n_embd, config.n_embd)
-        # causal mask to ensure that attention is only applied to the left in the input sequence
-        # self.register_buffer("mask", torch.tril(torch.ones(config.block_size, config.block_size))
-        #                              .view(1, 1, config.block_size, config.block_size))
-        self.register_buffer(
-            "mask",
-            torch.tril(torch.ones(config.block_size + 1, config.block_size + 1)).view(
-                1, 1, config.block_size + 1, config.block_size + 1
-            ),
-        )
-        self.n_head = config.n_head
-
-    def forward(self, x):
-        """
-        Apply attention
-        """
-        B, T, C = x.size()  # noqa: N806
-
-        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
-        k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
-        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
-        v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
-
-        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
-        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
-        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
-        att = func.softmax(att, dim=-1)
-        att = self.attn_drop(att)
-        y = att @ v  # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
-        y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-assemble all head outputs side by side
-
-        # output projection
-        y = self.resid_drop(self.proj(y))
-        return y
-
-
-class Block(nn.Module):
-    """an unassuming Transformer block"""
-
-    def __init__(self, config):
-        super().__init__()
-        self.ln1 = nn.LayerNorm(config.n_embd)
-        self.ln2 = nn.LayerNorm(config.n_embd)
-        self.attn = CausalSelfAttention(config)
-        self.mlp = nn.Sequential(
-            nn.Linear(config.n_embd, 4 * config.n_embd),
-            GELU(),
-            nn.Linear(4 * config.n_embd, config.n_embd),
-            nn.Dropout(config.resid_pdrop),
-        )
-
-    def forward(self, x):
-        """
-        :x: batch
-        """
-        x = x + self.attn(self.ln1(x))
-        x = x + self.mlp(self.ln2(x))
-        return x
-
-
-class StateReprModule(nn.Module):
-    """
-    Compute state for RL environment. Based on `DRR paper
-    <https://arxiv.org/pdf/1810.12027.pdf>`_
-
-    Computes State is a concatenation of user embedding,
-    weighted average pooling of `memory_size` latest relevant items
-    and their pairwise product.
-    """
-
-    def __init__(
-        self,
-        user_num,
-        item_num,
-        embedding_dim,
-        memory_size,
-    ):
-        super().__init__()
-        self.user_embeddings = nn.Embedding(user_num, embedding_dim)
-
-        self.item_embeddings = nn.Embedding(item_num + 1, embedding_dim, padding_idx=int(item_num))
-
-        self.drr_ave = torch.nn.Conv1d(in_channels=memory_size, out_channels=1, kernel_size=1)
-
-        self.linear = nn.Linear(3 * embedding_dim, embedding_dim)
-
-        self.initialize()
-
-    def initialize(self):
-        """weight init"""
-        nn.init.normal_(self.user_embeddings.weight, std=0.01)
-        self.item_embeddings.weight.data[-1].zero_()
-
-        nn.init.normal_(self.item_embeddings.weight, std=0.01)
-        nn.init.uniform_(self.drr_ave.weight)
-
-        self.drr_ave.bias.data.zero_()
-
-    def forward(self, user, memory):
-        """
-        :param user: user batch
-        :param memory: memory batch
-        :return: vector of dimension embedding_dim
-        """
-        user_embedding = self.user_embeddings(user.long()).squeeze(1)
-        item_embeddings = self.item_embeddings(memory.long())
-        drr_ave = self.drr_ave(item_embeddings).squeeze(1)
-        output = torch.cat((user_embedding, user_embedding * drr_ave, drr_ave), 1)
-        output = self.linear(output)
-
-        return output
-
-
-class GPT(nn.Module):
-    """the full GPT language model, with a context size of block_size"""
-
-    def __init__(self, config):
-        super().__init__()
-
-        self.config = config
-
-        self.user_num = config.user_num
-        self.memory_size = config.memory_size
-
-        self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size + 1, config.n_embd))
-        self.global_pos_emb = nn.Parameter(torch.zeros(1, config.max_timestep + 1, config.n_embd))
-        self.drop = nn.Dropout(config.embd_pdrop)
-
-        # transformer
-        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)])
-        # decoder head
-        self.ln_f = nn.LayerNorm(config.n_embd)
-        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
-
-        self.block_size = config.block_size
-        self.apply(self._init_weights)
-
-        logger.info(
-            "number of parameters: %e",
-            sum(p.numel() for p in self.parameters()),
-        )
-
-        self.state_repr = StateReprModule(
-            user_num=config.user_num,
-            item_num=config.vocab_size,
-            embedding_dim=config.n_embd,
-            memory_size=config.memory_size,
-        )
-
-        self.ret_emb = nn.Sequential(nn.Linear(1, config.n_embd), nn.Tanh())
-        self.action_embeddings = nn.Sequential(self.state_repr.item_embeddings, nn.Tanh())
-        nn.init.normal_(self.action_embeddings[0].weight, mean=0.0, std=0.02)
-
-    def get_block_size(self):
-        """
-        Return block_size
-        """
-        return self.block_size
-
-    @staticmethod
-    def _init_weights(module):
-        if isinstance(module, (nn.Linear, nn.Embedding)):
-            module.weight.data.normal_(mean=0.0, std=0.02)
-            if isinstance(module, nn.Linear) and module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.LayerNorm):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-
-    def configure_optimizers(self):
-        """
-        This long function is unfortunately doing something very simple and is being very defensive:
-        We are separating out all parameters of the model into two buckets: those that will experience
-        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).
-        We are then returning the PyTorch optimizer object.
-        """
-
-        # separate out all parameters to those that will and won't experience regularizing weight decay
-        decay = set()
-        no_decay = set()
-        whitelist_weight_modules = (
-            torch.nn.Linear,
-            torch.nn.Conv2d,
-            torch.nn.Conv1d,
-        )
-        blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
-        for mn, m in self.named_modules():
-            for pn, _ in m.named_parameters():
-                fpn = f"{mn}.{pn}" if mn else pn  # full param name
-
-                if pn.endswith("bias"):
-                    # all biases will not be decayed
-                    no_decay.add(fpn)
-                elif pn.endswith("weight") and isinstance(m, whitelist_weight_modules):
-                    # weights of whitelist modules will be weight decayed
-                    decay.add(fpn)
-                elif pn.endswith("weight") and isinstance(m, blacklist_weight_modules):
-                    # weights of blacklist modules will NOT be weight decayed
-                    no_decay.add(fpn)
-
-        # special case the position embedding parameter in the root GPT module as not decayed
-        no_decay.add("pos_emb")
-        no_decay.add("global_pos_emb")
-
-        # validate that we considered every parameter
-        param_dict = dict(self.named_parameters())
-        inter_params = decay & no_decay
-        union_params = decay | no_decay
-        assert len(inter_params) == 0, f"parameters {inter_params!s} made it into both decay/no_decay sets!"
-        assert (
-            len(param_dict.keys() - union_params) == 0
-        ), f"parameters {param_dict.keys() - union_params!s} were not separated into either decay/no_decay set!"
-
-        optim_groups = [
-            {
-                "params": [param_dict[pn] for pn in sorted(decay)],
-                "weight_decay": 0.1,
-            },
-            {
-                "params": [param_dict[pn] for pn in sorted(no_decay)],
-                "weight_decay": 0.0,
-            },
-        ]
-        return optim_groups
-
-    # state, action, and return
-    def forward(
-        self,
-        states,
-        actions,
-        rtgs,
-        timesteps,
-        users,
-    ):
-        """
-        :states: states batch, (batch, trajectory_len, 3)
-        :actions: actions batch, (batch, trajectory_len, 1)
-        :rtgs: rtgs batch, (batch, trajectory_len, 1)
-        :timesteps: timesteps batch, (batch, 1, 1)
-        :users:users batch,  (batch, 1)
-        """
-        inference = not self.training
-        state_embeddings = self.state_repr(
-            users.repeat((1, self.block_size // 3)).reshape(-1, 1),
-            states.reshape(-1, 3),
-        )
-
-        state_embeddings = state_embeddings.reshape(
-            states.shape[0], states.shape[1], self.config.n_embd
-        )  # (batch, block_size, n_embd)
-
-        if actions is not None:
-            rtg_embeddings = self.ret_emb(rtgs.type(torch.float32))
-            action_embeddings = self.action_embeddings(
-                actions.type(torch.long).squeeze(-1)
-            )  # (batch, block_size, n_embd)
-
-            token_embeddings = torch.zeros(
-                (
-                    states.shape[0],
-                    states.shape[1] * 3 - int(inference),
-                    self.config.n_embd,
-                ),
-                dtype=torch.float32,
-                device=state_embeddings.device,
-            )
-            token_embeddings[:, ::3, :] = rtg_embeddings
-            token_embeddings[:, 1::3, :] = state_embeddings
-            token_embeddings[:, 2::3, :] = action_embeddings[:, -states.shape[1] + int(inference) :, :]
-        else:
-            # only happens at very first timestep of evaluation
-            rtg_embeddings = self.ret_emb(rtgs.type(torch.float32))
-
-            token_embeddings = torch.zeros(
-                (states.shape[0], states.shape[1] * 2, self.config.n_embd),
-                dtype=torch.float32,
-                device=state_embeddings.device,
-            )
-            token_embeddings[:, ::2, :] = rtg_embeddings  # really just [:,0,:]
-            token_embeddings[:, 1::2, :] = state_embeddings  # really just [:,1,:]
-
-        batch_size = states.shape[0]
-        all_global_pos_emb = torch.repeat_interleave(
-            self.global_pos_emb, batch_size, dim=0
-        )  # batch_size, traj_length, n_embd
-
-        position_embeddings = (
-            torch.gather(
-                all_global_pos_emb,
-                1,
-                torch.repeat_interleave(timesteps, self.config.n_embd, dim=-1),
-            )
-            + self.pos_emb[:, : token_embeddings.shape[1], :]
-        )
-        x = self.drop(token_embeddings + position_embeddings)
-        x = self.blocks(x)
-        x = self.ln_f(x)
-        logits = self.head(x)
-
-        if actions is not None:
-            logits = logits[:, 1::3, :]  # only keep predictions from state_embeddings
-        elif actions is None:
-            logits = logits[:, 1:, :]
-
-        return logits
-
-    def predict(self, states, actions, rtgs, timesteps, users):
-        """
-        :states: states batch, (batch, block_size, 3)
-        :actions: actions batch, (batch, block_size, 1)
-        :rtgs: rtgs batch, (batch, block_size, 1)
-        :timesteps: timesteps batch, (batch, 1, 1)
-        :users:users batch,  (batch, 1)
-        """
-        logits, _ = self(
-            states=states.to(self.pos_emb.device),
-            actions=actions.to(self.pos_emb.device),
-            targets=None,
-            rtgs=rtgs.to(self.pos_emb.device),
-            timesteps=timesteps.to(self.pos_emb.device),
-            users=users.to(self.pos_emb.device),
-        )
-        logits = logits[:, -1, :]
-        actions = logits.argsort(dim=1, descending=True)
-        return actions