Evaluate NER predictions

NER eval.ipynb

ner_eval.py

from collections import Counter
def split_label(y):
    label = y.rsplit("-", 1)
    if len(label) == 2:
        return tuple(label)
    return (label[0], None)

def extract_entities(Y):
    i = 0
    while i < len(Y):
        y = Y[i]
        boundary, tag = split_label(y)
        start = i
        if boundary == "B":
            i += 1
            while i < len(Y) and split_label(Y[i])[0] == "I":
                i += 1
            yield (start, i, tag)
        else:
            i += 1

def seq_metrics(Y_true, Y_pred):
    true_counts = Counter()
    pred_counts = Counter()
    true_pred_counts = Counter()
    for y_true, y_pred in zip(Y_true, Y_pred):
        true_entities = set(extract_entities(y_true))
        pred_entities = set(extract_entities(y_pred))
        true_pred_entities = true_entities & pred_entities
        y_true_counts = Counter((x[2] for x in true_entities))
        y_pred_counts = Counter((x[2] for x in pred_entities))
        y_true_pred_counts = Counter((x[2] for x in true_pred_entities))
        true_counts += y_true_counts
        pred_counts += y_pred_counts
        true_pred_counts += y_true_pred_counts
    return true_counts, pred_counts, true_pred_counts

list(extract_entities(Y_pred[0]))  
seq_metrics(Y_true, Y_pred)