Fix kbcleaning cleaned text leakage

dataflow/operators/knowledge_cleaning/generate/kbc_text_cleaner.py

@@ -11,6 +11,76 @@
 from typing import Union
 
 import re
+
+
+_CLEANED_BLOCK_RE = re.compile(
+    r"<cleaned_start>\s*(.*?)\s*<cleaned_end>",
+    flags=re.IGNORECASE | re.DOTALL,
+)
+_ANSWER_BLOCK_RE = re.compile(
+    r"<answer>\s*(.*?)\s*</answer>",
+    flags=re.IGNORECASE | re.DOTALL,
+)
+_PROCESSING_STEPS_RE = re.compile(
+    r"\s*Processing Steps:\s*"
+    r"1\.\s*\[Tag Analysis\]\s*Classify markup tags\s*"
+    r"2\.\s*\[Reference Extraction\]\s*Isolate images/tables\s*"
+    r"3\.\s*\[Character Audit\]\s*Log special chars\s*"
+    r"4\.\s*\[Structure Check\]\s*Validate hierarchy\s*"
+    r"5\.\s*\[Final Output\]\s*Generate cleaned text.*$",
+    flags=re.IGNORECASE | re.DOTALL,
+)
+_PROCESSING_STEPS_WITHOUT_HEADER_RE = re.compile(
+    r"\s*1\.\s*\[Tag Analysis\]\s*Classify markup tags\s*"
+    r"2\.\s*\[Reference Extraction\]\s*Isolate images/tables\s*"
+    r"3\.\s*\[Character Audit\]\s*Log special chars\s*"
+    r"4\.\s*\[Structure Check\]\s*Validate hierarchy\s*"
+    r"5\.\s*\[Final Output\]\s*Generate cleaned text.*$",
+    flags=re.IGNORECASE | re.DOTALL,
+)
+_ZH_PROCESSING_STEPS_RE = re.compile(
+    r"\s*(?:处理步骤|处理流程|治理步骤)[:：]\s*"
+    r"1[.、]\s*\[(?:标签分析|标记分析)\].*?"
+    r"2[.、]\s*\[(?:引用提取|参考提取)\].*?"
+    r"3[.、]\s*\[(?:字符审核|字符审计)\].*?"
+    r"4[.、]\s*\[(?:结构检查|结构校验)\].*?"
+    r"5[.、]\s*\[(?:最终输出|最终结果)\].*$",
+    flags=re.DOTALL,
+)
+_ZH_PROCESSING_STEPS_WITHOUT_HEADER_RE = re.compile(
+    r"\s*1[.、]\s*\[(?:标签分析|标记分析)\].*?"
+    r"2[.、]\s*\[(?:引用提取|参考提取)\].*?"
+    r"3[.、]\s*\[(?:字符审核|字符审计)\].*?"
+    r"4[.、]\s*\[(?:结构检查|结构校验)\].*?"
+    r"5[.、]\s*\[(?:最终输出|最终结果)\].*$",
+    flags=re.DOTALL,
+)
+
+
+def extract_cleaned_text(text, post_process=None) -> str:
+    """Extract model output while dropping prompt-instruction leakage."""
+    text = "" if text is None else str(text)
+
+    cleaned_match = _CLEANED_BLOCK_RE.search(text)
+    if cleaned_match:
+        text = cleaned_match.group(1)
+    else:
+        answer_match = _ANSWER_BLOCK_RE.search(text)
+        if answer_match:
+            text = answer_match.group(1)
+        text = text.replace("<cleaned_start>", "").replace("<cleaned_end>", "")
+
+    text = _PROCESSING_STEPS_RE.sub("", text)
+    text = _PROCESSING_STEPS_WITHOUT_HEADER_RE.sub("", text)
+    text = _ZH_PROCESSING_STEPS_RE.sub("", text)
+    text = _ZH_PROCESSING_STEPS_WITHOUT_HEADER_RE.sub("", text)
+    text = text.strip()
+
+    if post_process:
+        text = post_process(text)
+    return text
+
+
 @prompt_restrict(
     KnowledgeCleanerPrompt       
 )
