Modify EL batching to doc-wise streaming approach

spacy/errors.py

@@ -950,7 +950,6 @@ class Errors(metaclass=ErrorsWithCodes):
              "case pass an empty list for the previously not specified argument to avoid this error.")
     E1043 = ("Expected None or a value in range [{range_start}, {range_end}] for entity linker threshold, but got "
              "{value}.")
-    E1044 = ("Expected `candidates_batch_size` to be >= 1, but got: {value}")
     E1045 = ("Encountered {parent} subclass without `{parent}.{method}` "
              "method in '{name}'. If you want to use this method, make "
              "sure it's overwritten on the subclass.")

spacy/kb/kb.pyx

@@ -1,14 +1,14 @@
 # cython: infer_types=True
 
 from pathlib import Path
-from typing import Iterable, Tuple, Union
+from typing import Iterable, Iterator, Tuple, Union
 
 from cymem.cymem cimport Pool
 
 from ..errors import Errors
-from ..tokens import Span, SpanGroup
+from ..tokens import SpanGroup
 from ..util import SimpleFrozenList
-from .candidate import Candidate
+from .candidate cimport Candidate
 
 
 cdef class KnowledgeBase:
@@ -19,6 +19,8 @@ cdef class KnowledgeBase:
 
     DOCS: https://spacy.io/api/kb
     """
+    CandidatesForMentionT = Iterable[Candidate]
+    CandidatesForDocT = Iterable[CandidatesForMentionT]
 
     def __init__(self, vocab: Vocab, entity_vector_length: int):
         """Create a KnowledgeBase."""
@@ -32,27 +34,14 @@ cdef class KnowledgeBase:
         self.entity_vector_length = entity_vector_length
         self.mem = Pool()
 
-    def get_candidates_batch(
-            self, mentions: SpanGroup
-    ) -> Iterable[Iterable[Candidate]]:
+    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.
         If no candidates are found for a given mention, an empty list is returned.
-        mentions (SpanGroup): Mentions for which to get candidates.
-        RETURNS (Iterable[Iterable[Candidate]]): Identified candidates.
-        """
-        return [self.get_candidates(span) for span in mentions]
-
-    def get_candidates(self, mention: Span) -> Iterable[Candidate]:
-        """
-        Return candidate entities for a specific 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.
-        If no candidate is found for the given mention, an empty list is returned.
-        mention (Span): Mention for which to get candidates.
-        RETURNS (Iterable[Candidate]): Identified candidates.
+        mentions (Iterator[SpanGroup]): Mentions for which to get candidates.
+        RETURNS (Iterator[Iterable[Iterable[Candidate]]]): Identified candidates per mention/doc/doc batch.
         """
         raise NotImplementedError(
             Errors.E1045.format(

spacy/kb/kb_in_memory.pyx

@@ -1,5 +1,5 @@
 # cython: infer_types=True
-from typing import Any, Callable, Dict, Iterable
+from typing import Any, Callable, Dict, Iterable, Iterator
 
 import srsly
 
@@ -12,7 +12,7 @@ from preshed.maps cimport PreshMap
 import warnings
 from pathlib import Path
 
-from ..tokens import Span
+from ..tokens import SpanGroup
 
 from ..typedefs cimport hash_t
 
@@ -255,8 +255,9 @@ cdef class InMemoryLookupKB(KnowledgeBase):
             alias_entry.probs = probs
             self._aliases_table[alias_index] = alias_entry
 
-    def get_candidates(self, mention: Span) -> Iterable[InMemoryCandidate]:
-        return self._get_alias_candidates(mention.text)  # type: ignore
+    def get_candidates(self, mentions: Iterator[SpanGroup]) -> Iterator[Iterable[Iterable[InMemoryCandidate]]]:
+        for mentions_for_doc in mentions:
+            yield [self._get_alias_candidates(span.text) for span in mentions_for_doc]
 
     def _get_alias_candidates(self, str alias) -> Iterable[InMemoryCandidate]:
         """

spacy/ml/models/entity_linker.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Callable, Iterable, List, Optional, Tuple
+from typing import Callable, Iterable, Iterator, List, Optional, Tuple
 
 from thinc.api import (
     Linear,
@@ -117,34 +117,20 @@ def empty_kb_factory(vocab: Vocab):
 
 
 @registry.misc("spacy.CandidateGenerator.v1")
-def create_candidates() -> Callable[[KnowledgeBase, Span], Iterable[Candidate]]:
-    return get_candidates
-
-
-@registry.misc("spacy.CandidateBatchGenerator.v1")
-def create_candidates_batch() -> Callable[
-    [KnowledgeBase, SpanGroup], Iterable[Iterable[Candidate]]
+def create_get_candidates() -> Callable[
+    [KnowledgeBase, Iterator[SpanGroup]],
+    Iterator[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)
+    return get_candidates
 
 
-def get_candidates_batch(
-    kb: KnowledgeBase, mentions: SpanGroup
-) -> Iterable[Iterable[Candidate]]:
+def get_candidates(
+    kb: KnowledgeBase, mentions: Iterator[SpanGroup]
+) -> Iterator[Iterable[Iterable[Candidate]]]:
     """
-    Return candidate entities for the given mentions and fetching appropriate entries from the index.
+    Return candidate entities for the given mentions from the KB.
     kb (KnowledgeBase): Knowledge base to query.
-    mentions (SpanGroup): Entity mentions for which to identify candidates.
-    RETURNS (Iterable[Iterable[Candidate]]): Identified candidates.
+    mentions (Iterator[SpanGroup]): Mentions per doc.
+    RETURNS (Iterator[Iterable[Iterable[Candidate]]]): Identified candidates per mentions in document/SpanGroup.
     """
-    return kb.get_candidates_batch(mentions)
+    return kb.get_candidates(mentions)