style: format code

Author: yangzhou23

tests/test_inbatch_sampling.py

@@ -6,7 +6,7 @@
 from torch_rechub.models.matching import DSSM
 from torch_rechub.trainers import MatchTrainer
 from torch_rechub.utils.data import MatchDataGenerator, df_to_dict
-from torch_rechub.utils.match import gather_inbatch_logits, generate_seq_feature_match, gen_model_input, inbatch_negative_sampling
+from torch_rechub.utils.match import gather_inbatch_logits, gen_model_input, generate_seq_feature_match, inbatch_negative_sampling
 
 
 def test_inbatch_negative_sampling_random_and_uniform():
@@ -32,13 +32,15 @@ def test_inbatch_negative_sampling_hard_negative():
 
 def _build_small_match_dataloader():
     n_users, n_items, n_samples = 12, 24, 80
-    data = pd.DataFrame(
-        {
-            "user_id": np.random.randint(0, n_users, n_samples),
-            "item_id": np.random.randint(0, n_items, n_samples),
-            "time": np.arange(n_samples),
-        }
-    )
+    data = pd.DataFrame({
+        "user_id": np.random.randint(0,
+                                     n_users,
+                                     n_samples),
+        "item_id": np.random.randint(0,
+                                     n_items,
+                                     n_samples),
+        "time": np.arange(n_samples),
+    })
     user_profile = pd.DataFrame({"user_id": np.arange(n_users)})
     item_profile = pd.DataFrame({"item_id": np.arange(n_items)})
 
@@ -48,8 +50,14 @@ def _build_small_match_dataloader():
     y_train = np.zeros(len(df_train))
 
     user_features = [
-        SparseFeature("user_id", n_users, embed_dim=8),
-        SequenceFeature("hist_item_id", n_items, embed_dim=8, pooling="mean", shared_with="item_id"),
+        SparseFeature("user_id",
+                      n_users,
+                      embed_dim=8),
+        SequenceFeature("hist_item_id",
+                        n_items,
+                        embed_dim=8,
+                        pooling="mean",
+                        shared_with="item_id"),
     ]
     item_features = [SparseFeature("item_id", n_items, embed_dim=8)]
 

torch_rechub/trainers/match_trainer.py

@@ -115,7 +115,7 @@ def train_one_epoch(self, data_loader, log_interval=10):
                 if user_embedding.dim() != 2 or item_embedding.dim() != 2:
                     raise ValueError(f"In-batch negative sampling requires 2D embeddings, got shapes {user_embedding.shape} and {item_embedding.shape}")
 
-                scores = torch.matmul(user_embedding, item_embedding.t()) # bs x bs
+                scores = torch.matmul(user_embedding, item_embedding.t())  # bs x bs
                 neg_indices = inbatch_negative_sampling(scores, neg_ratio=self.in_batch_neg_ratio, hard_negative=self.hard_negative, generator=self._sampler_generator)
                 logits = gather_inbatch_logits(scores, neg_indices)
                 if self.mode == 1:  # pair_wise

torch_rechub/utils/match.py

@@ -118,13 +118,13 @@ def inbatch_negative_sampling(scores, neg_ratio=None, hard_negative=False, gener
     Returns:
         torch.Tensor: sampled negative indices with shape (batch_size, neg_ratio).
     """
-    if scores.dim() != 2: # must be batch_size x batch_size
+    if scores.dim() != 2:  # must be batch_size x batch_size
         raise ValueError(f"inbatch_negative_sampling expects 2D scores, got shape {tuple(scores.shape)}")
     batch_size = scores.size(0)
     if batch_size <= 1:
         raise ValueError("In-batch negative sampling requires batch_size > 1")
 
-    max_neg = batch_size - 1 # each col can provide at most batch_size-1 negatives
+    max_neg = batch_size - 1  # each col can provide at most batch_size-1 negatives
     if neg_ratio is None or neg_ratio <= 0 or neg_ratio > max_neg:
         neg_ratio = max_neg
 
@@ -140,8 +140,8 @@ def inbatch_negative_sampling(scores, neg_ratio=None, hard_negative=False, gener
             topk = torch.topk(row_scores, k=neg_ratio).indices
             neg_indices[i] = topk
         else:
-            candidates = torch.cat([index_range[:i], index_range[i + 1 :]]) # all except i
-            perm = torch.randperm(candidates.size(0), device=device, generator=generator) # random negative sampling
+            candidates = torch.cat([index_range[:i], index_range[i + 1:]])  # all except i
+            perm = torch.randperm(candidates.size(0), device=device, generator=generator)  # random negative sampling
             neg_indices[i] = candidates[perm[:neg_ratio]]
 
     return neg_indices
@@ -157,8 +157,7 @@ def gather_inbatch_logits(scores, neg_indices):
     # positive: scores[i][i]
     positive_logits = torch.diagonal(scores).reshape(-1, 1)  # (B,1)
     # negatives: scores[i][neg_indices[i, j]]
-    negative_logits = scores[torch.arange(scores.size(0)).unsqueeze(1),
-                             neg_indices]  # (B,K)
+    negative_logits = scores[torch.arange(scores.size(0)).unsqueeze(1), neg_indices]  # (B,K)
     return torch.cat([positive_logits, negative_logits], dim=1)