|
import unittest |
|
import spacy |
|
from spacyEntityLinker.EntityLinker import EntityLinker |
|
|
|
|
|
class TestEntityLinker(unittest.TestCase): |
|
|
|
def __init__(self, arg, *args, **kwargs): |
|
super(TestEntityLinker, self).__init__(arg, *args, **kwargs) |
|
self.nlp = spacy.load('en_core_web_sm') |
|
|
|
def test_initialization(self): |
|
entityLinker = EntityLinker() |
|
|
|
self.nlp.add_pipe(entityLinker, last=True, name="entityLinker") |
|
|
|
doc = self.nlp( |
|
"Elon Musk was born in South Africa. Bill Gates and Steve Jobs come from in the United States") |
|
|
|
doc._.linkedEntities.pretty_print() |
|
doc._.linkedEntities.print_super_entities() |
|
for sent in doc.sents: |
|
sent._.linkedEntities.pretty_print() |
|
|
|
self.nlp.remove_pipe("entityLinker") |
|
|