@@ -27,7 +97,7 @@ def __init__(self, llm_serving: LLMServingABC, lang="en", prompt_template : Unio
         if prompt_template:
             self.prompt_template = prompt_template
         else:
-            self.prompt_template = KnowledgeCleanerPrompt()
+            self.prompt_template = KnowledgeCleanerPrompt(lang=lang)
 
     @staticmethod
     def get_desc(lang: str = "zh"):
@@ -123,15 +193,16 @@ def run(
         formatted_prompts = self._reformat_prompt(dataframe)
         cleaned = self.llm_serving.generate_from_input(formatted_prompts,"")
 
-        #for each in cleaned, only save the content in <cleaned_start> and <cleaned_end>
+        # Save only the final cleaned text, even if the model leaks prompt steps.
         cleaned_extracted = [
-            str(text).split('<cleaned_start>')[1].split('<cleaned_end>')[0].strip() 
-            if '<cleaned_start>' in str(text) and '<cleaned_end>' in str(text)
-            else str(text).strip()
+            extract_cleaned_text(
+                text,
+                getattr(self.prompt_template, "_post_process", None),
+            )
             for text in cleaned
         ]
         dataframe[self.output_key] = cleaned_extracted
         output_file = storage.write(dataframe)
         self.logger.info(f"Results saved to {output_file}")
 
-        return [output_key]
+        return [output_key]

dataflow/operators/knowledge_cleaning/generate/kbc_text_cleaner_batch.py

@@ -1,4 +1,5 @@
 from dataflow.prompts.kbcleaning import KnowledgeCleanerPrompt
+from dataflow.operators.knowledge_cleaning.generate.kbc_text_cleaner import extract_cleaned_text
 import pandas as pd
 from dataflow.utils.registry import OPERATOR_REGISTRY
 from dataflow import get_logger
@@ -125,11 +126,12 @@ def run(
                 raw_chunks, formatted_prompts = self._reformat_prompt_from_path(chunk_path)
                 cleaned = self.llm_serving.generate_from_input(formatted_prompts, "")
 
-                # for each in cleaned, only save the content in <cleaned_start> and <cleaned_end>
+                # Save only the final cleaned text, even if the model leaks prompt steps.
                 cleaned_extracted = [
-                    text.split('<cleaned_start>')[1].split('<cleaned_end>')[0].strip()
-                    if '<cleaned_start>' in str(text) and '<cleaned_end>' in str(text)
-                    else str(text).strip()
+                    extract_cleaned_text(
+                        text,
+                        getattr(self.prompts, "_post_process", None),
+                    )
                     for text in cleaned
                 ]
                 json_items=[{

dataflow/prompts/kbcleaning.py

@@ -177,35 +177,17 @@ def build_prompt(self, raw_content: str) -> str:
 """
 
       if self.lang == "en":
-         processing_steps = """
-Processing Steps:
-1. [Tag Analysis] Classify markup tags
-2. [Reference Extraction] Isolate images/tables
-3. [Character Audit] Log special chars
-4. [Structure Check] Validate hierarchy
-5. [Final Output] Generate cleaned text
-""".strip()
-         output_requirement = 'Response must contain ONLY cleaned text between <cleaned_start> and <cleaned_end>.'
+         output_requirement = 'Response must contain ONLY the final cleaned text between <cleaned_start> and <cleaned_end>. Do not include analysis, explanations, processing steps, or bullet lists outside the cleaned text.'
       else:
-         processing_steps = """
-处理步骤：
-1. [标签分析] 识别并分类所有标记标签
-2. [引用提取] 分离图片/表格/签名等引用内容
-3. [字符审核] 记录特殊字符变更
-4. [结构检查] 验证文本层级
-5. [最终输出] 生成清洗后文本
-""".strip()
          output_requirement = '响应必须只包含清洗后文本，以<cleaned_start>开头，<cleaned_end>结尾，无其他内容。'
 
       return f"""
 {self.prompt_header}
 
+{output_requirement}
+
 待清洗内容：
 {raw_content}
-
-{processing_steps}
-
-{output_requirement}
 """.strip()
 
    def _post_process(self, cleaned_text: str) -> str: