Update app.py

app.py

@@ -69,8 +69,8 @@ def train():
     optimizer.zero_grad()
     z = encoder(data.x, data.edge_index)
 
-    pos_edge = data.pos_edge_label_index
-    neg_edge = data.neg_edge_label_index
+    pos_edge = data.edge_label_index[:, data.edge_label == 1]
+    neg_edge = data.edge_label_index[:, data.edge_label == 0]
 
     pos_pred = predictor(z[pos_edge[0]], z[pos_edge[1]])
     neg_pred = predictor(z[neg_edge[0]], z[neg_edge[1]])
@@ -86,8 +86,10 @@ def train():
         encoder.eval()
         predictor.eval()
         z = encoder(test_data.x, test_data.edge_index)
-        pos_pred = predictor(z[test_data.pos_edge_label_index[0]], z[test_data.pos_edge_label_index[1]])
-        neg_pred = predictor(z[test_data.neg_edge_label_index[0]], z[test_data.neg_edge_label_index[1]])
+        pos_edge = test_data.edge_label_index[:, test_data.edge_label == 1]
+        neg_edge = test_data.edge_label_index[:, test_data.edge_label == 0]
+        pos_pred = predictor(z[pos_edge[0]], z[pos_edge[1]])
+        neg_pred = predictor(z[neg_edge[0]], z[neg_edge[1]])
         preds = torch.cat([pos_pred, neg_pred])
         labels = torch.cat([
             torch.ones(pos_pred.size(0)),
@@ -99,7 +101,7 @@ def train():
 
 def visualize_graph():
     temp_data = data.clone()
-    temp_data.edge_index = temp_data.pos_edge_label_index
+    temp_data.edge_index = data.edge_index  # Use actual edge index for visualization
     G = to_networkx(temp_data, to_undirected=True)
     plt.figure(figsize=(4, 4))
     nx.draw(G, with_labels=True, node_color='lightblue', edge_color='gray')
@@ -155,7 +157,7 @@ def predict_link(node_i, node_j):
         return []
 
     temp_data = data.clone()
-    temp_data.edge_index = temp_data.pos_edge_label_index
+    temp_data.edge_index = data.edge_index
     G = to_networkx(temp_data, to_undirected=True)
     neighbors = {node: list(G.neighbors(node)) for node in G.nodes}
     dfs = dfs_path(neighbors, int(node_i), int(node_j))