From cd98ab4e95b9d7e2de4f478587e6d6d6bffc9c95 Mon Sep 17 00:00:00 2001 From: Raphael Mitsch Date: Tue, 28 Feb 2023 13:49:52 +0100 Subject: [PATCH] Convert Candidate from Cython to Python class. --- setup.py | 1 - spacy/kb/__init__.py | 4 +- spacy/kb/candidate.pxd | 12 --- spacy/kb/candidate.py | 109 +++++++++++++++++++++ spacy/kb/candidate.pyx | 74 -------------- spacy/kb/kb_in_memory.pyx | 20 ++-- spacy/ml/models/entity_linker.py | 22 ++++- spacy/tests/pipeline/test_entity_linker.py | 4 +- 8 files changed, 147 insertions(+), 99 deletions(-) delete mode 100644 spacy/kb/candidate.pxd create mode 100644 spacy/kb/candidate.py delete mode 100644 spacy/kb/candidate.pyx diff --git a/setup.py b/setup.py index d5b82ec68..9b8897233 100755 --- a/setup.py +++ b/setup.py @@ -30,7 +30,6 @@ MOD_NAMES = [ "spacy.lexeme", "spacy.vocab", "spacy.attrs", - "spacy.kb.candidate", "spacy.kb.kb", "spacy.kb.kb_in_memory", "spacy.ml.tb_framework", diff --git a/spacy/kb/__init__.py b/spacy/kb/__init__.py index 1d70a9b34..6dd4a3222 100644 --- a/spacy/kb/__init__.py +++ b/spacy/kb/__init__.py @@ -1,3 +1,5 @@ from .kb import KnowledgeBase from .kb_in_memory import InMemoryLookupKB -from .candidate import Candidate, get_candidates, get_candidates_batch +from .candidate import Candidate + +__all__ = ["KnowledgeBase", "InMemoryLookupKB", "Candidate"] diff --git a/spacy/kb/candidate.pxd b/spacy/kb/candidate.pxd deleted file mode 100644 index 942ce9dd0..000000000 --- a/spacy/kb/candidate.pxd +++ /dev/null @@ -1,12 +0,0 @@ -from .kb cimport KnowledgeBase -from libcpp.vector cimport vector -from ..typedefs cimport hash_t - -# Object used by the Entity Linker that summarizes one entity-alias candidate combination. -cdef class Candidate: - cdef readonly KnowledgeBase kb - cdef hash_t entity_hash - cdef float entity_freq - cdef vector[float] entity_vector - cdef hash_t alias_hash - cdef float prior_prob diff --git a/spacy/kb/candidate.py b/spacy/kb/candidate.py new file mode 100644 index 000000000..190792fbe --- /dev/null +++ b/spacy/kb/candidate.py @@ -0,0 +1,109 @@ +import abc +from typing import List, Union, Callable + + +class BaseCandidate(abc.ABC): + """A `BaseCandidate` object refers to a textual mention (`alias`) that may or may not be resolved + to a specific `entity_id` from a Knowledge Base. This will be used as input for the entity_id linking + algorithm which will disambiguate the various candidates to the correct one. + Each candidate (alias, entity_id) pair is assigned a certain prior probability. + + DOCS: https://spacy.io/api/kb/#candidate-init + """ + + def __init__( + self, mention: str, entity_id: Union[int, str], entity_vector: List[float] + ): + """Create new instance of `Candidate`. Note: has to be a sub-class, otherwise error will be raised. + mention (str): Mention text for this candidate. + entity_id (Union[int, str]): Unique entity ID. + entity_vector (List[float]): Entity embedding. + """ + self._mention = mention + self._entity_id = entity_id + self._entity_vector = entity_vector + + @property + def entity(self) -> Union[int, str]: + """RETURNS (Union[int, str]): Entity ID.""" + return self._entity_id + + @property + @abc.abstractmethod + def entity_(self) -> str: + """RETURNS (str): Entity name.""" + + @property + def mention(self) -> str: + """RETURNS (str): Mention.""" + return self._mention + + @property + def entity_vector(self) -> List[float]: + """RETURNS (List[float]): Entity vector.""" + return self._entity_vector + + +class Candidate(BaseCandidate): + """`Candidate` for InMemoryLookupKBCandidate.""" + + def __init__( + self, + retrieve_string_from_hash: Callable[[int], str], + entity_hash: int, + entity_freq: int, + entity_vector: List[float], + alias_hash: int, + prior_prob: float, + ): + """ + retrieve_string_from_hash (Callable[[int], str]): Callable retrieveing entity name from provided entity/vocab + hash. + entity_hash (str): Hashed entity name /ID. + entity_freq (int): Entity frequency in KB corpus. + entity_vector (List[float]): Entity embedding. + alias_hash (int): Hashed alias. + prior_prob (float): Prior probability of entity for this mention - i.e. the probability that, independent of + the context, this mention resolves to this entity_id in the corpus used to build the knowledge base. In + cases in which this isn't always possible (e.g.: the corpus to analyse contains mentions that the KB corpus + doesn't) it might be better to eschew this information and always supply the same value. + """ + super().__init__( + mention=retrieve_string_from_hash(alias_hash), + entity_id=entity_hash, + entity_vector=entity_vector, + ) + self._retrieve_string_from_hash = retrieve_string_from_hash + self._entity_hash = entity_hash + self._entity_freq = entity_freq + self._alias_hash = alias_hash + self._prior_prob = prior_prob + + @property + def entity(self) -> int: + """RETURNS (int): hash of the entity_id's KB ID/name""" + return self._entity_hash + + @property + def entity_(self) -> str: + """RETURNS (str): ID/name of this entity_id in the KB""" + return self._retrieve_string_from_hash(self._entity_hash) + + @property + def alias(self) -> int: + """RETURNS (int): hash of the alias""" + return self._alias_hash + + @property + def alias_(self) -> str: + """RETURNS (str): ID of the original alias""" + return self._retrieve_string_from_hash(self._alias_hash) + + @property + def entity_freq(self) -> float: + return self._entity_freq + + @property + def prior_prob(self) -> float: + """RETURNS (List[float]): Entity vector.""" + return self._prior_prob diff --git a/spacy/kb/candidate.pyx b/spacy/kb/candidate.pyx deleted file mode 100644 index c89efeb03..000000000 --- a/spacy/kb/candidate.pyx +++ /dev/null @@ -1,74 +0,0 @@ -# cython: infer_types=True, profile=True - -from typing import Iterable -from .kb cimport KnowledgeBase -from ..tokens import Span - -cdef class Candidate: - """A `Candidate` object refers to a textual mention (`alias`) that may or may not be resolved - to a specific `entity` from a Knowledge Base. This will be used as input for the entity linking - algorithm which will disambiguate the various candidates to the correct one. - Each candidate (alias, entity) pair is assigned a certain prior probability. - - DOCS: https://spacy.io/api/kb/#candidate-init - """ - - def __init__(self, KnowledgeBase kb, entity_hash, entity_freq, entity_vector, alias_hash, prior_prob): - self.kb = kb - self.entity_hash = entity_hash - self.entity_freq = entity_freq - self.entity_vector = entity_vector - self.alias_hash = alias_hash - self.prior_prob = prior_prob - - @property - def entity(self) -> int: - """RETURNS (uint64): hash of the entity's KB ID/name""" - return self.entity_hash - - @property - def entity_(self) -> str: - """RETURNS (str): ID/name of this entity in the KB""" - return self.kb.vocab.strings[self.entity_hash] - - @property - def alias(self) -> int: - """RETURNS (uint64): hash of the alias""" - return self.alias_hash - - @property - def alias_(self) -> str: - """RETURNS (str): ID of the original alias""" - return self.kb.vocab.strings[self.alias_hash] - - @property - def entity_freq(self) -> float: - return self.entity_freq - - @property - def entity_vector(self) -> Iterable[float]: - return self.entity_vector - - @property - def prior_prob(self) -> float: - return self.prior_prob - - -def get_candidates(kb: KnowledgeBase, mention: Span) -> Iterable[Candidate]: - """ - Return candidate entities for a given mention and fetching appropriate entries from the index. - kb (KnowledgeBase): Knowledge base to query. - mention (Span): Entity mention for which to identify candidates. - RETURNS (Iterable[Candidate]): Identified candidates. - """ - return kb.get_candidates(mention) - - -def get_candidates_batch(kb: KnowledgeBase, mentions: Iterable[Span]) -> Iterable[Iterable[Candidate]]: - """ - Return candidate entities for the given mentions and fetching appropriate entries from the index. - kb (KnowledgeBase): Knowledge base to query. - mention (Iterable[Span]): Entity mentions for which to identify candidates. - RETURNS (Iterable[Iterable[Candidate]]): Identified candidates. - """ - return kb.get_candidates_batch(mentions) diff --git a/spacy/kb/kb_in_memory.pyx b/spacy/kb/kb_in_memory.pyx index edba523cf..0b8e3f2f4 100644 --- a/spacy/kb/kb_in_memory.pyx +++ b/spacy/kb/kb_in_memory.pyx @@ -238,14 +238,18 @@ cdef class InMemoryLookupKB(KnowledgeBase): alias_index = self._alias_index.get(alias_hash) alias_entry = self._aliases_table[alias_index] - return [Candidate(kb=self, - entity_hash=self._entries[entry_index].entity_hash, - entity_freq=self._entries[entry_index].freq, - entity_vector=self._vectors_table[self._entries[entry_index].vector_index], - alias_hash=alias_hash, - prior_prob=prior_prob) - for (entry_index, prior_prob) in zip(alias_entry.entry_indices, alias_entry.probs) - if entry_index != 0] + return [ + Candidate( + retrieve_string_from_hash=self.vocab.strings.__getitem__, + entity_hash=self._entries[entry_index].entity_hash, + entity_freq=self._entries[entry_index].freq, + entity_vector=self._vectors_table[self._entries[entry_index].vector_index], + alias_hash=alias_hash, + prior_prob=prior_prob + ) + for (entry_index, prior_prob) in zip(alias_entry.entry_indices, alias_entry.probs) + if entry_index != 0 + ] def get_vector(self, str entity): cdef hash_t entity_hash = self.vocab.strings[entity] diff --git a/spacy/ml/models/entity_linker.py b/spacy/ml/models/entity_linker.py index 299b6bb52..df42c09dc 100644 --- a/spacy/ml/models/entity_linker.py +++ b/spacy/ml/models/entity_linker.py @@ -6,7 +6,7 @@ from thinc.api import Model, Maxout, Linear, tuplify, Ragged from ...util import registry from ...kb import KnowledgeBase, InMemoryLookupKB -from ...kb import Candidate, get_candidates, get_candidates_batch +from ...kb import Candidate from ...vocab import Vocab from ...tokens import Span, Doc from ..extract_spans import extract_spans @@ -109,3 +109,23 @@ def create_candidates_batch() -> Callable[ [KnowledgeBase, Iterable[Span]], Iterable[Iterable[Candidate]] ]: return get_candidates_batch + + +def get_candidates(kb: KnowledgeBase, mention: Span) -> Iterable[Candidate]: + """ + Return candidate entities for a given mention and fetching appropriate entries from the index. + kb (KnowledgeBase): Knowledge base to query. + mention (Span): Entity mention for which to identify candidates. + RETURNS (Iterable[Candidate]): Identified candidates. + """ + return kb.get_candidates(mention) + + +def get_candidates_batch(kb: KnowledgeBase, mentions: Iterable[Span]) -> Iterable[Iterable[Candidate]]: + """ + Return candidate entities for the given mentions and fetching appropriate entries from the index. + kb (KnowledgeBase): Knowledge base to query. + mention (Iterable[Span]): Entity mentions for which to identify candidates. + RETURNS (Iterable[Iterable[Candidate]]): Identified candidates. + """ + return kb.get_candidates_batch(mentions) diff --git a/spacy/tests/pipeline/test_entity_linker.py b/spacy/tests/pipeline/test_entity_linker.py index 506530591..3933a1254 100644 --- a/spacy/tests/pipeline/test_entity_linker.py +++ b/spacy/tests/pipeline/test_entity_linker.py @@ -7,10 +7,10 @@ from thinc.types import Ragged from spacy import registry, util from spacy.attrs import ENT_KB_ID from spacy.compat import pickle -from spacy.kb import Candidate, InMemoryLookupKB, get_candidates, KnowledgeBase +from spacy.kb import Candidate, InMemoryLookupKB, KnowledgeBase from spacy.lang.en import English from spacy.ml import load_kb -from spacy.ml.models.entity_linker import build_span_maker +from spacy.ml.models.entity_linker import build_span_maker, get_candidates from spacy.pipeline import EntityLinker, TrainablePipe from spacy.pipeline.tok2vec import DEFAULT_TOK2VEC_MODEL from spacy.scorer import Scorer