feat: add pre main html

e06084 · e06084 · commit 00012524530e · 2025-11-26T16:52:36.000+08:00
diff --git a/examples/multi_extractor_compare.py b/examples/multi_extractor_compare.py
@@ -16,22 +16,24 @@ def all_extractor_comparison():
     config = {
         "use_preprocessed_html": True,          # 🔑 关键配置：启用预处理HTML模式
         "preprocessed_html_field": "llm_webkit_html"  # 指定预处理HTML字段名
+        "llm_response_html_field": "llm_response_html",  # 指定LLM响应HTML字段名
+        "typical_raw_tag_html_field": "typical_raw_tag_html",  # 指定典型原始标签HTML字段名
     }
 
     webkit_extractor = ExtractorFactory.create("llm-webkit", config=config)
     # 创建magic-extractor抽取器
-    magic_extractor = ExtractorFactory.create("magic-html")
-    # 创建trafilatura抽取器,抽取成markdown
-    trafilatura_extractor = ExtractorFactory.create("trafilatura")
-    # 创建trafilatura抽取器,抽取成txt
-    trafilatura_txt_extractor = ExtractorFactory.create("trafilatura_txt")
-    # 创建resiliparse抽取器
-    resiliparse_extractor = ExtractorFactory.create("resiliparse")
+    # magic_extractor = ExtractorFactory.create("magic-html")
+    # # 创建trafilatura抽取器,抽取成markdown
+    # trafilatura_extractor = ExtractorFactory.create("trafilatura")
+    # # 创建trafilatura抽取器,抽取成txt
+    # trafilatura_txt_extractor = ExtractorFactory.create("trafilatura_txt")
+    # # 创建resiliparse抽取器
+    # resiliparse_extractor = ExtractorFactory.create("resiliparse")
     
     # 运行对比
     evaluator = Evaluator()
