From da92fcfef92a748227c2d4b5c98a7ede999793a0 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 6 Jul 2026 18:47:37 +0000
Subject: [PATCH 1/2] [pre-commit.ci] pre-commit suggestions
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

updates:
- [github.com/psf/black-pre-commit-mirror: 26.3.1 → 26.5.1](https://github.com/psf/black-pre-commit-mirror/compare/26.3.1...26.5.1)
- [github.com/psf/black-pre-commit-mirror: 26.3.1 → 26.5.1](https://github.com/psf/black-pre-commit-mirror/compare/26.3.1...26.5.1)
- [github.com/pycqa/isort: 8.0.1 → 9.0.0a3](https://github.com/pycqa/isort/compare/8.0.1...9.0.0a3)
- [github.com/PyCQA/docformatter: v1.7.7 → v1.7.8](https://github.com/PyCQA/docformatter/compare/v1.7.7...v1.7.8)
---
 .pre-commit-config.yaml | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 1d974793..0ab841bb 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -12,23 +12,23 @@ repos:
       - id: check-toml
 
   - repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 26.3.1
+    rev: 26.5.1
     hooks:
       - id: black
 
   - repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 26.3.1
+    rev: 26.5.1
     hooks:
       - id: black-jupyter
 
   - repo: https://github.com/pycqa/isort
-    rev: 8.0.1
+    rev: 9.0.0a3
     hooks:
       - id: isort
         args: ["--profile", "black"]
 
   - repo: https://github.com/PyCQA/docformatter
-    rev: v1.7.7
+    rev: v1.7.8
     hooks:
       - id: docformatter
         language_version: python3.12

From 7be6313545b0c317ae9f46b13d022927307d3742 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 6 Jul 2026 18:47:47 +0000
Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 examples/tgn_attention_recommendation.py |  3 +-
 relbench/base/database.py                |  7 -----
 relbench/base/dataset.py                 |  3 --
 relbench/base/table.py                   |  4 ---
 relbench/base/task_autocomplete.py       | 38 ++++++++++++------------
 relbench/base/task_base.py               |  3 --
 relbench/base/task_recommendation.py     |  1 -
 relbench/datasets/__init__.py            |  3 --
 relbench/datasets/amazon.py              |  1 -
 relbench/datasets/arxiv.py               |  1 -
 relbench/datasets/dbinfer.py             |  1 -
 relbench/datasets/salt.py                |  1 -
 relbench/metrics.py                      | 12 ++++----
 relbench/modeling/graph.py               |  2 --
 relbench/modeling/utils.py               |  1 -
 relbench/tasks/__init__.py               |  3 --
 relbench/tasks/mimic.py                  |  3 +-
 relbench/tasks/tgb.py                    |  1 -
 18 files changed, 28 insertions(+), 60 deletions(-)

diff --git a/examples/tgn_attention_recommendation.py b/examples/tgn_attention_recommendation.py
index 4fc1bcec..f187116a 100644
--- a/examples/tgn_attention_recommendation.py
+++ b/examples/tgn_attention_recommendation.py
@@ -1,5 +1,4 @@
-"""
-Temporal recommendation baseline: TGN + attention (TransformerConv).
+"""Temporal recommendation baseline: TGN + attention (TransformerConv).
 
 This script ports the core TGB2 "TGN + GraphAttentionEmbedding" training recipe
 to RelBench's RecommendationTask interface.
diff --git a/relbench/base/database.py b/relbench/base/database.py
index 4bc1192b..3807ae55 100644
--- a/relbench/base/database.py
+++ b/relbench/base/database.py
@@ -15,7 +15,6 @@ class Database:
 
     def __init__(self, table_dict: Dict[str, Table]) -> None:
         r"""Creates a database from a dictionary of tables."""
-
         self.table_dict = table_dict
 
     def __repr__(self) -> str:
@@ -26,14 +25,12 @@ def save(self, path: Union[str, os.PathLike]) -> None:
 
         Simply saves each table individually with the table name as base name of file.
         """
-
         for name, table in self.table_dict.items():
             table.save(f"{path}/{name}.parquet")
 
     @classmethod
     def load(cls, path: Union[str, os.PathLike]) -> Self:
         r"""Load a database from a directory of tables in parquet files."""
-
         table_dict = {}
         for table_path in Path(path).glob("*.parquet"):
             table = Table.load(table_path)
@@ -45,7 +42,6 @@ def load(cls, path: Union[str, os.PathLike]) -> Self:
     @lru_cache(maxsize=None)
     def min_timestamp(self) -> pd.Timestamp:
         r"""Return the earliest timestamp in the database."""
-
         return min(
             table.min_timestamp
             for table in self.table_dict.values()
@@ -56,7 +52,6 @@ def min_timestamp(self) -> pd.Timestamp:
     @lru_cache(maxsize=None)
     def max_timestamp(self) -> pd.Timestamp:
         r"""Return the latest timestamp in the database."""
-
         return max(
             table.max_timestamp
             for table in self.table_dict.values()
@@ -65,7 +60,6 @@ def max_timestamp(self) -> pd.Timestamp:
 
     def upto(self, timestamp: pd.Timestamp) -> Self:
         r"""Return a database with all rows upto timestamp."""
-
         return Database(
             table_dict={
                 name: table.upto(timestamp) for name, table in self.table_dict.items()
@@ -74,7 +68,6 @@ def upto(self, timestamp: pd.Timestamp) -> Self:
 
     def from_(self, timestamp: pd.Timestamp) -> Self:
         r"""Return a database with all rows from timestamp."""
-
         return Database(
             table_dict={
                 name: table.from_(timestamp) for name, table in self.table_dict.items()
diff --git a/relbench/base/dataset.py b/relbench/base/dataset.py
index dc929d3c..c169eedc 100644
--- a/relbench/base/dataset.py
+++ b/relbench/base/dataset.py
@@ -42,7 +42,6 @@ def __init__(
                 the cached file. If None, we will not use cached file and re-process
                 everything from scratch without saving the cache.
         """
-
         self.cache_dir = cache_dir
 
         self.target_col = None
@@ -91,7 +90,6 @@ def get_db(self, upto_test_timestamp=True) -> Database:
 
         `upto_test_timestamp` is True by default to prevent test leakage.
         """
-
         db_path = f"{self.cache_dir}/db"
         if self.cache_dir and Path(db_path).exists() and any(Path(db_path).iterdir()):
             print(f"Loading Database object from {db_path}...")
@@ -141,7 +139,6 @@ def get_modified_db(self, db) -> Database:
         Returns:
             Database: The modified database object.
         """
-
         # Remove the target column from the source entity table
         # Ensure the entity table has a primary key if not add one
         # Remove any other columns marked for removal
diff --git a/relbench/base/table.py b/relbench/base/table.py
index f62886f0..27c39f79 100644
--- a/relbench/base/table.py
+++ b/relbench/base/table.py
@@ -109,7 +109,6 @@ def upto(self, timestamp: pd.Timestamp) -> Self:
 
         Table without time_col are returned as is.
         """
-
         if self.time_col is None:
             return self
 
@@ -125,7 +124,6 @@ def from_(self, timestamp: pd.Timestamp) -> Self:
 
         Table without time_col are returned as is.
         """
-
         if self.time_col is None:
             return self
 
@@ -140,7 +138,6 @@ def from_(self, timestamp: pd.Timestamp) -> Self:
     @lru_cache(maxsize=None)
     def min_timestamp(self) -> pd.Timestamp:
         r"""Return the earliest time in the table."""
-
         if self.time_col is None:
             raise ValueError("Table has no time column.")
 
@@ -150,7 +147,6 @@ def min_timestamp(self) -> pd.Timestamp:
     @lru_cache(maxsize=None)
     def max_timestamp(self) -> pd.Timestamp:
         r"""Return the latest time in the table."""
-
         if self.time_col is None:
             raise ValueError("Table has no time column.")
 
diff --git a/relbench/base/task_autocomplete.py b/relbench/base/task_autocomplete.py
index 59497921..f262680d 100644
--- a/relbench/base/task_autocomplete.py
+++ b/relbench/base/task_autocomplete.py
@@ -28,21 +28,22 @@
 
 
 class AutoCompleteTask(EntityTask):
-    r"""Auto complete column task on a dataset. Predict all values in the target column.
-
-    The task is constructed by specifying the entity table, entity column, time column, and target column.
-    The target column is removed from the entity table and saved to `db.table_dict[entity_table].removed_cols`,
-    which is used to construct the table for the predict column task.
-
-    The entity table needs to have a time column by which the data is split into training and validation set.
-
-    Args:
-        dataset: The dataset object.
-        task_type: The type of the task.
-        entity_table: The name of the entity table.
-        target_col: The name of the target column to be predicted.
-        cache_dir: The directory to cache the task tables.
-        remove_columns: List of columns, table pairs to remove from the graph.
+    r"""Auto complete column task on a dataset.
+
+    Predict all values in the target column.
+        The task is constructed by specifying the entity table, entity column, time column, and target column.
+        The target column is removed from the entity table and saved to `db.table_dict[entity_table].removed_cols`,
+        which is used to construct the table for the predict column task.
+
+        The entity table needs to have a time column by which the data is split into training and validation set.
+
+        Args:
+            dataset: The dataset object.
+            task_type: The type of the task.
+            entity_table: The name of the entity table.
+            target_col: The name of the target column to be predicted.
+            cache_dir: The directory to cache the task tables.
+            remove_columns: List of columns, table pairs to remove from the graph.
     """
 
     timedelta = pd.Timedelta(seconds=1)
@@ -111,11 +112,10 @@ def filter_dangling_entities(self, table: Table) -> Table:
     def _get_table(self, split: str) -> Table:
         r"""Helper function to get a table for a split.
 
-        This function overrides the `_get_table` method in `EntityTask`.
-        Because we predict all values in the target column, we only look at the min and max timestamp
-        for each split and take all rows in the table between them.
+        This function overrides the `_get_table` method in `EntityTask`. Because we
+        predict all values in the target column, we only look at the min and max
+        timestamp for each split and take all rows in the table between them.
         """
-
         db = self.dataset.get_db(upto_test_timestamp=split != "test")
 
         if split == "train":
diff --git a/relbench/base/task_base.py b/relbench/base/task_base.py
index 4e880a53..0e79e402 100644
--- a/relbench/base/task_base.py
+++ b/relbench/base/task_base.py
@@ -96,12 +96,10 @@ def make_table(
         To be implemented by subclass. The table rows need not be ordered
         deterministically.
         """
-
         raise NotImplementedError
 
     def _get_table(self, split: str) -> Table:
         r"""Helper function to get a table for a split."""
-
         db = self.dataset.get_db(upto_test_timestamp=split != "test")
 
         if split == "train":
@@ -169,7 +167,6 @@ def get_table(self, split, mask_input_cols=None):
 
         The table is cached in memory.
         """
-
         if mask_input_cols is None:
             mask_input_cols = split == "test"
 
diff --git a/relbench/base/task_recommendation.py b/relbench/base/task_recommendation.py
index d9aa062b..c1041bd9 100644
--- a/relbench/base/task_recommendation.py
+++ b/relbench/base/task_recommendation.py
@@ -113,7 +113,6 @@ def stats(self) -> Dict[str, Dict[str, int]]:
         r"""Get train / val / test table statistics for each timestamp and the whole
         table, including number of unique source entities, number of unique destination
         entities, number of destination entities and number of rows."""
-
         res = {}
         for split in ["train", "val", "test"]:
             split_stats = {}
diff --git a/relbench/datasets/__init__.py b/relbench/datasets/__init__.py
index 91dcf8f5..9374317e 100644
--- a/relbench/datasets/__init__.py
+++ b/relbench/datasets/__init__.py
@@ -54,7 +54,6 @@ def register_dataset(
     `cache_dir` is added to kwargs by default. If you want to override it, you
     can pass `cache_dir` as a keyword argument in `kwargs`.
     """
-
     cache_dir = f"{get_relbench_cache_dir()}/{name}"
     kwargs = {"cache_dir": cache_dir, **kwargs}
     dataset_registry[name] = (cls, args, kwargs)
@@ -71,7 +70,6 @@ def download_dataset(name: str) -> None:
     The downloaded database will be automatically picked up by the dataset object, when
     `dataset.get_db()` is called.
     """
-
     if name == "rel-mimic":
         from relbench.datasets.mimic import verify_mimic_access
 
@@ -104,7 +102,6 @@ def get_dataset(name: str, download=True) -> Dataset:
     raw files, the cache will be used. `download=True` will verify that the
     cached database matches the RelBench version even in this case.
     """
-
     if download:
         download_dataset(name)
 
diff --git a/relbench/datasets/amazon.py b/relbench/datasets/amazon.py
index 459e9459..d7881811 100644
--- a/relbench/datasets/amazon.py
+++ b/relbench/datasets/amazon.py
@@ -32,7 +32,6 @@ def __init__(
 
     def make_db(self) -> Database:
         r"""Process the raw files into a database."""
-
         ### product table ###
 
         url_key = self._category_to_url_key[self.category]
diff --git a/relbench/datasets/arxiv.py b/relbench/datasets/arxiv.py
index feb25400..1ae13e90 100644
--- a/relbench/datasets/arxiv.py
+++ b/relbench/datasets/arxiv.py
@@ -13,7 +13,6 @@ class ArxivDataset(Dataset):
 
     def make_db(self) -> Database:
         r"""Process the raw files into a database."""
-
         url = (
             "https://www.dropbox.com/scl/fi/tjj6r1fqikt4j0rz4qomu/db.zip?rlkey=1ykfkp8pj3hu6n4utz8g9dkx2&st"
             "=azmm56dc&dl=1"
diff --git a/relbench/datasets/dbinfer.py b/relbench/datasets/dbinfer.py
index c757b59b..3fd7eebc 100644
--- a/relbench/datasets/dbinfer.py
+++ b/relbench/datasets/dbinfer.py
@@ -80,7 +80,6 @@ def make_db(self) -> Database:
 
     def get_db(self, upto_test_timestamp: bool = True) -> Database:
         """DBInfer datasets are static, so never trim by timestamp."""
-
         return super().get_db(upto_test_timestamp=False)
 
 
diff --git a/relbench/datasets/salt.py b/relbench/datasets/salt.py
index 2416e131..006bc9c3 100644
--- a/relbench/datasets/salt.py
+++ b/relbench/datasets/salt.py
@@ -29,7 +29,6 @@ def __init__(self, **kwargs):
 
     def make_db(self) -> Database:
         r"""Process the raw files into a database."""
-
         try:
             _split = "train+test"
             salesdocument = load_dataset(
diff --git a/relbench/metrics.py b/relbench/metrics.py
index e976c098..104beacc 100644
--- a/relbench/metrics.py
+++ b/relbench/metrics.py
@@ -149,13 +149,13 @@ def multiclass_f1(true: NDArray[np.int_], pred: NDArray[np.int_]) -> float:
     if pred.ndim > 1:
         pred = pred.argmax(axis=1)
     return skm.f1_score(true, pred, average="micro")
-
-
 ####### Link prediction metrics
-"""All link prediction metrics take two arguments
-    - pred_isin: Numpy boolean array of size (num_src_nodes, eval_k)
-    - dst_count: Numpy integer array of size (num_src_nodes, ), storing
-        the number of destination nodes attached to each source node.
+\
+"""All link prediction metrics take two arguments.
+
+- pred_isin: Numpy boolean array of size (num_src_nodes, eval_k)
+- dst_count: Numpy integer array of size (num_src_nodes, ), storing
+    the number of destination nodes attached to each source node.
 """
 
 
diff --git a/relbench/modeling/graph.py b/relbench/modeling/graph.py
index 98aa9dff..1e696ef2 100644
--- a/relbench/modeling/graph.py
+++ b/relbench/modeling/graph.py
@@ -149,7 +149,6 @@ def get_node_train_table_input(
     task: EntityTask,
 ) -> NodeTrainTableInput:
     r"""Get the training table input for node prediction."""
-
     nodes = torch.from_numpy(table.df[task.entity_col].astype(int).values)
 
     time: Optional[Tensor] = None
@@ -201,7 +200,6 @@ def get_link_train_table_input(
     task: RecommendationTask,
 ) -> LinkTrainTableInput:
     r"""Get the training table input for link prediction."""
-
     src_node_idx: Tensor = torch.from_numpy(
         table.df[task.src_entity_col].astype(int).values
     )
diff --git a/relbench/modeling/utils.py b/relbench/modeling/utils.py
index f21f9297..1575d7a5 100644
--- a/relbench/modeling/utils.py
+++ b/relbench/modeling/utils.py
@@ -47,7 +47,6 @@ def get_stype_proposal(db: Database) -> Dict[str, Dict[str, stype]]:
         Dict[str, Dict[str, Any]]: A dictionary mapping table name into
             :obj:`col_to_stype` (mapping column names into inferred stypes).
     """
-
     inferred_col_to_stype_dict = {}
     for table_name, table in db.table_dict.items():
         df = table.df
diff --git a/relbench/tasks/__init__.py b/relbench/tasks/__init__.py
index cb643d95..e8528422 100644
--- a/relbench/tasks/__init__.py
+++ b/relbench/tasks/__init__.py
@@ -58,7 +58,6 @@ def register_task(
     `cache_dir` is added to kwargs by default. If you want to override it, you
     can pass `cache_dir` as a keyword argument in `kwargs`.
     """
-
     cache_dir = f"{get_relbench_cache_dir()}/{dataset_name}/tasks/{task_name}"
     kwargs = {"cache_dir": cache_dir, **kwargs}
     task_registry[dataset_name][task_name] = (cls, args, kwargs)
@@ -75,7 +74,6 @@ def download_task(dataset_name: str, task_name: str) -> None:
     The downloaded task tables will be automatically picked up by the task object, when
     `task.get_table(split)` is called.
     """
-
     DOWNLOAD_REGISTRY.fetch(
         f"{dataset_name}/tasks/{task_name}.zip",
         processor=pooch.Unzip(extract_dir="."),
@@ -104,7 +102,6 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask:
     scratch, the cache will be used. `download=True` will verify that the
     cached task tables matches the RelBench version even in this case.
     """
-
     if download:
         download_task(dataset_name, task_name)
     dataset = get_dataset(dataset_name, download=download)
diff --git a/relbench/tasks/mimic.py b/relbench/tasks/mimic.py
index 86b57dc1..5a8504cb 100644
--- a/relbench/tasks/mimic.py
+++ b/relbench/tasks/mimic.py
@@ -9,7 +9,8 @@
 
 
 class ICULengthOfStayTask(EntityTask):
-    r"""Binary classification: Predict if ICU length of stay is ≥ 3 days during the first ICU stay."""
+    r"""Binary classification: Predict if ICU length of stay is ≥ 3 days during the
+    first ICU stay."""
 
     task_type = TaskType.BINARY_CLASSIFICATION
     entity_col = "subject_id"
diff --git a/relbench/tasks/tgb.py b/relbench/tasks/tgb.py
index 240757d7..4ea5668a 100644
--- a/relbench/tasks/tgb.py
+++ b/relbench/tasks/tgb.py
@@ -32,7 +32,6 @@ def _tgb_eval_hits_and_mrr(
     k_value: int,
 ) -> dict[str, float]:
     r"""Compute one-vs-many Hits@k and MRR with tie-aware ranking."""
-
     y_pred_pos = np.asarray(y_pred_pos).reshape(-1, 1)
     y_pred_neg = np.asarray(y_pred_neg)
     if y_pred_neg.ndim != 2 or y_pred_neg.shape[0] != y_pred_pos.shape[0]: