fix tasks/avito entity_table name (#269)

A tiny fix on entity_table name.
snap-stanford · Oct 28, 2024 · 00fb9d2 · 00fb9d2
1 parent d60904e
commit 00fb9d2
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/relbench/tasks/avito.py b/relbench/tasks/avito.py
@@ -60,7 +60,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab
 
         return Table(
             df=df,
-            fkey_col_to_pkey_table={"AdID": "entity_table"},
+            fkey_col_to_pkey_table={self.entity_col: self.entity_table},
             pkey_col=None,
             time_col="timestamp",
         )
@@ -108,7 +108,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab
 
         return Table(
             df=df,
-            fkey_col_to_pkey_table={"UserID": "entity_table"},
+            fkey_col_to_pkey_table={self.entity_col: self.entity_table},
             pkey_col=None,
             time_col="timestamp",
         )
@@ -165,7 +165,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab
 
         return Table(
             df=df,
-            fkey_col_to_pkey_table={"UserID": "entity_table"},
+            fkey_col_to_pkey_table={self.entity_col: self.entity_table},
             pkey_col=None,
             time_col="timestamp",
         )