diff --git a/flair/datasets/entity_linking.py b/flair/datasets/entity_linking.py index bc97444617..330785d349 100644 --- a/flair/datasets/entity_linking.py +++ b/flair/datasets/entity_linking.py @@ -41,7 +41,7 @@ def __init__( self._idx_to_candidates = {candidate.concept_id: candidate for candidate in candidates} # one name can map to multiple concepts - self._text_to_index: Dict[str, List] = {} + self._text_to_index: Dict[str, List[str]] = {} for candidate in candidates: for text in [candidate.concept_name, *candidate.synonyms]: if text not in self._text_to_index: @@ -54,7 +54,7 @@ def database_name(self) -> str: return self._dataset_name @property - def text_to_index(self) -> Dict[str, List]: + def text_to_index(self) -> Dict[str, List[str]]: return self._text_to_index @property