From 56bba0baa00f56d50a9c243d9aa87c41096e04cc Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Wed, 5 Nov 2025 19:48:15 +0800 Subject: [PATCH 01/12] docs: add schema_guided_config --- graphgen/configs/schema_guided_config.yaml | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 graphgen/configs/schema_guided_config.yaml diff --git a/graphgen/configs/schema_guided_config.yaml b/graphgen/configs/schema_guided_config.yaml new file mode 100644 index 00000000..6510a91d --- /dev/null +++ b/graphgen/configs/schema_guided_config.yaml @@ -0,0 +1,22 @@ +read: + input_file: resources/input_examples/jsonl_demo.jsonl # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples +split: + chunk_size: 1024 # chunk size for text splitting + chunk_overlap: 100 # chunk overlap for text splitting +search: # web search configuration + enabled: false # whether to enable web search + search_types: ["google"] # search engine types, support: google, bing, uniprot, wikipedia +quiz_and_judge: # quiz and test whether the LLM masters the knowledge points + enabled: true + quiz_samples: 2 # number of quiz samples to generate + re_judge: false # whether to re-judge the existing quiz samples +partition: # graph partition configuration + method: ece # ece is a custom partition method based on comprehension loss + method_params: + max_units_per_community: 20 # max nodes and edges per community + min_units_per_community: 5 # min nodes and edges per community + max_tokens_per_community: 10240 # max tokens per community + unit_sampling: max_loss # unit sampling strategy, support: random, max_loss, min_loss +generate: + mode: aggregated # atomic, aggregated, multi_hop, cot, vqa + data_format: ChatML # Alpaca, Sharegpt, ChatML From b7af2e4988d8709ce62070557cdb464eb8063892 Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Thu, 6 Nov 2025 11:42:57 +0800 Subject: [PATCH 02/12] feat: add schema_guided extraction prompt template --- graphgen/templates/__init__.py | 1 + graphgen/templates/extraction/__init__.py | 1 + .../extraction/schema_guided_extraction.py | 56 +++++++++++++++++++ 3 files changed, 58 insertions(+) create mode 100644 graphgen/templates/extraction/__init__.py create mode 100644 graphgen/templates/extraction/schema_guided_extraction.py diff --git a/graphgen/templates/__init__.py b/graphgen/templates/__init__.py index ea28c4d0..0940e910 100644 --- a/graphgen/templates/__init__.py +++ b/graphgen/templates/__init__.py @@ -1,5 +1,6 @@ from .coreference_resolution import COREFERENCE_RESOLUTION_PROMPT from .description_rephrasing import DESCRIPTION_REPHRASING_PROMPT +from .extraction import SCHEMA_GUIDED_EXTRACTION_PROMPT from .generation import ( AGGREGATED_GENERATION_PROMPT, ATOMIC_GENERATION_PROMPT, diff --git a/graphgen/templates/extraction/__init__.py b/graphgen/templates/extraction/__init__.py new file mode 100644 index 00000000..70f1369a --- /dev/null +++ b/graphgen/templates/extraction/__init__.py @@ -0,0 +1 @@ +from .schema_guided_extraction import SCHEMA_GUIDED_EXTRACTION_PROMPT diff --git a/graphgen/templates/extraction/schema_guided_extraction.py b/graphgen/templates/extraction/schema_guided_extraction.py new file mode 100644 index 00000000..997bd7fe --- /dev/null +++ b/graphgen/templates/extraction/schema_guided_extraction.py @@ -0,0 +1,56 @@ +TEMPLATE_EN = """You are an expert at extracting information from text based on a given schema. +Extract relevant information about {field} from a given contract document according to the provided schema. + +Instructions: +1. Carefully read the entire document provided at the end of this prompt. +2. Extract the relevant information. +3. Present your findings in JSON format as specified below. + +Important Notes: +- Extract only relevant information. +- Consider the context of the entire document when determining relevance. +- Do not be verbose, only respond with the correct format and information. +- Some docs may have multiple relevant excerpts -- include all that apply. +- Some questions may have no relevant excerpts -- just return ["N/A"]. +- Do not include additional JSON keys beyond the ones listed here. +- Do not include the same key multiple times in the JSON. +- Use English for your response. + +Expected JSON keys and explanation of what they are: +{schema_explanation} + +{examples} + +Document to extract from: +{text} +""" + +TEMPLATE_ZH = """你是一个擅长根据给定的模式从文本中提取信息的专家。 +根据提供的模式，从合同文件中提取与{field}相关的信息。 +操作说明： +1. 仔细阅读本提示末尾提供的整份文件。 +2. 提取相关信息。 +3. 按照下面指定的JSON格式呈现你的发现。 + +重要注意事项： +- 仅提取相关信息。 +- 在确定相关性时，考虑整份文件的上下文。 +- 不要冗长，只需以正确的格式和信息进行回应。 +- 有些文件可能有多个相关摘录——请包含所有适用的内容。 +- 有些问题可能没有相关摘录——只需返回["N/A"]。 +- 不要在JSON中包含除列出的键之外的其他键。 +- 不要多次包含同一个键。 +- 使用中文回答。 + +预期的JSON键及其说明： +{schema_explanation} + +{examples} +要提取的文件： +{text} +""" + +SCHEMA_GUIDED_EXTRACTION_PROMPT = { + "en": TEMPLATE_EN, + "zh": TEMPLATE_ZH, +} From c5407b5a72fd52fb84dbb504218b89e661a46958 Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Thu, 6 Nov 2025 11:59:26 +0800 Subject: [PATCH 03/12] feat: add schema_guided_extraction config --- graphgen/configs/schema_guided_config.yaml | 22 ++++------------------ 1 file changed, 4 insertions(+), 18 deletions(-) diff --git a/graphgen/configs/schema_guided_config.yaml b/graphgen/configs/schema_guided_config.yaml index 6510a91d..5e71fa3d 100644 --- a/graphgen/configs/schema_guided_config.yaml +++ b/graphgen/configs/schema_guided_config.yaml @@ -1,22 +1,8 @@ read: input_file: resources/input_examples/jsonl_demo.jsonl # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples split: - chunk_size: 1024 # chunk size for text splitting + chunk_size: 10240 # chunk size for text splitting chunk_overlap: 100 # chunk overlap for text splitting -search: # web search configuration - enabled: false # whether to enable web search - search_types: ["google"] # search engine types, support: google, bing, uniprot, wikipedia -quiz_and_judge: # quiz and test whether the LLM masters the knowledge points - enabled: true - quiz_samples: 2 # number of quiz samples to generate - re_judge: false # whether to re-judge the existing quiz samples -partition: # graph partition configuration - method: ece # ece is a custom partition method based on comprehension loss - method_params: - max_units_per_community: 20 # max nodes and edges per community - min_units_per_community: 5 # min nodes and edges per community - max_tokens_per_community: 10240 # max tokens per community - unit_sampling: max_loss # unit sampling strategy, support: random, max_loss, min_loss -generate: - mode: aggregated # atomic, aggregated, multi_hop, cot, vqa - data_format: ChatML # Alpaca, Sharegpt, ChatML +extract: + method: schema_guided # extraction method, support: schema_guided + schema_file: resources/schemas/legal_contract.json # schema file path for schema_guided method From 33e1b271279e97629442c4edf2db408f92b6f34a Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Thu, 6 Nov 2025 12:02:22 +0800 Subject: [PATCH 04/12] feat: add extract_schema_guided.sh --- scripts/extract/extract_schema_guided.sh | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 scripts/extract/extract_schema_guided.sh diff --git a/scripts/extract/extract_schema_guided.sh b/scripts/extract/extract_schema_guided.sh new file mode 100644 index 00000000..a5d0564a --- /dev/null +++ b/scripts/extract/extract_schema_guided.sh @@ -0,0 +1,3 @@ +python3 -m graphgen.generate \ +--config_file graphgen/configs/schema_guided_config.yaml \ +--output_dir cache/ From 9de5b2f16c69e87600e39e73de75a587b8336cb5 Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Thu, 6 Nov 2025 12:25:32 +0800 Subject: [PATCH 05/12] wip: add schema_guided_extractor --- graphgen/bases/__init__.py | 1 + graphgen/bases/base_extractor.py | 22 +++++++ graphgen/models/extractor/__init__.py | 1 + .../extractor/key_information_extractor.py | 1 + .../extractor/schema_guided_extractor.py | 41 +++++++++++++ graphgen/operators/extract/__init__.py | 0 graphgen/operators/extract/extract.py | 47 +++++++++++++++ .../extraction/schemas/legal_contract.json | 58 +++++++++++++++++++ 8 files changed, 171 insertions(+) create mode 100644 graphgen/bases/base_extractor.py create mode 100644 graphgen/models/extractor/__init__.py create mode 100644 graphgen/models/extractor/key_information_extractor.py create mode 100644 graphgen/models/extractor/schema_guided_extractor.py create mode 100644 graphgen/operators/extract/__init__.py create mode 100644 graphgen/operators/extract/extract.py create mode 100644 graphgen/templates/extraction/schemas/legal_contract.json diff --git a/graphgen/bases/__init__.py b/graphgen/bases/__init__.py index ed452628..55048d6c 100644 --- a/graphgen/bases/__init__.py +++ b/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 diff --git a/graphgen/bases/base_extractor.py b/graphgen/bases/base_extractor.py new file mode 100644 index 00000000..45f6f55c --- /dev/null +++ b/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""" diff --git a/graphgen/models/extractor/__init__.py b/graphgen/models/extractor/__init__.py new file mode 100644 index 00000000..49db9fc2 --- /dev/null +++ b/graphgen/models/extractor/__init__.py @@ -0,0 +1 @@ +from .schema_guided_extractor import SchemaGuidedExtractor diff --git a/graphgen/models/extractor/key_information_extractor.py b/graphgen/models/extractor/key_information_extractor.py new file mode 100644 index 00000000..1a894fee --- /dev/null +++ b/graphgen/models/extractor/key_information_extractor.py @@ -0,0 +1 @@ +# TODO: text2json diff --git a/graphgen/models/extractor/schema_guided_extractor.py b/graphgen/models/extractor/schema_guided_extractor.py new file mode 100644 index 00000000..687574ed --- /dev/null +++ b/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 diff --git a/graphgen/operators/extract/__init__.py b/graphgen/operators/extract/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/graphgen/operators/extract/extract.py b/graphgen/operators/extract/extract.py new file mode 100644 index 00000000..533fd81e --- /dev/null +++ b/graphgen/operators/extract/extract.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 [] diff --git a/graphgen/templates/extraction/schemas/legal_contract.json b/graphgen/templates/extraction/schemas/legal_contract.json new file mode 100644 index 00000000..f10dee71 --- /dev/null +++ b/graphgen/templates/extraction/schemas/legal_contract.json @@ -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"] +} From 0cefddfe106c7faf8b9f376bbeb87113faa002af Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Thu, 6 Nov 2025 20:48:21 +0800 Subject: [PATCH 06/12] wip: add extract_info --- graphgen/bases/base_storage.py | 3 ++ graphgen/configs/schema_guided_config.yaml | 19 ++++---- graphgen/graphgen.py | 17 +++++++ .../extractor/schema_guided_extractor.py | 4 +- graphgen/models/storage/json_storage.py | 3 ++ graphgen/operators/__init__.py | 1 + graphgen/operators/extract/__init__.py | 1 + graphgen/operators/extract/extract.py | 47 ------------------- graphgen/operators/extract/extract_info.py | 47 +++++++++++++++++++ 9 files changed, 85 insertions(+), 57 deletions(-) delete mode 100644 graphgen/operators/extract/extract.py create mode 100644 graphgen/operators/extract/extract_info.py diff --git a/graphgen/bases/base_storage.py b/graphgen/bases/base_storage.py index f82e6f64..c8d515a3 100644 --- a/graphgen/bases/base_storage.py +++ b/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 diff --git a/graphgen/configs/schema_guided_config.yaml b/graphgen/configs/schema_guided_config.yaml index 5e71fa3d..9c5a1459 100644 --- a/graphgen/configs/schema_guided_config.yaml +++ b/graphgen/configs/schema_guided_config.yaml @@ -1,8 +1,11 @@ -read: - input_file: resources/input_examples/jsonl_demo.jsonl # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples -split: - chunk_size: 10240 # chunk size for text splitting - chunk_overlap: 100 # chunk overlap for text splitting -extract: - method: schema_guided # extraction method, support: schema_guided - schema_file: resources/schemas/legal_contract.json # schema file path for schema_guided method +pipeline: + - name: insert + 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: 10240 # chunk size for text splitting + chunk_overlap: 100 # chunk overlap for text splitting + + - name: extract + params: + method: schema_guided # extraction method, support: schema_guided + schema_file: resources/schemas/legal_contract.json # schema file path for schema_guided method diff --git a/graphgen/graphgen.py b/graphgen/graphgen.py index 5fa9f791..c8b76b5f 100644 --- a/graphgen/graphgen.py +++ b/graphgen/graphgen.py @@ -19,6 +19,7 @@ from graphgen.operators import ( build_kg, chunk_documents, + extract_info, generate_qas, init_llm, judge_statement, @@ -240,6 +241,22 @@ async def partition(self, partition_config: Dict): await self.partition_storage.upsert(batches) return batches + @op("extract", deps=["insert"]) + @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 + print(results) + @op("generate", deps=["insert", "partition"]) @async_to_sync_method async def generate(self, generate_config: Dict): diff --git a/graphgen/models/extractor/schema_guided_extractor.py b/graphgen/models/extractor/schema_guided_extractor.py index 687574ed..86e28694 100644 --- a/graphgen/models/extractor/schema_guided_extractor.py +++ b/graphgen/models/extractor/schema_guided_extractor.py @@ -37,5 +37,5 @@ def __init__(self, llm_client: BaseLLMWrapper, schema: dict): def build_prompt(self, text: str) -> str: pass - def extract(self, text_or_documents: str) -> dict: - pass + async def extract(self, chunk: dict) -> dict: + print(chunk) diff --git a/graphgen/models/storage/json_storage.py b/graphgen/models/storage/json_storage.py index 171eb988..c2801e0f 100644 --- a/graphgen/models/storage/json_storage.py +++ b/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} diff --git a/graphgen/operators/__init__.py b/graphgen/operators/__init__.py index ace334d6..a9ce24cd 100644 --- a/graphgen/operators/__init__.py +++ b/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 diff --git a/graphgen/operators/extract/__init__.py b/graphgen/operators/extract/__init__.py index e69de29b..ec576cb6 100644 --- a/graphgen/operators/extract/__init__.py +++ b/graphgen/operators/extract/__init__.py @@ -0,0 +1 @@ +from .extract_info import extract_info diff --git a/graphgen/operators/extract/extract.py b/graphgen/operators/extract/extract.py deleted file mode 100644 index 533fd81e..00000000 --- a/graphgen/operators/extract/extract.py +++ /dev/null @@ -1,47 +0,0 @@ -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 [] diff --git a/graphgen/operators/extract/extract_info.py b/graphgen/operators/extract/extract_info.py new file mode 100644 index 00000000..814884ec --- /dev/null +++ b/graphgen/operators/extract/extract_info.py @@ -0,0 +1,47 @@ +from typing import List + +import gradio as gr + +from graphgen.bases import BaseKVStorage, BaseLLMWrapper +from graphgen.bases.datatypes import Chunk +from graphgen.models.extractor import SchemaGuidedExtractor +from graphgen.utils import logger, run_concurrent + + +async def extract_info( + llm_client: BaseLLMWrapper, + chunk_storage: BaseKVStorage, + extract_config: dict, + progress_bar: gr.Progress = None, +): + """ + Extract information from chunks + :param llm_client: LLM client + :param chunk_storage: storage for chunks + :param extract_config + :param progress_bar + :return: extracted information + """ + + method = extract_config.get("method") + if method == "schema_guided": + schema = extract_config.get("schema") + extractor = SchemaGuidedExtractor(llm_client, schema) + else: + raise ValueError(f"Unsupported extraction method: {method}") + + chunks = await chunk_storage.get_all() + chunks = [{k: v} for k, v in chunks.items()] + logger.info(f"Start extracting information from {len(chunks)} chunks") + + results = await run_concurrent( + extractor.extract, + chunks, + desc="Extracting information", + unit="chunk", + progress_bar=progress_bar, + ) + + # TODO: 对results合并，去重 + + return [] From 2f1442b08345674ea6934d80c633fee9c318aea2 Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Fri, 7 Nov 2025 11:50:44 +0800 Subject: [PATCH 07/12] wip --- graphgen/graphgen.py | 3 +++ scripts/extract/extract_schema_guided.sh | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/graphgen/graphgen.py b/graphgen/graphgen.py index c8b76b5f..df10a872 100644 --- a/graphgen/graphgen.py +++ b/graphgen/graphgen.py @@ -295,3 +295,6 @@ async def clear(self): # TODO: add data filtering step here in the future # graph_gen.filter(filter_config=config["filter"]) + + +# TODO: 把insert拆成两个: read + build_kg，这样更合理 diff --git a/scripts/extract/extract_schema_guided.sh b/scripts/extract/extract_schema_guided.sh index a5d0564a..fe59aab8 100644 --- a/scripts/extract/extract_schema_guided.sh +++ b/scripts/extract/extract_schema_guided.sh @@ -1,3 +1,3 @@ -python3 -m graphgen.generate \ +python3 -m graphgen.run \ --config_file graphgen/configs/schema_guided_config.yaml \ --output_dir cache/ From f89a3203026528ef2ef2c9d0b069e7ec1d5dcae9 Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Fri, 7 Nov 2025 15:40:18 +0800 Subject: [PATCH 08/12] wip --- graphgen/bases/base_extractor.py | 2 +- graphgen/configs/schema_guided_config.yaml | 8 ++-- graphgen/graphgen.py | 2 +- .../extractor/schema_guided_extractor.py | 41 ++++++++++++++++++- graphgen/operators/extract/extract_info.py | 7 +++- .../extraction/schema_guided_extraction.py | 18 +++++++- graphgen/utils/__init__.py | 7 +++- graphgen/utils/hash.py | 5 +++ 8 files changed, 77 insertions(+), 13 deletions(-) diff --git a/graphgen/bases/base_extractor.py b/graphgen/bases/base_extractor.py index 45f6f55c..054e6e3a 100644 --- a/graphgen/bases/base_extractor.py +++ b/graphgen/bases/base_extractor.py @@ -14,7 +14,7 @@ def __init__(self, llm_client: BaseLLMWrapper): self.llm_client = llm_client @abstractmethod - def extract(self, text_or_documents: str) -> Any: + async def extract(self, chunk: dict) -> Any: """Extract information from the given text""" @abstractmethod diff --git a/graphgen/configs/schema_guided_config.yaml b/graphgen/configs/schema_guided_config.yaml index 9c5a1459..8452338a 100644 --- a/graphgen/configs/schema_guided_config.yaml +++ b/graphgen/configs/schema_guided_config.yaml @@ -1,11 +1,11 @@ pipeline: - - name: insert + - 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: 10240 # chunk size for text splitting + input_file: resources/input_examples/extract_demo.txt # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples + chunk_size: 20480 # chunk size for text splitting chunk_overlap: 100 # chunk overlap for text splitting - name: extract params: method: schema_guided # extraction method, support: schema_guided - schema_file: resources/schemas/legal_contract.json # schema file path for schema_guided method + schema_file: graphgen/templates/extraction/schemas/legal_contract.json # schema file path for schema_guided method diff --git a/graphgen/graphgen.py b/graphgen/graphgen.py index 1fe9ad92..98f7b75f 100644 --- a/graphgen/graphgen.py +++ b/graphgen/graphgen.py @@ -249,7 +249,7 @@ async def partition(self, partition_config: Dict): await self.partition_storage.upsert(batches) return batches - @op("extract", deps=["insert"]) + @op("extract", deps=["read"]) @async_to_sync_method async def extract(self, extract_config: Dict): logger.info("Extracting information from given chunks...") diff --git a/graphgen/models/extractor/schema_guided_extractor.py b/graphgen/models/extractor/schema_guided_extractor.py index 86e28694..6aa49f2a 100644 --- a/graphgen/models/extractor/schema_guided_extractor.py +++ b/graphgen/models/extractor/schema_guided_extractor.py @@ -1,4 +1,8 @@ +import json + from graphgen.bases import BaseExtractor, BaseLLMWrapper +from graphgen.templates import SCHEMA_GUIDED_EXTRACTION_PROMPT +from graphgen.utils import compute_dict_hash, detect_main_language class SchemaGuidedExtractor(BaseExtractor): @@ -33,9 +37,42 @@ class SchemaGuidedExtractor(BaseExtractor): 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: - pass + 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: - print(chunk) + 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): + return {} + main_keys_info = {key: extracted_info[key] for key in self.required_keys} + return {compute_dict_hash(main_keys_info): extracted_info} + except json.JSONDecodeError: + return {} + + # async def merge_extractions(self): diff --git a/graphgen/operators/extract/extract_info.py b/graphgen/operators/extract/extract_info.py index 814884ec..38483ea3 100644 --- a/graphgen/operators/extract/extract_info.py +++ b/graphgen/operators/extract/extract_info.py @@ -1,4 +1,4 @@ -from typing import List +import json import gradio as gr @@ -25,7 +25,9 @@ async def extract_info( method = extract_config.get("method") if method == "schema_guided": - schema = extract_config.get("schema") + schema_file = extract_config.get("schema_file") + with open(schema_file, "r", encoding="utf-8") as f: + schema = json.load(f) extractor = SchemaGuidedExtractor(llm_client, schema) else: raise ValueError(f"Unsupported extraction method: {method}") @@ -41,6 +43,7 @@ async def extract_info( unit="chunk", progress_bar=progress_bar, ) + print(results) # TODO: 对results合并，去重 diff --git a/graphgen/templates/extraction/schema_guided_extraction.py b/graphgen/templates/extraction/schema_guided_extraction.py index 997bd7fe..710900e5 100644 --- a/graphgen/templates/extraction/schema_guided_extraction.py +++ b/graphgen/templates/extraction/schema_guided_extraction.py @@ -11,7 +11,7 @@ - Consider the context of the entire document when determining relevance. - Do not be verbose, only respond with the correct format and information. - Some docs may have multiple relevant excerpts -- include all that apply. -- Some questions may have no relevant excerpts -- just return ["N/A"]. +- Some questions may have no relevant excerpts -- just return "". - Do not include additional JSON keys beyond the ones listed here. - Do not include the same key multiple times in the JSON. - Use English for your response. @@ -19,6 +19,13 @@ Expected JSON keys and explanation of what they are: {schema_explanation} +Expected format: +{{ + "key1": "value1", + "key2": "value2", + ... +}} + {examples} Document to extract from: @@ -37,7 +44,7 @@ - 在确定相关性时，考虑整份文件的上下文。 - 不要冗长，只需以正确的格式和信息进行回应。 - 有些文件可能有多个相关摘录——请包含所有适用的内容。 -- 有些问题可能没有相关摘录——只需返回["N/A"]。 +- 有些问题可能没有相关摘录——只需返回""。 - 不要在JSON中包含除列出的键之外的其他键。 - 不要多次包含同一个键。 - 使用中文回答。 @@ -45,6 +52,13 @@ 预期的JSON键及其说明： {schema_explanation} +预期格式： +{{ + "key1": "value1", + "key2": "value2", + ... +}} + {examples} 要提取的文件： {text} diff --git a/graphgen/utils/__init__.py b/graphgen/utils/__init__.py index eaf86762..d3e6df7b 100644 --- a/graphgen/utils/__init__.py +++ b/graphgen/utils/__init__.py @@ -9,7 +9,12 @@ split_string_by_multi_markers, write_json, ) -from .hash import compute_args_hash, compute_content_hash, compute_mm_hash +from .hash import ( + compute_args_hash, + compute_content_hash, + compute_dict_hash, + compute_mm_hash, +) from .help_nltk import NLTKHelper from .log import logger, parse_log, set_logger from .loop import create_event_loop diff --git a/graphgen/utils/hash.py b/graphgen/utils/hash.py index 59812e60..04ba96e7 100644 --- a/graphgen/utils/hash.py +++ b/graphgen/utils/hash.py @@ -21,3 +21,8 @@ def compute_mm_hash(item, prefix: str = ""): else: content = str(item) return prefix + md5(content.encode()).hexdigest() + + +def compute_dict_hash(d: dict, prefix: str = ""): + items = tuple(sorted(d.items())) + return prefix + md5(str(items).encode()).hexdigest() From c4f69b5616d95f7ee3382c36a4a81b815326c1ef Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Fri, 7 Nov 2025 18:22:58 +0800 Subject: [PATCH 09/12] fix: fix txt_reader, not splitting lines --- graphgen/models/reader/txt_reader.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/graphgen/models/reader/txt_reader.py b/graphgen/models/reader/txt_reader.py index e0a9e5c0..ec2ff747 100644 --- a/graphgen/models/reader/txt_reader.py +++ b/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) From 912508cee7e2202319d534375b3b29490549e989 Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Fri, 7 Nov 2025 18:23:45 +0800 Subject: [PATCH 10/12] docs: update example data --- .../templates/extraction/schemas/legal_contract.json | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/graphgen/templates/extraction/schemas/legal_contract.json b/graphgen/templates/extraction/schemas/legal_contract.json index f10dee71..7ed15090 100644 --- a/graphgen/templates/extraction/schemas/legal_contract.json +++ b/graphgen/templates/extraction/schemas/legal_contract.json @@ -3,11 +3,6 @@ "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"}, @@ -33,6 +28,11 @@ "items": {"type": "string"}, "description": "The start date of the lease." }, + "end_date": { + "type": "array", + "items": {"type": "string"}, + "description": "The end date of the lease." + }, "term_of_payment": { "type": "array", "items": {"type": "string"}, From 3f79260cab33ddd4ac5c6ffab0b4136f0bf113a8 Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Fri, 7 Nov 2025 18:38:27 +0800 Subject: [PATCH 11/12] feat: complete extract_info pipeline --- graphgen/configs/aggregated_config.yaml | 7 +++- graphgen/configs/atomic_config.yaml | 3 ++ graphgen/configs/cot_config.yaml | 7 +++- graphgen/configs/multi_hop_config.yaml | 3 ++ graphgen/configs/schema_guided_config.yaml | 8 +++- graphgen/configs/vqa_config.yaml | 7 +++- graphgen/graphgen.py | 39 +++++++++++++++---- .../extractor/schema_guided_extractor.py | 29 ++++++++++++-- graphgen/operators/extract/extract_info.py | 9 ++--- graphgen/operators/split/split_chunks.py | 12 +++--- 10 files changed, 94 insertions(+), 30 deletions(-) diff --git a/graphgen/configs/aggregated_config.yaml b/graphgen/configs/aggregated_config.yaml index 25ea691e..cffdffd8 100644 --- a/graphgen/configs/aggregated_config.yaml +++ b/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 diff --git a/graphgen/configs/atomic_config.yaml b/graphgen/configs/atomic_config.yaml index 94481c50..be109457 100644 --- a/graphgen/configs/atomic_config.yaml +++ b/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 diff --git a/graphgen/configs/cot_config.yaml b/graphgen/configs/cot_config.yaml index f7d2b735..7197f73a 100644 --- a/graphgen/configs/cot_config.yaml +++ b/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 diff --git a/graphgen/configs/multi_hop_config.yaml b/graphgen/configs/multi_hop_config.yaml index 3d00cc29..a0b75767 100644 --- a/graphgen/configs/multi_hop_config.yaml +++ b/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 diff --git a/graphgen/configs/schema_guided_config.yaml b/graphgen/configs/schema_guided_config.yaml index 8452338a..3944b326 100644 --- a/graphgen/configs/schema_guided_config.yaml +++ b/graphgen/configs/schema_guided_config.yaml @@ -2,8 +2,12 @@ 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 - chunk_size: 20480 # chunk size for text splitting - chunk_overlap: 100 # chunk overlap for text splitting + + - name: chunk + params: + chunk_size: 20480 + chunk_overlap: 2000 + separators: [] - name: extract params: diff --git a/graphgen/configs/vqa_config.yaml b/graphgen/configs/vqa_config.yaml index fb61cc52..d89800eb 100644 --- a/graphgen/configs/vqa_config.yaml +++ b/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 diff --git a/graphgen/graphgen.py b/graphgen/graphgen.py index 98f7b75f..dcdc535f 100644 --- a/graphgen/graphgen.py +++ b/graphgen/graphgen.py @@ -71,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" ) @@ -81,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 @@ -104,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( @@ -127,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): """ @@ -162,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( @@ -249,7 +268,7 @@ async def partition(self, partition_config: Dict): await self.partition_storage.upsert(batches) return batches - @op("extract", deps=["read"]) + @op("extract", deps=["chunk"]) @async_to_sync_method async def extract(self, extract_config: Dict): logger.info("Extracting information from given chunks...") @@ -263,7 +282,11 @@ async def extract(self, extract_config: Dict): if not results: logger.warning("No information extracted") return - print(results) + + 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 diff --git a/graphgen/models/extractor/schema_guided_extractor.py b/graphgen/models/extractor/schema_guided_extractor.py index 6aa49f2a..b7f15365 100644 --- a/graphgen/models/extractor/schema_guided_extractor.py +++ b/graphgen/models/extractor/schema_guided_extractor.py @@ -1,8 +1,9 @@ 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 +from graphgen.utils import compute_dict_hash, detect_main_language, logger class SchemaGuidedExtractor(BaseExtractor): @@ -69,10 +70,32 @@ async def extract(self, chunk: dict) -> dict: 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} - return {compute_dict_hash(main_keys_info): extracted_info} + 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): + 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]}{v}" + return merged diff --git a/graphgen/operators/extract/extract_info.py b/graphgen/operators/extract/extract_info.py index 38483ea3..98d8e98a 100644 --- a/graphgen/operators/extract/extract_info.py +++ b/graphgen/operators/extract/extract_info.py @@ -3,7 +3,6 @@ import gradio as gr from graphgen.bases import BaseKVStorage, BaseLLMWrapper -from graphgen.bases.datatypes import Chunk from graphgen.models.extractor import SchemaGuidedExtractor from graphgen.utils import logger, run_concurrent @@ -34,7 +33,7 @@ async def extract_info( chunks = await chunk_storage.get_all() chunks = [{k: v} for k, v in chunks.items()] - logger.info(f"Start extracting information from {len(chunks)} chunks") + logger.info("Start extracting information from %d chunks", len(chunks)) results = await run_concurrent( extractor.extract, @@ -43,8 +42,6 @@ async def extract_info( unit="chunk", progress_bar=progress_bar, ) - print(results) - # TODO: 对results合并，去重 - - return [] + results = await extractor.merge_extractions(results) + return results diff --git a/graphgen/operators/split/split_chunks.py b/graphgen/operators/split/split_chunks.py index e400ea63..f88b90e2 100644 --- a/graphgen/operators/split/split_chunks.py +++ b/graphgen/operators/split/split_chunks.py @@ -31,16 +31,18 @@ def split_chunks(text: str, language: str = "en", **kwargs) -> list: f"Unsupported language: {language}. " f"Supported languages are: {list(_MAPPING.keys())}" ) - splitter = _get_splitter(language, frozenset(kwargs.items())) + frozen_kwargs = frozenset( + (k, tuple(v) if isinstance(v, list) else v) for k, v in kwargs.items() + ) + splitter = _get_splitter(language, frozen_kwargs) return splitter.split_text(text) async def chunk_documents( new_docs: dict, - chunk_size: int = 1024, - chunk_overlap: int = 100, tokenizer_instance: Tokenizer = None, progress_bar=None, + **kwargs, ) -> dict: inserting_chunks = {} cur_index = 1 @@ -51,11 +53,11 @@ async def chunk_documents( doc_type = doc.get("type") if doc_type == "text": doc_language = detect_main_language(doc["content"]) + text_chunks = split_chunks( doc["content"], language=doc_language, - chunk_size=chunk_size, - chunk_overlap=chunk_overlap, + **kwargs, ) chunks = { From 2b6934c9c786f9f4b4c9c54b1de1ffac21b2a179 Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Fri, 7 Nov 2025 18:41:01 +0800 Subject: [PATCH 12/12] fix: update webui --- webui/app.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/webui/app.py b/webui/app.py index 843486fa..2e74e203 100644 --- a/webui/app.py +++ b/webui/app.py @@ -103,6 +103,11 @@ def sum_tokens(client): "name": "read", "params": { "input_file": params.upload_file, + }, + }, + { + "name": "chunk", + "params": { "chunk_size": params.chunk_size, "chunk_overlap": params.chunk_overlap, },