improve typing and run black

flairNLP · Feb 2, 2024 · 170a31b · 170a31b
1 parent 960af8b
commit 170a31b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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