[Feature]: schema guided extraction

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
+    async def extract(self, chunk: dict) -> 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/bases/base_storage.py

@@ -45,6 +45,9 @@ async def get_by_ids(
     ) -> list[Union[T, None]]:
         raise NotImplementedError
 
+    async def get_all(self) -> dict[str, T]:
+        raise NotImplementedError
+
     async def filter_keys(self, data: list[str]) -> set[str]:
         """return un-exist keys"""
         raise NotImplementedError

graphgen/configs/aggregated_config.yaml

@@ -2,8 +2,11 @@ pipeline:
   - name: read
     params:
       input_file: resources/input_examples/jsonl_demo.jsonl # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples
-      chunk_size: 1024 # chunk size for text splitting
-      chunk_overlap: 100 # chunk overlap for text splitting
+
+  - name: chunk
+    params:
+        chunk_size: 1024 # chunk size for text splitting
+        chunk_overlap: 100 # chunk overlap for text splitting
 
   - name: build_kg
 

graphgen/configs/atomic_config.yaml

@@ -2,6 +2,9 @@ pipeline:
   - name: read
     params:
       input_file: resources/input_examples/json_demo.json # input file path, support json, jsonl, txt, csv, pdf. See resources/input_examples for examples
+
+  - name: chunk
+    params:
       chunk_size: 1024 # chunk size for text splitting
       chunk_overlap: 100 # chunk overlap for text splitting
 

graphgen/configs/cot_config.yaml

@@ -2,8 +2,11 @@ pipeline:
   - name: read
     params:
       input_file: resources/input_examples/txt_demo.txt  # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples
-      chunk_size: 1024 # chunk size for text splitting
-      chunk_overlap: 100 # chunk overlap for text splitting
+
+  - name: chunk
+    params:
+        chunk_size: 1024 # chunk size for text splitting
+        chunk_overlap: 100 # chunk overlap for text splitting
 
   - name: build_kg
 

graphgen/configs/multi_hop_config.yaml

@@ -2,6 +2,9 @@ pipeline:
   - name: read
     params:
       input_file: resources/input_examples/csv_demo.csv # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples
+
+  - name: chunk
+    params:
       chunk_size: 1024 # chunk size for text splitting
       chunk_overlap: 100 # chunk overlap for text splitting
 

graphgen/configs/schema_guided_config.yaml

@@ -0,0 +1,15 @@
+pipeline:
+  - name: read
+    params:
+      input_file: resources/input_examples/extract_demo.txt # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples
+
+  - name: chunk
+    params:
+      chunk_size: 20480
+      chunk_overlap: 2000
+      separators: []
+
+  - name: extract
+    params:
+      method: schema_guided # extraction method, support: schema_guided
+      schema_file: graphgen/templates/extraction/schemas/legal_contract.json # schema file path for schema_guided method

graphgen/configs/vqa_config.yaml

@@ -2,8 +2,11 @@ pipeline:
   - name: read
     params:
       input_file: resources/input_examples/vqa_demo.json # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples
-      chunk_size: 1024 # chunk size for text splitting
-      chunk_overlap: 100 # chunk overlap for text splitting
+
+  - name: chunk
+    params:
+        chunk_size: 1024 # chunk size for text splitting
+        chunk_overlap: 100 # chunk overlap for text splitting
 
   - name: build_kg
 

graphgen/graphgen.py

