snap-stanford · rishabh-ranjan · Jul 1, 2025 · Jul 1, 2025 · Jul 1, 2025 · Jul 1, 2025
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -28,7 +28,7 @@ repos:
         args: ["--profile", "black"]
 
   - repo: https://github.com/PyCQA/docformatter
-    rev: v1.7.5
+    rev: v1.7.7
     hooks:
       - id: docformatter
         additional_dependencies: [tomli]

diff --git a/relbench/metrics.py b/relbench/metrics.py
@@ -182,6 +182,7 @@ def link_prediction_map(
     maps = (precision_mat * pred_isin).sum(axis=1) / clipped_dst_count
     return maps.mean()
 
+
 def link_prediction_ndcg(
     pred_isin: NDArray[np.int_],
     dst_count: NDArray[np.int_],
@@ -190,24 +191,29 @@ def link_prediction_ndcg(
     eval_k = pred_isin.shape[1]
 
     # Compute the discounted multiplier (1 / log2(i + 2) for i = 0, ..., k-1)
-    discounted_multiplier = np.concatenate((
-        np.zeros(1),
-        1 / np.log2(np.arange(1, eval_k + 1) + 1)
-    ))
+    discounted_multiplier = np.concatenate(
+        (np.zeros(1), 1 / np.log2(np.arange(1, eval_k + 1) + 1))
+    )
 
     # Compute Discounted Cumulative Gain (DCG)
-    discounted_cumulative_gain = (pred_isin * discounted_multiplier[1:eval_k + 1]).sum(axis=1)
+    discounted_cumulative_gain = (
+        pred_isin * discounted_multiplier[1 : eval_k + 1]
+    ).sum(axis=1)
 
     # Clip dst_count to the range [0, eval_k]
     clipped_dst_count = np.clip(dst_count, 0, eval_k)
 
     # Compute Ideal Discounted Cumulative Gain (IDCG)
     ideal_discounted_multiplier_cumsum = np.cumsum(discounted_multiplier)
-    ideal_discounted_cumulative_gain = ideal_discounted_multiplier_cumsum[clipped_dst_count]
+    ideal_discounted_cumulative_gain = ideal_discounted_multiplier_cumsum[
+        clipped_dst_count
+    ]
 
     # Avoid division by zero
-    ideal_discounted_cumulative_gain = np.clip(ideal_discounted_cumulative_gain, 1e-10, None)
+    ideal_discounted_cumulative_gain = np.clip(
+        ideal_discounted_cumulative_gain, 1e-10, None
+    )
 
     # Compute NDCG
     ndcg_scores = discounted_cumulative_gain / ideal_discounted_cumulative_gain
-    return ndcg_scores.mean()
+    return ndcg_scores.mean()
diff --git a/test/test_metrics.py b/test/test_metrics.py
@@ -2,9 +2,9 @@
 
 from relbench.metrics import (
     link_prediction_map,
+    link_prediction_ndcg,
     link_prediction_precision,
     link_prediction_recall,
-    link_prediction_ndcg,
 )
 
 

diff --git a/tutorials/train_model.ipynb b/tutorials/train_model.ipynb
@@ -621,16 +621,14 @@
     "\n",
     "\n",
     "class GloveTextEmbedding:\n",
-    "    def __init__(self, device: Optional[torch.device\n",
-    "                                       ] = None):\n",
+    "    def __init__(self, device: Optional[torch.device] = None):\n",
     "        self.model = SentenceTransformer(\n",
     "            \"sentence-transformers/average_word_embeddings_glove.6B.300d\",\n",
     "            device=device,\n",
     "        )\n",
     "\n",
     "    def __call__(self, sentences: List[str]) -> Tensor:\n",
-    "        return torch.from_numpy(self.model.encode(sentences))\n",
-    "\n"
+    "        return torch.from_numpy(self.model.encode(sentences))"
    ]
   },
   {