Fix evaluation bug

Aliang-CN · Aug 18, 2019 · 7482e12 · 7482e12
1 parent acda000
commit 7482e12
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/src/main.py b/src/main.py
@@ -75,25 +75,29 @@ def get_score(local_model, node1, node2):
         vector2 = local_model[node2]
         return np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2))
     except Exception as e:
-        print(e)
+        pass
 
 
 def evaluate(model, true_edges, false_edges):
     true_list = list()
     prediction_list = list()
+    true_num = 0
     for edge in true_edges:
         tmp_score = get_score(model, str(edge[0]), str(edge[1]))
-        true_list.append(1)
-        prediction_list.append(tmp_score)
+        if tmp_score is not None:
+            true_list.append(1)
+            prediction_list.append(tmp_score)
+            true_num += 1
 
     for edge in false_edges:
         tmp_score = get_score(model, str(edge[0]), str(edge[1]))
-        true_list.append(0)
-        prediction_list.append(tmp_score)
+        if tmp_score is not None:
+            true_list.append(0)
+            prediction_list.append(tmp_score)
 
     sorted_pred = prediction_list[:]
     sorted_pred.sort()
-    threshold = sorted_pred[-len(true_edges)]
+    threshold = sorted_pred[-true_num]
 
     y_pred = np.zeros(len(prediction_list), dtype=np.int32)
     for i in range(len(prediction_list)):