Merge pull request #12 from e06084/main

feat: add save_dataset_with_extraction

Author: e06084

examples/basic_usage.py

@@ -699,11 +699,76 @@ def hello_world():
     print("✅ LLM-WebKit 6项指标评测完成！")
 
 
+def demo_dataset_with_extraction():
+    """演示保存带有抽取内容的数据集"""
+    print("=== 演示：保存带有抽取内容的数据集 ===")
+    
+    from webmainbench import DataLoader, DataSaver, Evaluator, ExtractorFactory
+    from pathlib import Path
+    
+    # 从文件加载数据集
+    data_dir = Path("data")
+    dataset_path = data_dir / "sample_dataset.jsonl"
+    # dataset_path = "/Users/chupei/Downloads/WebMainBench_dataset_merge_2549.jsonl"
+    
+    print(f"📂 从文件加载数据集: {dataset_path}")
+    dataset = DataLoader.load_jsonl(dataset_path, include_results=False)
+    dataset.name = "WebMainBench_with_extraction"
+    dataset.description = "演示抽取内容保存的测试数据集"
+    
+    print(f"📊 加载数据集完成，包含 {len(dataset.samples)} 个样本")
+    
+    # 创建抽取器并运行评测
+    try:
+        extractor = ExtractorFactory.create("llm-webkit", config={"model_path": "/Users/chupei/model/checkpoint-3296"})
+        print(f"🤖 使用抽取器: {extractor.name}")
+    except Exception as e:
+        print(f"⚠️ LLM-WebKit抽取器创建失败，使用mock抽取器: {e}")
+        extractor = ExtractorFactory.create("mock")
+    
+    # 运行评测
+    evaluator = Evaluator()
+    result = evaluator.evaluate(dataset, extractor)
+    
+    print(f"✅ 评测完成，总体得分: {result.overall_metrics.get('overall', 0):.4f}")
+    
+    # 保存带有抽取内容的数据集
+    results_dir = Path("results")
+    enriched_dataset_path = results_dir / f"{dataset.name}_with_{extractor.name}_extraction.jsonl"
+    
+    DataSaver.save_dataset_with_extraction(
+        results=result,
+        dataset=dataset, 
+        file_path=enriched_dataset_path,
+        extractor_name=extractor.name
+    )
+    
+    print(f"💾 已保存带有抽取内容的数据集到: {enriched_dataset_path}")
+    
+    # 保存评测结果和摘要报告
+    evaluation_results_path = results_dir / f"{dataset.name}_{extractor.name}_evaluation_results.json"
+    summary_report_path = results_dir / f"{dataset.name}_{extractor.name}_evaluation_report.csv"
+    
+    DataSaver.save_evaluation_results(result, evaluation_results_path)
+    DataSaver.save_summary_report(result, summary_report_path)
+    
+    print(f"📊 已保存评测结果到: {evaluation_results_path}")
+    print(f"📈 已保存摘要报告到: {summary_report_path}")
+    
+    # 显示保存的字段信息
+    print("\n📋 保存的新字段包括:")
+    print(f"  - {extractor.name}_content: 抽取的内容")
+    print(f"  - {extractor.name}_content_list: 抽取的结构化内容列表")
+    print(f"  - {extractor.name}_success: 抽取是否成功")
+    print(f"  - {extractor.name}_time: 抽取耗时")
+    print(f"  - {extractor.name}_*_score: 各项指标分数")
+
 if __name__ == "__main__":
     try:
         demo_basic_mock_evaluation()
         demo_llm_webkit_evaluation()  # 使用LLM-WebKit评测示例
         demo_extractor_comparison()
+        demo_dataset_with_extraction()  # 演示保存带有抽取内容的数据集
         print("\n✅ 示例运行完成！")
 
     except Exception as e:

results/WebMainBench_with_extraction_llm-webkit_evaluation_report.csv

@@ -0,0 +1,2 @@
+extractor,dataset,total_samples,success_rate,overall,code_edit,formula_edit,table_TEDS,table_edit,text_edit
+llm-webkit,WebMainBench_with_extraction,4,1.0,0.8597,0.6215,1.0,1.0,0.9073,0.7695