Upload node2vec_model.py with huggingface_hub

node2vec_model.py

@@ -13,11 +13,9 @@ from torch_geometric.data.data import DataEdgeAttr, DataTensorAttr
 from torch_geometric.data.storage import BaseStorage, EdgeStorage, GlobalStorage
 from torch_geometric.nn import Node2Vec
 
-import os
+from paths import EMB_PATH, GRAPH_PATH
 
-DATA_DIR = os.environ.get("DATA_DIR", ".")
-
-GRAPH = f"{DATA_DIR}/edge_ML_expected_ge5_pyg.pt"
+GRAPH = GRAPH_PATH
 
 
 def load_graph(path: str = GRAPH) -> Data:
@@ -56,7 +54,7 @@ def main() -> None:
     ap.add_argument("--num-workers", type=int, default=4)
     ap.add_argument("--epochs", type=int, default=0, help="0 = smoke check only")
     ap.add_argument("--steps", type=int, default=5, help="steps for the smoke check")
-    ap.add_argument("--out", default=f"{DATA_DIR}/edge_ML_expected_ge5_n2v.pt")
+    ap.add_argument("--out", default=EMB_PATH)
     args = ap.parse_args()
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")