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))