A script to generate precision scoring for ranked embedding matches from BioBert.

bert-scorer.py

import pickle

ref = pickle.load(open('/path/to/biobert_results/gold-reference.pkl', 'rb'))
res = pickle.load(open('/path/to/biobert_results/results.pkl', 'rb'))

scores = []
for res_key in res.keys():
    if res_key == 'text':
        continue
    precisions = []

    cleaned_res = [x[0] for x in res[res_key]]

    # p@1
    precisions.append(1 if len(set(ref[res_key]).intersection(cleaned_res[:1])) > 0 else 0)

    # p@3
    precisions.append(1 if len(set(ref[res_key]).intersection(cleaned_res[:3])) > 0 else 0)

    # p@5
    precisions.append(1 if len(set(ref[res_key]).intersection(cleaned_res[:5])) > 0 else 0)

    # p@10
    precisions.append(1 if len(set(ref[res_key]).intersection(cleaned_res[:10])) > 0 else 0)

    scores.append(precisions)

print('p@1 =', sum([x[0] for x in scores])/len(scores))
print('p@3 =', sum([x[1] for x in scores])/len(scores))
print('p@5 =', sum([x[2] for x in scores])/len(scores))
print('p@10 =', sum([x[3] for x in scores])/len(scores))