[pre-commit.ci] pre-commit suggestions

.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

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.

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()

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

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.")
 

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":

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"
 

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 = {}

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)
 

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]

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"

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)
 
 

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(

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.
 """
 
 

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
     )

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

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)

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"

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]: