finalize desgin

Author: thunderock

graph_ml/models/node2vec.py

@@ -1,5 +1,4 @@
-import numpy as np
-
+import torch
 
 from ..utils import config
 
@@ -26,11 +25,16 @@ def __init__(
     def fit(self, epochs=1, learning_rate=0.1, batch_size=128):
         return self._fit(epochs, learning_rate, batch_size)
 
-    def transform(self, nodes=None, type_=np.ndarray):
+    def transform(self, nodes=None, type_=torch.Tensor):
         return self._transform(nodes, type_)
 
     def fit_transform(
-        self, epochs=1, learning_rate=0.1, batch_size=128, nodes=None, type_=np.ndarray
+        self,
+        epochs=1,
+        learning_rate=0.1,
+        batch_size=128,
+        nodes=None,
+        type_=torch.Tensor,
     ):
         self.fit(epochs, learning_rate, batch_size)
         return self.transform(nodes, type_)

graph_ml/models/torch_node2vec.py

@@ -2,7 +2,6 @@
 from torch_geometric.nn import Node2Vec as PyGNode2Vec
 import torch
 from ..models.node2vec import Node2Vec
-from ..utils import torch_utils
 
 
 class TorchNode2Vec(Node2Vec):
@@ -21,7 +20,7 @@ def __init__(self, **params):
     @property
     def edge_index(self):
         # should not be called too often, no caching here
-        return torch_utils.adj_list_to_edge_index(self.adj_list)
+        return self.adj_list.nonzero().t().contiguous()
 
     def _fit(self, epochs, learning_rate, batch_size, shuffle=True):
         # TODO (ashutosh): check if training two times works
@@ -47,7 +46,7 @@ def _fit(self, epochs, learning_rate, batch_size, shuffle=True):
                 print(f"Epoch: {epoch}, Loss: {total_loss[epoch]}")
         return self
 
-    def _transform(self, nodes=None, type_=np.ndarray):
+    def _transform(self, nodes=None, type_=torch.Tensor):
         self.model.eval()
         if nodes is None:
             nodes = torch.arange(self.num_nodes, device=self.device)