-    extractors = [webkit_extractor, magic_extractor, trafilatura_extractor,trafilatura_txt_extractor, resiliparse_extractor]
-    # extractors = [webkit_extractor]
+    # extractors = [webkit_extractor, magic_extractor, trafilatura_extractor,trafilatura_txt_extractor, resiliparse_extractor]
+    extractors = [webkit_extractor]
 
     
     results = evaluator.compare_extractors(
diff --git a/scripts/process_dataset.py b/scripts/process_dataset.py
@@ -0,0 +1,275 @@
+#!/usr/bin/env python3
+"""
+为数据集添加 llm_webkit430_main_html 字段
+
+根据 llm_webkit_extractor.py 的逻辑，从现有字段构建 main_html
+
+使用方法:
+    python scripts/process_dataset.py data/WebMainBench_7887_within_formula_code.jsonl
+"""
+
+import json
+import argparse
+from pathlib import Path
+
+
+def process_single_item(data: dict, verbose: bool = False) -> dict:
+    """
+    为单个JSON对象添加 llm_webkit430_main_html 字段
+    
+    参考 llm_webkit_extractor.py:665-670 的逻辑
+    """
+    try:
+        # 导入必要的模块
+        from llm_web_kit.input.pre_data_json import PreDataJson, PreDataJsonKey
+        from llm_web_kit.main_html_parser.parser.tag_mapping import MapItemToHtmlTagsParser
+        
+        # 从数据中获取字段
+        typical_raw_tag_html = data.get('llm_webkit_html', '')  # 预处理HTML
+        html = data.get('llm_webkit_html', '')  # 预处理HTML
+        llm_response = data.get('llm_response_html', '')  # LLM响应HTML
+        
+        if verbose:
+            print(f"\n📝 字段信息:")
+            print(f"  llm_webkit_html 长度: {len(typical_raw_tag_html)}")
+            print(f"  llm_response_html 长度: {len(llm_response)}")
+        
+        # 检查必要字段
+        if not typical_raw_tag_html:
+            if verbose:
+                print("  ⚠️  llm_webkit_html 字段为空，跳过")
+            data['llm_webkit430_main_html'] = ""
+            return data
+        
+        # 构建 pre_data（参考 llm_webkit_extractor.py:665）
+        pre_data = {
+            'typical_raw_tag_html': typical_raw_tag_html,
+            'typical_raw_html': html,
+            'llm_response': llm_response,
+            'html_source': typical_raw_tag_html
+        }
+        
+        if verbose:
+            print(f"  构建 pre_data 完成")
+        
+        # 转换为 PreDataJson 对象
+        pre_data = PreDataJson(pre_data)
+        
+        if verbose:
+            print(f"  PreDataJson 创建完成")
+        
+        # 映射 - 使用 MapItemToHtmlTagsParser
+        parser = MapItemToHtmlTagsParser({})
+        pre_data = parser.parse_single(pre_data)
+        
+        if verbose:
+            print(f"  映射解析完成")
+        
+        # 提取 main_html
+        main_html = pre_data.get(PreDataJsonKey.TYPICAL_MAIN_HTML, "")
+        
+        if verbose:
+            print(f"  main_html 长度: {len(main_html)}")
+        
+        # 添加新字段
+        data['llm_webkit430_main_html'] = main_html
+        
+        return data
+        
+    except ImportError as e:
+        if verbose:
+            print(f"\n❌ 导入错误: {e}")
+            print("   请确保安装了 llm_web_kit: pip install llm-webkit")
+        data['llm_webkit430_main_html'] = ""
+        return data
+    except Exception as e:
+        if verbose:
+            import traceback
+            print(f"\n⚠️  处理失败: {e}")
+            print(f"   错误详情: {traceback.format_exc()}")
+        # 失败时添加空字段
+        data['llm_webkit430_main_html'] = ""
+        return data
+
+
+def process_dataset(input_file: str, output_file: str = None, verbose: bool = False, test_first: int = None):
+    """
+    处理整个数据集
+    
+    Args:
+        input_file: 输入JSONL文件路径
+        output_file: 输出JSONL文件路径（默认为输入文件名_with_main_html.jsonl）
+        verbose: 是否显示详细信息
+        test_first: 仅处理前N条数据（用于测试）
+    """
+    input_path = Path(input_file)
+    
+    if not input_path.exists():
+        print(f"❌ 文件不存在: {input_file}")
+        return
+    
+    # 确定输出文件名
+    if output_file is None:
+        output_file = str(input_path.parent / f"{input_path.stem}_with_main_html.jsonl")
+    
+    print(f"📄 输入文件: {input_file}")
+    print(f"📄 输出文件: {output_file}")
+    if test_first:
+        print(f"🧪 测试模式: 仅处理前 {test_first} 条数据")
+    
+    # 检查依赖
+    print("\n🔍 检查依赖...")
+    try:
+        from llm_web_kit.input.pre_data_json import PreDataJson, PreDataJsonKey
+        from llm_web_kit.main_html_parser.parser.tag_mapping import MapItemToHtmlTagsParser
+        print("✅ llm_web_kit 模块可用")
+    except ImportError as e:
+        print(f"❌ llm_web_kit 模块未安装: {e}")
+        print("   请运行: pip install llm-webkit")
+        return
+    
+    # 统计信息
+    total = 0
+    success = 0
+    failed = 0
+    
+    # 先统计总行数（用于进度条）
+    print("\n📊 统计总行数...")
+    with open(input_file, 'r', encoding='utf-8') as f:
+        total_lines = sum(1 for _ in f)
+    
+    if test_first:
+        total_lines = min(total_lines, test_first)
+    
+    print(f"📦 总共 {total_lines:,} 条数据\n")
+    
+    # 处理数据
+    print("🔄 开始处理...\n")
+    try:
+        with open(input_file, 'r', encoding='utf-8') as fin, \
+             open(output_file, 'w', encoding='utf-8') as fout:
+            
+            for idx, line in enumerate(fin, 1):
+                # 测试模式：只处理前N条
+                if test_first and idx > test_first:
+                    break
+                
+                if not line.strip():
+                    continue
+                
+                try:
+                    # 解析JSON
+                    data = json.loads(line)
+                    total += 1
+                    
+                    # 显示进度（每100条显示一次）
+                    if total % 100 == 0:
+                        print(f"  处理进度: {total}/{total_lines} ({total/total_lines*100:.1f}%)")
+                    
+                    # 处理单条数据
+                    if verbose and idx <= 3:
+                        print(f"\n处理第 {idx} 条数据...")
+                    
+                    processed_data = process_single_item(data, verbose=(verbose and idx <= 3))
+                    
+                    # 检查是否成功添加字段
+                    if processed_data.get('llm_webkit430_main_html'):
+                        success += 1
+                    else:
+                        failed += 1
+                    
+                    # 写入输出文件
+                    fout.write(json.dumps(processed_data, ensure_ascii=False) + '\n')
+                    
+                except json.JSONDecodeError as e:
+                    print(f"\n⚠️  行 {idx} JSON解析错误: {e}")
+                    failed += 1
+                    # 写入原始行
+                    fout.write(line)
+                except Exception as e:
+                    print(f"\n❌ 行 {idx} 处理错误: {e}")
+                    if verbose:
+                        import traceback
+                        print(traceback.format_exc())
+                    failed += 1
+                    # 写入原始数据
+                    try:
+                        data['llm_webkit430_main_html'] = ""
+                        fout.write(json.dumps(data, ensure_ascii=False) + '\n')
+                    except:
+                        fout.write(line)
+    
+    except Exception as e:
+        print(f"\n❌ 处理过程中发生严重错误: {e}")
+        import traceback
+        print(traceback.format_exc())
+        return
+    
+    # 输出统计信息
+    print("\n" + "="*60)
+    print("✅ 处理完成！")
+    print("="*60)
+    print(f"总处理数: {total:,}")
+    print(f"成功: {success:,} ({success/total*100:.1f}%)" if total > 0 else "成功: 0")
+    print(f"失败: {failed:,} ({failed/total*100:.1f}%)" if total > 0 else "失败: 0")
+    print(f"\n输出文件: {output_file}")
+    print("="*60)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='为数据集添加 llm_webkit430_main_html 字段',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog='''
+示例:
+  # 基本使用
+  python scripts/process_dataset.py data/WebMainBench_7887_within_formula_code.jsonl
+  
+  # 指定输出文件
+  python scripts/process_dataset.py data/WebMainBench_7887_within_formula_code.jsonl \\
+    --output data/WebMainBench_7887_with_main_html.jsonl
+  
+  # 测试前10条数据
+  python scripts/process_dataset.py data/WebMainBench_7887_within_formula_code.jsonl \\
+    --test-first 10 --verbose
+  
+  # 详细模式（显示前3条的处理细节）
+  python scripts/process_dataset.py data/WebMainBench_7887_within_formula_code.jsonl \\
+    --verbose
+        '''
+    )
+    
+    parser.add_argument(
+        'input_file',
+        help='输入JSONL文件路径'
+    )
+    
+    parser.add_argument(
+        '--output',
+        '-o',
+        help='输出JSONL文件路径（默认：输入文件名_with_main_html.jsonl）'
+    )
+    
+    parser.add_argument(
+        '--verbose',
+        '-v',
+        action='store_true',
+        help='显示详细处理信息（仅显示前3条）'
+    )
+    
+    parser.add_argument(
+        '--test-first',
+        '-t',
+        type=int,
+        help='仅处理前N条数据（用于测试）'
+    )
+    
+    args = parser.parse_args()
+    
+    # 处理数据集
+    process_dataset(args.input_file, args.output, args.verbose, args.test_first)
+
+
+if __name__ == '__main__':
+    main()
+
diff --git a/webmainbench/data/dataset.py b/webmainbench/data/dataset.py
@@ -30,6 +30,8 @@ class DataSample:
     tags: Optional[List[str]] = None
     llm_webkit_md: Optional[str] = None
     llm_webkit_html: Optional[str] = None  # 预处理HTML字段
