From 1c76ca247c2b74c3d1c6e20ed9bda55307a51148 Mon Sep 17 00:00:00 2001 From: Leon-L20 Date: Mon, 19 Oct 2026 02:44:44 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20=E9=80=9A=E7=94=A8=E7=89=88=E6=9C=AC?= =?UTF-8?q?=E5=8C=96=20prompt=20profile=E3=80=81=E8=BE=93=E5=87=BA?= =?UTF-8?q?=E9=97=A8=E6=B3=A8=E5=86=8C=E8=A1=A8=E3=80=81=E7=94=9F=E6=88=90?= =?UTF-8?q?=E5=B9=B6=E5=8F=91=E4=B8=8A=E9=99=90=E4=B8=8E=20env=20=E9=97=A8?= =?UTF-8?q?=E6=8E=A7=20LLM=20=E5=AE=A2=E6=88=B7=E7=AB=AF=E8=A1=A5=E4=B8=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - prompt profile：BaseGenerator 新增 TEMPLATE_KEY/apply_prompt_profile/template， 全部生成器 build_prompt 覆写改为实例方法；prompt_profiles.py 结构校验 + fail-closed；内置 example.v1 通用 profile（内容与官方模板逐字一致） - 输出门注册表：gates.py register_gate/GATES + 内置 non_empty_answer 通用门， GenerateService 增加 prompt_profile/output_gate/llm_concurrency 选项 - 生成并发：llm_concurrency.py 事件循环内有界并发包装 - LLM 客户端 4 个 env 门控补丁：GRAPHGEN_EMPTY_RETRY 空响应有界重试、 GRAPHGEN_EMPTY_EXTRACTION_PREFILL 空抽取 assistant prefill 重试、 GRAPHGEN_DISABLE_THINKING logprob 调用关思考、可选 assistant_prefill 参数； GRAPHGEN_LLM_ACTOR_MAX_CONCURRENCY 提升 Ray actor 并发；logger 空上下文回退 - 未设置任何 env / 选项时行为与官方默认逐字一致；不改动包身份与版本号 Co-Authored-By: Claude Opus 5.5 --- MANIFEST.in | 1 + graphgen/bases/base_generator.py | 36 ++++- graphgen/common/init_llm.py | 20 ++- .../extractor/schema_guided_extractor.py | 15 ++ .../models/generator/aggregated_generator.py | 15 +- graphgen/models/generator/atomic_generator.py | 8 +- graphgen/models/generator/cot_generator.py | 12 +- .../generator/fill_in_blank_generator.py | 6 +- .../masked_fill_in_blank_generator.py | 8 +- .../generator/multi_answer_generator.py | 6 +- .../generator/multi_choice_generator.py | 6 +- .../models/generator/multi_hop_generator.py | 8 +- graphgen/models/generator/quiz_generator.py | 3 +- .../models/generator/true_false_generator.py | 6 +- graphgen/models/generator/vqa_generator.py | 8 +- graphgen/models/llm/api/openai_client.py | 75 +++++++--- graphgen/operators/generate/gates.py | 48 +++++++ .../operators/generate/generate_service.py | 136 +++++++++++------- .../operators/generate/llm_concurrency.py | 31 ++++ .../profiles/example.v1/aggregated.json | 12 ++ .../templates/profiles/example.v1/atomic.json | 4 + .../templates/profiles/example.v1/cot.json | 10 ++ .../profiles/example.v1/fill_in_blank.json | 4 + .../profiles/example.v1/multi_answer.json | 4 + .../profiles/example.v1/multi_choice.json | 4 + .../profiles/example.v1/multi_hop.json | 4 + .../profiles/example.v1/true_false.json | 4 + .../templates/profiles/example.v1/vqa.json | 4 + graphgen/templates/prompt_profiles.py | 119 +++++++++++++++ graphgen/utils/log.py | 6 +- setup.py | 7 +- tests/unit_tests/__init__.py | 0 .../llm/api/test_openai_client_patches.py | 128 +++++++++++++++++ tests/unit_tests/operators/__init__.py | 0 .../unit_tests/operators/generate/__init__.py | 0 .../operators/generate/test_gates.py | 34 +++++ .../generate/test_llm_concurrency.py | 34 +++++ tests/unit_tests/templates/__init__.py | 0 .../templates/test_prompt_profiles.py | 86 +++++++++++ tests/unit_tests/utils/__init__.py | 0 tests/unit_tests/utils/test_log.py | 13 ++ 41 files changed, 791 insertions(+), 134 deletions(-) create mode 100644 graphgen/operators/generate/gates.py create mode 100644 graphgen/operators/generate/llm_concurrency.py create mode 100644 graphgen/templates/profiles/example.v1/aggregated.json create mode 100644 graphgen/templates/profiles/example.v1/atomic.json create mode 100644 graphgen/templates/profiles/example.v1/cot.json create mode 100644 graphgen/templates/profiles/example.v1/fill_in_blank.json create mode 100644 graphgen/templates/profiles/example.v1/multi_answer.json create mode 100644 graphgen/templates/profiles/example.v1/multi_choice.json create mode 100644 graphgen/templates/profiles/example.v1/multi_hop.json create mode 100644 graphgen/templates/profiles/example.v1/true_false.json create mode 100644 graphgen/templates/profiles/example.v1/vqa.json create mode 100644 graphgen/templates/prompt_profiles.py create mode 100644 tests/unit_tests/__init__.py create mode 100644 tests/unit_tests/models/llm/api/test_openai_client_patches.py create mode 100644 tests/unit_tests/operators/__init__.py create mode 100644 tests/unit_tests/operators/generate/__init__.py create mode 100644 tests/unit_tests/operators/generate/test_gates.py create mode 100644 tests/unit_tests/operators/generate/test_llm_concurrency.py create mode 100644 tests/unit_tests/templates/__init__.py create mode 100644 tests/unit_tests/templates/test_prompt_profiles.py create mode 100644 tests/unit_tests/utils/__init__.py create mode 100644 tests/unit_tests/utils/test_log.py diff --git a/MANIFEST.in b/MANIFEST.in index d1a9003b..0b3c24dc 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,3 +1,4 @@ include requirements.txt include graphgen/configs/*.yaml recursive-include graphgen/configs *.yaml +recursive-include graphgen/templates/profiles *.json diff --git a/graphgen/bases/base_generator.py b/graphgen/bases/base_generator.py index 42e55180..dc3870da 100644 --- a/graphgen/bases/base_generator.py +++ b/graphgen/bases/base_generator.py @@ -1,23 +1,45 @@ from abc import ABC, abstractmethod -from typing import Any +from typing import Any, Optional from graphgen.bases.base_llm_wrapper import BaseLLMWrapper class BaseGenerator(ABC): - """ - Generate QAs based on given prompts. - """ + """Generate QAs based on given prompts.""" + + TEMPLATE_KEY: Optional[str] = None def __init__(self, llm_client: BaseLLMWrapper): self.llm_client = llm_client + self._profile_templates: Optional[dict] = None + + def apply_prompt_profile(self, profile: str) -> None: + """Attach a versioned prompt profile to this generator instance.""" + from graphgen.templates.prompt_profiles import load_profile_template + + if self.TEMPLATE_KEY is None: + raise ValueError( + f"prompt_profile_unsupported_generator:{type(self).__name__}" + ) + self._profile_templates = load_profile_template(profile, self.TEMPLATE_KEY) + + def template(self, language: str, key: Optional[str] = None) -> str: + """Resolve the effective template for this generator.""" + from graphgen.templates.prompt_profiles import OFFICIAL_TEMPLATES + + if self.TEMPLATE_KEY is None: + raise ValueError(f"prompt_template_key_missing:{type(self).__name__}") + if self._profile_templates is not None: + node = self._profile_templates[language] + else: + node = OFFICIAL_TEMPLATES[self.TEMPLATE_KEY][language] + return node[key] if key is not None else node - @staticmethod @abstractmethod def build_prompt( - batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] + self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] ) -> str: - """Build prompt for LLM based on the given batch""" + """Build prompt for LLM based on the given batch.""" @staticmethod @abstractmethod diff --git a/graphgen/common/init_llm.py b/graphgen/common/init_llm.py index 56bffedf..17e8e626 100644 --- a/graphgen/common/init_llm.py +++ b/graphgen/common/init_llm.py @@ -129,19 +129,27 @@ def create_llm( ) actor_name = f"Actor_LLM_{model_type}" + # Ray async actors default to max_concurrency=1: all workers' LLM + # requests serialize inside a single actor, which becomes the + # generation-throughput bottleneck. GRAPHGEN_LLM_ACTOR_MAX_CONCURRENCY + # raises it (async methods are concurrency-safe); unset keeps the + # official default behavior. + max_concurrency = int(os.environ.get("GRAPHGEN_LLM_ACTOR_MAX_CONCURRENCY", "1")) + actor_options: Dict[str, Any] = { + "name": actor_name, + "num_gpus": float(config.pop("num_gpus", 0)), + "get_if_exists": True, + } + if max_concurrency > 1: + actor_options["max_concurrency"] = max_concurrency try: actor_handle = ray.get_actor(actor_name) print(f"Using existing Ray actor: {actor_name}") except ValueError: print(f"Creating Ray actor for LLM {model_type} with backend {backend}.") - num_gpus = float(config.pop("num_gpus", 0)) actor_handle = ( ray.remote(LLMServiceActor) - .options( - name=actor_name, - num_gpus=num_gpus, - get_if_exists=True, - ) + .options(**actor_options) .remote(backend, config) ) diff --git a/graphgen/models/extractor/schema_guided_extractor.py b/graphgen/models/extractor/schema_guided_extractor.py index dee6e3e3..89f7fc7d 100644 --- a/graphgen/models/extractor/schema_guided_extractor.py +++ b/graphgen/models/extractor/schema_guided_extractor.py @@ -1,4 +1,5 @@ import json +import os from graphgen.bases import BaseExtractor, BaseLLMWrapper, Chunk from graphgen.templates import SCHEMA_GUIDED_EXTRACTION_PROMPT @@ -63,6 +64,20 @@ async def extract(self, chunk: Chunk) -> dict: prompt = self.build_prompt(text) response = await self.llm_client.generate_answer(prompt) + + # Strict-prompt extraction backends occasionally hit EOS as the first + # token. When GRAPHGEN_EMPTY_EXTRACTION_PREFILL=1 and the first + # response is empty, retry once with a deterministic assistant prefill + # that seeds the opening of the expected JSON output; non-empty + # responses are never retried or modified. + if ( + not response.strip() + and os.environ.get("GRAPHGEN_EMPTY_EXTRACTION_PREFILL", "0") == "1" + ): + logger.warning("Empty extraction response; retrying with assistant prefill") + response = await self.llm_client.generate_answer( + prompt, assistant_prefill="{" + ) try: extracted_info = json.loads(response) # Ensure all required keys are present diff --git a/graphgen/models/generator/aggregated_generator.py b/graphgen/models/generator/aggregated_generator.py index 3f223325..0f4a4f00 100644 --- a/graphgen/models/generator/aggregated_generator.py +++ b/graphgen/models/generator/aggregated_generator.py @@ -2,7 +2,6 @@ from typing import Any, Optional from graphgen.bases import BaseGenerator -from graphgen.templates import AGGREGATED_GENERATION_PROMPT from graphgen.utils import detect_main_language, logger @@ -14,9 +13,10 @@ class AggregatedGenerator(BaseGenerator): 2. question generation: Generate relevant questions based on the rephrased text. """ - @staticmethod + TEMPLATE_KEY = "aggregated" + def build_prompt( - batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] + self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] ) -> str: """ Build prompts for REPHRASE. @@ -51,7 +51,7 @@ def build_prompt( # for index, text in enumerate(original_text) # ] # ) - prompt = AGGREGATED_GENERATION_PROMPT[language]["ANSWER_REPHRASING"].format( + prompt = self.template(language, "ANSWER_REPHRASING").format( entities=entities_str, relationships=relations_str ) return prompt @@ -73,17 +73,14 @@ def parse_rephrased_text(response: str) -> Optional[str]: return None return rephrased_text.strip('"').strip("'") - @staticmethod - def _build_prompt_for_question_generation(answer: str) -> str: + def _build_prompt_for_question_generation(self, answer: str) -> str: """ Build prompts for QUESTION GENERATION. :param answer: :return: """ language = detect_main_language(answer) - prompt = AGGREGATED_GENERATION_PROMPT[language]["QUESTION_GENERATION"].format( - answer=answer - ) + prompt = self.template(language, "QUESTION_GENERATION").format(answer=answer) return prompt @staticmethod diff --git a/graphgen/models/generator/atomic_generator.py b/graphgen/models/generator/atomic_generator.py index d045b0da..6f9d5e67 100644 --- a/graphgen/models/generator/atomic_generator.py +++ b/graphgen/models/generator/atomic_generator.py @@ -2,14 +2,14 @@ from typing import Any from graphgen.bases import BaseGenerator -from graphgen.templates import ATOMIC_GENERATION_PROMPT from graphgen.utils import detect_main_language, logger class AtomicGenerator(BaseGenerator): - @staticmethod + TEMPLATE_KEY = "atomic" + def build_prompt( - batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] + self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] ) -> str: nodes, edges = batch context = "" @@ -19,7 +19,7 @@ def build_prompt( context += f"- {edge[0]} - {edge[1]}: {edge[2]['description']}\n" language = detect_main_language(context) - prompt = ATOMIC_GENERATION_PROMPT[language].format(context=context) + prompt = self.template(language).format(context=context) return prompt @staticmethod diff --git a/graphgen/models/generator/cot_generator.py b/graphgen/models/generator/cot_generator.py index 88d04324..dfd779f2 100644 --- a/graphgen/models/generator/cot_generator.py +++ b/graphgen/models/generator/cot_generator.py @@ -2,14 +2,14 @@ from typing import Any from graphgen.bases import BaseGenerator -from graphgen.templates import COT_GENERATION_PROMPT from graphgen.utils import detect_main_language, logger class CoTGenerator(BaseGenerator): - @staticmethod + TEMPLATE_KEY = "cot" + def build_prompt( - batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] + self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] ) -> str: """ Build prompts for COT Template Design. @@ -30,13 +30,13 @@ def build_prompt( ] ) language = detect_main_language(entities_str + relationships_str) - prompt = COT_GENERATION_PROMPT[language]["COT_TEMPLATE_DESIGN"].format( + prompt = self.template(language, "COT_TEMPLATE_DESIGN").format( entities=entities_str, relationships=relationships_str ) return prompt - @staticmethod def build_prompt_for_cot_generation( + self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]], question: str, reasoning_path: str, @@ -58,7 +58,7 @@ def build_prompt_for_cot_generation( ] ) language = detect_main_language(entities_str + relationships_str) - prompt = COT_GENERATION_PROMPT[language]["COT_GENERATION"].format( + prompt = self.template(language, "COT_GENERATION").format( entities=entities_str, relationships=relationships_str, question=question, diff --git a/graphgen/models/generator/fill_in_blank_generator.py b/graphgen/models/generator/fill_in_blank_generator.py index a26daf3e..53de8f60 100644 --- a/graphgen/models/generator/fill_in_blank_generator.py +++ b/graphgen/models/generator/fill_in_blank_generator.py @@ -2,11 +2,12 @@ from typing import Any from graphgen.bases import BaseGenerator -from graphgen.templates import FILL_IN_BLANK_GENERATION_PROMPT from graphgen.utils import detect_main_language, logger class FillInBlankGenerator(BaseGenerator): + TEMPLATE_KEY = "fill_in_blank" + def __init__(self, llm_client, num_of_questions) -> None: super().__init__(llm_client) self.num_of_questions = num_of_questions @@ -72,7 +73,6 @@ def parse_response(response: str) -> list[dict]: return qa_pairs - # pylint: disable=W0221 def build_prompt( self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] ) -> str: @@ -92,7 +92,7 @@ def build_prompt( ) context = entities_str + "\n" + relationships_str language = detect_main_language(entities_str + relationships_str) - prompt = FILL_IN_BLANK_GENERATION_PROMPT[language].format( + prompt = self.template(language).format( context=context, num_of_questions=self.num_of_questions, ) diff --git a/graphgen/models/generator/masked_fill_in_blank_generator.py b/graphgen/models/generator/masked_fill_in_blank_generator.py index 355cadcf..8b0619fd 100644 --- a/graphgen/models/generator/masked_fill_in_blank_generator.py +++ b/graphgen/models/generator/masked_fill_in_blank_generator.py @@ -3,7 +3,6 @@ from typing import Any, Optional from graphgen.bases import BaseGenerator -from graphgen.templates import AGGREGATED_GENERATION_PROMPT from graphgen.utils import detect_main_language, logger random.seed(42) @@ -16,9 +15,10 @@ class MaskedFillInBlankGenerator(BaseGenerator): 2. mask: Randomly select a node from the input nodes, and then mask the name of the node in the rephrased text. """ - @staticmethod + TEMPLATE_KEY = "masked_fill_in_blank" + def build_prompt( - batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] + self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] ) -> str: """ Build prompts for REPHRASE. @@ -53,7 +53,7 @@ def build_prompt( # for index, text in enumerate(original_text) # ] # ) - prompt = AGGREGATED_GENERATION_PROMPT[language]["ANSWER_REPHRASING"].format( + prompt = self.template(language, "ANSWER_REPHRASING").format( entities=entities_str, relationships=relations_str ) return prompt diff --git a/graphgen/models/generator/multi_answer_generator.py b/graphgen/models/generator/multi_answer_generator.py index a341a4fd..df9b3c5d 100644 --- a/graphgen/models/generator/multi_answer_generator.py +++ b/graphgen/models/generator/multi_answer_generator.py @@ -2,11 +2,12 @@ from typing import Any from graphgen.bases import BaseGenerator -from graphgen.templates import MAQ_GENERATION_PROMPT from graphgen.utils import detect_main_language, logger class MultiAnswerGenerator(BaseGenerator): + TEMPLATE_KEY = "multi_answer" + def __init__(self, llm_client, num_of_questions) -> None: super().__init__(llm_client) self.num_of_questions = num_of_questions @@ -93,7 +94,6 @@ def parse_response(response: str) -> list[dict]: return qa_pairs - # pylint: disable=W0221 def build_prompt( self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] ) -> str: @@ -113,7 +113,7 @@ def build_prompt( ) context = entities_str + "\n" + relationships_str language = detect_main_language(entities_str + relationships_str) - prompt = MAQ_GENERATION_PROMPT[language].format( + prompt = self.template(language).format( context=context, num_of_questions=self.num_of_questions, ) diff --git a/graphgen/models/generator/multi_choice_generator.py b/graphgen/models/generator/multi_choice_generator.py index 0c48b76d..0853faf2 100644 --- a/graphgen/models/generator/multi_choice_generator.py +++ b/graphgen/models/generator/multi_choice_generator.py @@ -2,11 +2,12 @@ from typing import Any from graphgen.bases import BaseGenerator -from graphgen.templates import MCQ_GENERATION_PROMPT from graphgen.utils import detect_main_language, logger class MultiChoiceGenerator(BaseGenerator): + TEMPLATE_KEY = "multi_choice" + def __init__(self, llm_client, num_of_questions) -> None: super().__init__(llm_client) self.num_of_questions = num_of_questions @@ -91,7 +92,6 @@ def parse_response(response: str) -> list[dict]: return qa_pairs - # pylint: disable=W0221 def build_prompt( self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] ) -> str: @@ -111,7 +111,7 @@ def build_prompt( ) context = entities_str + "\n" + relationships_str language = detect_main_language(entities_str + relationships_str) - prompt = MCQ_GENERATION_PROMPT[language].format( + prompt = self.template(language).format( context=context, num_of_questions=self.num_of_questions, ) diff --git a/graphgen/models/generator/multi_hop_generator.py b/graphgen/models/generator/multi_hop_generator.py index a19082b9..eb654af4 100644 --- a/graphgen/models/generator/multi_hop_generator.py +++ b/graphgen/models/generator/multi_hop_generator.py @@ -2,14 +2,14 @@ from typing import Any from graphgen.bases import BaseGenerator -from graphgen.templates import MULTI_HOP_GENERATION_PROMPT from graphgen.utils import detect_main_language, logger class MultiHopGenerator(BaseGenerator): - @staticmethod + TEMPLATE_KEY = "multi_hop" + def build_prompt( - batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] + self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] ) -> str: nodes, edges = batch entities_str = "\n".join( @@ -26,7 +26,7 @@ def build_prompt( ] ) language = detect_main_language(entities_str + relationships_str) - prompt = MULTI_HOP_GENERATION_PROMPT[language].format( + prompt = self.template(language).format( entities=entities_str, relationships=relationships_str ) return prompt diff --git a/graphgen/models/generator/quiz_generator.py b/graphgen/models/generator/quiz_generator.py index 2b1cf74a..073988ff 100644 --- a/graphgen/models/generator/quiz_generator.py +++ b/graphgen/models/generator/quiz_generator.py @@ -10,9 +10,8 @@ class QuizGenerator(BaseGenerator): Quiz Generator rephrases given descriptions to create quiz questions. """ - @staticmethod def build_prompt( - batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] + self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] ) -> str: """ Build prompt for rephrasing the description. diff --git a/graphgen/models/generator/true_false_generator.py b/graphgen/models/generator/true_false_generator.py index 1a1fa0d3..4312319f 100644 --- a/graphgen/models/generator/true_false_generator.py +++ b/graphgen/models/generator/true_false_generator.py @@ -2,11 +2,12 @@ from typing import Any from graphgen.bases import BaseGenerator -from graphgen.templates import TF_GENERATION_PROMPT from graphgen.utils import detect_main_language, logger class TrueFalseGenerator(BaseGenerator): + TEMPLATE_KEY = "true_false" + def __init__(self, llm_client, num_of_questions) -> None: super().__init__(llm_client) self.num_of_questions = num_of_questions @@ -64,7 +65,6 @@ def parse_response(response: str) -> list[dict]: return qa_pairs - # pylint: disable=W0221 def build_prompt( self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] ) -> str: @@ -84,7 +84,7 @@ def build_prompt( ) context = entities_str + "\n" + relationships_str language = detect_main_language(entities_str + relationships_str) - prompt = TF_GENERATION_PROMPT[language].format( + prompt = self.template(language).format( context=context, num_of_questions=self.num_of_questions, ) diff --git a/graphgen/models/generator/vqa_generator.py b/graphgen/models/generator/vqa_generator.py index 723bd2a6..f597bb23 100644 --- a/graphgen/models/generator/vqa_generator.py +++ b/graphgen/models/generator/vqa_generator.py @@ -3,14 +3,14 @@ from typing import Any from graphgen.bases import BaseGenerator -from graphgen.templates import VQA_GENERATION_PROMPT from graphgen.utils import detect_main_language, logger class VQAGenerator(BaseGenerator): - @staticmethod + TEMPLATE_KEY = "vqa" + def build_prompt( - batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] + self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] ) -> str: nodes, edges = batch entities_str = "\n".join( @@ -27,7 +27,7 @@ def build_prompt( ] ) language = detect_main_language(entities_str + relationships_str) - prompt = VQA_GENERATION_PROMPT[language].format( + prompt = self.template(language).format( entities=entities_str, relationships=relationships_str ) return prompt diff --git a/graphgen/models/llm/api/openai_client.py b/graphgen/models/llm/api/openai_client.py index 532b981c..eb553226 100644 --- a/graphgen/models/llm/api/openai_client.py +++ b/graphgen/models/llm/api/openai_client.py @@ -1,4 +1,5 @@ import math +import os from typing import Any, Dict, List, Optional import openai @@ -13,6 +14,7 @@ from graphgen.bases.base_llm_wrapper import BaseLLMWrapper from graphgen.bases.datatypes import Token from graphgen.models.llm.limitter import RPM, TPM +from graphgen.utils import logger def get_top_response_tokens(response: openai.ChatCompletion) -> List[Token]: @@ -86,7 +88,9 @@ def __post_init__(self): else: raise ValueError(f"Unsupported backend {self.backend}. Use 'openai_api' or 'azure_openai_api'.") - def _pre_generate(self, text: str, history: List[str]) -> Dict: + def _pre_generate( + self, text: str, history: List[str], assistant_prefill: Optional[str] = None + ) -> Dict: kwargs = { "temperature": self.temperature, "top_p": self.top_p, @@ -106,6 +110,11 @@ def _pre_generate(self, text: str, history: List[str]) -> Dict: assert len(history) % 2 == 0, "History should have even number of elements." messages = history + messages + # Optional deterministic assistant prefill for strict-prompt backends + # that occasionally emit EOS as the first token. + if assistant_prefill is not None: + messages.append({"role": "assistant", "content": assistant_prefill}) + kwargs["messages"] = messages return kwargs @@ -130,6 +139,16 @@ async def generate_topk_per_token( # Limit max_tokens to 1 to avoid long completions kwargs["max_tokens"] = 1 + # Thinking models spend the first tokens inside content, which + # makes next-token yes/no logprobs meaningless for judge/loss calls. + # GRAPHGEN_DISABLE_THINKING=1 asks the backend to disable thinking for + # these logprob-based calls only. + if os.environ.get("GRAPHGEN_DISABLE_THINKING", "0") == "1": + kwargs["extra_body"] = { + **kwargs.get("extra_body", {}), + "chat_template_kwargs": {"enable_thinking": False}, + } + completion = await self.client.chat.completions.create( # pylint: disable=E1125 model=self.model, **kwargs ) @@ -149,31 +168,53 @@ async def generate_answer( self, text: str, history: Optional[List[str]] = None, + assistant_prefill: Optional[str] = None, **extra: Any, ) -> str: - kwargs = self._pre_generate(text, history) + kwargs = self._pre_generate(text, history, assistant_prefill) prompt_tokens = 0 for message in kwargs["messages"]: prompt_tokens += len(self.tokenizer.encode(message["content"])) estimated_tokens = prompt_tokens + kwargs["max_tokens"] - if self.request_limit: - await self.rpm.wait(silent=True) - await self.tpm.wait(estimated_tokens, silent=True) - - completion = await self.client.chat.completions.create( # pylint: disable=E1125 - model=self.model, **kwargs - ) - if hasattr(completion, "usage"): - self.token_usage.append( - { - "prompt_tokens": completion.usage.prompt_tokens, - "completion_tokens": completion.usage.completion_tokens, - "total_tokens": completion.usage.total_tokens, - } + # Bounded retry on empty responses. Thinking backends occasionally + # return only content that is stripped to nothing downstream; + # retry instead of propagating an empty string that would fail + # parsing. GRAPHGEN_EMPTY_RETRY=0 keeps the official single-attempt + # behavior. + try: + empty_retries = max(0, int(os.environ.get("GRAPHGEN_EMPTY_RETRY", "0"))) + except ValueError: + empty_retries = 0 + + content = "" + for attempt in range(empty_retries + 1): + if self.request_limit: + await self.rpm.wait(silent=True) + await self.tpm.wait(estimated_tokens, silent=True) + + completion = await self.client.chat.completions.create( # pylint: disable=E1125 + model=self.model, **kwargs + ) + if hasattr(completion, "usage"): + self.token_usage.append( + { + "prompt_tokens": completion.usage.prompt_tokens, + "completion_tokens": completion.usage.completion_tokens, + "total_tokens": completion.usage.total_tokens, + } + ) + content = self.filter_think_tags(completion.choices[0].message.content) + if content.strip(): + break + logger.warning( + "Empty LLM response (attempt %d/%d, model=%s)", + attempt + 1, + empty_retries + 1, + self.model, ) - return self.filter_think_tags(completion.choices[0].message.content) + return content async def generate_inputs_prob( self, text: str, history: Optional[List[str]] = None, **extra: Any diff --git a/graphgen/operators/generate/gates.py b/graphgen/operators/generate/gates.py new file mode 100644 index 00000000..49841916 --- /dev/null +++ b/graphgen/operators/generate/gates.py @@ -0,0 +1,48 @@ +"""Deterministic output gates for generated QA pairs. + +A gate is a pure, synchronous filter applied to the parsed QA list of one +generate call before results reach the native cache. Gates never modify or +fabricate model output; they only drop pairs that violate a declared +deterministic contract. Gates are registered by name via +:func:`register_gate` and selected by name through ``GenerateService``. +""" + +from __future__ import annotations + +from typing import Any, Callable + +GATES: dict[str, Callable[[Any], list[dict]]] = {} + + +def register_gate(name: str) -> Callable: + """Register a gate function under a unique name.""" + + def _decorator(func: Callable[[Any], list[dict]]) -> Callable[[Any], list[dict]]: + if name in GATES: + raise ValueError(f"output_gate_already_registered:{name}") + GATES[name] = func + return func + + return _decorator + + +@register_gate("non_empty_answer") +def non_empty_answer(qa_pairs: Any) -> list[dict]: + """Generic built-in gate: keep only pairs with a non-empty answer.""" + if not isinstance(qa_pairs, list): + return [] + return [ + pair + for pair in qa_pairs + if isinstance(pair, dict) + and isinstance(pair.get("question"), str) + and isinstance(pair.get("answer"), str) + and pair["answer"].strip() + ] + + +def apply_output_gate(name: str, qa_pairs: Any) -> list[dict]: + """Apply a registered gate by name; fail-closed on unknown gates.""" + if name not in GATES: + raise ValueError(f"output_gate_not_registered:{name}") + return GATES[name](qa_pairs) diff --git a/graphgen/operators/generate/generate_service.py b/graphgen/operators/generate/generate_service.py index 18ce1d43..fe6a6058 100644 --- a/graphgen/operators/generate/generate_service.py +++ b/graphgen/operators/generate/generate_service.py @@ -6,9 +6,83 @@ from graphgen.utils import logger, run_concurrent +def _gated_generate(generate, gate): + """Wrap an async generate so its parsed QA list passes through ``gate``.""" + + async def _wrapper(batch: list) -> list[dict]: + return gate(await generate(batch)) + + return _wrapper + + +#: method -> (generator class name, default num_of_questions or None) +_GENERATOR_SPECS = { + "atomic": ("AtomicGenerator", None), + "aggregated": ("AggregatedGenerator", None), + "multi_hop": ("MultiHopGenerator", None), + "cot": ("CoTGenerator", None), + "vqa": ("VQAGenerator", None), + "multi_choice": ("MultiChoiceGenerator", 5), + "multi_answer": ("MultiAnswerGenerator", 3), + "fill_in_blank": ("FillInBlankGenerator", 5), + "masked_fill_in_blank": ("MaskedFillInBlankGenerator", None), + "true_false": ("TrueFalseGenerator", 5), +} + + +def _create_generator(llm_client: BaseLLMWrapper, method: str, generate_kwargs: dict): + """Instantiate the generator selected by ``method``.""" + if method not in _GENERATOR_SPECS: + raise ValueError(f"Unsupported generation mode: {method}") + class_name, default_num_of_questions = _GENERATOR_SPECS[method] + from graphgen import models as graphgen_models + + generator_cls = getattr(graphgen_models, class_name) + if default_num_of_questions is None: + return generator_cls(llm_client) + return generator_cls( + llm_client, + num_of_questions=generate_kwargs.get( + "num_of_questions", default_num_of_questions + ), + ) + + +def _configure_generator( + generator, + prompt_profile: str | None, + output_gate: str | None, + llm_concurrency: int | None, +) -> None: + """Apply optional profile / concurrency / gate configuration.""" + if prompt_profile: + generator.apply_prompt_profile(prompt_profile) + if llm_concurrency is not None: + from graphgen.operators.generate.llm_concurrency import ( + bound_llm_concurrency, + ) + + generator.generate = bound_llm_concurrency( + generator.generate, llm_concurrency + ) + if output_gate: + from graphgen.operators.generate.gates import GATES + + if output_gate not in GATES: + raise ValueError(f"output_gate_not_registered:{output_gate}") + generator.generate = _gated_generate(generator.generate, GATES[output_gate]) + + class GenerateService(BaseOperator): """ Generate question-answer pairs based on nodes and edges. + + Optional additions (all default to official behavior when unset): + - ``prompt_profile``: versioned prompt override loaded via + graphgen.templates.prompt_profiles; + - ``output_gate``: deterministic QA filter registered in + graphgen.operators.generate.gates; + - ``llm_concurrency``: cap on in-flight LLM requests per event loop. """ def __init__( @@ -17,6 +91,9 @@ def __init__( kv_backend: str = "rocksdb", method: str = "aggregated", data_format: str = "ChatML", + prompt_profile: str | None = None, + output_gate: str | None = None, + llm_concurrency: int | None = None, **generate_kwargs, ): super().__init__( @@ -29,61 +106,10 @@ def __init__( self.method = method self.data_format = data_format - - if self.method == "atomic": - from graphgen.models import AtomicGenerator - - self.generator = AtomicGenerator(self.llm_client) - elif self.method == "aggregated": - from graphgen.models import AggregatedGenerator - - self.generator = AggregatedGenerator(self.llm_client) - elif self.method == "multi_hop": - from graphgen.models import MultiHopGenerator - - self.generator = MultiHopGenerator(self.llm_client) - elif self.method == "cot": - from graphgen.models import CoTGenerator - - self.generator = CoTGenerator(self.llm_client) - elif self.method == "vqa": - from graphgen.models import VQAGenerator - - self.generator = VQAGenerator(self.llm_client) - elif self.method == "multi_choice": - from graphgen.models import MultiChoiceGenerator - - self.generator = MultiChoiceGenerator( - self.llm_client, - num_of_questions=generate_kwargs.get("num_of_questions", 5), - ) - elif self.method == "multi_answer": - from graphgen.models import MultiAnswerGenerator - - self.generator = MultiAnswerGenerator( - self.llm_client, - num_of_questions=generate_kwargs.get("num_of_questions", 3), - ) - elif self.method == "fill_in_blank": - from graphgen.models import FillInBlankGenerator - - self.generator = FillInBlankGenerator( - self.llm_client, - num_of_questions=generate_kwargs.get("num_of_questions", 5), - ) - elif self.method == "masked_fill_in_blank": - from graphgen.models import MaskedFillInBlankGenerator - - self.generator = MaskedFillInBlankGenerator(self.llm_client) - elif self.method == "true_false": - from graphgen.models import TrueFalseGenerator - - self.generator = TrueFalseGenerator( - self.llm_client, - num_of_questions=generate_kwargs.get("num_of_questions", 5), - ) - else: - raise ValueError(f"Unsupported generation mode: {method}") + self.generator = _create_generator(self.llm_client, method, generate_kwargs) + _configure_generator( + self.generator, prompt_profile, output_gate, llm_concurrency + ) def process(self, batch: list) -> Tuple[list, dict]: """ diff --git a/graphgen/operators/generate/llm_concurrency.py b/graphgen/operators/generate/llm_concurrency.py new file mode 100644 index 00000000..87bec16b --- /dev/null +++ b/graphgen/operators/generate/llm_concurrency.py @@ -0,0 +1,31 @@ +"""Bound in-flight LLM requests of a generator inside one event loop. + +Official ``run_concurrent`` has no cross-batch concurrency cap; a busy LLM +endpoint can deadlock under high fan-out. This helper caps concurrent +``generator.generate`` calls with an asyncio.Semaphore created lazily per +running loop (Ray actors create a fresh loop per batch). +""" + +from __future__ import annotations + +import asyncio +from typing import Any + + +def bound_llm_concurrency(generate: Any, limit: int) -> Any: + """Wrap an async ``generate(batch)`` so at most ``limit`` run in-flight.""" + if limit < 1: + raise ValueError("llm_concurrency_must_be_positive") + state: dict[str, Any] = {} + + async def _gated_generate(batch: Any) -> Any: + loop = asyncio.get_running_loop() + semaphore = state.get("semaphore") + if semaphore is None or state.get("loop") is not loop: + semaphore = asyncio.Semaphore(limit) + state["semaphore"] = semaphore + state["loop"] = loop + async with semaphore: + return await generate(batch) + + return _gated_generate diff --git a/graphgen/templates/profiles/example.v1/aggregated.json b/graphgen/templates/profiles/example.v1/aggregated.json new file mode 100644 index 00000000..aa6073ca --- /dev/null +++ b/graphgen/templates/profiles/example.v1/aggregated.json @@ -0,0 +1,12 @@ +{ + "en": { + "ANSWER_REPHRASING": "---Role---\nYou are an NLP expert responsible for generating a logically structured and coherent rephrased version of the TEXT based on ENTITIES and RELATIONSHIPS provided below.\nUse English as output language.\n\n---Goal---\nTo generate a version of the text that is rephrased and conveys the same meaning as the original entity and relationship descriptions, while:\n1. Following a clear logical flow and structure\n2. Establishing proper cause-and-effect relationships\n3. Ensuring temporal and sequential consistency\n4. Creating smooth transitions between ideas using conjunctions and appropriate linking words like \"firstly,\" \"however,\" \"therefore,\" etc.\n\n---Instructions---\n1. Analyze the provided ENTITIES and RELATIONSHIPS carefully to identify:\n - Key concepts and their hierarchies\n - Temporal sequences and chronological order\n - Cause-and-effect relationships\n - Dependencies between different elements\n\n2. Organize the information in a logical sequence by:\n - Starting with foundational concepts\n - Building up to more complex relationships\n - Grouping related ideas together\n - Creating clear transitions between sections\n\n3. Rephrase the text while maintaining:\n - Logical flow and progression\n - Clear connections between ideas\n - Proper context and background\n - Coherent narrative structure\n\n4. Review and refine the text to ensure:\n - Logical consistency throughout\n - Clear cause-and-effect relationships\n\n**Attention: Please directly provide the rephrased text without any additional content or analysis.**\n\n################\n-ENTITIES-\n################\n{entities}\n\n################\n-RELATIONSHIPS-\n################\n{relationships}\n\n\n################\nPlease directly output the coherent rephrased text below, without any additional content.\n\nOutput format:\nrephrased_text_here\n\nRephrased Text:\n", + "ANSWER_REPHRASING_CONTEXT": "---Role---\nYou are an NLP expert responsible for generating a logically structured and coherent rephrased version of the TEXT based on ENTITIES and RELATIONSHIPS provided below. You may refer to the original text to assist in generating the rephrased version, but ensure that the final output text meets the requirements.\nUse English as output language.\n\n---Goal---\nTo generate a version of the text that is rephrased and conveys the same meaning as the original entity and relationship descriptions, while:\n1. Following a clear logical flow and structure\n2. Establishing proper cause-and-effect relationships\n3. Ensuring temporal and sequential consistency\n4. Creating smooth transitions between ideas using conjunctions and appropriate linking words like \"firstly,\" \"however,\" \"therefore,\" etc.\n\n---Instructions---\n1. Analyze the provided ENTITIES and RELATIONSHIPS carefully to identify:\n - Key concepts and their hierarchies\n - Temporal sequences and chronological order\n - Cause-and-effect relationships\n - Dependencies between different elements\n\n2. Organize the information in a logical sequence by:\n - Starting with foundational concepts\n - Building up to more complex relationships\n - Grouping related ideas together\n - Creating clear transitions between sections\n\n3. Rephrase the text while maintaining:\n - Logical flow and progression\n - Clear connections between ideas\n - Proper context and background\n - Coherent narrative structure\n\n4. Review and refine the text to ensure:\n - Logical consistency throughout\n - Clear cause-and-effect relationships\n\n################\n-ORIGINAL TEXT-\n################\n{original_text}\n\n################\n-ENTITIES-\n################\n{entities}\n\n################\n-RELATIONSHIPS-\n################\n{relationships}\n\n\n################\nPlease directly output the coherent rephrased text below, without any additional content.\n\nOutput format:\nrephrased_text_here\n\nRephrased Text:\n", + "QUESTION_GENERATION": "The answer to a question is provided. Please generate a question that corresponds to the answer.\n\nThe answer for which a question needs to be generated is as follows:\n{answer}\n\nPlease note the following requirements:\n1. Only output one question text without any additional explanations or analysis.\n2. Do not repeat the content of the answer or any fragments of it.\n3. The question must be independently understandable and fully match the answer.\n\nOutput format:\nquestion_text\n\nQuestion:\n" + }, + "zh": { + "ANSWER_REPHRASING": "---角色---\n你是一位NLP专家，负责根据下面提供的实体和关系生成逻辑结构清晰且连贯的文本重述版本。\n使用中文作为输出语言。\n\n---目标---\n生成文本的重述版本，使其传达与原始实体和关系描述相同的含义，同时：\n1. 遵循清晰的逻辑流和结构\n2. 建立适当的因果关系\n3. 确保时间和顺序的一致性\n4. 使用连词和适当的连接词(如\"首先\"、\"然而\"、\"因此\"等)创造流畅的过渡\n\n---说明---\n1. 仔细分析提供的实体和关系，以识别：\n - 关键概念及其层级关系\n - 时间序列和时间顺序\n - 因果关系\n - 不同元素之间的依赖关系\n2. 通过以下方式将信息组织成逻辑顺序：\n - 从基础概念开始\n - 逐步建立更复杂的关系\n - 将相关的想法分组在一起\n - 在各部分之间创建清晰的过渡\n3. 重述文本时保持：\n - 逻辑流畅\n - 概念之间的清晰联系\n - 适当的上下文和背景\n - 连贯的叙述结构\n4. 检查和完善文本以确保：\n - 整体逻辑一致性\n - 清晰的因果关系\n\n**注意： 请你直接给出重述文本，不要输出任何额外的内容，也不要进行任何分析。**\n\n################\n-实体-\n################\n{entities}\n\n################\n-关系-\n################\n{relationships}\n\n\n################\n请在下方直接输出连贯的重述文本，不要输出任何额外的内容。\n\n输出格式：\nrephrased_text_here\n\n重述文本:\n", + "ANSWER_REPHRASING_CONTEXT": "---角色---\n你是一位NLP专家，负责根据下面提供的实体和关系生成逻辑结构清晰且连贯的文本重述版本。你可以参考原始文本辅助生成，但需要确保最终输出的文本符合要求。\n使用中文作为输出语言。\n\n---目标---\n生成文本的重述版本，使其传达与原始实体和关系描述相同的含义，同时：\n1. 遵循清晰的逻辑流和结构\n2. 建立适当的因果关系\n3. 确保时间和顺序的一致性\n4. 使用连词和适当的连接词(如\"首先\"、\"然而\"、\"因此\"等)创造流畅的过渡\n\n---说明---\n1. 仔细分析提供的实体和关系，以识别：\n - 关键概念及其层级关系\n - 时间序列和时间顺序\n - 因果关系\n - 不同元素之间的依赖关系\n2. 通过以下方式将信息组织成逻辑顺序：\n - 从基础概念开始\n - 逐步建立更复杂的关系\n - 将相关的想法分组在一起\n - 在各部分之间创建清晰的过渡\n3. 重述文本时保持：\n - 逻辑流畅\n - 概念之间的清晰联系\n - 适当的上下文和背景\n - 连贯的叙述结构\n4. 检查和完善文本以确保：\n - 整体逻辑一致性\n - 清晰的因果关系\n\n################\n-原始文本-\n################\n{original_text}\n\n################\n-实体-\n################\n{entities}\n\n################\n-关系-\n################\n{relationships}\n\n\n################\n请在下方直接输出连贯的重述文本，不要输出任何额外的内容。\n\n输出格式：\nrephrased_text_here\n\n重述文本:\n", + "QUESTION_GENERATION": "下面提供了一个问题的答案，请生成一个与答案对应的问题。\n\n需要生成问题的答案如下：\n{answer}\n\n请注意下列要求：\n1. 仅输出一个问题文本，不得包含任何额外说明或分析\n2. 不得重复答案内容或其中任何片段\n3. 问题必须可独立理解且与答案完全匹配\n\n输出格式：\nquestion_text\n\n问题:\n" + } +} diff --git a/graphgen/templates/profiles/example.v1/atomic.json b/graphgen/templates/profiles/example.v1/atomic.json new file mode 100644 index 00000000..f90db6ac --- /dev/null +++ b/graphgen/templates/profiles/example.v1/atomic.json @@ -0,0 +1,4 @@ +{ + "en": "You are given a text passage. Your task is to generate a question and answer (QA) pair based on the content of that text.\n\nPlease note the following requirements:\n1. Output only one QA pair without any additional explanations or analysis.\n2. Do not repeat the content of the answer or any part of it.\n3. The answer should be accurate and directly derived from the text. Make sure the QA pair is relevant to the main theme or important details of the given text.\n\nOutput format:\nquestion_text\nanswer_text\n\nFor example:\nWhat is the effect of overexpressing the BG1 gene on grain size and development?\nOverexpression of the BG1 gene leads to significantly increased grain size, demonstrating its role in grain development.\n\nHere is the text passage you need to generate a QA pair for:\n{context}\n\nOutput:\n", + "zh": "给定一个文本段落。你的任务是根据该文本的内容生成一个问答（QA）对。\n\n请注意下列要求：\n1. 仅输出一个问答（QA）对，不得包含任何额外说明或分析\n2. 不得重复答案内容或其中任何片段\n3. 答案应准确且直接从文本中得出。确保QA对与给定文本的主题或重要细节相关。\n\n输出格式如下：\nquestion_text\nanswer_text\n\n例如：\n过表达BG1基因对谷粒大小和发育有什么影响？\nBG1基因的过表达显著增加了谷粒大小，表明其在谷物发育中的作用。\n\n以下是你需要为其生成QA对的文本段落：\n{context}\n\n输出：\n" +} diff --git a/graphgen/templates/profiles/example.v1/cot.json b/graphgen/templates/profiles/example.v1/cot.json new file mode 100644 index 00000000..4db643a6 --- /dev/null +++ b/graphgen/templates/profiles/example.v1/cot.json @@ -0,0 +1,10 @@ +{ + "en": { + "COT_GENERATION": "Given the raw knowledge graph information and the provided reasoning-path, produce one Chain-of-Thought (CoT) sample that strictly follows the template and can be directly used for downstream training or inference.\nCoT (Chain-of-Thought) means that when answering a complex question, the intermediate reasoning steps are explicitly written out one by one, making the reasoning process transparent and traceable instead of giving only the final answer.\n\n-Input Format-\n[Entities:]:\n(ENTITY_NAME: ENTITY_DESCRIPTION)\n...\n\n[Relationships:]:\n(ENTITY_SOURCE)-[RELATIONSHIP_DESCRIPTION]->(ENTITY_TARGET)\n...\n\n[Question and Reasoning Path:]:\n(QUESTION)\n(REASONING_PATH)\n\n-Output Requirements-\n1. Each step completes a single, indivisible sub-task and is naturally connected, avoiding abrupt transition words.\n2. Use English.\n3. Do not use ordered lists or numbering.\n4. Do not generate extraneous information, just provide the answer.\n\n-Real Data-\nInput:\n[Entities:]:\n{entities}\n\n[Relationships:]:\n{relationships}\n\n[Question:]:\n{question}\n\n[Reasoning_Template:]:\n{reasoning_template}\n\nOutput:\n", + "COT_TEMPLATE_DESIGN": "You are a “meta-reasoning architect”. Your task is NOT to answer the question, but to design a reusable, generalizable CoT reasoning-path template based solely on the names and descriptions of entities and relationships in the provided knowledge graph.\n\n- Steps -\n1. Entity Recognition\n- Accurately recognize entity information in the [Entities:] section, including entity names and descriptions.\n- The general formats for entity information are:\n(ENTITY_NAME: ENTITY_DESCRIPTION)\n\n2. Relationship Recognition\n- Accurately recognize relationship information in the [Relationships:] section, including source_entity_name, target_entity_name, and relationship descriptions.\n- The general formats for relationship information are:\n(SOURCE_ENTITY_NAME)-[RELATIONSHIP_DESCRIPTION]->(TARGET_ENTITY_NAME)\n\n3. Graph Structure Understanding\n- Correctly associate the source entity name in the relationship information with the entity information.\n- Reconstruct the graph structure based on the provided relationship information.\n\n4. Question Design\n- Design a question around the \"core theme\" expressed by the knowledge graph.\n- The question must be verifiable directly within the graph through entities, relationships, or attributes; avoid subjective judgments.\n- The question should allow the model to think sufficiently, fully utilizing the entities and relationships in the graph, avoiding overly simple or irrelevant questions.\n\n5. Reasoning-Path Design\n- Output a **blueprint that any later model can directly execute**.\n- Keep steps minimal: each step solves one indivisible sub-problem.\n\n\n- Constraints -\n1. Do NOT describe your thinking; output only the reasoning-path design.\n2. If the provided descriptions are contradictory, resolve conflicts and provide a single coherent logic.\n3. Avoid using stop words and overly common words.\n4. Do not include specific numerical values or conclusions, and DO NOT describing meaningless operations like \"Identify the entity\" or \"Identify the relationship\".\n5. Use English as the output language.\n6. The output format is:\nquestion text\nreasoning path design text\n\nPlease summarize the information expressed by the knowledge graph based on the following [Entities:] and [Relationships:] provided.\n\n- Real Data -\nInput:\n[Entities:]:\n{entities}\n\n[Relationships:]:\n{relationships}\n\nOutput:\n" + }, + "zh": { + "COT_GENERATION": "根据给定的知识图谱原始信息及已生成的推理路径，产出一条符合模板要求、可直接用于下游训练或推理的 CoT 数据。CoT（Chain-of-Thought，思维链）指在回答复杂问题时，把中间推理步骤一步一步显式写出来，使推理过程透明、可追溯，而不是直接给出最终答案。\n\n-输入格式-\n[Entities:]\n(实体名:实体描述)\n...\n\n[Relationships:]\n(来源实体)-[关系描述]->(目标实体)\n...\n\n[Question and Reasoning Path:]\n(问题)\n(推理路径)\n\n-输出要求-\n1. 每一步只完成一个不可分割的子任务，并用自然语言衔接，但是要避免生硬的连接词。\n2. 使用中文。\n3. 不要使用有序列表或编号。\n4. 请直接给出答案，不要生成无关信息。\n\n-真实数据-\n输入:\n[Entities:]:\n{entities}\n\n[Relationships:]:\n{relationships}\n\n[Question:]:\n{question}\n\n[Reasoning_Template:]:\n{reasoning_template}\n\n输出：\n\n", + "COT_TEMPLATE_DESIGN": "你是一位“元推理架构师”。你的任务不是回答问题，而是根据给定的知识图谱中的实体和关系的名称以及描述信息，设计一条可复用、可泛化的 CoT 推理路径模板。\n-步骤-\n1. 实体识别\n- 准确地识别[Entities:]章节中的实体信息，包括实体名、实体描述信息。\n- 实体信息的一般格式为:\n(实体名:实体描述)\n\n2. 关系识别\n- 准确地识别[Relationships:]章节中的关系信息，包括来源实体名、目标实体名、关系描述信息。\n- 关系信息的一般格式为:\n(来源实体名)-[关系描述]->(目标实体名)\n\n3. 图结构理解\n- 正确地将关系信息中的来源实体名与实体信息关联。\n- 根据提供的关系信息还原出图结构。\n\n4. 问题设计\n- 围绕知识图谱所表达的“核心主题”设计一个问题。\n- 问题必须能在图谱内部通过实体、关系或属性直接验证；避免主观判断。\n- 问题应该能够模型足够的思考，充分利用图谱中的实体和关系，避免过于简单或无关的问题。\n\n5. 推理路径生成\n- 根据问题设计一个**可被后续模型直接执行的推理蓝图**。\n- 保持步骤最小化：每一步只解决一个“不可分割”的子问题。\n\n-约束条件-\n1. 不要在回答中描述你的思考过程，直接给出回复，只给出问题和推理路径设计，不要生成无关信息。\n2. 如果提供的描述信息相互矛盾，请解决矛盾并提供一个单一、连贯的逻辑。\n3. 避免使用停用词和过于常见的词汇。\n4. 不要出现具体数值或结论，不要出现“识别实体”、“识别关系”这类无意义的操作描述。\n5. 使用中文作为输出语言。\n6. 输出格式为：\n问题文本\n推理路径设计文本\n\n-真实数据-\n输入:\n[Entities:]:\n{entities}\n\n[Relationships:]:\n{relationships}\n\n输出:\n" + } +} diff --git a/graphgen/templates/profiles/example.v1/fill_in_blank.json b/graphgen/templates/profiles/example.v1/fill_in_blank.json new file mode 100644 index 00000000..b07684ef --- /dev/null +++ b/graphgen/templates/profiles/example.v1/fill_in_blank.json @@ -0,0 +1,4 @@ +{ + "zh": "请根据上下文资料生成独立的知识问答填空题。填空题的答案必须能在原文中直接找到。\n\n生成要求：\n1. **语言一致性**：若上下文资料为中文，则生成中文问题；若为英文，则生成英文问题\n2. **数量**：每个上下文资料生成{num_of_questions}个填空题\n3. **独立性**：每个问题必须完整独立，不依赖其他问题\n4. **准确性**：正确答案必须能从原文直接得出\n5. **占位符格式**：使用________（四个下划线）作为填空占位符\n\n输出格式：\n\n\n问题文本（使用________作为占位符）\n正确答案文本（多个空用逗号分隔）\n\n\n\n示例（根据iPad Air 2生成2题）：\n\n\niPad Air 2 是由________制造的？\n美国苹果公司（Apple）\n\n\niPad Air 2 的发布日期是________，上市日期是________。\n2014年10月16日，2014年10月22日\n\n\n\n\n上下文资料：\n{context}\n\n请为以下资料生成{num_of_questions}个填空题：\n", + "en": "Generate independent fill-in-the-blank questions based on the provided context. Answers must be directly derivable from the text.\n\nRequirements:\n1. **Language Consistency**: Generate in the same language as the context (Chinese/English)\n2. **Quantity**: Generate {num_of_questions} questions per context\n3. **Independence**: Each question must be self-contained\n4. **Accuracy**: Correct answer must be directly found in the source text\n5. **Placeholder Format**: Use ________ (four underscores) as the blank placeholder\n\nOutput Format:\n\n\nQuestion text (use ________ as placeholder)\nCorrect answer text (separate multiple blanks with commas)\n\n\n\nExample (2 questions):\n\n\nThe iPad Air 2 was manufactured by ________?\nApple Inc.\n\n\nThe iPad Air 2 was released on ________ and launched on ________.\nOctober 16, 2014, October 22, 2014\n\n\n\nContext:\n{context}\n\nPlease generate {num_of_questions} fill-in-the-blank questions for the following context:\n" +} diff --git a/graphgen/templates/profiles/example.v1/multi_answer.json b/graphgen/templates/profiles/example.v1/multi_answer.json new file mode 100644 index 00000000..85d1a186 --- /dev/null +++ b/graphgen/templates/profiles/example.v1/multi_answer.json @@ -0,0 +1,4 @@ +{ + "zh": "请根据上下文资料生成独立的知识问答不定项选择题，每个选择题包含四个选项，其中有若干个正确答案（至少一个），其他为干扰项。\n\n生成要求：\n1. **语言一致性**：若上下文资料为中文，则生成中文问题；若为英文，则生成英文问题\n2. **数量**：每个上下文资料生成{num_of_questions}个选择题\n3. **独立性**：每个问题必须完整独立，不依赖其他问题\n4. **准确性**：正确答案必须能从原文直接得出，干扰项需合理且有区分度\n5. **答案格式**：当有多个正确答案时，用逗号分隔选项字母，如\"A, B, C\"\n\n输出格式：\n\n\n问题文本\nA. 选项A文本\nB. 选项B文本\nC. 选项C文本\nD. 选项D文本\n正确答案选项字母（多个答案用逗号分隔）\n\n\n\n示例（根据iPad Air 2生成2题）：\n\n\niPad Air 2的发布年份是？\nA. 2012年\nB. 2014年\nC. 2015年\nD. 2017年\nB\n\n\n以下哪些是 iPad Air 2 的特点？\nA. Touch ID指纹识别功能\nB. A8X高效处理器\nC. 十百万像素前置相机\nD. 八百万像素后置相机镜头\nA, B, D\n\n\n\n\n上下文资料：\n{context}\n\n请为以下资料生成{num_of_questions}个不定项选择题：\n", + "en": "Generate independent multiple-select knowledge questions based on the provided context. Each question should contain four options with one or more correct answers and distractors.\n\nRequirements:\n1. **Language Consistency**: Generate in the same language as the context (Chinese/English)\n2. **Quantity**: Generate {num_of_questions} questions per context\n3. **Independence**: Each question must be self-contained\n4. **Accuracy**: Correct answer(s) must be derivable from text, distractors should be plausible\n5. **Answer Format**: For multiple correct answers, separate option letters with commas, e.g., \"A, B, C\"\n\nOutput Format:\n\n\nQuestion text\nA. Option A text\nB. Option B text\nC. Option C text\nD. Option D text\nCorrect option letter(s) (separate multiple answers with commas)\n\n\n\nExample (2 questions):\n\n\nWhat are the features of iPad Air 2?\nA. Touch ID fingerprint recognition\nB. A8X processor\nC. Ten-megapixel front camera\nD. Eight-megapixel rear camera\nA, B, D\n\n\nWhen was iPad Air 2 discontinued?\nA. March 21, 2016\nB. March 21, 2017\nC. October 22, 2017\nD. October 16, 2016\nB\n\n\n\nContext:\n{context}\n\nPlease generate {num_of_questions} multiple-select questions for the following context:\n" +} diff --git a/graphgen/templates/profiles/example.v1/multi_choice.json b/graphgen/templates/profiles/example.v1/multi_choice.json new file mode 100644 index 00000000..67775afa --- /dev/null +++ b/graphgen/templates/profiles/example.v1/multi_choice.json @@ -0,0 +1,4 @@ +{ + "zh": "请根据上下文资料生成独立的知识问答单选题，每个选择题包含四个选项，其中仅有一个正确答案，其他三个为干扰项。\n\n生成要求：\n1. **语言一致性**：若上下文资料为中文，则生成中文问题；若为英文，则生成英文问题\n2. **数量**：每个上下文资料生成{num_of_questions}个选择题\n3. **独立性**：每个问题必须完整独立，不依赖其他问题\n4. **准确性**：正确答案必须能从原文直接得出，干扰项需合理且有区分度\n\n输出格式：\n\n\n问题文本\nA. 选项A文本\nB. 选项B文本\nC. 选项C文本\nD. 选项D文本\n正确答案选项字母\n\n\n\n示例（根据iPad Air 2生成2题）：\n\n\niPad Air 2的发布年份是？\nA. 2012年\nB. 2014年\nC. 2015年\nD. 2017年\nB\n\n\niPad Air 2搭载的处理器型号是？\nA. A8\nB. A9X\nC. A8X\nD. A10\nC\n\n\n\n\n上下文资料：\n{context}\n\n请为以下资料生成{num_of_questions}个选择题：\n", + "en": "Generate independent multiple-choice questions based on the provided context. Each question should contain four options with only one correct answer and three distractors.\n\nRequirements:\n1. **Language Consistency**: Generate in the same language as the context (Chinese/English)\n2. **Quantity**: Generate {num_of_questions} questions per context\n3. **Independence**: Each question must be self-contained\n4. **Accuracy**: Correct answer must be derivable from text, distractors should be plausible\n\nOutput Format:\n\n\nQuestion text\nA. Option A text\nB. Option B text\nC. Option C text\nD. Option D text\nCorrect option letter\n\n\n\nExample (2 questions):\n\n\nWhat year was the iPad Air 2 released?\nA. 2012\nB. 2014\nC. 2015\nD. 2017\nB\n\n\nWhich processor does iPad Air 2 use?\nA. A8\nB. A9X\nC. A8X\nD. A10\nC\n\n\n\nContext:\n{context}\n\nPlease generate {num_of_questions} questions for the following context:\n" +} diff --git a/graphgen/templates/profiles/example.v1/multi_hop.json b/graphgen/templates/profiles/example.v1/multi_hop.json new file mode 100644 index 00000000..256041a1 --- /dev/null +++ b/graphgen/templates/profiles/example.v1/multi_hop.json @@ -0,0 +1,4 @@ +{ + "en": "Please generate a multi-hop reasoning question and answer based on the following knowledge subgraph. You will be provided with a knowledge subgraph that contains a series of entities, relations, and facts.\nYour task is to generate a question-answer (QA) pair where the question requires multiple steps of reasoning to answer. The answer to the question should be inferred from the given knowledge subgraph. Ensure that the question is of moderate difficulty and requires multiple steps of reasoning to answer.\n\nPlease note the following requirements:\n1. Output only one QA pair without any additional explanations or analysis.\n2. Do not repeat the content of the answer or any part of it. Do not directly copy the example question and answer.\n3. The answer should be accurate and directly derived from the text. Make sure the QA pair is relevant to the main theme or important details of the given text.\n\nFor example:\nInput:\n--Entities--\n1. Apple\n2. Fruit\n3. Vitamin C\n--Relations--\n1. Apple-Fruit: Apple is a type of fruit\n2. Fruit-Vitamin C: Fruits are rich in Vitamin C\n\nOutput:\nWhat substance, obtained by eating apples, helps maintain health?\nVitamin C\n\nReal input:\n--Entities--\n{entities}\n--Relations--\n{relationships}\n\nOutput:\n", + "zh": "请基于以下知识子图生成多跳推理问题和答案。你将获得一个知识子图，其中包含一系列实体、关系和事实。\n你的任务是生成一个问答对，其中问题需要经过多次推理才能回答。问题的答案应该是从给定的知识子图中推断出来的。确保问题的难度适中，需要多步推理才能回答。\n\n请注意下列要求：\n1. 仅输出一个问答（QA）对，不得包含任何额外说明或分析\n2. 不得重复答案内容或其中任何片段，不要直接复制示例问题和答案\n3. 答案应准确且直接从文本中得出。确保QA对与给定文本的主题或重要细节相关。\n\n输出格式：\nquestion_text\nanswer_text\n\n例如：\n输入为：\n--实体--\n1. 苹果\n2. 水果\n3. 维生素C\n--关系--\n1. 苹果-水果：苹果是一种水果\n2. 水果-维生素C：水果中富含维生素C\n\n输出：\n通过吃苹果补充的什么物质，有助于维持健康？\n维生素C\n\n真实输入如下：\n--实体--\n{entities}\n--关系--\n{relationships}\n\n输出：\n" +} diff --git a/graphgen/templates/profiles/example.v1/true_false.json b/graphgen/templates/profiles/example.v1/true_false.json new file mode 100644 index 00000000..9974d5ce --- /dev/null +++ b/graphgen/templates/profiles/example.v1/true_false.json @@ -0,0 +1,4 @@ +{ + "zh": "请根据上下文资料生成独立的知识判断题，每个判断题包含一个陈述句，答案只能是正确(True)或错误(False)。\n\n生成要求：\n1. **语言一致性**：若上下文资料为中文，则生成中文问题；若为英文，则生成英文问题\n2. **数量**：每个上下文资料生成{num_of_questions}个判断题\n3. **独立性**：每个问题必须完整独立，不依赖其他问题\n4. **准确性**：正确答案必须能从原文直接得出，陈述需有明确的判断依据\n\n输出格式：\n\n\n陈述句文本\nTrue或False\n\n\n\n示例（根据iPad Air 2生成2题）：\n\n\niPad Air 2于2014年发布。\nTrue\n\n\niPad Air 2搭载的是A10处理器。\nFalse\n\n\n\n\n上下文资料：\n{context}\n\n请为以下资料生成{num_of_questions}个判断题：\n", + "en": "Generate independent true/false questions based on the provided context. Each question should be a factual statement that can be clearly determined as true or false.\n\nRequirements:\n1. **Language Consistency**: Generate in the same language as the context (Chinese/English)\n2. **Quantity**: Generate {num_of_questions} true/false questions per context\n3. **Independence**: Each question must be self-contained\n4. **Accuracy**: Correct answer must be directly derivable from the text with clear evidence\n\nOutput Format:\n\n\nStatement text\nTrue or False\n\n\n\nExample (2 questions):\n\n\nThe iPad Air 2 was released in 2014.\nTrue\n\n\nThe iPad Air 2 uses an A10 processor.\nTrue\nFalse\nFalse\n\n\n\nContext:\n{context}\n\nPlease generate {num_of_questions} true/false questions for the following context:\n" +} diff --git a/graphgen/templates/profiles/example.v1/vqa.json b/graphgen/templates/profiles/example.v1/vqa.json new file mode 100644 index 00000000..6f98dd54 --- /dev/null +++ b/graphgen/templates/profiles/example.v1/vqa.json @@ -0,0 +1,4 @@ +{ + "en": "You are a senior VQA data engineer. Your task is to generate logically coherent, verifiable and non-hallucinated question-answer pairs for the given multi-modal samples.\nUse English as the output language.\n\n---Objectives---\nCreate multiple sets of VQA question-answer pairs that satisfy the following:\n1. Only ask about objectively existing facts in the given data, avoiding subjective or ambiguous questions.\n2. Ensure that each question has a clear and verifiable answer, avoiding questions with no answer or uncertainty.\n3. Questions should cover various aspects of both image and text content, ensuring diversity and comprehensiveness.\n4. Avoid repetitive questions, ensuring that each question is unique and meaningful.\n5. Use clear and concise language, avoiding complex or ambiguous wording.\n\n---Instructions---\n1. Carefully analyze the provided entities and relationships to identify:\n - Key concepts and their hierarchical relationships\n - Temporal sequences and time order\n - Cause-and-effect relationships\n - Dependencies between different elements\n2. Organize the information into a logical sequence by:\n - Starting with foundational concepts\n - Gradually building up to more complex relationships\n - Grouping related ideas together\n - Creating clear transitions between sections\n3. Maintain the following when generating question-answer pairs:\n - Logical flow\n - Clear connections between concepts\n - Appropriate context and background\n - Coherent narrative structure\n4. Review and refine the question-answer pairs to ensure:\n - Overall logical consistency\n - Clear cause-and-effect relationships\n\n################\n-Entities-\n################\n{entities}\n################\n-Relationships-\n################\n{relationships}\n################\n\nPlease directly output the generated questions and answers, do not directly copy the example questions and answers, and do not provide irrelevant information.\n\nHere is the response format you should follow:\nquestion1\nanswer1\nquestion2\nanswer2\n\nOutput:\n", + "zh": "---角色---\n你是一位资深 VQA 数据工程师。你需要为给定的多模态样本生成逻辑连贯、可验证、无幻觉的问答对。\n使用中文作为输出语言。\n\n---目标---\n创建多组 VQA 问答对，满足：\n1. 仅询问给定数据中客观存在的事实，避免主观或模糊的问题。\n2. 确保每个问题都有明确且可验证的答案，避免无答案或不确定的问题。\n3. 问题应涵盖图像和文本内容的各个方面，确保多样性和全面性。\n4. 避免重复问题，确保每个问题都是独特且有意义的。\n5. 使用清晰简洁的语言，避免复杂或含糊的措辞。\n\n---说明---\n1. 仔细分析提供的实体和关系，以识别：\n - 关键概念及其层级关系\n - 时间序列和时间顺序\n - 因果关系\n - 不同元素之间的依赖关系\n2. 通过以下方式将信息组织成逻辑顺序：\n - 从基础概念开始\n - 逐步建立更复杂的关系\n - 将相关的想法分组在一起\n - 在各部分之间创建清晰的过渡\n3. 生成问答对时保持：\n - 逻辑流畅\n - 概念之间的清晰联系\n - 适当的上下文和背景\n - 连贯的叙述结构\n4. 检查和完善问答对以确保：\n - 整体逻辑一致性\n - 清晰的因果关系\n\n################\n-实体-\n################\n{entities}\n\n################\n-关系-\n################\n{relationships}\n################\n\n请直接输出生成的问题和答案，不要直接复制示例问题和答案，也不要提供无关信息。\n以下是你应遵循的响应格式：\nquestion1\nanswer1\nquestion2\nanswer2\n\n输出：\n" +} diff --git a/graphgen/templates/prompt_profiles.py b/graphgen/templates/prompt_profiles.py new file mode 100644 index 00000000..a3ba6db7 --- /dev/null +++ b/graphgen/templates/prompt_profiles.py @@ -0,0 +1,119 @@ +"""Versioned prompt profiles: config-driven prompt customization for generators. + +Official GraphGen hard-codes all prompts as Python constants under +``graphgen.templates``; this module adds a generic, fail-closed override +channel: + +- a *profile* is a directory containing one JSON file per generator + ``TEMPLATE_KEY`` (e.g. ``atomic.json``); +- each JSON file must mirror the exact key structure of the corresponding + official template (same languages, same nested keys, leaf values are str); +- ``profile`` values are either built-in profile ids resolved against + ``graphgen/templates/profiles/`` or explicit filesystem paths (containing + a path separator). + +Default behaviour is untouched: when no profile is applied every generator +falls back to the official template constant byte-for-byte. +""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +from typing import Any + +from graphgen.templates import ( + AGGREGATED_GENERATION_PROMPT, + ATOMIC_GENERATION_PROMPT, + COT_GENERATION_PROMPT, + FILL_IN_BLANK_GENERATION_PROMPT, + MAQ_GENERATION_PROMPT, + MCQ_GENERATION_PROMPT, + MULTI_HOP_GENERATION_PROMPT, + TF_GENERATION_PROMPT, + VQA_GENERATION_PROMPT, +) + +BUILTIN_PROFILE_ROOT = Path(__file__).resolve().parent / "profiles" + +#: template_key -> official template constant (structure validation source). +OFFICIAL_TEMPLATES: dict[str, Any] = { + "atomic": ATOMIC_GENERATION_PROMPT, + "multi_hop": MULTI_HOP_GENERATION_PROMPT, + "aggregated": AGGREGATED_GENERATION_PROMPT, + "cot": COT_GENERATION_PROMPT, + "multi_choice": MCQ_GENERATION_PROMPT, + "multi_answer": MAQ_GENERATION_PROMPT, + "fill_in_blank": FILL_IN_BLANK_GENERATION_PROMPT, + "masked_fill_in_blank": AGGREGATED_GENERATION_PROMPT, + "true_false": TF_GENERATION_PROMPT, + "vqa": VQA_GENERATION_PROMPT, +} + + +def resolve_profile_dir(profile: str) -> Path: + """Resolve a profile id / path to its directory; fail-closed if missing.""" + if not profile or not isinstance(profile, str): + raise ValueError("prompt_profile_must_be_non_empty_string") + if "/" in profile or "\\" in profile: + profile_dir = Path(profile) + else: + profile_dir = BUILTIN_PROFILE_ROOT / profile + if not profile_dir.is_dir(): + raise ValueError(f"prompt_profile_not_found:{profile}") + return profile_dir + + +def _validate_structure(official: Any, override: Any, path: str) -> None: + """Override must mirror the official template's key structure exactly.""" + if isinstance(official, dict): + if not isinstance(override, dict): + raise ValueError(f"prompt_profile_structure_mismatch:{path}") + if set(override.keys()) != set(official.keys()): + raise ValueError( + f"prompt_profile_structure_mismatch:{path}:" + f"keys={sorted(map(str, override.keys()))}" + ) + for key in official: + _validate_structure(official[key], override[key], f"{path}.{key}") + return + if not isinstance(override, str) or not override.strip(): + raise ValueError(f"prompt_profile_structure_mismatch:{path}:non_string_leaf") + + +def load_profile_template(profile: str, template_key: str) -> dict[str, Any]: + """Load one generator's template override from a profile. + + :param profile: profile id (builtin) or filesystem path. + :param template_key: generator TEMPLATE_KEY, must exist in OFFICIAL_TEMPLATES. + :return: parsed override dict (same structure as the official template). + """ + if template_key not in OFFICIAL_TEMPLATES: + raise ValueError(f"prompt_profile_unknown_template_key:{template_key}") + profile_file = resolve_profile_dir(profile) / f"{template_key}.json" + if not profile_file.is_file(): + raise ValueError(f"prompt_profile_template_missing:{profile}/{template_key}") + try: + override = json.loads(profile_file.read_text(encoding="utf-8")) + except json.JSONDecodeError as error: + raise ValueError(f"prompt_profile_json_invalid:{profile_file}:{error}") from error + _validate_structure(OFFICIAL_TEMPLATES[template_key], override, template_key) + return override + + +def profile_template_sha256s(profile: str, template_keys: list[str]) -> dict[str, str]: + """Lineage fingerprint: SHA-256 of each override file actually in use.""" + profile_dir = resolve_profile_dir(profile) + digests: dict[str, str] = {} + for template_key in template_keys: + profile_file = profile_dir / f"{template_key}.json" + if not profile_file.is_file(): + continue + digests[template_key] = hashlib.sha256(profile_file.read_bytes()).hexdigest() + return digests + + +def apply_prompt_profile(generator: Any, profile: str) -> None: + """Attach profile templates to a generator instance (BaseGenerator API).""" + generator.apply_prompt_profile(profile) diff --git a/graphgen/utils/log.py b/graphgen/utils/log.py index e29e994e..abaa587b 100644 --- a/graphgen/utils/log.py +++ b/graphgen/utils/log.py @@ -62,9 +62,9 @@ def set_logger( def get_current_logger() -> logging.Logger: - current_logger = CURRENT_LOGGER_VAR.get() - if not current_logger: - raise RuntimeError("No logger is set in the current context.") + current_logger = CURRENT_LOGGER_VAR.get(None) + if current_logger is None: + return logging.getLogger("GraphGen") return current_logger diff --git a/setup.py b/setup.py index 2b30f103..0bed8d24 100644 --- a/setup.py +++ b/setup.py @@ -45,7 +45,12 @@ def read_requirements(): author="open-sciencelab", author_email="open-sciencelab@pjlab.org.cn", packages=find_packages(exclude=["models"]), - package_data={"GraphGen": ["configs/*"]}, + package_data={ + "graphgen": [ + "configs/*", + "templates/profiles/*/*.json", + ], + }, include_package_data=True, install_requires=install_packages, classifiers=[ diff --git a/tests/unit_tests/__init__.py b/tests/unit_tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/models/llm/api/test_openai_client_patches.py b/tests/unit_tests/models/llm/api/test_openai_client_patches.py new file mode 100644 index 00000000..4468d793 --- /dev/null +++ b/tests/unit_tests/models/llm/api/test_openai_client_patches.py @@ -0,0 +1,128 @@ +"""Tests for env-gated OpenAI client patches (retry / prefill / thinking).""" + +import asyncio + +from graphgen.models.llm.api.openai_client import OpenAIClient + + +def _client(monkeypatch, **overrides): + monkeypatch.setenv("TOKENIZER_MODEL", "cl100k_base") + client = OpenAIClient.__new__(OpenAIClient) + for key, value in { + "model": "fake-model", + "api_key": "k", + "api_version": None, + "base_url": None, + "json_mode": False, + "seed": None, + "topk_per_token": 5, + "request_limit": False, + "system_prompt": None, + }.items(): + setattr(client, key, value) + for key, value in overrides.items(): + setattr(client, key, value) + + class _Tok: + def encode(self, text): + return list(text) + + client.tokenizer = _Tok() + client.temperature = 0 + client.top_p = 1 + client.max_tokens = 16 + client.token_usage = [] + + class _RPM: + async def wait(self, silent=True): + return None + + class _TPM: + async def wait(self, tokens, silent=True): + return None + + client.rpm = _RPM() + client.tpm = _TPM() + return client + + +def _stub_completions(client, contents): + calls = {"count": 0} + + class _Usage: + prompt_tokens = 1 + completion_tokens = 1 + total_tokens = 2 + + class _Msg: + content = None + + class _Logprobs: + content = [] + + class _Choice: + message = _Msg() + logprobs = _Logprobs() + + class _Completion: + usage = _Usage() + choices = [_Choice()] + + async def _create(**kwargs): + idx = min(calls["count"], len(contents) - 1) + _Choice.message.content = contents[idx] + calls["count"] += 1 + calls.setdefault("kwargs", []).append(kwargs) + return _Completion() + + class _StubCompletions: + create = staticmethod(_create) + + class _StubChat: + completions = _StubCompletions() + + class _StubClient: + chat = _StubChat() + + client.client = _StubClient() + return calls + + +def test_empty_response_retry(monkeypatch): + monkeypatch.setenv("GRAPHGEN_EMPTY_RETRY", "2") + client = _client(monkeypatch) + calls = _stub_completions(client, ["", "", "real answer"]) + result = asyncio.run(client.generate_answer("hello")) + assert result == "real answer" + assert calls["count"] == 3 + + +def test_empty_response_no_retry_by_default(monkeypatch): + monkeypatch.delenv("GRAPHGEN_EMPTY_RETRY", raising=False) + client = _client(monkeypatch) + calls = _stub_completions(client, ["", "unused"]) + result = asyncio.run(client.generate_answer("hello")) + assert result == "" + assert calls["count"] == 1 + + +def test_assistant_prefill_added_last(monkeypatch): + client = _client(monkeypatch) + calls = _stub_completions(client, ["ok"]) + asyncio.run(client.generate_answer("hello", assistant_prefill="{")) + messages = calls["kwargs"][0]["messages"] + assert messages[-1] == {"role": "assistant", "content": "{"} + + +def test_disable_thinking_env_gates_extra_body(monkeypatch): + monkeypatch.setenv("GRAPHGEN_DISABLE_THINKING", "1") + client = _client(monkeypatch) + calls = _stub_completions(client, ["yes"]) + asyncio.run(client.generate_topk_per_token("hello")) + extra_body = calls["kwargs"][0]["extra_body"] + assert extra_body["chat_template_kwargs"] == {"enable_thinking": False} + + monkeypatch.delenv("GRAPHGEN_DISABLE_THINKING", raising=False) + calls = _stub_completions(client, ["yes"]) + asyncio.run(client.generate_topk_per_token("hello")) + assert "extra_body" not in calls["kwargs"][0] diff --git a/tests/unit_tests/operators/__init__.py b/tests/unit_tests/operators/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/operators/generate/__init__.py b/tests/unit_tests/operators/generate/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/operators/generate/test_gates.py b/tests/unit_tests/operators/generate/test_gates.py new file mode 100644 index 00000000..5e46a02c --- /dev/null +++ b/tests/unit_tests/operators/generate/test_gates.py @@ -0,0 +1,34 @@ +"""Tests for the generic output-gate registry.""" + +import pytest + +from graphgen.operators.generate.gates import GATES, apply_output_gate, register_gate + + +def test_builtin_gate_keeps_only_non_empty_answers(): + gate = GATES["non_empty_answer"] + pairs = [ + {"question": "q1", "answer": "a1"}, + {"question": "q2", "answer": " "}, + {"question": "q3"}, + "not-a-dict", + ] + assert gate(pairs) == [{"question": "q1", "answer": "a1"}] + assert gate(None) == [] + + +def test_apply_unknown_gate_fails_closed(): + with pytest.raises(ValueError, match="output_gate_not_registered"): + apply_output_gate("no_such_gate", []) + + +def test_register_duplicate_gate_fails_closed(): + @register_gate("dup_gate_test") + def _gate(qa_pairs): + return [] + + with pytest.raises(ValueError, match="output_gate_already_registered"): + + @register_gate("dup_gate_test") + def _other(qa_pairs): + return [] diff --git a/tests/unit_tests/operators/generate/test_llm_concurrency.py b/tests/unit_tests/operators/generate/test_llm_concurrency.py new file mode 100644 index 00000000..23e5f7c1 --- /dev/null +++ b/tests/unit_tests/operators/generate/test_llm_concurrency.py @@ -0,0 +1,34 @@ +"""Tests for bounded generator LLM concurrency.""" + +import asyncio + +import pytest + +from graphgen.operators.generate.llm_concurrency import bound_llm_concurrency + + +async def _fake_generate(batch, *, tracker=None, delay=0.01): + tracker["inflight"] += 1 + tracker["peak"] = max(tracker["peak"], tracker["inflight"]) + await asyncio.sleep(delay) + tracker["inflight"] -= 1 + return [{"question": batch, "answer": "a"}] + + +@pytest.mark.parametrize("limit", [1, 3]) +def test_concurrency_is_bounded(limit): + tracker = {"inflight": 0, "peak": 0} + + async def run(): + generate = bound_llm_concurrency( + lambda batch: _fake_generate(batch, tracker=tracker), limit + ) + await asyncio.gather(*[generate(f"b{i}") for i in range(10)]) + + asyncio.run(run()) + assert tracker["peak"] <= limit + + +def test_non_positive_limit_fails_closed(): + with pytest.raises(ValueError, match="llm_concurrency_must_be_positive"): + bound_llm_concurrency(_fake_generate, 0) diff --git a/tests/unit_tests/templates/__init__.py b/tests/unit_tests/templates/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/templates/test_prompt_profiles.py b/tests/unit_tests/templates/test_prompt_profiles.py new file mode 100644 index 00000000..185406c7 --- /dev/null +++ b/tests/unit_tests/templates/test_prompt_profiles.py @@ -0,0 +1,86 @@ +"""Tests for the versioned prompt-profile mechanism.""" + +import json +from pathlib import Path + +import pytest + +from graphgen.models.generator.atomic_generator import AtomicGenerator +from graphgen.templates import ATOMIC_GENERATION_PROMPT +from graphgen.templates.prompt_profiles import ( + BUILTIN_PROFILE_ROOT, + load_profile_template, + profile_template_sha256s, + resolve_profile_dir, +) + + +class _FakeLLM: + tokenizer = None + + +def test_builtin_example_profile_matches_official_structure(): + override = load_profile_template("example.v1", "atomic") + assert override == ATOMIC_GENERATION_PROMPT + + +def test_unknown_profile_fails_closed(): + with pytest.raises(ValueError, match="prompt_profile_not_found"): + resolve_profile_dir("no_such_profile") + + +def test_missing_template_key_fails_closed(tmp_path: Path): + profile_dir = tmp_path / "my.profile" + profile_dir.mkdir() + with pytest.raises(ValueError, match="prompt_profile_template_missing"): + load_profile_template(str(profile_dir), "atomic") + + +def test_invalid_json_fails_closed(tmp_path: Path): + profile_dir = tmp_path / "my.profile" + profile_dir.mkdir() + (profile_dir / "atomic.json").write_text("{not json", encoding="utf-8") + with pytest.raises(ValueError, match="prompt_profile_json_invalid"): + load_profile_template(str(profile_dir), "atomic") + + +def test_structure_mismatch_fails_closed(tmp_path: Path): + profile_dir = tmp_path / "my.profile" + profile_dir.mkdir() + (profile_dir / "atomic.json").write_text( + json.dumps({"xx": {}, "zh": {}}), encoding="utf-8" + ) + with pytest.raises(ValueError, match="prompt_profile_structure_mismatch"): + load_profile_template(str(profile_dir), "atomic") + + +def test_unknown_template_key_fails_closed(tmp_path: Path): + with pytest.raises(ValueError, match="prompt_profile_unknown_template_key"): + load_profile_template(str(tmp_path), "no_such_key") + + +def test_profile_template_sha256s(tmp_path: Path): + profile_dir = tmp_path / "my.profile" + profile_dir.mkdir() + (profile_dir / "atomic.json").write_text("{}", encoding="utf-8") + digests = profile_template_sha256s(str(profile_dir), ["atomic", "cot"]) + assert set(digests) == {"atomic"} + + +def test_generator_profile_overrides_are_instance_local(): + generator_a = AtomicGenerator(_FakeLLM()) + generator_b = AtomicGenerator(_FakeLLM()) + generator_a.apply_prompt_profile("example.v1") + batch = ( + [("n1", {"description": "alpha"})], + [("n1", "n2", {"description": "beta"})], + ) + prompt_a = generator_a.build_prompt(batch) + prompt_b = generator_b.build_prompt(batch) + assert prompt_a == prompt_b + language = "en" + assert prompt_a.startswith(ATOMIC_GENERATION_PROMPT[language].split("{")[0]) + + +def test_builtin_profile_root_exists(): + assert (BUILTIN_PROFILE_ROOT / "example.v1" / "atomic.json").is_file() diff --git a/tests/unit_tests/utils/__init__.py b/tests/unit_tests/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/utils/test_log.py b/tests/unit_tests/utils/test_log.py new file mode 100644 index 00000000..661a3e9c --- /dev/null +++ b/tests/unit_tests/utils/test_log.py @@ -0,0 +1,13 @@ +"""Tests for the logger context fallback.""" + +import logging + +from graphgen.utils.log import get_current_logger + + +def test_get_current_logger_falls_back_without_context(): + # Outside a Ray worker context no logger is set; the helper must return a + # usable fallback instead of raising. + logger = get_current_logger() + assert isinstance(logger, logging.Logger) + assert logger.name == "GraphGen"