diff --git a/spacy/kb/candidate.pyx b/spacy/kb/candidate.pyx index 5ad52618a..613b70483 100644 --- a/spacy/kb/candidate.pyx +++ b/spacy/kb/candidate.pyx @@ -1,6 +1,6 @@ # cython: infer_types=True, profile=True -from typing import Iterable, Generator +from typing import Iterable, Generator, Iterator from .kb cimport KnowledgeBase from ..tokens import Span @@ -66,7 +66,7 @@ def get_candidates(kb: KnowledgeBase, mention: Span) -> Iterable[Candidate]: def get_candidates_all( kb: KnowledgeBase, mentions: Generator[Iterable[Span], None, None] -) -> Generator[Iterable[Iterable[Candidate]], None, None]: +) -> Iterator[Iterable[Iterable[Candidate]]]: """ Return candidate entities for the given mentions and fetching appropriate entries from the index. kb (KnowledgeBase): Knowledge base to query. diff --git a/spacy/ml/models/entity_linker.py b/spacy/ml/models/entity_linker.py index 455bcc3b1..4045a3206 100644 --- a/spacy/ml/models/entity_linker.py +++ b/spacy/ml/models/entity_linker.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Optional, Callable, Iterable, List, Tuple, Generator +from typing import Optional, Callable, Iterable, List, Tuple, Generator, Iterator from thinc.types import Floats2d from thinc.api import chain, list2ragged, reduce_mean, residual from thinc.api import Model, Maxout, Linear, tuplify, Ragged @@ -108,6 +108,6 @@ def create_candidates() -> Callable[[KnowledgeBase, Span], Iterable[Candidate]]: @registry.misc("spacy.CandidateAllGenerator.v1") def create_candidates_all() -> Callable[ [KnowledgeBase, Generator[Iterable[Span], None, None]], - Generator[Iterable[Iterable[Candidate]], None, None], + Iterator[Iterable[Iterable[Candidate]]], ]: return get_candidates_all diff --git a/spacy/pipeline/entity_linker.py b/spacy/pipeline/entity_linker.py index eb546b3a0..439c3af59 100644 --- a/spacy/pipeline/entity_linker.py +++ b/spacy/pipeline/entity_linker.py @@ -1,4 +1,14 @@ -from typing import Optional, Iterable, Callable, Dict, Union, List, Any, Generator +from typing import ( + Optional, + Iterable, + Callable, + Dict, + Union, + List, + Any, + Generator, + Iterator, +) from thinc.types import Floats2d from pathlib import Path from itertools import islice @@ -79,7 +89,7 @@ def make_entity_linker( get_candidates: Callable[[KnowledgeBase, Span], Iterable[Candidate]], get_candidates_all: Callable[ [KnowledgeBase, Generator[Iterable[Span], None, None]], - Generator[Iterable[Iterable[Candidate]], None, None], + Iterator[Iterable[Iterable[Candidate]]], ], overwrite: bool, scorer: Optional[Callable], @@ -102,7 +112,7 @@ def make_entity_linker( get_candidates_all ( Callable[ [KnowledgeBase, Generator[Iterable[Span], None, None]], - Generator[Iterable[Iterable[Candidate]], None, None] + Iterator[Iterable[Iterable[Candidate]]] ]): Function that produces a list of candidates per document, given a certain knowledge base and several textual documents with textual mentions. scorer (Optional[Callable]): The scoring method. @@ -180,7 +190,7 @@ class EntityLinker(TrainablePipe): get_candidates: Callable[[KnowledgeBase, Span], Iterable[Candidate]], get_candidates_all: Callable[ [KnowledgeBase, Generator[Iterable[Span], None, None]], - Generator[Iterable[Iterable[Candidate]], None, None], + Iterator[Iterable[Iterable[Candidate]]], ], overwrite: bool = BACKWARD_OVERWRITE, scorer: Optional[Callable] = entity_linker_score, @@ -204,7 +214,7 @@ class EntityLinker(TrainablePipe): get_candidates_all ( Callable[ [KnowledgeBase, Generator[Iterable[Span], None, None]], - Generator[Iterable[Iterable[Candidate]], None, None] + Iterator[Iterable[Iterable[Candidate]]] ]): Function that produces a list of candidates per document, given a certain knowledge base and several textual documents with textual mentions. scorer (Optional[Callable]): The scoring method. Defaults to Scorer.score_links. diff --git a/spacy/tests/pipeline/test_entity_linker.py b/spacy/tests/pipeline/test_entity_linker.py index 877a4c5ce..86956c2cc 100644 --- a/spacy/tests/pipeline/test_entity_linker.py +++ b/spacy/tests/pipeline/test_entity_linker.py @@ -1,4 +1,4 @@ -from typing import Callable, Iterable, Dict, Any, Generator +from typing import Callable, Iterable, Dict, Any, Generator, Iterator import pytest from numpy.testing import assert_equal @@ -515,7 +515,7 @@ def test_el_pipe_configuration(nlp): @registry.misc("spacy.LowercaseCandidateAllGenerator.v1") def create_candidates_batch() -> Callable[ [InMemoryLookupKB, Generator[Iterable["Span"], None, None]], - Generator[Iterable[Iterable[Candidate]], None, None], + Iterator[Iterable[Iterable[Candidate]]], ]: return get_lowercased_candidates_all