预测公式先经过正则，再用LLM修正

Author: 1041206149

examples/multi_extractor_compare.py

@@ -8,7 +8,7 @@ def all_extractor_comparison():
     print("\n=== 多抽取器对比演示 ===\n")
 
     # 创建数据集
-    dataset_path = Path("../data/test_math.jsonl")
+    dataset_path = Path("../data/WebMainBench_llm-webkit_v1_WebMainBench_7887_within_formula.jsonl")
     dataset = DataLoader.load_jsonl(dataset_path)
 
     # 创建webkit抽取器
@@ -26,8 +26,8 @@ def all_extractor_comparison():
 
     # 运行对比
     evaluator = Evaluator()
-    extractors = [webkit_extractor, magic_extractor, trafilatura_extractor, resiliparse_extractor]
-    # extractors = [webkit_extractor]
+    # extractors = [webkit_extractor, magic_extractor, trafilatura_extractor, resiliparse_extractor]
+    extractors = [webkit_extractor]
 
 
     results = evaluator.compare_extractors(

webmainbench/data/saver.py

@@ -303,15 +303,13 @@ def save_dataset_with_extraction(results: Union["EvaluationResult", Dict[str, An
 
                     # 解析预测值（predicted）
                     predicted_content = extraction_result.get('extracted_content', '')
-                    # TODO: 这里可以根据需要选择不同的解析方法
                     predicted_parts = BaseMetric._extract_from_markdown(predicted_content, field_name="llm_webkit_md")  # 关键：解析预测内容
                     for part_type in ['code', 'formula', 'table', 'text']:
                         sample_dict[f'{current_extractor_name}_predicted_{part_type}'] = predicted_parts.get(part_type, '')
 
             # 解析真实值（groundtruth）- 只需要解析一次
             if extractor_names:  # 只有当存在extractor时才解析
                 groundtruth_content = sample_dict.get('groundtruth_content', '')
-                # TODO: 这里可以根据需要选择不同的解析方法
                 groundtruth_parts = BaseMetric._extract_from_markdown(groundtruth_content, field_name="groundtruth_content")  # 关键：解析真实内容
                 for part_type in ['code', 'formula', 'table', 'text']:
                     # 使用第一个extractor的名字作为前缀，或者使用通用前缀

webmainbench/metrics/base.py

@@ -293,43 +293,55 @@ def _extract_from_markdown(text: str, field_name: str = None) -> Dict[str, str]:
             if code_content.strip():
                 code_parts.append(code_content)
 
-        # 提取公式 - 根据字段类型决定使用API还是正则
+        # 提取公式 - 新的两步处理逻辑
         formula_parts = []
 
-        # 如果是groundtruth_content，使用正则提取公式
-        if field_name == "llm_webkit_md":
-            print(f"[DEBUG] 检测到groundtruth内容，使用正则提取公式")
-            # 统一的公式提取模式
-            latex_patterns = [
-                r'(?<!\\)\$\$(.*?)(?<!\\)\$\$',  # 行间 $$...$$
-                r'(?<!\\)\\\[(.*?)(?<!\\)\\\]',  # 行间 \[...\]
-                r'(?<!\\)\$(.*?)(?<!\\)\$',  # 行内 $...$
-                r'(?<!\\)\\\((.*?)(?<!\\)\\\)',  # 行内 \(...\)
-            ]
-
-            for pattern in latex_patterns:
-                for match in re.finditer(pattern, text, re.DOTALL):
-                    formula_full = match.group(0)
-                    formula_content = match.group(1)
-                    extracted_segments.append(formula_full)
-                    if formula_content.strip():
-                        formula_parts.append(formula_content.strip())
+        # 第一步：先用正则提取公式
+        regex_formulas = []
+        latex_patterns = [
+            r'(?<!\\)\$\$(.*?)(?<!\\)\$\$',  # 行间 $$...$$
+            r'(?<!\\)\\\[(.*?)(?<!\\)\\\]',  # 行间 \[...\]
+            r'(?<!\\)\$(.*?)(?<!\\)\$',  # 行内 $...$
+            r'(?<!\\)\\\((.*?)(?<!\\)\\\)',  # 行内 \(...\)
+        ]
+
+        for pattern in latex_patterns:
+            for match in re.finditer(pattern, text, re.DOTALL):
+                formula_full = match.group(0)
+                formula_content = match.group(1)
+                extracted_segments.append(formula_full)
+                if formula_content.strip():
+                    regex_formulas.append(formula_content.strip())
+
+        # 第二步：根据字段类型决定是否需要API修正
+        if field_name == "groundtruth_content":
+            print(f"[DEBUG] 检测到groundtruth内容，仅使用正则提取公式")
+            formula_parts = regex_formulas
         else:
-            # 其他内容使用API提取公式
-            cache_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '.cache')
-            os.makedirs(cache_dir, exist_ok=True)
-
-            # 使用文本哈希作为缓存文件名
-            text_hash = hashlib.md5(text.encode('utf-8')).hexdigest()
-            cache_file = os.path.join(cache_dir, f'formula_cache_{text_hash}.json')
-
-            # 使用LLM API提取公式
-            try:
-                from .formula_extractor import extract_formulas_with_llm
-                formula_parts = extract_formulas_with_llm(text, cache_file)
-                print(f"[DEBUG] 公式提取成功，提取到 {len(formula_parts)} 个公式")
-            except Exception as e:
-                print(f"[DEBUG] 公式提取失败: {type(e).__name__}: {e}")
+            print(f"[DEBUG] 检测到llm_webkit_md内容，使用正则+API修正模式")
+            # 对于llm_webkit_md，将正则结果传递给API进行修正
+            if regex_formulas:
+                # 将正则提取的公式作为输入传递给API
+                regex_formulas_text = '\n'.join(regex_formulas)
+                print(f"[DEBUG] 正则提取到 {len(regex_formulas)} 个公式，准备API修正")
+
+                cache_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '.cache')
+                os.makedirs(cache_dir, exist_ok=True)
+
+                # 使用正则结果的哈希作为缓存文件名
+                text_hash = hashlib.md5(regex_formulas_text.encode('utf-8')).hexdigest()
+                cache_file = os.path.join(cache_dir, f'formula_correction_cache_{text_hash}.json')
+
+                try:
+                    from .formula_extractor import correct_formulas_with_llm
+                    corrected_formulas = correct_formulas_with_llm(regex_formulas, cache_file)
+                    formula_parts = corrected_formulas
+                    print(f"[DEBUG] API修正成功，最终得到 {len(formula_parts)} 个公式")
+                except Exception as e:
+                    print(f"[DEBUG] API修正失败: {type(e).__name__}: {e}，使用正则结果")
+                    formula_parts = regex_formulas
+            else:
+                print(f"[DEBUG] 正则未提取到公式，跳过API修正")
                 formula_parts = []
 
         # 提取表格

webmainbench/metrics/formula_extractor.py

@@ -1 +1 @@
-# webmainbench/metrics/formula_extractor.pyimport jsonimport osfrom openai import OpenAIdef extract_formulas_with_llm(text, cache_file=None):    """使用LLM API提取文本中的数学公式"""    # 预检查：如果$符号数量<2，直接返回空列表    dollar_count = text.count('$')    if dollar_count < 2:        print(f"[DEBUG] 输入文本$符号数量为{dollar_count}，小于2，跳过API调用")        return []    # 检查缓存    if cache_file and os.path.exists(cache_file):        try:            with open(cache_file, 'r', encoding='utf-8') as f:                cached_result = json.load(f)                print(f"[DEBUG] 从缓存加载公式: {len(cached_result)} 个")                return cached_result        except Exception as e:            print(f"[DEBUG] 缓存读取失败: {e}")    # API配置    client = OpenAI(        base_url="http://35.220.164.252:3888/v1/",        api_key="sk-PZgDr7sZdt77805Cg8s5ZB9QnGMGke61ovYnHYcHKIYVGHNA"    )    PROMPT = '''任务：请作为信息抽取专家，精确提取所提供 Markdown 文本中的所有数学公式，并按要求输出。### 公式格式说明Markdown 中数学公式通常包括以下两类：- **行内公式（Inline）**：由单个美元符号 `$...$` 包裹，例如：`$E = mc^2$`- **行间公式（Block）**：  - 双美元符号包裹：`$$...$$`，例如：`$$\sum_{i=1}^n i = \frac{n(n+1)}{2}$$`  - 数学代码块：以 ```` ```math ```` 开头和结尾的代码块### 提取要求1. **精准提取**：仅提取正确标记的公式部分，排除普通文本、代码（除非是数学代码块）、注释或无关内容。2. **保持原貌**：提取的公式必须与原文完全一致，不得修改、简化或转译。3. **LaTeX 公式**：若识别到 LaTeX 格式的公式（包括 LaTeX 环境或命令），也应原样提取。4. **区分货币与公式**：避免将美元货币金额（如 `$3.99`）误提取为公式，需结合上下文判断是否为数学表达式。### 输出格式- 提取所有识别到的公式，按出现顺序逐行输出原始字符串。- 每个公式以独立行形式呈现，不附加任何额外信息。- 若无公式，则不返回任何内容。[输入文本开始]'''    try:        print(f"[DEBUG] 开始调用 OpenAI API...")        response = client.chat.completions.create(            model="deepseek-chat",            temperature=0,            messages=[                {"role": "user", "content": PROMPT + f"{text}" + '''[输入文本结束]---请根据上述要求，仅输出提取后的公式内容或空字符串。---请注意：- 绝对不要对公式内容做任何修改或解释。- 确保不遗漏任何符合要求的公式，也不添加非公式文本。Optimized for: clarity, precision, context-awareness, and strict formatting compliance.'''}            ]        )        result_text = response.choices[0].message.content.strip()        print(f"[DEBUG] API 返回原始结果: {repr(result_text)}")        # 解析返回的公式        if not result_text:            formulas = []        else:            formulas = [line.strip() for line in result_text.split('\n') if line.strip()]        print(f"[DEBUG] 解析后的公式列表: {formulas}")        # 保存缓存        if cache_file:            try:                os.makedirs(os.path.dirname(cache_file), exist_ok=True)                with open(cache_file, 'w', encoding='utf-8') as f:                    json.dump(formulas, f, ensure_ascii=False, indent=2)                print(f"[DEBUG] 结果已缓存到: {cache_file}")            except Exception as e:                print(f"[DEBUG] 缓存保存失败: {e}")        return formulas    except Exception as e:        print(f"[DEBUG] API 调用异常: {type(e).__name__}: {e}")        raise e
+# webmainbench/metrics/formula_extractor.pyimport jsonimport osfrom openai import OpenAIdef correct_formulas_with_llm(regex_formulas, cache_file=None):    """使用LLM API修正正则提取的公式"""    if not regex_formulas:        print(f"[DEBUG] 输入公式列表为空，跳过API修正")        return []    # 检查缓存    if cache_file and os.path.exists(cache_file):        try:            with open(cache_file, 'r', encoding='utf-8') as f:                cached_result = json.load(f)                print(f"[DEBUG] 从缓存加载修正结果: {len(cached_result)} 个")                return cached_result        except Exception as e:            print(f"[DEBUG] 缓存读取失败: {e}")    # API配置    client = OpenAI(        base_url="",        api_key=""    )    # 将正则提取的公式转换为文本    formulas_text = '\n'.join(regex_formulas)    CORRECTION_PROMPT = '''任务：请从以下正则表达式提取的内容中，识别并保留真正的LaTeX数学公式，剔除货币形式的内容。    ### 识别规则    **真正的数学公式**（保留）：    - 包含数学符号：+ - × ÷ = < > ≤ ≥ ± ∞ ∑ ∫ ∂ √ ^ _ { } 等    - 包含希腊字母：α β γ δ θ λ μ π σ ω 等    - 包含LaTeX命令：\\frac \\sum \\int \\sqrt \\alpha \\beta \\sin \\cos 等    - 包含数学表达式：变量、函数、方程等    **货币形式内容**（剔除）：    - 仅包含数字、逗号、小数点的价格：如 1,150.00    - 纯粹的金额数值：如 25.99、1,200、5.50    - 不包含任何数学运算符或数学符号的数字    ### 处理要求    1. **严格区分**：只保留真正的数学公式，剔除所有货币价格    2. **格式标准化**：统一公式格式，确保LaTeX语法正确    3. **保持原意**：不修改数学公式内容        ### 输出格式    - 每个有效的数学公式独占一行    - 只输出公式内容，不包含$符号或其他包装    - 如果输入不是有效的数学公式（如货币），则输出<空>    - 按原顺序输出保留的公式    ### 示例 1 (剔除后有有效公式)    输入：1,150.00 → 剔除（货币）    输入：x^2 + y^2 = r^2 → 保留（数学公式）    输入：25.99 → 剔除（货币）    输入：\\frac{a}{b} + c → 保留（数学公式）        ### 示例 2 (剔除后无有效公式)    输入：1,150.00 → 剔除（货币）    输入：25.99 → 剔除（货币）        输出：<空>        注意，输出结果中不要添加任何解释！。    [输入内容列表开始]'''    try:        print(f"[DEBUG] 开始调用 OpenAI API 进行公式修正...")        response = client.chat.completions.create(            model="deepseek-chat",            temperature=0,            messages=[                {"role": "user", "content": CORRECTION_PROMPT + f"\n{formulas_text}\n" + '''[输入内容列表结束]---请按要求识别并输出真正的数学公式，剔除货币形式的内容。---'''}            ]        )        result_text = response.choices[0].message.content.strip()        print(f"[DEBUG] API 返回修正结果: {repr(result_text)}")        # 检测返回内容是否包含"空"字 - 如果包含则整个结果为空        if '空' in result_text:            print(f"[DEBUG] 检测到API返回包含'空'字，将整个结果设置为空列表")            corrected_formulas = []        elif not result_text:            corrected_formulas = []        else:            # 正常解析返回的公式列表            corrected_formulas = [line.strip() for line in result_text.split('\n') if line.strip()]        print(f"[DEBUG] 修正后的公式列表: {corrected_formulas}")        # 保存缓存        if cache_file:            try:                os.makedirs(os.path.dirname(cache_file), exist_ok=True)                with open(cache_file, 'w', encoding='utf-8') as f:                    json.dump(corrected_formulas, f, ensure_ascii=False, indent=2)                print(f"[DEBUG] 修正结果已缓存到: {cache_file}")            except Exception as e:                print(f"[DEBUG] 缓存保存失败: {e}")        return corrected_formulas    except Exception as e:        print(f"[DEBUG] API 修正异常: {type(e).__name__}: {e}")        print(f"[DEBUG] 回退到原始正则结果")        return regex_formulas