def test_pipeline():
base_url = Path(__file__).parent
wiki_subfolder = "wiki_test"
sample = {
"test_doc": ["the brown fox jumped over the lazy dog", [[10, 3]]]
}
config = {
"mode": "eval",
"model_path": f"{base_url}/{wiki_subfolder}/generated/model",
}
md = MentionDetection(base_url, wiki_subfolder)
tagger = Cmns(base_url, wiki_subfolder, n=5)
model = EntityDisambiguation(base_url, wiki_subfolder, config)
mentions_dataset, total_mentions = md.format_spans(sample)
predictions, _ = model.predict(mentions_dataset)
results = process_results(mentions_dataset,
predictions,
sample,
include_offset=False)
gold_truth = {"test_doc": [(10, 3, "Fox", "fox", -1, "NULL", 0.0)]}
return results == gold_truth
from flair.models import SequenceTagger
from REL.mention_detection import MentionDetection
from REL.entity_disambiguation import EntityDisambiguation
from time import time
base_url = "C:/Users/mickv/desktop/data_back/"
flair.device = torch.device('cuda:0')
mention_detection = MentionDetection(base_url, wiki_version)
# Alternatively use Flair NER tagger.
tagger_ner = SequenceTagger.load("ner-fast")
start = time()
mentions_dataset, n_mentions = mention_detection.find_mentions(
docs, tagger_ner)
print('MD took: {}'.format(time() - start))
# 3. Load model.
config = {
"mode": "eval",
"model_path": "{}/{}/generated/model".format(base_url, wiki_version),
}
model = EntityDisambiguation(base_url, wiki_version, config)
# 4. Entity disambiguation.
start = time()
predictions, timing = model.predict(mentions_dataset)
print('ED took: {}'.format(time() - start))