+    llm_response_html: Optional[str] = None  # LLM响应HTML字段
+    typical_raw_tag_html: Optional[str] = None  # 典型原始标签HTML字段
     main_html: Optional[str] = None  # 主要HTML内容字段
     
     # Extracted results (populated during evaluation)
diff --git a/webmainbench/extractors/llm_webkit_extractor.py b/webmainbench/extractors/llm_webkit_extractor.py
@@ -30,6 +30,8 @@ class LLMInferenceConfig:
     enforce_eager: bool = True      # 使用eager模式
     use_preprocessed_html: bool = False  # 是否使用预处理的HTML（跳过HTML简化步骤）
     preprocessed_html_field: str = "llm_webkit_html"  # 预处理HTML字段名
+    llm_response_html_field: str = "llm_response_html"  # LLM响应HTML字段名
+    typical_raw_tag_html_field: str = "typical_raw_tag_html"  # 典型原始标签HTML字段名
 
 
 class TokenState(Enum):
@@ -654,11 +656,20 @@ def _extract_content(self, html: str, url: str = None) -> ExtractionResult:
         start_time = time.time()
         
         try:
-            # 检查是否使用预处理的HTML（跳过HTML简化步骤）
+            # 检查是否使用预处理的HTML（跳过HTML简化步骤）d
             if self.inference_config.use_preprocessed_html:
                 # 传入的html已经是预处理的内容（由Evaluator从指定字段提取），直接用作main_html
                 print(f"📥 使用预处理HTML，跳过HTML简化步骤")
-                content, content_list = self._extract_content_from_main_html(html, url)
+                llm_response = self.inference_config.llm_response_html_field
+                typical_raw_tag_html = self.inference_config.typical_raw_tag_html_field
+                pre_data = {'typical_raw_tag_html': typical_raw_tag_html, 'typical_raw_html': html, 'llm_response': llm_response, 'html_source': typical_raw_tag_html}
+                pre_data = self._PreDataJson(pre_data)
+                # 映射
+                parser = self._MapItemToHtmlTagsParser({})
+                pre_data = parser.parse_single(pre_data)
+                main_html = pre_data.get(self._PreDataJsonKey.TYPICAL_MAIN_HTML, "")
+
+                content, content_list = self._extract_content_from_main_html(main_html, url)
                 
                 extraction_time = time.time() - start_time
                 
