Add biomedical entity normalization

flair/data.py

@@ -336,11 +336,13 @@ def get_metadata(self, key: str) -> typing.Any:
     def has_metadata(self, key: str) -> bool:
         return key in self._metadata
 
-    def add_label(self, typename: str, value: str, score: float = 1.0):
+    def add_label(self, typename: str, value_or_label: Union[str, Label], score: float = 1.0):
+        label = value_or_label if isinstance(value_or_label, Label) else Label(self, value_or_label, score)
+
         if typename not in self.annotation_layers:
-            self.annotation_layers[typename] = [Label(self, value, score)]
+            self.annotation_layers[typename] = [label]
         else:
-            self.annotation_layers[typename].append(Label(self, value, score))
+            self.annotation_layers[typename].append(label)
 
         return self
 
@@ -431,6 +433,81 @@ def __len__(self):
         raise NotImplementedError
 
 
+class EntityLinkingLabel(Label):
+    """
+    Label class models entity linking annotations. Each entity linking label has a data point it refers
+    to as well as the identifier and name of the concept / entity from a knowledge base or ontology.
+
+    Optionally, additional concepts identifier and the database name can be provided.
+    """
+
+    def __init__(
+        self,
+        data_point: DataPoint,
+        concept_id: str,
+        concept_name: str,
+        score: float = 1.0,
+        additional_ids: Optional[Union[List[str], str]] = None,
+        database: Optional[str] = None,
+    ):
+        """
+        Initializes the label instance.
+
+        :param data_point: Data point / span the label refers to
+        :param concept_id: Identifier of the entity / concept from the knowledge base / ontology
+        :param concept_name: (Canonical) name of the entity / concept from the knowledge base / ontology
+        :param score: Matching score of the entity / concept according to the entity mention
+        :param additional_ids: List of additional identifiers for the concept / entity in the KB / ontology
+        :param database: Name of the knowlege base / ontology
+        """
+        super().__init__(data_point, concept_id, score)
+        self.concept_name = concept_name
+        self.database = database
+
+        if isinstance(additional_ids, str):
+            additional_ids = [additional_ids]
+        self.additional_ids = additional_ids
+
+    def spawn(self, value: str, score: float = 1.0):
+        return EntityLinkingLabel(
+            data_point=self.data_point,
+            concept_id=value,
+            score=score,
+            concept_name=self.concept_name,
+            additional_ids=self.additional_ids,
+            database=self.database,
+        )
+
+    def __str__(self):
+        return (
+            f"{self.data_point.unlabeled_identifier}{flair._arrow} "
+            f"{self.concept_name} - {self.database}:{self._value} ({round(self._score, 4)})"
+        )
+
+    def __repr__(self):
+        return (
+            f"{self.data_point.unlabeled_identifier}{flair._arrow} "
+            f"{self.concept_name} - {self.database}:{self._value} ({round(self._score, 4)})"
+        )
+
+    def __len__(self):
+        return len(self.data_point)
+
+    def __eq__(self, other):
+        return (
+            self.value == other.value
+            and self.data_point == other.data_point
+            and self.concept_name == other.concept_name
+            and self.identifier == other.identifier
+            and self.database == other.database
+            and self.score == other.score
+        )
+
+    @property
+    def identifier(self):
+        return f"{self.value}"
+
+
 DT = typing.TypeVar("DT", bound=DataPoint)
 DT2 = typing.TypeVar("DT2", bound=DataPoint)
 

flair/datasets/__init__.py

@@ -37,6 +37,8 @@
     CLL,
     CRAFT,
     CRAFT_V4,
+    CTD_CHEMICAL_DICTIONARY,
+    CTD_DISEASE_DICTIONARY,
     DECA,
     FSU,
     GELLUS,
@@ -90,6 +92,8 @@
     LOCTEXT,
     MIRNA,
     NCBI_DISEASE,
+    NCBI_GENE_HUMAN_DICTIONARY,
+    NCBI_TAXONOMY_DICTIONARY,
     OSIRIS,
     PDR,
     S800,
@@ -386,6 +390,10 @@
     "LINNEAUS",
     "LOCTEXT",
     "MIRNA",
+    "NEL_NCBI_HUMAN_GENE_DICT",
+    "NEL_NCBI_TAXONOMY_DICT",
+    "NEL_CTD_CHEMICAL_DICT",
+    "NEL_CTD_DISEASE_DICT",
     "NCBI_DISEASE",
     "ONTONOTES",
     "OSIRIS",