Fix remaining review items: make_ffn docstring, N812 violation, padding test

- Add all parameters and Returns to make_ffn docstring (feldlime)
- Replace import torch.nn.functional as F with direct imports in
  lightning.py to eliminate N812 violation; remove N812 from setup.cfg
- Add padding-zeroed assertion to test_determinism_and_padding_masking
  so the test matches its docstring claim (copilot)

Author: TOPAPEC

rectools/fast_transformers/unisrec/lightning.py

@@ -5,7 +5,7 @@
 
 import pytorch_lightning as pl
 import torch
-import torch.nn.functional as F
+from torch.nn.functional import binary_cross_entropy_with_logits, cross_entropy
 from torch.optim.lr_scheduler import LambdaLR
 
 from .net import UniSRecNet
@@ -114,7 +114,7 @@ def _full_softmax_loss(self, hidden: torch.Tensor, labels: torch.Tensor) -> torc
 
         targets = labels.clone()
         targets[targets == 0] = -100
-        return F.cross_entropy(
+        return cross_entropy(
             logits.view(-1, logits.size(-1)),
             targets.view(-1),
             ignore_index=-100,
@@ -125,7 +125,7 @@ def _sampled_softmax_loss(self, logits: torch.Tensor, mask: torch.Tensor) -> tor
         logits = logits.clone()
         logits[:, :, [0, 1]] = logits[:, :, [1, 0]]
         targets = mask.long()  # 1 where non-padding, 0 where padding
-        return F.cross_entropy(
+        return cross_entropy(
             logits.view(-1, logits.size(-1)),
             targets.view(-1),
             ignore_index=0,
@@ -134,7 +134,7 @@ def _sampled_softmax_loss(self, logits: torch.Tensor, mask: torch.Tensor) -> tor
     def _bce_loss(self, logits: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
         target = torch.zeros_like(logits)
         target[:, :, 0] = 1.0
-        loss = F.binary_cross_entropy_with_logits(logits, target, reduction="none")
+        loss = binary_cross_entropy_with_logits(logits, target, reduction="none")
         loss = loss.mean(-1) * mask
         return loss.sum() / mask.sum().clamp(min=1)
 

rectools/fast_transformers/unisrec/net.py

@@ -44,8 +44,19 @@ def make_ffn(n_factors: int, ffn_type: str, expansion: int, dropout: float) -> n
 
     Parameters
     ----------
+    n_factors : int
+        Input and output dimension.
     ffn_type : ``"conv1d"`` | ``"linear_gelu"`` | ``"linear_relu"``
-    expansion : hidden-dim multiplier (e.g. 1 or 4).
+        Type of feed-forward block.
+    expansion : int
+        Hidden-dimension multiplier (e.g. 1 or 4).
+    dropout : float
+        Dropout rate applied inside the block.
+
+    Returns
+    -------
+    nn.Module
+        A feed-forward network module.
     """
     if ffn_type == "conv1d":
         return FeedForwardConv1d(n_factors, dropout)

setup.cfg

@@ -50,7 +50,6 @@ per-file-ignores =
     rectools/dataset/torch_datasets.py: D101,D102
     rectools/models/implicit_als.py: N806
     rectools/fast_transformers/net.py: N806
-    rectools/fast_transformers/unisrec/lightning.py: N812
     rectools/fast_transformers/unisrec/net.py: N806
 
 [mypy]

tests/fast_transformers/test_unisrec_net.py

@@ -92,4 +92,7 @@ def test_determinism_and_padding_masking(self, net: UniSRecNet) -> None:
         with torch.no_grad():
             e_a = net.encode_last(x_a)
             e_b = net.encode_last(x_b)
+            h_a = net(x_a)
         torch.testing.assert_close(e_a, e_b)
+        # Padding positions (first 3 columns) should be zeroed in full output
+        torch.testing.assert_close(h_a[:, :3, :], torch.zeros_like(h_a[:, :3, :]))