wip: add schema_guided_extractor

Author: ChenZiHong-Gavin

graphgen/bases/__init__.py

@@ -1,3 +1,4 @@
+from .base_extractor import BaseExtractor
 from .base_generator import BaseGenerator
 from .base_kg_builder import BaseKGBuilder
 from .base_llm_wrapper import BaseLLMWrapper

graphgen/bases/base_extractor.py

@@ -0,0 +1,22 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from graphgen.bases.base_llm_wrapper import BaseLLMWrapper
+
+
+class BaseExtractor(ABC):
+    """
+    Extract information from given text.
+
+    """
+
+    def __init__(self, llm_client: BaseLLMWrapper):
+        self.llm_client = llm_client
+
+    @abstractmethod
+    def extract(self, text_or_documents: str) -> Any:
+        """Extract information from the given text"""
+
+    @abstractmethod
+    def build_prompt(self, text: str) -> str:
+        """Build prompt for LLM based on the given text"""

graphgen/models/extractor/__init__.py

@@ -0,0 +1 @@
+from .schema_guided_extractor import SchemaGuidedExtractor

graphgen/models/extractor/key_information_extractor.py

@@ -0,0 +1 @@
+# TODO: text2json

graphgen/models/extractor/schema_guided_extractor.py

@@ -0,0 +1,41 @@
+from graphgen.bases import BaseExtractor, BaseLLMWrapper
+
+
+class SchemaGuidedExtractor(BaseExtractor):
+    """
+    Use JSON/YAML Schema or Pydantic Model to guide the LLM to extract structured information from text.
+
+    Usage example:
+        schema = {
+                "type": "legal contract",
+                "description": "A legal contract for leasing property.",
+                "properties": {
+                    "end_date": {"type": "string", "description": "The end date of the lease."},
+                    "leased_space": {"type": "string", "description": "Description of the space that is being leased."},
+                    "lessee": {"type": "string", "description": "The lessee's name (and possibly address)."},
+                    "lessor": {"type": "string", "description": "The lessor's name (and possibly address)."},
+                    "signing_date": {"type": "string", "description": "The date the contract was signed."},
+                    "start_date": {"type": "string", "description": "The start date of the lease."},
+                    "term_of_payment": {"type": "string", "description": "Description of the payment terms."},
+                    "designated_use": {"type": "string",
+                    "description": "Description of the designated use of the property being leased."},
+                    "extension_period": {"type": "string",
+                    "description": "Description of the extension options for the lease."},
+                    "expiration_date_of_lease": {"type": "string", "description": "The expiration data of the lease."}
+                },
+                "required": ["lessee", "lessor", "start_date", "end_date"]
+            }
+        extractor = SchemaGuidedExtractor(llm_client, schema)
+        result = extractor.extract(text)
+
+    """
+
+    def __init__(self, llm_client: BaseLLMWrapper, schema: dict):
+        super().__init__(llm_client)
+        self.schema = schema
+
+    def build_prompt(self, text: str) -> str:
+        pass
+
+    def extract(self, text_or_documents: str) -> dict:
+        pass

graphgen/operators/extract/__init__.py

@@ -0,0 +1,47 @@
+from typing import List
+
+import gradio as gr
+
+from graphgen.bases import BaseLLMWrapper
+from graphgen.bases.datatypes import Chunk
+from graphgen.models.extractor import SchemaGuidedExtractor
+from graphgen.utils import logger, run_concurrent
+
+
+async def extract(
+    llm_client: BaseLLMWrapper,
+    chunks: List[Chunk],
+    generation_config: dict,
+    progress_bar: gr.Progress = None,
+):
+    """
+    Extract information from chunks
+    :param llm_client: LLM client
+    :param chunks
+    :param generation_config
+    :param progress_bar
+    :return: extracted information
+    """
+
+    method = generation_config.get("method")
+    if method == "schema_guided":
+        schema = generation_config.get("schema")
+        extractor = SchemaGuidedExtractor(llm_client, schema)
+        print(extractor)
+    else:
+        raise ValueError(f"Unsupported extraction method: {method}")
+
+    logger.info("[Extraction] method: %s, chunks: %d", method, len(chunks))
+
+    # results = await run_concurrent(
+    #     extractor.extract,
+    #     [chunk.content for chunk in chunks],
+    #     desc="Extracting information",
+    #     unit="chunk",
+    #     progress_bar=progress_bar,
+    # )
+    #
+    # # TODO: 对results合并，去重
+    # return results
+
+    return []

graphgen/operators/extract/extract.py

@@ -0,0 +1,58 @@
+{
+  "type": "object",
+  "name": "legal contract",
+  "description": "A legal contract for leasing property.",
+  "properties": {
+      "end_date": {
+          "type": "array",
+          "items": {"type": "string"},
+          "description": "The end date of the lease."
+      },
+      "leased_space": {
+          "type": "array",
+          "items": {"type": "string"},
+          "description": "Description of the space that is being leased."
+      },
+      "lessee": {
+          "type": "array",
+          "items": {"type": "string"},
+          "description": "The lessee's name (and possibly address)."
+      },
+      "lessor": {
+          "type": "array",
+          "items": {"type": "string"},
+          "description": "The lessor's name (and possibly address)."
+      },
+      "signing_date": {
+          "type": "array",
+          "items": {"type": "string"},
+          "description": "The date the contract was signed."
+      },
+      "start_date": {
+          "type": "array",
+          "items": {"type": "string"},
+          "description": "The start date of the lease."
+      },
+      "term_of_payment": {
+          "type": "array",
+          "items": {"type": "string"},
+          "description": "Description of the payment terms."
+      },
+      "designated_use": {
+          "type": "array",
+          "items": {"type": "string"},
+          "description": "Designated use of the property being leased."
+      },
+      "extension_period": {
+          "type": "array",
+          "items": {"type": "string"},
+          "description": "Description of the extension options for the lease."
+      },
+      "expiration_date_of_lease": {
+          "type": "array",
+          "items": {"type": "string"},
+          "description": "The expiration date of the lease."
+      }
+  },
+    "required": ["lessee", "lessor", "start_date", "end_date"]
+}