Fixed arvix import (#151)

* Fixed arvix * Reran tox * Restore cpp change --------- Co-authored-by: dsarda <[email protected]> Co-authored-by: Devesh <[email protected]>
marius-team · Nov 25, 2023 · 2f27ffe · 2f27ffe
1 parent 483277a
commit 2f27ffe
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/src/python/tools/preprocess/datasets/ogbn_arxiv.py b/src/python/tools/preprocess/datasets/ogbn_arxiv.py
@@ -81,7 +81,6 @@ def preprocess(
             output_dir=self.output_directory,
             train_edges=self.input_edge_list_file,
             num_partitions=num_partitions,
-            columns=[0, 1],
             src_column=0,
             dst_column=1,
             remap_ids=remap_ids,

diff --git a/src/python/tools/preprocess/datasets/twitter.py b/src/python/tools/preprocess/datasets/twitter.py
@@ -28,7 +28,14 @@ def download(self, overwrite=False):
             archive_path = download_url(self.dataset_url, self.output_directory, overwrite)
             extract_file(archive_path, remove_input=True)
 
-    def preprocess(self, num_partitions=1, remap_ids=True, splits=[0.9, 0.05, 0.05], sequential_train_nodes=False):
+    def preprocess(
+        self,
+        num_partitions=1,
+        remap_ids=True,
+        splits=[0.9, 0.05, 0.05],
+        sequential_train_nodes=False,
+        partitioned_eval=False,
+    ):
         converter = TorchEdgeListConverter(
             output_dir=self.output_directory,
             train_edges=self.input_edges,
@@ -38,6 +45,7 @@ def preprocess(self, num_partitions=1, remap_ids=True, splits=[0.9, 0.05, 0.05],
             num_partitions=num_partitions,
             splits=splits,
             remap_ids=remap_ids,
+            partitioned_evaluation=partitioned_eval,
         )
 
         return converter.convert()