kapipe 0.0.1__py3-none-any.whl

kapipe/__init__.py

@@ -0,0 +1,39 @@
+# Types defined
+from . import datatypes
+
+# Top-level modules
+from . import systems
+
+# Submodules for `systems`
+from . import models
+from . import decoders
+from . import prompt_processors
+from . import misc
+
+# Top modules for training and evaluation
+from . import trainers
+
+# Evaluation metrics
+from . import evaluation
+
+# Retrievers for few-shot in-context learning
+from . import demonstration_retrievers
+
+from . import utils
+
+
+from .pipeline import Pipeline
+
+
+__version__ = "0.0.1"
+
+
+def load(identifier, gpu_map=None):
+    ka = Pipeline(identifier=identifier, gpu_map=gpu_map)
+    return ka
+
+
+def blank(gpu_map=None):
+    ka = Pipeline(identifier=None, gpu_map=gpu_map)
+    return ka
+

kapipe/datatypes/__init__.py

@@ -0,0 +1,11 @@
+from .datatypes import DocKey
+from .datatypes import Document
+from .datatypes import Mention
+from .datatypes import Entity
+from .datatypes import Triple
+
+from .datatypes import EntityPage
+from .datatypes import CandEntKeyInfo
+from .datatypes import EntDoc
+
+from .datatypes import DemoKeyInfo

kapipe/datatypes/datatypes.py

@@ -0,0 +1,69 @@
+from typing import Any, TypeAlias
+
+
+DocKey : TypeAlias = str
+
+Document : TypeAlias = dict[str, Any]
+# {
+#     "doc_key": str,
+#     "sentences": list[str],
+#     "mentions": list[Mention],
+#     "entities": list[Entity],
+#     "relations": list[Triple]
+# }
+
+Mention : TypeAlias = dict[str, Any]
+# {
+#     "span": tuple[int],
+#     "name": str,
+#     "entity_type": str
+# }
+
+Entity : TypeAlias = dict[str, Any]
+# {
+#     "mention_indices": list[int],
+#     "entity_type": str,
+#     "entity_id": str    
+# }
+
+Triple : TypeAlias = dict[str, Any]
+# {
+#     "arg1": int,
+#     "relation": str,
+#     "arg2": int
+# }
+
+##########
+# For candidate entities in ED
+##########
+
+EntityPage : TypeAlias = dict[str, Any]
+# {
+#     "entity_id": str,
+#     "canonical_name": str,
+#     "synonyms": list[str],
+#     "description": str
+# }
+
+CandEntKeyInfo : TypeAlias = dict[str, Any]
+# {
+#     "entity_id": str,
+#     "score": float
+# }
+
+EntDoc: TypeAlias = dict[str, str]
+# {
+#     "entity_id": str,
+#     "canonical_name": str,
+#     "text": str
+# }
+
+##########
+# For in-context learning
+##########
+
+DemoKeyInfo : TypeAlias = dict[str, Any]
+# {
+#     "doc_key": str,
+#     "score": float
+# }

kapipe/decoders/__init__.py

@@ -0,0 +1,6 @@
+# NER
+from .spanbaseddecoder import SpanBasedDecoder
+
+# Entity Linking
+
+# DocRE

kapipe/decoders/spanbaseddecoder.py

