Modify EL batching to doc-wise streaming approach

spacy/cli/templates/quickstart_training.jinja

@@ -238,7 +238,7 @@ grad_factor = 1.0
 {% if "entity_linker" in components -%}
 [components.entity_linker]
 factory = "entity_linker"
-get_candidates = {"@misc":"spacy.CandidateGenerator.v1"}
+get_candidates = {"@misc":"spacy.CandidateGenerator.v2"}
 incl_context = true
 incl_prior = true
 
@@ -517,7 +517,7 @@ width = ${components.tok2vec.model.encode.width}
 {% if "entity_linker" in components -%}
 [components.entity_linker]
 factory = "entity_linker"
-get_candidates = {"@misc":"spacy.CandidateGenerator.v1"}
+get_candidates = {"@misc":"spacy.CandidateGenerator.v2"}
 incl_context = true
 incl_prior = true
 

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,15 @@ 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.
+        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 (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,
@@ -21,6 +21,9 @@
 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(
@@ -117,34 +120,38 @@ def empty_kb_factory(vocab: Vocab):
 
 
 @registry.misc("spacy.CandidateGenerator.v1")
-def create_candidates() -> Callable[[KnowledgeBase, Span], Iterable[Candidate]]:
+def create_get_candidates() -> Callable[[KnowledgeBase, Span], Iterable[Candidate]]:
     return get_candidates
 
 
-@registry.misc("spacy.CandidateBatchGenerator.v1")
-def create_candidates_batch() -> Callable[
-    [KnowledgeBase, SpanGroup], Iterable[Iterable[Candidate]]
+@registry.misc("spacy.CandidateGenerator.v2")
+def create_get_candidates_v2() -> Callable[
+    [KnowledgeBase, Iterator[SpanGroup]], Iterator[CandidatesForDocT]
 ]:
-    return get_candidates_batch
+    return get_candidates_v2
 
 
 def get_candidates(kb: KnowledgeBase, mention: Span) -> Iterable[Candidate]:
     """
-    Return candidate entities for a given mention and fetching appropriate entries from the index.
+    Return candidate entities for the given mention from the KB.
     kb (KnowledgeBase): Knowledge base to query.
-    mention (Span): Entity mention for which to identify candidates.
-    RETURNS (Iterable[Candidate]): Identified candidates.
+    mention (Span): Entity mention.
+    RETURNS (Iterable[Candidate]): Identified candidates for specified mention.
     """
-    return kb.get_candidates(mention)
+    cands_per_doc = next(
+        get_candidates_v2(kb, iter([SpanGroup(mention.doc, spans=[mention])]))
+    )
+    assert isinstance(cands_per_doc, list)
+    return next(cands_per_doc[0])
 
 
-def get_candidates_batch(
-    kb: KnowledgeBase, mentions: SpanGroup
-) -> Iterable[Iterable[Candidate]]:
+def get_candidates_v2(
+    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)