diff --git a/examples/relational/relbench/README.md b/examples/relational/relbench/README.md index 953ef92df0..3e360880d6 100644 --- a/examples/relational/relbench/README.md +++ b/examples/relational/relbench/README.md @@ -12,41 +12,41 @@ python main.py --dataset=rel-f1 --task=driver-dnf - `--num_estimators`: Number of estimators (default: `1`). - `--context_size`: Maximum number of context rows (default: `10_000`). -- `--context_mode`: Context selection strategy, either `latest` or `random` (default: `latest`). +- `--context_mode`: Context selection strategy, either `last` or `rnd` (default: `last`). - `--num_neighbors`: The number of neighbors to sample for each hop (default: `[16, 16]`). -- `--text`: Include text embeddings from the task entity table. -- `--text_dim`: Final text embedding dimension after PCA (default: `64`). +- `--num_lags`: The number of lag targets to use as additional input (default: `20`). +- `--text_dim`: Text embedding dimension after PCA for the entity table (default: `0`). ## Results -Optimal hyperparameters found via the [RelArena-α](../../../benchmark/relational/relarena) evaluation protocol, using 8 estimators: - -| Dataset | Task | Val AUROC | Test AUROC | Size | Mode | Neighbors | Text Model | Text Dim | -| :----------- | :---------------- | --------: | ---------: | ----: | :------- | ---------: | :---------- | -------: | -| `rel-amazon` | `user-churn` | 0.6935 | 0.6947 | `20K` | `latest` | `[8, 8]` | | | -| `rel-amazon` | `item-churn` | 0.8215 | 0.8254 | `20K` | `latest` | `[32, 32]` | | | -| `rel-avito` | `user-visits` | 0.6956 | 0.6711 | `20K` | `latest` | `[32, 32]` | | | -| `rel-avito` | `user-clicks` | 0.6369 | 0.6824 | `20K` | `latest` | `[32, 32]` | | | -| `rel-event` | `user-repeat` | 0.7330 | 0.8090 | `1K` | `random` | `[4, 4]` | | | -| `rel-event` | `user-ignore` | 0.8467 | 0.8380 | `2K` | `latest` | `[16, 16]` | | | -| `rel-f1` | `driver-dnf` | 0.7937 | 0.8418 | `2K` | `latest` | `[4, 4]` | | | -| `rel-f1` | `driver-top3` | 0.8816 | 0.9077 | `1K` | `random` | `[4, 4]` | | | -| `rel-hm` | `user-churn` | 0.7036 | 0.6992 | `20K` | `latest` | `[1, 1]` | | | -| `rel-stack` | `user-engagement` | 0.9001 | 0.9019 | `20K` | `random` | `[8, 8]` | `mini`[^1] | `32` | -| `rel-stack` | `user-badge` | 0.8932 | 0.8809 | `20K` | `random` | `[8, 8]` | `mpnet`[^2] | `64` | -| `rel-trial` | `study-outcome` | 0.6636 | 0.7393 | `20K` | `latest` | `[8, 8]` | `mini`[^1] | `32` | - -| Dataset | Task | Val nMAE | Test nMAE | Size | Mode | Neighbors | Text Model | Text Dim | -| :----------- | :---------------- | -------: | --------: | ----: | :------- | ---------: | :--------- | -------: | -| `rel-amazon` | `user-ltv` | 21.60 | 25.01 | `20K` | `random` | `[64, 64]` | | | -| `rel-amazon` | `item-ltv` | 7.38 | 7.56 | `20K` | `random` | `[]` | | | -| `rel-avito` | `ad-ctr` | 32.40 | 32.95 | `2K` | `latest` | `[64, 64]` | | | -| `rel-event` | `user-attendance` | 31.25 | 31.70 | `20K` | `random` | `[64, 64]` | | | -| `rel-f1` | `driver-position` | 44.63 | 43.55 | `20K` | `latest` | `[8, 8]` | | | -| `rel-hm` | `item-sales` | 9.81 | 8.00 | `20K` | `random` | `[]` | | | -| `rel-stack` | `post-votes` | 11.36 | 12.49 | `20K` | `random` | `[]` | | | -| `rel-trial` | `study-adverse` | 10.21 | 9.42 | `20K` | `latest` | `[]` | `mini`[^1] | `32` | -| `rel-trial` | `site-success` | 87.90 | 84.31 | `10K` | `latest` | `[8, 8]` | `mini`[^1] | `32` | +Optimal hyperparameters found via the [RelArena-α](../../../benchmark/relational/relarena) evaluation protocol, using 8 estimators and 20 lag targets: + +| Dataset | Task | Val AUROC | Test AUROC | Size | Mode | Neighbors | Text Dim | Text Model | +| :----------- | :---------------- | --------: | ---------: | ----: | :----- | ---------: | -------: | :---------- | +| `rel-amazon` | `user-churn` | 69.35 | 69.47 | `20K` | `last` | `[8, 8]` | `0` | | +| `rel-amazon` | `item-churn` | 82.15 | 82.54 | `20K` | `last` | `[32, 32]` | `0` | | +| `rel-avito` | `user-visits` | 69.56 | 67.11 | `20K` | `last` | `[32, 32]` | `0` | | +| `rel-avito` | `user-clicks` | 63.69 | 68.24 | `20K` | `last` | `[32, 32]` | `0` | | +| `rel-event` | `user-repeat` | 73.30 | 80.90 | `1K` | `rnd` | `[4, 4]` | `0` | | +| `rel-event` | `user-ignore` | 84.67 | 83.80 | `2K` | `last` | `[16, 16]` | `0` | | +| `rel-f1` | `driver-dnf` | 79.37 | 84.18 | `2K` | `last` | `[4, 4]` | `0` | | +| `rel-f1` | `driver-top3` | 88.16 | 90.77 | `1K` | `rnd` | `[4, 4]` | `0` | | +| `rel-hm` | `user-churn` | 70.36 | 69.92 | `20K` | `last` | `[1, 1]` | `0` | | +| `rel-stack` | `user-engagement` | 90.01 | 90.19 | `20K` | `rnd` | `[8, 8]` | `32` | `mini`[^1] | +| `rel-stack` | `user-badge` | 89.32 | 88.09 | `20K` | `rnd` | `[8, 8]` | `64` | `mpnet`[^2] | +| `rel-trial` | `study-outcome` | 66.36 | 73.93 | `20K` | `last` | `[8, 8]` | `32` | `mini`[^1] | + +| Dataset | Task | Val nMAE | Test nMAE | Size | Mode | Neighbors | Text Dim | Text Model | +| :----------- | :---------------- | -------: | --------: | ----: | :----- | ---------: | -------: | :--------- | +| `rel-amazon` | `user-ltv` | 21.60 | 25.01 | `20K` | `rnd` | `[64, 64]` | `0` | | +| `rel-amazon` | `item-ltv` | 7.38 | 7.56 | `20K` | `rnd` | `[]` | `0` | | +| `rel-avito` | `ad-ctr` | 32.40 | 32.95 | `2K` | `last` | `[64, 64]` | `0` | | +| `rel-event` | `user-attendance` | 31.25 | 31.70 | `20K` | `rnd` | `[64, 64]` | `0` | | +| `rel-f1` | `driver-position` | 44.63 | 43.55 | `20K` | `last` | `[8, 8]` | `0` | | +| `rel-hm` | `item-sales` | 9.81 | 8.00 | `20K` | `rnd` | `[]` | `0` | | +| `rel-stack` | `post-votes` | 11.36 | 12.49 | `20K` | `rnd` | `[]` | `0` | | +| `rel-trial` | `study-adverse` | 10.21 | 9.42 | `20K` | `last` | `[]` | `32` | `mini`[^1] | +| `rel-trial` | `site-success` | 87.90 | 84.31 | `10K` | `last` | `[8, 8]` | `32` | `mini`[^1] | [^1]: `sentence-transformers/all-MiniLM-L6-v2` diff --git a/examples/relational/relbench/lag_target.py b/examples/relational/relbench/lag_target.py new file mode 100644 index 0000000000..86052451d0 --- /dev/null +++ b/examples/relational/relbench/lag_target.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pandas as pd +import relbench + + +def add_lag_target( + df: pd.DataFrame, + task: relbench.base.EntityTask, + num_lags: int, +) -> pd.DataFrame: + r"""Add lag target features to the task table.""" + df = df.copy() + right = df.sort_values(task.time_col) + left = pd.DataFrame( + { + task.entity_col: df[task.entity_col], + "__lookup_time__": df[task.time_col], + "__row__": range(len(df)), + } + ) + for lag in range(1, num_lags + 1): + column = f"{task.target_col}_lag_{lag}" + if task.task_type == relbench.base.TaskType.REGRESSION: + df[column] = float("NaN") + else: + df[column] = pd.Series(pd.NA, index=df.index, dtype="object") + + merged = pd.merge_asof( + left.sort_values("__lookup_time__"), + right, + by=task.entity_col, + left_on="__lookup_time__", + right_on=task.time_col, + direction="backward", + allow_exact_matches=False, + ) + mask = merged[task.time_col].notna() + rows = merged.loc[mask, "__row__"].to_numpy() + df.loc[rows, column] = merged.loc[mask, task.target_col].to_numpy() + + # Find the next lag strictly before the previous observation. + left = merged.loc[mask, [task.entity_col, "__row__", task.time_col]] + left = left.rename(columns={task.time_col: "__lookup_time__"}) + + return df diff --git a/examples/relational/relbench/main.py b/examples/relational/relbench/main.py index 722e32240f..d7a8a6ca65 100644 --- a/examples/relational/relbench/main.py +++ b/examples/relational/relbench/main.py @@ -9,114 +9,24 @@ import relbench import torch import torchmetrics +from lag_target import add_lag_target +from text import embed_text from tqdm import tqdm import sdm -import sdm.processing as sp - - -def embed_entity_text( - entity: sdm.TableTensor, - task_ids: pd.Series, - pkey: str, - *, - model_name: str, - device: torch.device, -) -> tuple[torch.Tensor, tuple[str, ...], torch.Tensor] | None: - """Embed task entities and return CPU embeddings.""" - if len(entity.text_columns) == 0: - return None - entity_ids = entity.select_columns(pkey).to_pandas()[pkey] - in_task = torch.from_numpy( - entity_ids.isin(pd.unique(task_ids)).to_numpy(dtype=bool) - ) - text = entity.select_stypes(sdm.Stype.text)[in_task] - if text.size(0) == 0: - return None - encoder = sp.SentenceTransformer( - model_name=model_name, - batch_size=32, - max_seq_length=2048, - ) - row_chunk = 1024 - parts = [] - columns: tuple[str, ...] = () - with torch.inference_mode(): - for start in range(0, text.size(0), row_chunk): - embedded = encoder(text[start : start + row_chunk].to(device)) - columns = embedded.numerical_columns - parts.append(embedded.numerical.cpu()) - del encoder, embedded, text - if device.type == "cuda": - torch.cuda.empty_cache() - return torch.cat(parts, dim=0), columns, in_task - - -def attach_text( - entity: sdm.TableTensor, - embeddings: torch.Tensor, - columns: tuple[str, ...], - in_task: torch.Tensor, - context_ids: pd.Series, - pkey: str, - text_dim: int, -) -> sdm.TableTensor: - """Project embeddings with a PCA fit on the context entities.""" - entity_ids = entity.select_columns(pkey).to_pandas()[pkey] - in_context = torch.from_numpy( - entity_ids.isin(pd.unique(context_ids)).to_numpy(dtype=bool) - ) - pca = sp.PCA(text_dim) - pca.fit( - sdm.TableTensor( - columns={"numerical": columns}, - numerical=embeddings[in_context[in_task]], - ) - ) - reduced = pca.transform( - sdm.TableTensor( - columns={"numerical": columns}, - numerical=embeddings, - ) - ) - values = torch.full( - (entity.size(0), reduced.numerical.size(-1)), - torch.nan, - dtype=reduced.numerical.dtype, - ) - values[in_task] = reduced.numerical - return cast( - sdm.TableTensor, - torch.cat( - [ - entity.drop_stypes(sdm.Stype.text), - sdm.TableTensor( - columns={ - "numerical": reduced.numerical_columns, - }, - numerical=values, - ), - ], - dim=-1, - ), - ) - parser = argparse.ArgumentParser() parser.add_argument("--dataset", type=str, required=True) parser.add_argument("--task", type=str, required=True) parser.add_argument("--context_size", type=int, default=10_000) -parser.add_argument( - "--context_mode", choices=("latest", "random"), default="latest" -) +parser.add_argument("--context_mode", choices=("last", "rnd"), default="last") parser.add_argument("--batch_size", type=int, default=10_000) parser.add_argument("--max_test_steps", type=int, default=None) parser.add_argument("--num_neighbors", type=int, nargs="*", default=[16, 16]) parser.add_argument("--num_estimators", type=int, default=1) parser.add_argument("--num_lags", type=int, default=20) parser.add_argument("--seed", type=int, default=0) -parser.add_argument("--text", action="store_true") -parser.add_argument("--text_dim", type=int, default=64) +parser.add_argument("--text_dim", type=int, default=0) args = parser.parse_args() torch.manual_seed(args.seed) @@ -126,25 +36,26 @@ def attach_text( dataset = relbench.load_dataset(args.dataset) task = dataset.load_task(args.task) db = task.get_db(upto_test_timestamp=False) -tables = { - name: sdm.TableTensor.from_pandas( - df=table.df, - stypes=sdm.infer_stypes( - table.df.head(10_000), - overrides={ - cast(str, table.pkey_col): "id", - **dict.fromkeys(table.fkey_col_to_pkey_table, "id"), - }, - text=( - "infer" if args.text and name == task.entity_table else "drop" - ), - unsupported="drop", - ), - ) - for name, table in db.table_dict.items() -} data = sdm.RelationalData( - tables=tables, + tables={ + name: sdm.TableTensor.from_pandas( + df=table.df, + stypes=sdm.infer_stypes( + table.df.head(10_000), + overrides={ + cast(str, table.pkey_col): "id", + **dict.fromkeys(table.fkey_col_to_pkey_table, "id"), + }, + text=( + "infer" + if args.text_dim > 0 and name == task.entity_table + else "drop" + ), + unsupported="drop", + ), + ) + for name, table in db.table_dict.items() + }, relationships=[ { "left_table": left_table, @@ -156,11 +67,6 @@ def attach_text( for left_column, right_table in table.fkey_col_to_pkey_table.items() ], ) -time_columns = { - name: table.time_col - for name, table in db.table_dict.items() - if table.time_col is not None -} # Collect Task Table ########################################################## if task.task_type not in { @@ -175,39 +81,7 @@ def attach_text( for split in ["train", "val", "test"] ] df = pd.concat(dfs, ignore_index=True) - -# Add lag target features to task table: -right = df.sort_values(task.time_col) -left = pd.DataFrame( - { - task.entity_col: df[task.entity_col], - "__lookup_time__": df[task.time_col], - "__row__": range(len(df)), - } -) -for lag in range(1, args.num_lags + 1): - column = f"{task.target_col}_lag_{lag}" - if task.task_type == relbench.base.TaskType.REGRESSION: - df[column] = float("NaN") - else: - df[column] = pd.Series(pd.NA, index=df.index, dtype="object") - - merged = pd.merge_asof( - left.sort_values("__lookup_time__"), - right, - by=task.entity_col, - left_on="__lookup_time__", - right_on=task.time_col, - direction="backward", - allow_exact_matches=False, - ) - mask = merged[task.time_col].notna() - rows = merged.loc[mask, "__row__"].to_numpy() - df.loc[rows, column] = merged.loc[mask, task.target_col].to_numpy() - - # Find the next lag strictly before the previous observation. - left = merged.loc[mask, [task.entity_col, "__row__", task.time_col]] - left = left.rename(columns={task.time_col: "__lookup_time__"}) +df = add_lag_target(df, task, args.num_lags) if task.task_type == relbench.base.TaskType.REGRESSION: target_stype = "numerical" @@ -226,12 +100,26 @@ def attach_text( }, }, ) -context, query = task_table.split([len(dfs[0]) + len(dfs[1]), len(dfs[2])]) +if args.text_dim > 0: # Embed entity table text: + data.tables[task.entity_table] = embed_text( # type: ignore + table=data.tables[task.entity_table], + text_model="sentence-transformers/all-MiniLM-L6-v2", + text_dim=args.text_dim, + mask=torch.from_numpy( + db.table_dict[task.entity_table] + .df[db.table_dict[task.entity_table].pkey_col] + .isin(df[task.entity_col]) + .to_numpy() + ), + device=device, + ) + +context, query = task_table.split([len(dfs[0]) + len(dfs[1]), len(dfs[2])]) num_estimators = args.num_estimators if len(context) > args.context_size: # Sample different context per estimator: width = args.context_size * num_estimators - if args.context_mode == "latest" and len(context) > width: + if args.context_mode == "last" and len(context) > width: order = context.datetime.reshape(-1).argsort(descending=True) context = context[order[:width]] repeats = math.ceil(width / len(context)) @@ -240,37 +128,17 @@ def attach_text( if num_estimators > 1: context = context.unflatten(0, (num_estimators, args.context_size)) query = query.expand(num_estimators, *query.size()) - num_estimators = None + num_estimators = None # Context provides ensemble dimension. -# Execute Model ############################################################### -if args.text: - entity_pkey = cast(str, db.table_dict[task.entity_table].pkey_col) - entity = tables[task.entity_table] - encoded = embed_entity_text( - entity=entity, - task_ids=df[task.entity_col], - pkey=entity_pkey, - model_name="sentence-transformers/all-MiniLM-L6-v2", - device=device, - ) - if encoded is not None: - embeddings, columns, in_task = encoded - tables[task.entity_table] = attach_text( - entity=entity, - embeddings=embeddings, - columns=columns, - in_task=in_task, - context_ids=pd.concat( - [dfs[0][task.entity_col], dfs[1][task.entity_col]], - ignore_index=True, - ), - pkey=entity_pkey, - text_dim=args.text_dim, - ) -sampler = data.sampler(time_columns=time_columns) -model = sdm.models.KumoRelational(device=device) -recipe = model.default_recipe() +# Create Subgraph Sampler ##################################################### +sampler = data.sampler( + time_columns={ + name: table.time_col + for name, table in db.table_dict.items() + if table.time_col is not None + }, +) kwargs: dict[str, Any] = { "task_link": { "task_column": task.entity_col, @@ -281,6 +149,9 @@ def attach_text( "task_time_column": task.time_col, } +# Execute Model ############################################################### +model = sdm.models.KumoRelational(device=device) + context, related_tables = sampler(context, **kwargs).to(device) with torch.amp.autocast(device.type, torch.float16, enabled=True): model.fit( @@ -288,7 +159,6 @@ def attach_text( y=context[task.target_col], related_tables=related_tables, num_estimators=num_estimators, - recipe=recipe, ) if task.task_type == relbench.base.TaskType.REGRESSION: @@ -319,8 +189,8 @@ def attach_text( metric.update(pred, target) if task.task_type == relbench.base.TaskType.REGRESSION: - print(f"nMAE: {metric.compute() / task.nmae_std:.4f}") + print(f"nMAE: {100 * metric.compute() / task.nmae_std:.2f}") elif task.task_type == relbench.base.TaskType.BINARY_CLASSIFICATION: - print(f"AUROC: {metric.compute():.4f}") + print(f"AUROC: {100 * metric.compute():.2f}") else: - print(f"ACC: {metric.compute():.4f}") + print(f"ACC: {100 * metric.compute():.2f}") diff --git a/examples/relational/relbench/text.py b/examples/relational/relbench/text.py new file mode 100644 index 0000000000..419c4b31df --- /dev/null +++ b/examples/relational/relbench/text.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import time +from typing import cast + +import torch +from torch import Tensor + +import sdm +import sdm.processing as sp + + +@torch.inference_mode() +def embed_text( + table: sdm.TableTensor, + text_model: str, + text_dim: int, + mask: Tensor | None = None, + device: torch.device | str | None = None, +) -> sdm.TableTensor: + r"""Embed all text columns of all masked rows.""" + if len(table.text_columns) == 0: + return table + + text = table.select_stypes("text") + if mask is not None: + text = text[mask] + text = cast(sdm.TableTensor, text.to(device)) + + processor = sp.Sequential( + sp.SentenceTransformer( + text_model, + batch_size=32, + max_seq_length=2048, + ), + sp.PCA(text_dim), + ) + + print(f"Embedding {text.size(-1):,} text columns... ", end="", flush=True) + t = time.perf_counter() + out = processor.fit_transform(text).to(table.device) + print(f"Done! [{time.perf_counter() - t:.2f}s]") + + if mask is not None: + emb = out.numerical.new_full( + (table.size(0), out.size(-1)), + fill_value=torch.nan, + ) + emb[mask] = out.numerical + out = sdm.TableTensor(columns=out.columns, numerical=emb) + + return cast( + sdm.TableTensor, + torch.cat([table.drop_stypes("text"), out], dim=-1), + )