@@ -0,0 +1,97 @@
+import numpy as np
+
+
+class SpanBasedDecoder:
+    """A span-based decoder for NER. This decoder applies constraints (for either Flat NER or Nested NER) to given scored spans and outputs a list of mentions.
+    """
+
+    def __init__(
+        self,
+        allow_nested_entities
+    ):
+        """
+        Parameters
+        ----------
+        allow_nested_entities : bool
+        """
+        self.allow_nested_entities = allow_nested_entities
+
+    def decode(
+        self,
+        spans,
+        words
+    ):
+        """
+        Parameters
+        ----------
+        spans : list[tuple[int, int, str, float]]
+        words : list[str]
+
+        Results
+        -------
+        list[Mention]
+        """
+        mentions = [] # list[Mention]
+
+        # Sort the candidate spans based on their scores
+        spans = sorted(spans, key=lambda x: -x[-1])
+
+        # Select spans
+        n_words = len(words)
+        self.check_matrix = np.zeros((n_words, n_words)) # Used in Flat NER
+        self.check_set = set() # Uased in Nested NER
+        for span in spans:
+            begin_token_index, end_token_index, etype, _ = span
+            name = " ".join(words[begin_token_index: end_token_index + 1])
+            if self.is_violation(
+                begin_token_index=begin_token_index,
+                end_token_index=end_token_index
+            ):
+                continue
+            mentions.append({
+                "span": (begin_token_index, end_token_index),
+                "name": name,
+                "entity_type": etype,
+            })
+            self.check_matrix[begin_token_index: end_token_index + 1] = 1
+            self.check_set.add((begin_token_index, end_token_index))
+
+        # Sort mentions based on the positions
+        mentions = sorted(mentions, key=lambda m: m["span"])
+
+        return mentions
+
+    def is_violation(self, begin_token_index, end_token_index):
+        """
+        Parameters
+        ----------
+        begin_token_index : int
+        end_token_index : int
+
+        Returns
+        -------
+        bool
+        """
+        if not self.allow_nested_entities:
+            # Flat NER
+            if (
+                self.check_matrix[begin_token_index: end_token_index + 1].sum()
+                > 0
+            ):
+                return True
+            else:
+                return False
+        else:
+            # Nested NER
+            for begin_token_j, end_token_j in self.check_set:
+                if (
+                    (begin_token_index < begin_token_j
+                     <= end_token_index < end_token_j)
+                    or
+                    (begin_token_j < begin_token_index
+                     <= end_token_j < end_token_index)
+                ):
+                    return True
+            return False
+
+

kapipe/demonstration_retrievers/__init__.py

@@ -0,0 +1 @@
+from .demonstration_retrievers import DemonstrationRetriever

kapipe/demonstration_retrievers/demonstration_retrievers.py

@@ -0,0 +1,85 @@
+import random
+
+from .. import utils
+
+
+class DemonstrationRetriever:
+    """A demonstration (i.e., a few-shot examplers) retriever for LLM-based in-context learning methods. This retriever returns the few-shot examplers for a given document.
+    """
+
+    def __init__(
+        self,
+        path_demonstration_pool,
+        method,
+        task="docre"
+    ):
+        """
+        Parameters
+        ----------
+        path_demonstration_pool : str
+        method : str
+        task : str, optional
+            by default "docre"
+        """
+        assert method in ["first", "random"]
+
+        self.path_demonstration_pool = path_demonstration_pool
+        self.method = method
+        self.task = task
+
+        # Get a map from DocKey to Document
+        self.demonstration_pool = {
+            demo_doc["doc_key"]: demo_doc
+            for demo_doc in utils.read_json(self.path_demonstration_pool)
+        }
+
+        # Get a pool of candidate document keys for retrieval
+        self.doc_keys = list(self.demonstration_pool.keys())
+
+    def retrieve(self, document, top_k, doc_keys=None):
+        """
+        Parameters
+        ----------
+        document : Document
+        top_k : int
+        doc_keys : list[str] | None, optional
+            by default None
+
+        Returns
+        -------
+        dict[str, str | list[dict[str, str | float]]]
+        """
+        # `doc_keys` can be used for retrieval on limited candidates
+        if doc_keys is None:
+            doc_keys = self.doc_keys
+
+        # list[dict[str, str | float]]
+        if self.method == "first":
+            demonstrations_for_doc = doc_keys[:top_k]
+            demonstrations_for_doc = [
+                {
+                    "doc_key": key,
+                    "score": 1.0
+                }
+                for key in demonstrations_for_doc
+            ]
+        elif self.method == "random":
+            demonstrations_for_doc = random.sample(doc_keys, top_k)
+            demonstrations_for_doc = [
+                {
+                    "doc_key": key,
+                    "score": 1.0
+                }
+                for key in demonstrations_for_doc
+            ]
+        else:
+            raise Exception(f"Invalid method: {self.method}")
+
+        # dict[str, str | list[dict[str, str | float]]]
+        demonstrations_for_doc = {
+            "doc_key": document["doc_key"],
+            "demonstrations": demonstrations_for_doc,
+        }
+
+        return demonstrations_for_doc
+

kapipe/evaluation/__init__.py

@@ -0,0 +1,3 @@
+from . import ner
+from . import ed
+from . import docre

kapipe/evaluation/docre/__init__.py

@@ -0,0 +1,4 @@
+from .fscore import fscore
+
+from .docred_official import to_official
+from .docred_official import official_evaluate