@@ -18,6 +18,7 @@
 from graphgen.operators import (
     build_kg,
     chunk_documents,
+    extract_info,
     generate_qas,
     init_llm,
     judge_statement,
@@ -70,6 +71,7 @@ def __init__(
         self.search_storage: JsonKVStorage = JsonKVStorage(
             self.working_dir, namespace="search"
         )
+
         self.rephrase_storage: JsonKVStorage = JsonKVStorage(
             self.working_dir, namespace="rephrase"
         )
@@ -80,6 +82,10 @@ def __init__(
             os.path.join(self.working_dir, "data", "graphgen", f"{self.unique_id}"),
             namespace="qa",
         )
+        self.extract_storage: JsonKVStorage = JsonKVStorage(
+            os.path.join(self.working_dir, "data", "graphgen", f"{self.unique_id}"),
+            namespace="extraction",
+        )
 
         # webui
         self.progress_bar: gr.Progress = progress_bar
@@ -103,16 +109,30 @@ async def read(self, read_config: Dict):
         _add_doc_keys = await self.full_docs_storage.filter_keys(list(new_docs.keys()))
         new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys}
 
+        if len(new_docs) == 0:
+            logger.warning("All documents are already in the storage")
+            return
+
+        await self.full_docs_storage.upsert(new_docs)
+        await self.full_docs_storage.index_done_callback()
+
+    @op("chunk", deps=["read"])
+    @async_to_sync_method
+    async def chunk(self, chunk_config: Dict):
+        """
+        chunk documents into smaller pieces from full_docs_storage if not already present
+        """
+
+        new_docs = await self.meta_storage.get_new_data(self.full_docs_storage)
         if len(new_docs) == 0:
             logger.warning("All documents are already in the storage")
             return
 
         inserting_chunks = await chunk_documents(
             new_docs,
-            read_config["chunk_size"],
-            read_config["chunk_overlap"],
             self.tokenizer_instance,
             self.progress_bar,
+            **chunk_config,
         )
 
         _add_chunk_keys = await self.chunks_storage.filter_keys(
@@ -126,12 +146,12 @@ async def read(self, read_config: Dict):
             logger.warning("All chunks are already in the storage")
             return
 
-        await self.full_docs_storage.upsert(new_docs)
-        await self.full_docs_storage.index_done_callback()
         await self.chunks_storage.upsert(inserting_chunks)
         await self.chunks_storage.index_done_callback()
+        await self.meta_storage.mark_done(self.full_docs_storage)
+        await self.meta_storage.index_done_callback()
 
-    @op("build_kg", deps=["read"])
+    @op("build_kg", deps=["chunk"])
     @async_to_sync_method
     async def build_kg(self):
         """
@@ -161,7 +181,7 @@ async def build_kg(self):
 
         return _add_entities_and_relations
 
-    @op("search", deps=["read"])
+    @op("search", deps=["chunk"])
     @async_to_sync_method
     async def search(self, search_config: Dict):
         logger.info(
@@ -248,6 +268,26 @@ async def partition(self, partition_config: Dict):
         await self.partition_storage.upsert(batches)
         return batches
 
+    @op("extract", deps=["chunk"])
+    @async_to_sync_method
+    async def extract(self, extract_config: Dict):
+        logger.info("Extracting information from given chunks...")
+
+        results = await extract_info(
+            self.synthesizer_llm_client,
+            self.chunks_storage,
+            extract_config,
+            progress_bar=self.progress_bar,
+        )
+        if not results:
+            logger.warning("No information extracted")
+            return
+
+        await self.extract_storage.upsert(results)
+        await self.extract_storage.index_done_callback()
+        await self.meta_storage.mark_done(self.chunks_storage)
+        await self.meta_storage.index_done_callback()
+
     @op("generate", deps=["partition"])
     @async_to_sync_method
     async def generate(self, generate_config: Dict):

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,101 @@
+import json
+from typing import Dict, List
+
+from graphgen.bases import BaseExtractor, BaseLLMWrapper
+from graphgen.templates import SCHEMA_GUIDED_EXTRACTION_PROMPT
+from graphgen.utils import compute_dict_hash, detect_main_language, logger
+
+
+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
+        self.required_keys = self.schema.get("required")
+        if not self.required_keys:
+            # If no required keys are specified, use all keys from the schema as default
+            self.required_keys = list(self.schema.get("properties", {}).keys())
+
+    def build_prompt(self, text: str) -> str:
+        schema_explanation = ""
+        for field, details in self.schema.get("properties", {}).items():
+            description = details.get("description", "No description provided.")
+            schema_explanation += f'- "{field}": {description}\n'
+
+        lang = detect_main_language(text)
+
+        prompt = SCHEMA_GUIDED_EXTRACTION_PROMPT[lang].format(
+            field=self.schema.get("name", "the document"),
+            schema_explanation=schema_explanation,
+            examples="",
+            text=text,
+        )
+        return prompt
+
+    async def extract(self, chunk: dict) -> dict:
+        text = chunk.get("text", "")
+        prompt = self.build_prompt(text)
+        response = await self.llm_client.generate_answer(prompt)
+        try:
+            extracted_info = json.loads(response)
+            # Ensure all required keys are present
+            for key in self.required_keys:
+                if key not in extracted_info:
+                    extracted_info[key] = ""
+            if any(extracted_info[key] == "" for key in self.required_keys):
+                logger.debug("Missing required keys in extraction: %s", extracted_info)
+                return {}
+            main_keys_info = {key: extracted_info[key] for key in self.required_keys}
+            logger.debug("Extracted info: %s", extracted_info)
+            return {compute_dict_hash(main_keys_info, prefix="extract"): extracted_info}
+        except json.JSONDecodeError:
+            logger.error("Failed to parse extraction response: %s", response)
+            return {}
+
+    async def merge_extractions(
+        self, extraction_list: List[Dict[str, dict]]
+    ) -> Dict[str, dict]:
+        """
+        Merge multiple extraction results based on their hashes.
+        :param extraction_list: List of extraction results, each is a dict with hash as key and record as value.
+        :return: Merged extraction results.
+        """
+        merged: Dict[str, dict] = {}
+        for ext in extraction_list:
+            for h, rec in ext.items():
+                if h not in merged:
+                    merged[h] = rec.copy()
+                else:
+                    for k, v in rec.items():
+                        if k not in merged[h] or merged[h][k] == v:
+                            merged[h][k] = v
+                        else:
+                            merged[h][k] = f"{merged[h][k]}<SEP>{v}"
+        return merged

graphgen/models/reader/txt_reader.py

@@ -5,10 +5,6 @@
 
 class TXTReader(BaseReader):
     def read(self, file_path: str) -> List[Dict[str, Any]]:
-        docs = []
         with open(file_path, "r", encoding="utf-8") as f:
-            for line in f:
-                line = line.strip()
-                if line:
-                    docs.append({self.text_column: line})
+            docs = [{"type": "text", self.text_column: f.read()}]
         return self.filter(docs)

graphgen/models/storage/json_storage.py

@@ -39,6 +39,9 @@ async def get_by_ids(self, ids, fields=None) -> list:
             for id in ids
         ]
 
+    async def get_all(self) -> dict[str, str]:
+        return self._data
+
     async def filter_keys(self, data: list[str]) -> set[str]:
         return {s for s in data if s not in self._data}
 

graphgen/operators/__init__.py

@@ -1,4 +1,5 @@
 from .build_kg import build_kg
+from .extract import extract_info
 from .generate import generate_qas
 from .init import init_llm
 from .judge import judge_statement

graphgen/operators/extract/__init__.py

@@ -0,0 +1 @@
+from .extract_info import extract_info