Formatting

Author: TOPAPEC

rectools/fast_transformers/__init__.py

@@ -1,13 +1,13 @@
 """Fast Transformers: flat sequential recommenders without ItemNet hierarchy."""
 
-from .gpu_data import build_sequences, align_embeddings, GPUBatchDataset, make_dataloader
+from .gpu_data import GPUBatchDataset, align_embeddings, build_sequences, make_dataloader
 from .lightning_wrap import FlatSASRecLightning
 from .model import FlatSASRecConfig, FlatSASRecModel
 from .net import FlatSASRec, SASRecBlock
 from .ranking import rank_topk
-from .unisrec_net import UniSRec, FeedForward
 from .unisrec_lightning import UniSRecLightning
 from .unisrec_model import UniSRecModel
+from .unisrec_net import FeedForward, UniSRec
 
 __all__ = [
     "build_sequences",

rectools/fast_transformers/gpu_data.py

@@ -3,7 +3,8 @@
 import typing as tp
 
 import torch
-from torch.utils.data import Dataset as TorchDataset, DataLoader
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset as TorchDataset
 
 
 def build_sequences(
@@ -52,7 +53,9 @@ def build_sequences(
     if total_elements > 0:
         user_indices = torch.repeat_interleave(torch.arange(n_users, device=device), effective_lens)
         cumsum = effective_lens.cumsum(0)
-        offsets = torch.arange(total_elements, device=device) - torch.repeat_interleave(cumsum - effective_lens, effective_lens)
+        offsets = torch.arange(total_elements, device=device) - torch.repeat_interleave(
+            cumsum - effective_lens, effective_lens
+        )
 
         x_src = torch.repeat_interleave(ends - capped_lens, effective_lens) + offsets
         y_src = x_src + 1

rectools/fast_transformers/lightning_wrap.py

@@ -2,8 +2,8 @@
 
 import typing as tp
 
-import torch
 import pytorch_lightning as pl
+import torch
 from torch import nn
 
 from .net import FlatSASRec
@@ -47,7 +47,9 @@ def training_step(self, batch: tp.Dict[str, torch.Tensor], batch_idx: int) -> to
         if self.loss_name == "softmax":
             # logits: (B, L, n_items) — full catalog
             # targets need to be 0-indexed item ids (subtract 1 since item ids start from 1)
-            targets = y - 1  # shift to 0-based for CrossEntropyLoss; padding (0) becomes -1 -> ignore_index=0 won't work
+            targets = (
+                y - 1
+            )  # shift to 0-based for CrossEntropyLoss; padding (0) becomes -1 -> ignore_index=0 won't work
             # Actually, we set ignore_index=0 but padding maps to -1.
             # Let's use a different approach: set padding targets to 0 and use ignore_index=0
             targets = y.clone()

rectools/fast_transformers/model.py

@@ -2,18 +2,15 @@
 
 import typing as tp
 
-import numpy as np
 import pandas as pd
-import torch
 import pytorch_lightning as pl
+import torch
 from scipy import sparse
 
-from rectools import Columns
 from rectools.dataset import Dataset
-from rectools.dataset.identifiers import IdMap
 from rectools.models.base import InternalRecoTriplet, ModelBase, ModelConfig
-from rectools.models.nn.transformers.sasrec import SASRecDataPreparator
 from rectools.models.nn.transformers.negative_sampler import CatalogUniformSampler
+from rectools.models.nn.transformers.sasrec import SASRecDataPreparator
 from rectools.types import InternalIdsArray
 from rectools.utils.config import BaseConfig
 
@@ -157,10 +154,6 @@ def _fit(self, dataset: Dataset, *args: tp.Any, **kwargs: tp.Any) -> None:
         dp.process_dataset_train(dataset)
         self._data_preparator = dp
 
-        n_items = dp.item_id_map.size  # includes extra tokens (padding)
-        # item ids in the preparator go from 0 (padding) to n_items-1
-        # FlatSASRec expects n_items = max real item count (embedding table = n_items+1 with padding at 0)
-        # The preparator's item_id_map.size includes the padding token, so real items = size - 1
         n_real_items = dp.item_id_map.size - dp.n_item_extra_tokens
 
         net = FlatSASRec(
@@ -242,7 +235,6 @@ def _recommend_u2i(
         sorted_item_ids_to_recommend: tp.Optional[InternalIdsArray],
     ) -> InternalRecoTriplet:
         assert self._data_preparator is not None
-        device = next(self._net.parameters()).device  # type: ignore
 
         user_embs = self._get_user_embeddings(dataset)  # (n_users, D)
         item_embs = self._get_item_embeddings()  # (n_items, D)
@@ -278,7 +270,9 @@ def _recommend_u2i(
             whitelist = wl[(wl >= 0) & (wl < item_embs.shape[0])]
 
         u_ids, i_ids, scores = rank_topk(
-            user_embs, item_embs, k,
+            user_embs,
+            item_embs,
+            k,
             filter_csr=filter_csr,
             whitelist=whitelist,
             batch_size=self.recommend_batch_size,
@@ -298,7 +292,6 @@ def _recommend_i2i(
         sorted_item_ids_to_recommend: tp.Optional[InternalIdsArray],
     ) -> InternalRecoTriplet:
         assert self._data_preparator is not None and self._net is not None
-        device = next(self._net.parameters()).device
 
         item_embs = self._get_item_embeddings()  # (n_items, D)
         n_extra = self._data_preparator.n_item_extra_tokens
@@ -313,7 +306,9 @@ def _recommend_i2i(
             whitelist = wl[(wl >= 0) & (wl < item_embs.shape[0])]
 
         t_ids, i_ids, scores = rank_topk(
-            target_embs, item_embs, k,
+            target_embs,
+            item_embs,
+            k,
             whitelist=whitelist,
             batch_size=self.recommend_batch_size,
         )

rectools/fast_transformers/net.py

@@ -127,19 +127,7 @@ def encode_last(self, x: torch.Tensor) -> torch.Tensor:
         Tensor (B, D)
         """
         h = self.encode(x)  # (B, L, D)
-        # Find last non-padding position per row
-        non_pad = (x != self.PADDING_IDX)  # (B, L)
-        # lengths: number of non-pad tokens
-        lengths = non_pad.sum(dim=1)  # (B,)
-        # Clamp to at least 1 to avoid index -1 for fully-padded rows
-        last_idx = (lengths - 1).clamp(min=0)
-        # We use left-padding, so last non-pad is at position (L - 1) if any token exists
-        # Actually with left padding, non-pad tokens are at the end, so the last position is L-1
-        # But let's compute correctly: the last non-pad index
-        # With left-padding: first non-pad is at L - length, last non-pad is at L - 1
-        B = x.shape[0]
-        last_pos = x.shape[1] - 1  # last position is always the last for left-padded sequences
-        return h[:, last_pos, :]  # (B, D)
+        return h[:, -1, :]  # left-padded: last position is always rightmost
 
     def all_item_embeddings(self) -> torch.Tensor:
         """

rectools/fast_transformers/unisrec_lightning.py

@@ -3,9 +3,9 @@
 import math
 import typing as tp
 
+import pytorch_lightning as pl
 import torch
 import torch.nn.functional as F
-import pytorch_lightning as pl
 from torch.optim.lr_scheduler import LambdaLR
 
 from .unisrec_net import UniSRec
@@ -63,23 +63,29 @@ def _get_all_embs(self) -> torch.Tensor:
         return self.net.project_all()
 
     def _get_pos_neg_logits(
-        self, hidden: torch.Tensor, labels: torch.Tensor, negatives: torch.Tensor,
+        self,
+        hidden: torch.Tensor,
+        labels: torch.Tensor,
+        negatives: torch.Tensor,
     ) -> torch.Tensor:
         """Compute (B, L, 1+N) logits where index 0 = positive."""
         emb_pos = self._get_item_embs(labels)
         logits_pos = (hidden * emb_pos).sum(dim=-1)
 
         emb_neg = self._get_item_embs(negatives)
         logits_neg = torch.matmul(
-            hidden.unsqueeze(2), emb_neg.transpose(2, 3),
+            hidden.unsqueeze(2),
+            emb_neg.transpose(2, 3),
         ).squeeze(2)
 
         return torch.cat([logits_pos.unsqueeze(-1), logits_neg], dim=-1)
 
     # ── losses ──
 
     def _calc_loss(
-        self, hidden: torch.Tensor, batch: tp.Dict[str, torch.Tensor],
+        self,
+        hidden: torch.Tensor,
+        batch: tp.Dict[str, torch.Tensor],
     ) -> torch.Tensor:
         labels = batch["y"]
         has_neg = "negatives" in batch
@@ -114,7 +120,9 @@ def _full_softmax_loss(self, hidden: torch.Tensor, labels: torch.Tensor) -> torc
         targets = labels.clone()
         targets[targets == 0] = -100
         return F.cross_entropy(
-            logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100,
+            logits.view(-1, logits.size(-1)),
+            targets.view(-1),
+            ignore_index=-100,
         )
 
     def _sampled_softmax_loss(self, logits: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
@@ -123,7 +131,9 @@ def _sampled_softmax_loss(self, logits: torch.Tensor, mask: torch.Tensor) -> tor
         logits[:, :, [0, 1]] = logits[:, :, [1, 0]]
         targets = mask.long()  # 1 where non-padding, 0 where padding
         return F.cross_entropy(
-            logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=0,
+            logits.view(-1, logits.size(-1)),
+            targets.view(-1),
+            ignore_index=0,
         )
 
     def _bce_loss(self, logits: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:

rectools/fast_transformers/unisrec_model.py

@@ -3,13 +3,13 @@
 import typing as tp
 from pathlib import Path
 
-import torch
 import pytorch_lightning as pl
+import torch
 from pytorch_lightning.callbacks import EarlyStopping
 
+from .gpu_data import align_embeddings, build_sequences, make_dataloader
+from .unisrec_lightning import SUPPORTED_LOSSES, SUPPORTED_OPTIMIZERS, SUPPORTED_SCHEDULERS, UniSRecLightning
 from .unisrec_net import UniSRec
-from .unisrec_lightning import UniSRecLightning, SUPPORTED_LOSSES, SUPPORTED_OPTIMIZERS, SUPPORTED_SCHEDULERS
-from .gpu_data import build_sequences, align_embeddings, make_dataloader
 
 
 class UniSRecModel:
@@ -143,7 +143,12 @@ def _make_trainer(self, max_epochs: int, val_dl: tp.Any = None) -> pl.Trainer:
         )
 
     def _make_lightning(
-        self, net: UniSRec, param_groups: tp.List[tp.Dict], use_id: bool, max_epochs: int, train_dl: tp.Any,
+        self,
+        net: UniSRec,
+        param_groups: tp.List[tp.Dict],
+        use_id: bool,
+        max_epochs: int,
+        train_dl: tp.Any,
     ) -> UniSRecLightning:
         total_steps = len(train_dl) * max_epochs if self.scheduler else None
         return UniSRecLightning(
@@ -172,16 +177,22 @@ def _phase2_params(self, net: UniSRec) -> tp.List[tp.Dict[str, tp.Any]]:
                 {"params": [net.whitening_bias], "lr": self.phase2_lr * 10.0, "weight_decay": 0.0},
             ]
             if net.head is not None:
-                groups.append({
-                    "params": list(net.head.parameters()),
-                    "lr": self.phase2_lr * self.lr_head,
-                    "weight_decay": self.weight_decay,
-                })
+                groups.append(
+                    {
+                        "params": list(net.head.parameters()),
+                        "lr": self.phase2_lr * self.lr_head,
+                        "weight_decay": self.weight_decay,
+                    }
+                )
         else:
             groups = [
                 {"params": list(net.bn_input.parameters()), "lr": self.phase2_lr, "weight_decay": 0.0},
                 {"params": list(net.bn_score.parameters()), "lr": self.phase2_lr, "weight_decay": 0.0},
-                {"params": list(net.head.parameters()), "lr": self.phase2_lr * self.lr_head, "weight_decay": self.weight_decay},
+                {
+                    "params": list(net.head.parameters()),
+                    "lr": self.phase2_lr * self.lr_head,
+                    "weight_decay": self.weight_decay,
+                },
             ]
         return groups
 
@@ -198,21 +209,27 @@ def _phase3_params(self, net: UniSRec) -> tp.List[tp.Dict[str, tp.Any]]:
             ]
         head: tp.List[tp.Dict[str, tp.Any]] = []
         if net.head is not None:
-            head = [{"params": list(net.head.parameters()), "lr": self.phase3_lr * self.lr_head, "weight_decay": self.weight_decay}]
+            head = [
+                {
+                    "params": list(net.head.parameters()),
+                    "lr": self.phase3_lr * self.lr_head,
+                    "weight_decay": self.weight_decay,
+                }
+            ]
         transformer = [
             {"params": list(net.pos_emb.parameters()), "lr": self.phase3_lr * self.lr_transformer, "weight_decay": 0.0},
             {
                 "params": (
-                    [p for l in net.attention_layers for p in l.parameters()]
-                    + [p for l in net.forward_layers for p in l.parameters()]
+                    [p for layer in net.attention_layers for p in layer.parameters()]
+                    + [p for layer in net.forward_layers for p in layer.parameters()]
                 ),
                 "lr": self.phase3_lr * self.lr_transformer,
                 "weight_decay": self.weight_decay,
             },
             {
                 "params": (
-                    [p for l in net.attention_layernorms for p in l.parameters()]
-                    + [p for l in net.forward_layernorms for p in l.parameters()]
+                    [p for layer in net.attention_layernorms for p in layer.parameters()]
+                    + [p for layer in net.forward_layernorms for p in layer.parameters()]
                     + list(net.last_layernorm.parameters())
                 ),
                 "lr": self.phase3_lr,
@@ -246,7 +263,9 @@ def fit(
         self
         """
         x, y, unique_items, unique_users = build_sequences(
-            user_ids, item_ids, timestamps,
+            user_ids,
+            item_ids,
+            timestamps,
             max_len=self.session_max_len,
             min_interactions=self.train_min_user_interactions,
         )
@@ -303,12 +322,15 @@ def _run_phase(param_groups: tp.List[tp.Dict], use_id: bool, max_epochs: int) ->
 
     def save_checkpoint(self, path: tp.Union[str, Path]) -> None:
         assert self._net is not None
-        torch.save({
-            "net": self._net.state_dict(),
-            "unique_items": self._unique_items,
-            "unique_users": self._unique_users,
-            "n_items": len(self._unique_items),
-        }, path)
+        torch.save(
+            {
+                "net": self._net.state_dict(),
+                "unique_items": self._unique_items,
+                "unique_users": self._unique_users,
+                "n_items": len(self._unique_items),
+            },
+            path,
+        )
 
     def load_checkpoint(self, path: tp.Union[str, Path], device: str = "cuda") -> None:
         ckpt = torch.load(path, map_location=device, weights_only=False)

rectools/fast_transformers/unisrec_net.py

@@ -51,12 +51,17 @@ def make_ffn(n_factors: int, ffn_type: str, expansion: int, dropout: float) -> n
     hidden = n_factors * expansion
     if ffn_type == "linear_gelu":
         return nn.Sequential(
-            nn.Linear(n_factors, hidden), nn.GELU(), nn.Dropout(dropout),
-            nn.Linear(hidden, n_factors), nn.Dropout(dropout),
+            nn.Linear(n_factors, hidden),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden, n_factors),
+            nn.Dropout(dropout),
         )
     if ffn_type == "linear_relu":
         return nn.Sequential(
-            nn.Linear(n_factors, hidden), nn.ReLU(), nn.Dropout(dropout),
+            nn.Linear(n_factors, hidden),
+            nn.ReLU(),
+            nn.Dropout(dropout),
             nn.Linear(hidden, n_factors),
         )
     raise ValueError(f"Unknown ffn_type: {ffn_type}. Choose from: conv1d, linear_gelu, linear_relu")
@@ -238,8 +243,10 @@ def project_all(self) -> torch.Tensor:
     @property
     def transformer_params(self) -> tp.List[nn.Parameter]:
         modules = (
-            list(self.attention_layernorms) + list(self.attention_layers)
-            + list(self.forward_layernorms) + list(self.forward_layers)
+            list(self.attention_layernorms)
+            + list(self.attention_layers)
+            + list(self.forward_layernorms)
+            + list(self.forward_layers)
             + [self.last_layernorm, self.pos_emb]
         )
         return [p for m in modules for p in m.parameters()]
@@ -272,9 +279,9 @@ def _encode(self, seqs: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
         seqs = seqs + self.pos_emb(positions)
         seqs = self.emb_dropout(seqs)
 
-        pad_mask = (input_ids == self.PADDING_IDX)               # (B, L)
-        pad_mask_3d = pad_mask.unsqueeze(-1)                       # (B, L, 1)
-        seqs = seqs.masked_fill(pad_mask_3d, 0.0)                 # zero out padding
+        pad_mask = input_ids == self.PADDING_IDX  # (B, L)
+        pad_mask_3d = pad_mask.unsqueeze(-1)  # (B, L, 1)
+        seqs = seqs.masked_fill(pad_mask_3d, 0.0)  # zero out padding
 
         attn_mask = self._causal_mask(L, seqs.device)
         key_padding_mask = pad_mask
@@ -284,7 +291,9 @@ def _encode(self, seqs: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
             # Zero padding in Q/K/V so NaN can never appear in dot-products
             normed = normed.masked_fill(pad_mask_3d, 0.0)
             mha_out, _ = self.attention_layers[i](
-                normed, normed, normed,
+                normed,
+                normed,
+                normed,
                 attn_mask=attn_mask,
                 key_padding_mask=key_padding_mask,
                 need_weights=False,