generate dataframe from ground truth

usc-isi-i2 · Jul 23, 2018 · 8d3b231 · 8d3b231
1 parent f7a8d07
commit 8d3b231
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/rltk/evaluation/ground_truth.py b/rltk/evaluation/ground_truth.py
@@ -2,6 +2,7 @@
 import heapq
 import random
 from typing import Callable
+import pandas as pd
 
 from rltk.utils import get_record_pairs
 from rltk.io.reader import GroundTruthReader
@@ -199,4 +200,11 @@ def train_test_split(self, test_ratio: float = 0.2, random_seed: int = None):
             else:
                 train_gt.add_ground_truth(id1, id2, label)
 
-        return train_gt, test_gt
+        return train_gt, test_gt
+
+    def generate_dataframe(self, **kwargs):
+        columns = ['id1', 'id2', 'label']
+        table = []
+        for id1, id2, label in self:
+            table.append([id1, id2, label])
+        return pd.DataFrame(table, columns=columns, **kwargs)