Add annotation dataclasses to pkg_types.py

pkg_api/core/pkg_types.py

@@ -1,4 +1,11 @@
 """PKG package types."""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Union
+
 from rfc3987 import match
 
 SPARQLQuery = str
@@ -9,3 +16,77 @@ def __new__(cls, *args, **kwargs):
         """Creates a new URI."""
         assert match(args[0], rule="IRI"), f"Invalid URI: {args[0]}"
         return super().__new__(cls, *args, **kwargs)
+
+
+@dataclass
+class Concept:
+    """Class representing a SKOS concept."""
+
+    description: str
+    related_entities: List[URI] = field(default_factory=list)
+    broader_entities: List[URI] = field(default_factory=list)
+    narrower_entities: List[URI] = field(default_factory=list)
+
+
+@dataclass
+class TripleElement:
+    """Class representing a subject, predicate or object element.
+
+    Attributes:
+        reference: Raw string reference of the element.
+        value: URI, Concept or literal value of the element.
+    """
+
+    reference: str
+    value: Union[URI, Concept, str, None] = field(default=None)
+
+    @staticmethod
+    def from_value(value: Union[URI, Concept, str]) -> TripleElement:
+        """Creates a TripleElement from a value.
+
+        Args:
+            value: URI, Concept or literal value.
+
+        Returns:
+            TripleElement.
+        """
+        if isinstance(value, URI):
+            return TripleElement("", value)
+        elif isinstance(value, Concept):
+            return TripleElement(value.description, value)
+        return TripleElement(value, value)
+
+
+@dataclass
+class Triple:
+    """Class representing a subject, predicate, object triple."""
+
+    subject: Optional[TripleElement] = None
+    predicate: Optional[TripleElement] = None
+    object: Optional[TripleElement] = None
+
+
+@dataclass
+class Preference:
+    """Class representing a preference.
+
+    Note: In the current version of the PKG API, topic refers to the object of
+    a triple.
+    """
+
+    topic: TripleElement
+    weight: float
+
+
+@dataclass
+class PKGData:
+    """Class representing an annotated statement.
+
+    Annotations include a triple, a preference, and logging data.
+    """
+
+    id: uuid.UUID
+    statement: str
+    triple: Optional[Triple] = None
+    preference: Optional[Preference] = None
+    logging_data: Dict[str, Any] = field(default_factory=dict)

pkg_api/nl_to_pkg/annotators/annotator.py

@@ -5,12 +5,11 @@
 -1) in the query.
 """
 
-
 from abc import ABC, abstractmethod
 from typing import Tuple
 
-from pkg_api.core.annotation import PKGData
 from pkg_api.core.intents import Intent
+from pkg_api.core.pkg_types import PKGData
 from pkg_api.nl_to_pkg.llm.prompt import Prompt
 
 

pkg_api/nl_to_pkg/annotators/three_step_annotator.py

@@ -4,13 +4,12 @@
 with a triple and a preference using LLM.
 """
 
-
 import re
 import uuid
 from typing import Dict, Optional, Tuple
 
-from pkg_api.core.annotation import PKGData, Preference, Triple, TripleElement
 from pkg_api.core.intents import Intent
+from pkg_api.core.pkg_types import PKGData, Preference, Triple, TripleElement
 from pkg_api.nl_to_pkg.annotators.annotator import StatementAnnotator
 from pkg_api.nl_to_pkg.llm.llm_connector import LLMConnector
 from pkg_api.nl_to_pkg.llm.prompt import Prompt

pkg_api/nl_to_pkg/entity_linking/entity_linker.py

@@ -1,9 +1,8 @@
 """Abstract class for entity linking."""
 
-
 from abc import ABC, abstractmethod
 
-from pkg_api.core.annotation import PKGData
+from pkg_api.core.pkg_types import PKGData
 
 
 class EntityLinker(ABC):

pkg_api/nl_to_pkg/entity_linking/rel_entity_linking.py

@@ -4,8 +4,7 @@
 
 import requests
 
-from pkg_api.core.annotation import Concept, PKGData, TripleElement
-from pkg_api.core.pkg_types import URI
+from pkg_api.core.pkg_types import URI, Concept, PKGData, TripleElement
 from pkg_api.nl_to_pkg.entity_linking.entity_linker import EntityLinker
 
 _DEFAULT_API_URL = "https://rel.cs.ru.nl/api"

pkg_api/nl_to_pkg/entity_linking/spotlight_entity_linker.py

@@ -4,8 +4,7 @@
 
 import requests
 
-from pkg_api.core.annotation import Concept, PKGData, TripleElement
-from pkg_api.core.pkg_types import URI
+from pkg_api.core.pkg_types import URI, Concept, PKGData, TripleElement
 from pkg_api.nl_to_pkg.entity_linking.entity_linker import EntityLinker
 from pkg_api.util.load_config import load_yaml_config
 

pkg_api/nl_to_pkg/eval_nl_to_pkg.py

@@ -1,10 +1,13 @@
 """Evaluates the NL to PKG models."""
+
 import csv
-from typing import Dict, Any, List, Tuple
+from typing import Any, Dict, List, Tuple
+
 from sklearn.metrics import f1_score
 from tqdm import tqdm
-from pkg_api.core.annotation import PKGData
+
 from pkg_api.core.intents import Intent
