From d6c7636069297924bbe68d3278c872b2dc8c0ca3 Mon Sep 17 00:00:00 2001 From: Raphael Mitsch Date: Wed, 7 Feb 2024 10:01:21 +0100 Subject: [PATCH] Update docstring and type usage. --- spacy/kb/kb.pyx | 7 ++++--- spacy/ml/models/entity_linker.py | 6 ++++-- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/spacy/kb/kb.pyx b/spacy/kb/kb.pyx index 0e2ba7c65..2a12b204a 100644 --- a/spacy/kb/kb.pyx +++ b/spacy/kb/kb.pyx @@ -36,9 +36,10 @@ cdef class KnowledgeBase: def get_candidates(self, mentions: Iterator[SpanGroup]) -> Iterator[CandidatesForDocT]: """ - Return candidate entities for a specified Span mention. Each candidate defines at least the entity and the - entity's embedding vector. Depending on the KB implementation, further properties - such as the prior - probability of the specified mention text resolving to that entity - might be included. + Return candidate entities for the specified groups of mentions (as SpanGroup) per Doc. + Each candidate for a mention defines at least the entity and the entity's embedding vector. Depending on the KB + implementation, further properties - such as the prior probability of the specified mention text resolving to + that entity - might be included. If no candidates are found for a given mention, an empty list is returned. mentions (Iterator[SpanGroup]): Mentions for which to get candidates. RETURNS (Iterator[Iterable[Iterable[Candidate]]]): Identified candidates per mention/doc/doc batch. diff --git a/spacy/ml/models/entity_linker.py b/spacy/ml/models/entity_linker.py index d0465b618..f06876aa4 100644 --- a/spacy/ml/models/entity_linker.py +++ b/spacy/ml/models/entity_linker.py @@ -21,6 +21,9 @@ from ...util import registry from ...vocab import Vocab from ..extract_spans import extract_spans +CandidatesForMentionT = Iterable[Candidate] +CandidatesForDocT = Iterable[CandidatesForMentionT] + @registry.architectures("spacy.EntityLinker.v2") def build_nel_encoder( @@ -118,8 +121,7 @@ def empty_kb( @registry.misc("spacy.CandidateGenerator.v1") def create_get_candidates() -> Callable[ - [KnowledgeBase, Iterator[SpanGroup]], - Iterator[Iterable[Iterable[Candidate]]], + [KnowledgeBase, Iterator[SpanGroup]], Iterator[CandidatesForDocT] ]: return get_candidates