Skip to content

Commit b5c3641

Browse files
rusty1srusty1s
authored
Clean-up rel-bench example (#1326)
1 parent caa46a5 commit b5c3641

4 files changed

Lines changed: 189 additions & 216 deletions

File tree

‎examples/relational/relbench/README.md‎

Lines changed: 31 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -12,41 +12,41 @@ python main.py --dataset=rel-f1 --task=driver-dnf
1212

1313
- `--num_estimators`: Number of estimators (default: `1`).
1414
- `--context_size`: Maximum number of context rows (default: `10_000`).
15-
- `--context_mode`: Context selection strategy, either `latest` or `random` (default: `latest`).
15+
- `--context_mode`: Context selection strategy, either `last` or `rnd` (default: `last`).
1616
- `--num_neighbors`: The number of neighbors to sample for each hop (default: `[16, 16]`).
17-
- `--text`: Include text embeddings from the task entity table.
18-
- `--text_dim`: Final text embedding dimension after PCA (default: `64`).
17+
- `--num_lags`: The number of lag targets to use as additional input (default: `20`).
18+
- `--text_dim`: Text embedding dimension after PCA for the entity table (default: `0`).
1919

2020
## Results
2121

22-
Optimal hyperparameters found via the [RelArena-α](../../../benchmark/relational/relarena) evaluation protocol, using 8 estimators:
23-
24-
| Dataset | Task | Val AUROC | Test AUROC | Size | Mode | Neighbors | Text Model | Text Dim |
25-
| :----------- | :---------------- | --------: | ---------: | ----: | :------- | ---------: | :---------- | -------: |
26-
| `rel-amazon` | `user-churn` | 0.6935 | 0.6947 | `20K` | `latest` | `[8, 8]` | | |
27-
| `rel-amazon` | `item-churn` | 0.8215 | 0.8254 | `20K` | `latest` | `[32, 32]` | | |
28-
| `rel-avito` | `user-visits` | 0.6956 | 0.6711 | `20K` | `latest` | `[32, 32]` | | |
29-
| `rel-avito` | `user-clicks` | 0.6369 | 0.6824 | `20K` | `latest` | `[32, 32]` | | |
30-
| `rel-event` | `user-repeat` | 0.7330 | 0.8090 | `1K` | `random` | `[4, 4]` | | |
31-
| `rel-event` | `user-ignore` | 0.8467 | 0.8380 | `2K` | `latest` | `[16, 16]` | | |
32-
| `rel-f1` | `driver-dnf` | 0.7937 | 0.8418 | `2K` | `latest` | `[4, 4]` | | |
33-
| `rel-f1` | `driver-top3` | 0.8816 | 0.9077 | `1K` | `random` | `[4, 4]` | | |
34-
| `rel-hm` | `user-churn` | 0.7036 | 0.6992 | `20K` | `latest` | `[1, 1]` | | |
35-
| `rel-stack` | `user-engagement` | 0.9001 | 0.9019 | `20K` | `random` | `[8, 8]` | `mini`[^1] | `32` |
36-
| `rel-stack` | `user-badge` | 0.8932 | 0.8809 | `20K` | `random` | `[8, 8]` | `mpnet`[^2] | `64` |
37-
| `rel-trial` | `study-outcome` | 0.6636 | 0.7393 | `20K` | `latest` | `[8, 8]` | `mini`[^1] | `32` |
38-
39-
| Dataset | Task | Val nMAE | Test nMAE | Size | Mode | Neighbors | Text Model | Text Dim |
40-
| :----------- | :---------------- | -------: | --------: | ----: | :------- | ---------: | :--------- | -------: |
41-
| `rel-amazon` | `user-ltv` | 21.60 | 25.01 | `20K` | `random` | `[64, 64]` | | |
42-
| `rel-amazon` | `item-ltv` | 7.38 | 7.56 | `20K` | `random` | `[]` | | |
43-
| `rel-avito` | `ad-ctr` | 32.40 | 32.95 | `2K` | `latest` | `[64, 64]` | | |
44-
| `rel-event` | `user-attendance` | 31.25 | 31.70 | `20K` | `random` | `[64, 64]` | | |
45-
| `rel-f1` | `driver-position` | 44.63 | 43.55 | `20K` | `latest` | `[8, 8]` | | |
46-
| `rel-hm` | `item-sales` | 9.81 | 8.00 | `20K` | `random` | `[]` | | |
47-
| `rel-stack` | `post-votes` | 11.36 | 12.49 | `20K` | `random` | `[]` | | |
48-
| `rel-trial` | `study-adverse` | 10.21 | 9.42 | `20K` | `latest` | `[]` | `mini`[^1] | `32` |
49-
| `rel-trial` | `site-success` | 87.90 | 84.31 | `10K` | `latest` | `[8, 8]` | `mini`[^1] | `32` |
22+
Optimal hyperparameters found via the [RelArena-α](../../../benchmark/relational/relarena) evaluation protocol, using 8 estimators and 20 lag targets:
23+
24+
| Dataset | Task | Val AUROC | Test AUROC | Size | Mode | Neighbors | Text Dim | Text Model |
25+
| :----------- | :---------------- | --------: | ---------: | ----: | :----- | ---------: | -------: | :---------- |
26+
| `rel-amazon` | `user-churn` | 69.35 | 69.47 | `20K` | `last` | `[8, 8]` | `0` | |
27+
| `rel-amazon` | `item-churn` | 82.15 | 82.54 | `20K` | `last` | `[32, 32]` | `0` | |
28+
| `rel-avito` | `user-visits` | 69.56 | 67.11 | `20K` | `last` | `[32, 32]` | `0` | |
29+
| `rel-avito` | `user-clicks` | 63.69 | 68.24 | `20K` | `last` | `[32, 32]` | `0` | |
30+
| `rel-event` | `user-repeat` | 73.30 | 80.90 | `1K` | `rnd` | `[4, 4]` | `0` | |
31+
| `rel-event` | `user-ignore` | 84.67 | 83.80 | `2K` | `last` | `[16, 16]` | `0` | |
32+
| `rel-f1` | `driver-dnf` | 79.37 | 84.18 | `2K` | `last` | `[4, 4]` | `0` | |
33+
| `rel-f1` | `driver-top3` | 88.16 | 90.77 | `1K` | `rnd` | `[4, 4]` | `0` | |
34+
| `rel-hm` | `user-churn` | 70.36 | 69.92 | `20K` | `last` | `[1, 1]` | `0` | |
35+
| `rel-stack` | `user-engagement` | 90.01 | 90.19 | `20K` | `rnd` | `[8, 8]` | `32` | `mini`[^1] |
36+
| `rel-stack` | `user-badge` | 89.32 | 88.09 | `20K` | `rnd` | `[8, 8]` | `64` | `mpnet`[^2] |
37+
| `rel-trial` | `study-outcome` | 66.36 | 73.93 | `20K` | `last` | `[8, 8]` | `32` | `mini`[^1] |
38+
39+
| Dataset | Task | Val nMAE | Test nMAE | Size | Mode | Neighbors | Text Dim | Text Model |
40+
| :----------- | :---------------- | -------: | --------: | ----: | :----- | ---------: | -------: | :--------- |
41+
| `rel-amazon` | `user-ltv` | 21.60 | 25.01 | `20K` | `rnd` | `[64, 64]` | `0` | |
42+
| `rel-amazon` | `item-ltv` | 7.38 | 7.56 | `20K` | `rnd` | `[]` | `0` | |
43+
| `rel-avito` | `ad-ctr` | 32.40 | 32.95 | `2K` | `last` | `[64, 64]` | `0` | |
44+
| `rel-event` | `user-attendance` | 31.25 | 31.70 | `20K` | `rnd` | `[64, 64]` | `0` | |
45+
| `rel-f1` | `driver-position` | 44.63 | 43.55 | `20K` | `last` | `[8, 8]` | `0` | |
46+
| `rel-hm` | `item-sales` | 9.81 | 8.00 | `20K` | `rnd` | `[]` | `0` | |
47+
| `rel-stack` | `post-votes` | 11.36 | 12.49 | `20K` | `rnd` | `[]` | `0` | |
48+
| `rel-trial` | `study-adverse` | 10.21 | 9.42 | `20K` | `last` | `[]` | `32` | `mini`[^1] |
49+
| `rel-trial` | `site-success` | 87.90 | 84.31 | `10K` | `last` | `[8, 8]` | `32` | `mini`[^1] |
5050

5151
[^1]: `sentence-transformers/all-MiniLM-L6-v2`
5252

‎examples/relational/relbench/lag_target.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
import pandas as pd
5+
import relbench
6+
7+
8+
def add_lag_target(
9+
df: pd.DataFrame,
10+
task: relbench.base.EntityTask,
11+
num_lags: int,
12+
) -> pd.DataFrame:
13+
r"""Add lag target features to the task table."""
14+
df = df.copy()
15+
right = df.sort_values(task.time_col)
16+
left = pd.DataFrame(
17+
{
18+
task.entity_col: df[task.entity_col],
19+
"__lookup_time__": df[task.time_col],
20+
"__row__": range(len(df)),
21+
}
22+
)
23+
for lag in range(1, num_lags + 1):
24+
column = f"{task.target_col}_lag_{lag}"
25+
if task.task_type == relbench.base.TaskType.REGRESSION:
26+
df[column] = float("NaN")
27+
else:
28+
df[column] = pd.Series(pd.NA, index=df.index, dtype="object")
29+
30+
merged = pd.merge_asof(
31+
left.sort_values("__lookup_time__"),
32+
right,
33+
by=task.entity_col,
34+
left_on="__lookup_time__",
35+
right_on=task.time_col,
36+
direction="backward",
37+
allow_exact_matches=False,
38+
)
39+
mask = merged[task.time_col].notna()
40+
rows = merged.loc[mask, "__row__"].to_numpy()
41+
df.loc[rows, column] = merged.loc[mask, task.target_col].to_numpy()
42+
43+
# Find the next lag strictly before the previous observation.
44+
left = merged.loc[mask, [task.entity_col, "__row__", task.time_col]]
45+
left = left.rename(columns={task.time_col: "__lookup_time__"})
46+
47+
return df

0 commit comments

Comments
 (0)