+from pkg_api.core.pkg_types import PKGData
 from pkg_api.nl_to_pkg.annotators.three_step_annotator import (
     ThreeStepStatementAnnotator,
 )

pkg_api/nl_to_pkg/nl_to_pkg.py

@@ -2,8 +2,8 @@
 
 from typing import Tuple
 
-from pkg_api.core.annotation import PKGData
 from pkg_api.core.intents import Intent
+from pkg_api.core.pkg_types import PKGData
 from pkg_api.nl_to_pkg import EntityLinker, StatementAnnotator
 
 

pkg_api/pkg.py

@@ -24,9 +24,8 @@
 
 import pkg_api.utils as utils
 from pkg_api.connector import Connector, RDFStore
-from pkg_api.core.annotation import Concept, PKGData, Triple, TripleElement
 from pkg_api.core.namespaces import PKGPrefixes
-from pkg_api.core.pkg_types import URI
+from pkg_api.core.pkg_types import URI, Concept, PKGData, Triple, TripleElement
 from pkg_api.mapping_vocab import MappingVocab
 
 DEFAULT_VISUALIZATION_PATH = "data/pkg_visualizations"
@@ -279,9 +278,11 @@ def _parse_statement_node(
                 setattr(_triple, k, TripleElement.from_value(v))
 
         return PKGData(
-            id=uuid.UUID(f"{{{statement_node_id}}}")
-            if statement_node_id
-            else uuid.uuid1(),
+            id=(
+                uuid.UUID(f"{{{statement_node_id}}}")
+                if statement_node_id
+                else uuid.uuid1()
+            ),
             statement=statement_dict.get("statement"),
             triple=_triple,
             preference=None,

pkg_api/utils.py

@@ -11,9 +11,15 @@
 import re
 from typing import List, Optional, Union
 
-from pkg_api.core.annotation import Concept, PKGData, Triple, TripleElement
 from pkg_api.core.namespaces import PKGPrefixes
-from pkg_api.core.pkg_types import URI, SPARQLQuery
+from pkg_api.core.pkg_types import (
+    URI,
+    Concept,
+    PKGData,
+    SPARQLQuery,
+    Triple,
+    TripleElement,
+)
 
 
 def _clean_sparql_representation(sparql: str) -> str:

tests/nl_to_pkg/test_eval_nl_to_pkg.py

@@ -1,12 +1,13 @@
 """Tests eval_nl_to_pkg.py file."""
+
 import uuid
 from typing import Dict, List, Tuple
 from unittest.mock import MagicMock, patch
 
 import pytest
 
-from pkg_api.core.annotation import PKGData, Triple, TripleElement
 from pkg_api.core.intents import Intent
+from pkg_api.core.pkg_types import PKGData, Triple, TripleElement
 from pkg_api.nl_to_pkg.annotators.three_step_annotator import (
     ThreeStepStatementAnnotator,
 )

tests/nl_to_pkg/test_nl_to_pkg.py

@@ -5,8 +5,8 @@
 
 import pytest
 
-from pkg_api.core.annotation import PKGData, Preference, Triple, TripleElement
 from pkg_api.core.intents import Intent
+from pkg_api.core.pkg_types import PKGData, Preference, Triple, TripleElement
 from pkg_api.nl_to_pkg.nl_to_pkg import NLtoPKG
 
 

tests/nl_to_pkg/test_rel_entity_linker.py

@@ -1,11 +1,11 @@
 """Tests for REL entity linker."""
+
 import uuid
 from unittest.mock import Mock, patch
 
 import pytest
 
-from pkg_api.core.annotation import Concept, PKGData, Triple, TripleElement
-from pkg_api.core.pkg_types import URI
+from pkg_api.core.pkg_types import URI, Concept, PKGData, Triple, TripleElement
 from pkg_api.nl_to_pkg.entity_linking.rel_entity_linking import RELEntityLinker
 
 

tests/nl_to_pkg/test_spotlight_entity_linker.py

@@ -1,11 +1,11 @@
 """Tests the Spotlight entity linker class."""
+
 import uuid
 from unittest.mock import Mock, patch
 
 import pytest
 
-from pkg_api.core.annotation import Concept, PKGData, Triple, TripleElement
-from pkg_api.core.pkg_types import URI
+from pkg_api.core.pkg_types import URI, Concept, PKGData, Triple, TripleElement
 from pkg_api.nl_to_pkg.entity_linking.spotlight_entity_linker import (
     SpotlightEntityLinker,
 )

tests/nl_to_pkg/test_three_step_annotator.py

@@ -5,8 +5,8 @@
 
 import pytest
 
-from pkg_api.core.annotation import Preference, Triple, TripleElement
 from pkg_api.core.intents import Intent
+from pkg_api.core.pkg_types import Preference, Triple, TripleElement
 from pkg_api.nl_to_pkg.annotators.three_step_annotator import (
     ThreeStepStatementAnnotator,
 )

tests/pkg_api/test_pkg.py

@@ -1,18 +1,19 @@
 """Tests for the PKG module."""
+
 import re
 import uuid
 
 import pytest
 
 from pkg_api.connector import RDFStore
-from pkg_api.core.annotation import (
+from pkg_api.core.pkg_types import (
+    URI,
     Concept,
     PKGData,
     Preference,
     Triple,
     TripleElement,
 )
-from pkg_api.core.pkg_types import URI
 from pkg_api.pkg import PKG
 from pkg_api.utils import get_statement_node_id