masakhaner-zero-shot.py

import numpy as np
import sys

from flair.datasets import NER_MASAKHANE
from flair.models import SequenceTagger
from tabulate import tabulate

languages = ["amh", "hau", "ibo", "kin", "lug", "luo", "pcm", "swa", "wol", "yor"]

language_mapping = {}

label_name_map = {"DATE": "O"}

for language in languages:
    language_mapping[language] = NER_MASAKHANE(languages=language, version="v1",
                                               label_name_map=label_name_map)

model_names = sys.argv[1:]

print("Models:", model_names)

dev_table = []
test_table = []

for model_name in model_names:
    current_model = SequenceTagger.load(model_name)
    
    current_dev_entry = [model_name]
    current_test_entry = [model_name]
    
    for language in languages:
        current_corpus = language_mapping[language]
        
        dev_result = current_model.evaluate(current_corpus.dev, gold_label_type="ner",
                                            mini_batch_size=64).main_score
        test_result = current_model.evaluate(current_corpus.test, gold_label_type="ner",
                                             mini_batch_size=64).main_score
        
        dev_result = round(dev_result * 100, 1)
        test_result = round(test_result * 100, 1)
        current_dev_entry.append(dev_result)
        current_test_entry.append(test_result)
    
    dev_avg = round(np.mean(current_test_entry[1:]), 1)
    test_avg = round(np.mean(current_test_entry[1:]), 1)
    
    current_dev_entry.append(dev_avg)
    current_test_entry.append(test_avg)
    
    dev_table.append(current_dev_entry)
    test_table.append(current_test_entry)

# Calculate mean of language columns per language over all models
last_dev_row = ["Language Avg."]

for index, _ in enumerate(languages):
    dev_scores = []
    for row in dev_table:
        dev_scores.append(row[index + 1])
    last_dev_row.append(round(np.mean(dev_scores), 1))

last_dev_row.append(round(np.mean(last_dev_row[1:]), 1))
dev_table.append(last_dev_row)

# Same for test set
last_test_row = ["Language Avg."]

for index, _ in enumerate(languages):
    test_scores = []
    for row in test_table:
        test_scores.append(row[index + 1])
    last_test_row.append(round(np.mean(test_scores), 1))

last_test_row.append(round(np.mean(last_test_row[1:]), 1))
test_table.append(last_test_row)

# Final tables
headers = ["Model Name"] + languages + ["Avg."]

print("Development Results:")
print(tabulate(dev_table, headers=headers, tablefmt="github"))

print("")

print("Test Results:")
print(tabulate(test_table, headers=headers, tablefmt="github"))