update metrics

Hironsan · Jul 23, 2017 · 43229b1 · 43229b1
1 parent 806eaf6
commit 43229b1
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/anago/data/metrics.py b/anago/data/metrics.py
@@ -2,17 +2,16 @@
 from sklearn.metrics import classification_report
 
 
-def report(y_true, y_pred, index2chunk):
+def report(y_true, y_pred, entity_to_id):
     y_true = [y.argmax() for y in itertools.chain(*y_true)]
     y_pred = [y.argmax() for y in itertools.chain(*y_pred)]
 
-    tagset = set(index2chunk) - {'O', '<PAD>'}
+    tagset = set(entity_to_id) - {'O', '<PAD>'}
     tagset = sorted(tagset, key=lambda tag: tag.split('-', 1)[::-1])
-    class_indices = {cls: idx for idx, cls in enumerate(index2chunk)}
 
     print(classification_report(
         y_true,
         y_pred,
-        labels=[class_indices[cls] for cls in tagset],
+        labels=[entity_to_id[cls] for cls in tagset],
         target_names=tagset,
     ))