isage-refiner-benchmark 0.1.0.1__cp311-none-any.whl

isage_refiner_benchmark-0.1.0.1.dist-info/METADATA

@@ -0,0 +1,62 @@
+Metadata-Version: 2.4
+Name: isage-refiner-benchmark
+Version: 0.1.0.1
+Summary: Context compression and refiner algorithm benchmark framework for the SAGE ecosystem
+Author-email: IntelliStream Team <shuhao_zhang@hust.edu.cn>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/intellistream/sage-refiner-benchmark
+Project-URL: Documentation, https://github.com/intellistream/sage-refiner-benchmark#readme
+Project-URL: Repository, https://github.com/intellistream/sage-refiner-benchmark
+Project-URL: Issues, https://github.com/intellistream/sage-refiner-benchmark/issues
+Keywords: sage,benchmark,refiner,context-compression,evaluation,intellistream
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: isage-common
+Requires-Dist: isage-kernel
+Requires-Dist: isage-libs
+Requires-Dist: datasets>=2.14.0
+Provides-Extra: refiner
+Requires-Dist: jieba>=0.42; extra == "refiner"
+Requires-Dist: fuzzywuzzy>=0.18; extra == "refiner"
+Requires-Dist: python-Levenshtein>=0.12; extra == "refiner"
+Requires-Dist: rouge>=1.0.0; extra == "refiner"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
+Requires-Dist: black>=23.0.0; extra == "dev"
+Requires-Dist: ruff==0.14.6; extra == "dev"
+Requires-Dist: pre-commit>=3.0.0; extra == "dev"
+Requires-Dist: mypy>=1.0.0; extra == "dev"
+Dynamic: license-file
+
+# SAGE Refiner Benchmark
+
+**上下文压缩与 Refiner 算法性能评估框架** - 评测 LongRefiner、REFORM、Provence 等多种上下文压缩算法。
+
+**Context Compression & Refiner Benchmark Framework** - Evaluates multiple context compression algorithms including LongRefiner, REFORM, Provence, and more.
+
+## Install
+
+```bash
+pip install isage-refiner-benchmark
+```
+
+## Features
+
+- 📊 多种上下文压缩算法支持
+- ⚡ 标准化评估指标
+- 🔧 灵活的配置系统
+- 📈 详细的性能分析报告
+
+## Repository
+
+- GitHub: https://github.com/intellistream/sage-refiner-benchmark

isage_refiner_benchmark-0.1.0.1.dist-info/RECORD

@@ -0,0 +1,14 @@
+isage_refiner_benchmark-0.1.0.1.dist-info/licenses/LICENSE,sha256=vBNVIGkYYZY0B8f0Ui1ITYwRu7WNtSwyxvIAVGYS6jU,1075
+sage/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sage/benchmark_refiner/__init__.py,sha256=ps5uQv6Es7E0b53hUvpyOWGQE8b22z76vVPGBzIWanU,1812
+sage/benchmark_refiner/_version.py,sha256=jsH11gEb0e-Ciggsva1C1QDkTp82TcG33ZVQp2QhXB0,153
+sage/benchmark_refiner/batch.py,sha256=V_-b65827BauiGmYK97WVtYdVCdLEE7MMwlbEGXNIbc,5126
+sage/benchmark_refiner/constants.py,sha256=Rfs7JO34iqcqVqwADH1IMfnb_kbxixBQnLzwJknswp4,1467
+sage/benchmark_refiner/evaluator.py,sha256=cJExCQQDnc7zilJyBuLdZSIrNz1e-mldHD0UH3DOA_M,12846
+sage/benchmark_refiner/metrics.py,sha256=_HtI5ERi4P4pGVogQOFhLW_BBYPg171ThnMhGn7GdVE,7274
+sage/benchmark_refiner/promptor.py,sha256=rz6hhYTT0HaxtHQuFo-tOb0sZBL9B1CTPpoM__wMgWA,7690
+sage/benchmark_refiner/utils.py,sha256=lo7NBY_1keB6Mczel2RY1vgeVM3MXZetP_hj2nvgHS4,3044
+isage_refiner_benchmark-0.1.0.1.dist-info/METADATA,sha256=1arb06fh6yaHhP_HpEF7bh1iaYpKe55ZslVS-6jOKUA,2428
+isage_refiner_benchmark-0.1.0.1.dist-info/WHEEL,sha256=yk-B4c9kYsinhQ_MzhPAVcDm9mhkAVmdo0rg0jgFCmo,94
+isage_refiner_benchmark-0.1.0.1.dist-info/top_level.txt,sha256=hibFyzQHiLOMK68qL1OWsNKaXOmSXqZjeLTBem6Yy7I,5
+isage_refiner_benchmark-0.1.0.1.dist-info/RECORD,,

isage_refiner_benchmark-0.1.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.10.2)
+Root-Is-Purelib: true
+Tag: cp311-none-any
+

isage_refiner_benchmark-0.1.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 IntelliStream Team
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

isage_refiner_benchmark-0.1.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+sage

sage/benchmark_refiner/__init__.py

@@ -0,0 +1,65 @@
+"""
+Context Compression & Refiner Benchmark Module
+===============================================
+
+上下文压缩与 Refiner 算法性能评估组件。
+
+支持的算法/数据集：
+- LongBench (THUDM/LongBench) - 长文本理解基准
+- LongRefiner - 上下文压缩算法
+- REFORM - 检索优化
+- Provence - 上下文剪枝
+- 其他自定义 Refiner 实现
+
+组件：
+- LongBenchBatch: 从 THUDM/LongBench 加载数据
+- LongBenchPromptor: 生成 LongBench 官方 prompt
+- LongBenchEvaluator: LongBench 官方评估指标
+
+使用示例：
+    from sage.benchmark_refiner import (
+        LongBenchBatch,
+        LongBenchPromptor,
+        LongBenchEvaluator,
+    )
+
+    # 在 pipeline 中使用
+    env.from_batch(LongBenchBatch, config["source"])
+       .map(LongBenchPromptor, config["promptor"])
+       .map(OpenAIGenerator, config["generator"])
+       .map(LongBenchEvaluator, config["evaluate"])
+
+配置示例：
+    source:
+      hf_dataset_name: "THUDM/LongBench"
+      hf_dataset_config: "hotpotqa"
+      max_samples: 100
+
+    promptor:
+      max_input_tokens: 120000
+      is_chat_model: true
+      model_name_or_path: "Qwen/Qwen2.5-7B-Instruct"
+
+    evaluate:
+      longbench_e_buckets: false
+      output_path: "results/longbench_results.jsonl"
+      model_name: "Qwen/Qwen2.5-7B-Instruct"
+
+依赖安装：
+    pip install isage-refiner-benchmark[refiner]
+
+    # 可选依赖说明：
+    # - jieba: 中文分词（中文数据集评估需要）
+    # - fuzzywuzzy + python-Levenshtein: 代码相似度（lcc, repobench-p 需要）
+    # - rouge: ROUGE-L 分数（摘要任务需要）
+"""
+
+from .batch import LongBenchBatch
+from .evaluator import LongBenchEvaluator
+from .promptor import LongBenchPromptor
+
+__all__ = [
+    "LongBenchBatch",
+    "LongBenchPromptor",
+    "LongBenchEvaluator",
+]

sage/benchmark_refiner/_version.py

@@ -0,0 +1,4 @@
+"""Version information for sage-longbench-benchmark."""
+__version__ = "0.1.0.0"
+__author__ = "IntelliStream Team"
+__email__ = "shuhao_zhang@hust.edu.cn"

sage/benchmark_refiner/batch.py

@@ -0,0 +1,123 @@
+"""
+LongBench Batch Processing
+==========================
+
+LongBench 数据集批处理函数，从 THUDM/LongBench 加载数据并转换为 SAGE 标准格式。
+
+迁移自 sage-libs/foundation/io/batch.py，遵循 SAGE 架构：
+- benchmark 相关组件统一放在 sage-benchmark (L5)
+"""
+
+from typing import Any
+
+from sage.libs.foundation.io import HFDatasetBatch
+
+
+class LongBenchBatch(HFDatasetBatch):
+    """
+    LongBench 数据集批处理函数
+
+    专门用于 LongBench 长文本理解基准测试，字段映射：
+    - input → query（用户问题）
+    - context → context（长文本上下文，LongBench 自带，无需检索）
+    - answers → references（标准答案列表）
+    - all_classes → all_classes（分类任务的类别列表）
+    - length → length（原始文本长度，用于 LongBench-E 分桶评估）
+
+    **与 SAGE RAG Pipeline 的对齐说明**：
+
+    SAGE RAG 标准数据流:
+    - query: 用户问题
+    - references: 标准答案（评估用）
+    - retrieval_results: 检索到的文档 List[Dict]（Retriever 输出）
+    - refining_results: 压缩后的文档 List[str]（Refiner 输出）
+    - context: 上下文字符串或列表（Promptor 读取）
+    - generated: 生成的答案（Generator 输出）
+
+    LongBench 特殊处理:
+    - LongBench 自带 context，跳过 Retriever 阶段
+    - context 直接作为 `context` 字段供 Promptor 使用
+    - 同时设置 `retrieval_results` 为空列表（表示无检索）
+
+    Input: None (直接从HF数据集读取)
+    Output: SAGE RAGResponse 兼容格式 + LongBench 专用字段
+
+    Config Keys (继承自 HFDatasetBatch):
+        hf_dataset_name: str - 固定为 "THUDM/LongBench"
+        hf_dataset_config: str - 如 "multi_news", "hotpotqa", "multi_news_e" 等
+        hf_split: str - 默认 "test"
+        max_samples: int - 最大样本数限制
+
+    Output Fields (SAGE RAG 标准字段):
+        query: str - 用户问题（来自 LongBench input）
+        references: List[str] - 标准答案列表（来自 LongBench answers，评估用）
+        context: str - 长文本上下文（来自 LongBench context，供 Promptor 使用）
+        retrieval_results: List[Dict] - 空列表（LongBench 不走检索）
+
+    Output Fields (LongBench 专用字段):
+        all_classes: List[str] | None - 分类任务类别（trec, lsht 等）
+        length: int - 原始文本长度（LongBench-E 分桶评估用）
+        _dataset: str - 数据集名称（用于选择评估指标）
+        _is_longbench_e: bool - 是否是 LongBench-E 版本
+    """
+
+    def __init__(self, config: dict | None = None, **kwargs):
+        super().__init__(config, **kwargs)
+        # 解析数据集名称和是否为 LongBench-E
+        self._dataset_name = self._parse_dataset_name()
+        self._is_longbench_e = self._check_longbench_e()
+
+    def _parse_dataset_name(self) -> str:
+        """从 hf_dataset_config 解析数据集名称（去除 _e 后缀）"""
+        config_name = self.hf_config or ""
+        if config_name.endswith("_e"):
+            return config_name[:-2]  # 去除 _e 后缀
+        return config_name
+
+    def _check_longbench_e(self) -> bool:
+        """检查是否是 LongBench-E 版本"""
+        config_name = self.hf_config or ""
+        return config_name.endswith("_e")
+
+    def _build_iter(self):
+        """构建 LongBench 数据集迭代器，重写父类方法"""
+        try:
+            from datasets import load_dataset
+        except ImportError:
+            raise ImportError(
+                "datasets library is required for LongBenchBatch. "
+                "Install with: pip install datasets"
+            )
+
+        ds = load_dataset(self.hf_name, self.hf_config, split=self.hf_split, streaming=True)
+        for ex in ds:
+            if isinstance(ex, dict):
+                yield self._transform_example(ex)
+
+    def _transform_example(self, ex: dict[str, Any]) -> dict[str, Any]:
+        """
+        LongBench 字段映射到 SAGE RAG Pipeline 标准格式
+
+        LongBench 原始字段 → SAGE RAG 标准字段:
+        - input → query（用户问题）
+        - context → context（上下文，供 Promptor 使用）
+        - answers → references（标准答案，供 Evaluate 使用）
+
+        LongBench 专用字段保留:
+        - all_classes（分类任务类别）
+        - length（原始长度，LongBench-E 分桶用）
+        """
+        return {
+            # ========== SAGE RAG Pipeline 标准字段 ==========
+            "query": ex.get("input", ""),
+            "references": ex.get("answers") or [],
+            "context": ex.get("context", ""),
+            # 空列表表示跳过检索阶段（LongBench 自带 context）
+            "retrieval_results": [],
+            # ========== LongBench 专用字段 ==========
+            "all_classes": ex.get("all_classes"),
+            "length": ex.get("length", 0),
+            # ========== 内部元数据（下划线前缀，pipeline 流转用）==========
+            "_dataset": self._dataset_name,
+            "_is_longbench_e": self._is_longbench_e,
+        }

sage/benchmark_refiner/constants.py

@@ -0,0 +1,47 @@
+"""
+LongBench 常量定义
+
+数据集到评估指标的映射（照搬官方 eval.py）
+"""
+
+# 数据集到评估指标的映射（来自官方 eval.py）
+DATASET_TO_METRIC: dict[str, str] = {
+    # QA 任务 - F1 score
+    "narrativeqa": "qa_f1",
+    "qasper": "qa_f1",
+    "multifieldqa_en": "qa_f1",
+    "hotpotqa": "qa_f1",
+    "2wikimqa": "qa_f1",
+    "musique": "qa_f1",
+    "triviaqa": "qa_f1",
+    # 中文 QA（需要 jieba 分词）
+    "multifieldqa_zh": "qa_f1_zh",
+    # 摘要任务 - ROUGE score
+    "gov_report": "rouge",
+    "qmsum": "rouge",
+    "multi_news": "rouge",
+    "samsum": "rouge",
+    # 中文摘要（需要 jieba 分词）
+    "dureader": "rouge_zh",
+    "vcsum": "rouge_zh",
+    # 分类任务
+    "trec": "classification",
+    "lsht": "classification",
+    # 检索任务
+    "passage_retrieval_en": "retrieval",
+    "passage_retrieval_zh": "retrieval_zh",
+    "passage_count": "count",
+    # 代码任务（需要 fuzzywuzzy）
+    "lcc": "code_sim",
+    "repobench-p": "code_sim",
+}
+
+# 需要取第一行的数据集（照搬官方 eval.py）
+FIRST_LINE_DATASETS: set[str] = {"trec", "triviaqa", "samsum", "lsht"}
+
+# 不使用 chat template 的数据集（照搬官方 pred.py 注释）
+# chat models are better off without build prompts on these tasks
+NO_CHAT_DATASETS: set[str] = {"trec", "triviaqa", "samsum", "lsht", "lcc", "repobench-p"}
+
+# 支持的数据集列表
+SUPPORTED_DATASETS: set[str] = set(DATASET_TO_METRIC.keys())

sage/benchmark_refiner/evaluator.py

@@ -0,0 +1,339 @@
+"""
+LongBench Evaluator - 优化版
+
+改进点：
+1. 模型专用后处理（照搬官方 post_process）
+2. 结果持久化（JSONL 逐条写入）
+3. 压缩文本追踪
+4. 时间统计
+"""
+
+import json
+from pathlib import Path
+from typing import Any, Optional, TextIO
+
+from sage.common.core import StopSignal
+from sage.common.core.functions import MapFunction as MapOperator
+
+from .constants import DATASET_TO_METRIC, FIRST_LINE_DATASETS
+from .metrics import METRIC_FUNCTIONS
+from .utils import post_process
+
+
+class LongBenchEvaluator(MapOperator):
+    """
+    LongBench 专用评估器。
+
+    功能：
+    1. 根据数据集自动选择评估指标
+    2. 支持标准版单一分数和 LongBench-E 长度分桶
+    3. 集成所有 LongBench 指标函数
+    4. 预测结果后处理（特定数据集只取第一行）
+    5. 模型专用后处理（照搬官方 post_process）
+    6. 结果持久化（JSONL 格式）
+
+    输入数据格式（来自 Generator）：
+    {
+        "query": str,
+        "generated": str,          # 模型生成的答案
+        "references": List[str],   # 标准答案列表
+        "_dataset": str,           # 数据集名称
+        "all_classes": List[str],  # 分类任务类别（可选）
+        "length": int,             # 原始长度（LongBench-E 分桶用）
+    }
+
+    配置参数：
+        - longbench_e_buckets: bool - 是否输出 LongBench-E 分桶分数
+        - output_path: str | None - 结果保存路径（JSONL 格式）
+        - model_name: str | None - 模型名称（用于后处理）
+    """
+
+    def __init__(self, config: Optional[dict[str, Any]] = None, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.config = config or {}
+        self.longbench_e_buckets: bool = self.config.get("longbench_e_buckets", False)
+        self.model_name: str = self.config.get("model_name", "")
+
+        # 结果持久化
+        self._output_path: Optional[str] = self.config.get("output_path")
+        self._output_file: Optional[TextIO] = None
+        if self._output_path:
+            Path(self._output_path).parent.mkdir(parents=True, exist_ok=True)
+            self._output_file = open(self._output_path, "a", encoding="utf-8")
+
+        # 分数收集器（用于计算平均分）
+        self._scores: list[float] = []
+        self._dataset_scores: dict[str, list[float]] = {}
+
+        # LongBench-E 分桶分数
+        self._bucket_scores: dict[str, list[float]] = {
+            "0-4k": [],
+            "4-8k": [],
+            "8k+": [],
+        }
+
+        # 时间收集器
+        self._refine_times: list[float] = []
+        self._generate_times: list[float] = []
+        self._retrieve_times: list[float] = []
+
+    def _post_process_prediction(self, pred: str, dataset: str) -> str:
+        """预测结果后处理"""
+        # 1. 模型专用后处理（照搬官方）
+        if self.model_name:
+            pred = post_process(pred, self.model_name)
+
+        # 2. 数据集专用后处理：取第一行（照搬官方 eval.py）
+        if dataset in FIRST_LINE_DATASETS:
+            pred = pred.lstrip("\n").split("\n")[0]
+
+        return pred
+
+    def _get_length_bucket(self, length: int) -> str:
+        """根据长度获取分桶名称"""
+        if length < 4000:
+            return "0-4k"
+        elif length < 8000:
+            return "4-8k"
+        return "8k+"
+
+    def _compute_score(
+        self,
+        pred: str,
+        ground_truths: list[str],
+        dataset: str,
+        all_classes: Optional[list[str]] = None,
+    ) -> float:
+        """计算单个样本的分数"""
+        # 获取指标类型
+        metric_type = DATASET_TO_METRIC.get(dataset, "qa_f1")
+        metric_fn = METRIC_FUNCTIONS.get(metric_type)
+
+        if metric_fn is None:
+            self.logger.warning(f"Unknown metric type: {metric_type}")
+            return 0.0
+
+        # 后处理预测结果
+        pred = self._post_process_prediction(pred, dataset)
+
+        # 对所有参考答案计算分数，取最高（照搬官方 eval.py scorer）
+        best_score = 0.0
+        for ground_truth in ground_truths:
+            try:
+                score = metric_fn(pred, ground_truth, all_classes=all_classes or [])
+                best_score = max(best_score, score)
+            except Exception as e:
+                self.logger.warning(f"Error computing score for {dataset}: {e}")
+
+        return best_score
+
+    def _save_result(self, data: dict[str, Any], score: float, dataset: str) -> None:
+        """保存单条结果到 JSONL 文件"""
+        if not self._output_file:
+            return
+
+        result: dict[str, Any] = {
+            "pred": data.get("generated", ""),
+            "answers": data.get("references", []),
+            "score": score,
+            "dataset": dataset,
+            "all_classes": data.get("all_classes"),
+            "length": data.get("length", 0),
+        }
+
+        # 保存压缩后的上下文（如果存在）
+        refining_results = data.get("refining_results", [])
+        if refining_results:
+            if isinstance(refining_results, list):
+                result["compressed_text"] = "\n\n".join(refining_results)
+            else:
+                result["compressed_text"] = str(refining_results)
+
+        # 保存时间数据
+        times: dict[str, float] = {}
+        if "retrieve_time" in data:
+            times["retrieve"] = data["retrieve_time"]
+        if "refine_time" in data:
+            times["refine"] = data["refine_time"]
+        if "generate_time" in data:
+            times["generate"] = data["generate_time"]
+        if times:
+            result["times"] = times
+
+        self._output_file.write(json.dumps(result, ensure_ascii=False) + "\n")
+        self._output_file.flush()
+
+    def execute(self, data: Any) -> Any:
+        """执行评估"""
+        # Handle StopSignal - 输出汇总统计
+        if isinstance(data, StopSignal):
+            self._print_summary()
+            if self._output_file:
+                self._output_file.close()
+            return data
+
+        # 获取必要字段
+        dataset: str = data.get("_dataset", "unknown")
+        pred: str = data.get("generated", "")
+        references: list[str] = data.get("references", [])
+        all_classes: Optional[list[str]] = data.get("all_classes")
+        length: int = data.get("length", 0)
+
+        # 计算分数
+        score = self._compute_score(pred, references, dataset, all_classes)
+
+        # 分数 * 100（与原始 LongBench 一致）
+        score_percent = round(score * 100, 2)
+
+        # 收集分数
+        self._scores.append(score)
+        if dataset not in self._dataset_scores:
+            self._dataset_scores[dataset] = []
+        self._dataset_scores[dataset].append(score)
+
+        # LongBench-E 分桶
+        if self.longbench_e_buckets and length > 0:
+            bucket = self._get_length_bucket(length)
+            self._bucket_scores[bucket].append(score)
+
+        # 收集时间数据（由 MapOperator 自动添加）
+        if "refine_time" in data:
+            self._refine_times.append(data["refine_time"])
+        if "generate_time" in data:
+            self._generate_times.append(data["generate_time"])
+        if "retrieve_time" in data:
+            self._retrieve_times.append(data["retrieve_time"])
+
+        # 保存结果到 JSONL
+        self._save_result(data, score, dataset)
+
+        # 打印单个样本分数和时间
+        metric_type = DATASET_TO_METRIC.get(dataset, "qa_f1")
+        total_time = (
+            data.get("retrieve_time", 0) + data.get("refine_time", 0) + data.get("generate_time", 0)
+        )
+        time_str = f" (time={total_time:.3f}s)" if total_time > 0 else ""
+        print(f"\033[92m[LongBench {dataset}] {metric_type}: {score_percent}{time_str}\033[0m")
+
+        # 将分数添加到数据中
+        data["longbench_score"] = score
+        data["longbench_score_percent"] = score_percent
+        data["longbench_metric"] = metric_type
+
+        return data
+
+    def _print_summary(self) -> None:
+        """打印汇总统计"""
+        if not self._scores:
+            print("\n" + "=" * 80)
+            print("No LongBench samples processed")
+            print("=" * 80)
+            return
+
+        print("\n" + "=" * 80)
+        print(f"LONGBENCH EVALUATION SUMMARY ({len(self._scores)} samples)")
+        print("=" * 80)
+
+        # 总体平均分
+        avg_score = sum(self._scores) / len(self._scores) * 100
+        print(f"\033[92m[Overall Average Score]: {avg_score:.2f}\033[0m")
+
+        # 按数据集分组的平均分
+        if self._dataset_scores:
+            print("\n--- Per-Dataset Scores ---")
+            for dataset, scores in sorted(self._dataset_scores.items()):
+                avg = sum(scores) / len(scores) * 100
+                metric_type = DATASET_TO_METRIC.get(dataset, "qa_f1")
+                print(f"  {dataset} ({metric_type}): {avg:.2f} ({len(scores)} samples)")
+
+        # LongBench-E 分桶分数
+        if self.longbench_e_buckets:
+            print("\n--- LongBench-E Length Buckets ---")
+            for bucket, scores in self._bucket_scores.items():
+                if scores:
+                    avg = sum(scores) / len(scores) * 100
+                    print(f"  {bucket}: {avg:.2f} ({len(scores)} samples)")
+
+        # 时间统计
+        has_time_data = self._refine_times or self._generate_times or self._retrieve_times
+        if has_time_data:
+            print("\n--- Timing Statistics (seconds) ---")
+            if self._retrieve_times:
+                avg_retrieve = sum(self._retrieve_times) / len(self._retrieve_times)
+                total_retrieve = sum(self._retrieve_times)
+                print(
+                    f"  Retrieve: avg={avg_retrieve:.3f}s, total={total_retrieve:.2f}s ({len(self._retrieve_times)} samples)"
+                )
+            if self._refine_times:
+                avg_refine = sum(self._refine_times) / len(self._refine_times)
+                total_refine = sum(self._refine_times)
+                print(
+                    f"  Refine:   avg={avg_refine:.3f}s, total={total_refine:.2f}s ({len(self._refine_times)} samples)"
+                )
+            if self._generate_times:
+                avg_generate = sum(self._generate_times) / len(self._generate_times)
+                total_generate = sum(self._generate_times)
+                print(
+                    f"  Generate: avg={avg_generate:.3f}s, total={total_generate:.2f}s ({len(self._generate_times)} samples)"
+                )
+            # 总时间
+            total_time = (
+                sum(self._retrieve_times) + sum(self._refine_times) + sum(self._generate_times)
+            )
+            print(f"  \033[92mTotal Pipeline Time: {total_time:.2f}s\033[0m")
+
+        print("=" * 80 + "\n")
+
+    def get_results(self) -> dict[str, Any]:
+        """获取评估结果（用于程序化访问）"""
+        results: dict[str, Any] = {
+            "overall_score": (sum(self._scores) / len(self._scores) * 100 if self._scores else 0),
+            "sample_count": len(self._scores),
+            "per_dataset": {},
+            "timing": {},
+        }
+
+        for dataset, scores in self._dataset_scores.items():
+            results["per_dataset"][dataset] = {
+                "score": sum(scores) / len(scores) * 100 if scores else 0,
+                "count": len(scores),
+                "metric": DATASET_TO_METRIC.get(dataset, "qa_f1"),
+            }
+
+        if self.longbench_e_buckets:
+            results["buckets"] = {}
+            for bucket, scores in self._bucket_scores.items():
+                results["buckets"][bucket] = {
+                    "score": sum(scores) / len(scores) * 100 if scores else 0,
+                    "count": len(scores),
+                }
+
+        # 时间统计
+        if self._retrieve_times:
+            results["timing"]["retrieve"] = {
+                "avg": sum(self._retrieve_times) / len(self._retrieve_times),
+                "total": sum(self._retrieve_times),
+                "count": len(self._retrieve_times),
+            }
+        if self._refine_times:
+            results["timing"]["refine"] = {
+                "avg": sum(self._refine_times) / len(self._refine_times),
+                "total": sum(self._refine_times),
+                "count": len(self._refine_times),
+            }
+        if self._generate_times:
+            results["timing"]["generate"] = {
+                "avg": sum(self._generate_times) / len(self._generate_times),
+                "total": sum(self._generate_times),
+                "count": len(self._generate_times),
+            }
+
+        return results
+
+    def __del__(self) -> None:
+        """对象销毁时关闭文件"""
+        try:
+            if self._output_file:
+                self._output_file.close()
+        except Exception:
+            pass

sage/benchmark_refiner/metrics.py

@@ -0,0 +1,216 @@
+"""
+LongBench 官方评估指标函数
+
+直接照搬自: https://github.com/THUDM/LongBench/blob/main/metrics.py
+添加了类型注解和可选依赖处理。
+"""
+
+import re
+import string
+from collections import Counter
+from typing import List
+
+# 可选依赖
+try:
+    import jieba
+
+    HAS_JIEBA = True
+except ImportError:
+    HAS_JIEBA = False
+
+try:
+    from fuzzywuzzy import fuzz
+
+    HAS_FUZZYWUZZY = True
+except ImportError:
+    HAS_FUZZYWUZZY = False
+
+try:
+    from rouge import Rouge
+
+    HAS_ROUGE = True
+except ImportError:
+    HAS_ROUGE = False
+
+
+def normalize_answer(s: str) -> str:
+    """Lower text and remove punctuation, articles and extra whitespace."""
+
+    def remove_articles(text: str) -> str:
+        return re.sub(r"\b(a|an|the)\b", " ", text)
+
+    def white_space_fix(text: str) -> str:
+        return " ".join(text.split())
+
+    def remove_punc(text: str) -> str:
+        exclude = set(string.punctuation)
+        return "".join(ch for ch in text if ch not in exclude)
+
+    def lower(text: str) -> str:
+        return text.lower()
+
+    return white_space_fix(remove_articles(remove_punc(lower(s))))
+
+
+def normalize_zh_answer(s: str) -> str:
+    """Lower text and remove punctuation, extra whitespace (Chinese)."""
+
+    def white_space_fix(text: str) -> str:
+        return "".join(text.split())
+
+    def remove_punc(text: str) -> str:
+        cn_punctuation = (
+            "！？｡。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—''‛"
+            "„‟…‧﹏."
+        )
+        all_punctuation = set(string.punctuation + cn_punctuation)
+        return "".join(ch for ch in text if ch not in all_punctuation)
+
+    def lower(text: str) -> str:
+        return text.lower()
+
+    return white_space_fix(remove_punc(lower(s)))
+
+
+def f1_score(prediction: List[str], ground_truth: List[str], **kwargs) -> float:
+    """Token-level F1 score."""
+    common = Counter(prediction) & Counter(ground_truth)
+    num_same = sum(common.values())
+    if num_same == 0:
+        return 0.0
+    precision = 1.0 * num_same / len(prediction)
+    recall = 1.0 * num_same / len(ground_truth)
+    f1 = (2 * precision * recall) / (precision + recall)
+    return f1
+
+
+def qa_f1_score(prediction: str, ground_truth: str, **kwargs) -> float:
+    """QA F1 score (English)."""
+    normalized_prediction = normalize_answer(prediction)
+    normalized_ground_truth = normalize_answer(ground_truth)
+    prediction_tokens = normalized_prediction.split()
+    ground_truth_tokens = normalized_ground_truth.split()
+    return f1_score(prediction_tokens, ground_truth_tokens)
+
+
+def qa_f1_zh_score(prediction: str, ground_truth: str, **kwargs) -> float:
+    """QA F1 score (Chinese, requires jieba)."""
+    if not HAS_JIEBA:
+        raise ImportError(
+            "jieba is required for Chinese evaluation. Install with: pip install jieba"
+        )
+
+    prediction_tokens = list(jieba.cut(prediction, cut_all=False))
+    ground_truth_tokens = list(jieba.cut(ground_truth, cut_all=False))
+    prediction_tokens = [normalize_zh_answer(token) for token in prediction_tokens]
+    ground_truth_tokens = [normalize_zh_answer(token) for token in ground_truth_tokens]
+    prediction_tokens = [token for token in prediction_tokens if len(token) > 0]
+    ground_truth_tokens = [token for token in ground_truth_tokens if len(token) > 0]
+    return f1_score(prediction_tokens, ground_truth_tokens)
+
+
+def rouge_score(prediction: str, ground_truth: str, **kwargs) -> float:
+    """ROUGE-L F1 score."""
+    if not HAS_ROUGE:
+        raise ImportError("rouge is required. Install with: pip install rouge")
+
+    rouge = Rouge()
+    try:
+        scores = rouge.get_scores([prediction], [ground_truth], avg=True)
+        return scores["rouge-l"]["f"]
+    except Exception:
+        return 0.0
+
+
+def rouge_zh_score(prediction: str, ground_truth: str, **kwargs) -> float:
+    """ROUGE-L F1 score (Chinese, requires jieba)."""
+    if not HAS_JIEBA:
+        raise ImportError(
+            "jieba is required for Chinese evaluation. Install with: pip install jieba"
+        )
+
+    prediction = " ".join(list(jieba.cut(prediction, cut_all=False)))
+    ground_truth = " ".join(list(jieba.cut(ground_truth, cut_all=False)))
+    return rouge_score(prediction, ground_truth)
+
+
+def classification_score(prediction: str, ground_truth: str, **kwargs) -> float:
+    """Classification score with all_classes matching."""
+    all_classes = kwargs.get("all_classes", [])
+    if not all_classes:
+        return 0.0
+
+    em_match_list = []
+    for class_name in all_classes:
+        if class_name in prediction:
+            em_match_list.append(class_name)
+
+    for match_term in em_match_list.copy():
+        if match_term in ground_truth and match_term != ground_truth:
+            em_match_list.remove(match_term)
+
+    if ground_truth in em_match_list:
+        return 1.0 / len(em_match_list)
+    return 0.0
+
+
+def retrieval_score(prediction: str, ground_truth: str, **kwargs) -> float:
+    """Retrieval score (English)."""
+    pattern = r"Paragraph (\d+)"
+    matches = re.findall(pattern, ground_truth)
+    if not matches:
+        return 0.0
+    ground_truth_id = matches[0]
+    numbers = re.findall(r"\d+", prediction)
+    right_num = sum(1 for number in numbers if str(number) == str(ground_truth_id))
+    return 0.0 if len(numbers) == 0 else float(right_num / len(numbers))
+
+
+def retrieval_zh_score(prediction: str, ground_truth: str, **kwargs) -> float:
+    """Retrieval score (Chinese)."""
+    pattern = r"段落(\d+)"
+    matches = re.findall(pattern, ground_truth)
+    if not matches:
+        return 0.0
+    ground_truth_id = matches[0]
+    numbers = re.findall(r"\d+", prediction)
+    right_num = sum(1 for number in numbers if str(number) == str(ground_truth_id))
+    return 0.0 if len(numbers) == 0 else float(right_num / len(numbers))
+
+
+def count_score(prediction: str, ground_truth: str, **kwargs) -> float:
+    """Count score for passage_count task."""
+    numbers = re.findall(r"\d+", prediction)
+    right_num = sum(1 for number in numbers if str(number) == str(ground_truth))
+    return 0.0 if len(numbers) == 0 else float(right_num / len(numbers))
+
+
+def code_sim_score(prediction: str, ground_truth: str, **kwargs) -> float:
+    """Code similarity score (requires fuzzywuzzy)."""
+    if not HAS_FUZZYWUZZY:
+        raise ImportError(
+            "fuzzywuzzy is required for code evaluation. "
+            "Install with: pip install fuzzywuzzy python-Levenshtein"
+        )
+
+    all_lines = prediction.lstrip("\n").split("\n")
+    processed_prediction = ""
+    for line in all_lines:
+        if ("`" not in line) and ("#" not in line) and ("//" not in line):
+            processed_prediction = line
+            break
+    return fuzz.ratio(processed_prediction, ground_truth) / 100.0
+
+
+# 指标函数注册表（照搬官方 dataset2metric 映射）
+METRIC_FUNCTIONS = {
+    "qa_f1": qa_f1_score,
+    "qa_f1_zh": qa_f1_zh_score,
+    "rouge": rouge_score,
+    "rouge_zh": rouge_zh_score,
+    "classification": classification_score,
+    "retrieval": retrieval_score,
+    "retrieval_zh": retrieval_zh_score,
+    "count": count_score,
+    "code_sim": code_sim_score,
+}

sage/benchmark_refiner/promptor.py

@@ -0,0 +1,210 @@
+"""
+LongBench Promptor - 优化版
+
+改进点：
+1. 从 JSON 文件加载配置（而非硬编码）
+2. 支持模型专用 chat template（照搬官方 build_chat）
+3. 中间截断策略（照搬官方实现）
+4. 数据集验证
+"""
+
+import json
+from pathlib import Path
+from typing import Any, Optional
+
+from sage.common.core.functions import MapFunction as MapOperator
+
+from .constants import NO_CHAT_DATASETS, SUPPORTED_DATASETS
+from .utils import build_chat, truncate_middle
+
+
+class LongBenchPromptor(MapOperator):
+    """
+    LongBench 专用 Promptor。
+
+    功能：
+    1. 从 JSON 文件加载任务专用 prompt 模板
+    2. 使用 context 和 input (query) 填充模板
+    3. Token 级中间截断（当超过 max_input_tokens 时，保留首尾）
+    4. 按数据集决定是否使用 chat template（few-shot 和代码任务不使用）
+
+    配置参数：
+        - max_input_tokens: int | None - 最大输入 token 数，超过则中间截断
+        - is_chat_model: bool - 是否使用 chat template（默认 False）
+        - model_name_or_path: str | None - 模型路径，用于加载 tokenizer
+    """
+
+    # 配置文件缓存（类级别，避免重复加载）
+    _prompt_templates: Optional[dict[str, str]] = None
+    _max_gen_lengths: Optional[dict[str, int]] = None
+    _model_max_lengths: Optional[dict[str, int]] = None
+
+    @classmethod
+    def _load_configs(cls) -> None:
+        """延迟加载 JSON 配置文件"""
+        if cls._prompt_templates is None:
+            config_dir = Path(__file__).parent / "config"
+            with open(config_dir / "dataset2prompt.json", encoding="utf-8") as f:
+                cls._prompt_templates = json.load(f)
+            with open(config_dir / "dataset2maxlen.json", encoding="utf-8") as f:
+                cls._max_gen_lengths = json.load(f)
+            with open(config_dir / "model2maxlen.json", encoding="utf-8") as f:
+                cls._model_max_lengths = json.load(f)
+
+    def __init__(self, config: dict[str, Any], **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._load_configs()
+
+        self.config = config
+        self.max_input_tokens: Optional[int] = config.get("max_input_tokens")
+        self.is_chat_model: bool = config.get("is_chat_model", False)
+        self.model_name: str = config.get("model_name_or_path", "")
+
+        # 延迟加载 tokenizer
+        self._tokenizer: Optional[Any] = None
+
+    @property
+    def tokenizer(self) -> Optional[Any]:
+        """延迟加载 tokenizer"""
+        if self._tokenizer is None and self.model_name:
+            try:
+                from transformers import AutoTokenizer
+
+                self._tokenizer = AutoTokenizer.from_pretrained(
+                    self.model_name, trust_remote_code=True
+                )
+                self.logger.info(f"Loaded tokenizer from {self.model_name}")
+            except Exception as e:
+                self.logger.warning(f"Failed to load tokenizer: {e}")
+        return self._tokenizer
+
+    def execute(self, data: dict[str, Any]) -> list[Any]:
+        """
+        生成 LongBench 风格的 prompt。
+
+        输入格式（来自 LongBenchBatch 或经过 Refiner）：
+        {
+            "query": str,           # 用户问题（原 input 字段）
+            "context": str,         # 长文本上下文（原始）
+            "refining_results": List[str],  # 压缩后的上下文（可选，Refiner 输出）
+            "references": list,     # 标准答案
+            "_dataset": str,        # 数据集名称
+            ...
+        }
+
+        **上下文选择优先级**：
+        1. refining_results（如果存在且非空，来自 Refiner 压缩后的结果）
+        2. context（原始上下文，LongBench 自带）
+
+        输出格式：
+        [original_data, prompt_string]
+
+        Args:
+            data: 包含 query, context, _dataset 等字段的字典
+
+        Returns:
+            [原始数据, prompt 字符串] 列表
+        """
+        dataset = data.get("_dataset", "")
+        query = data.get("query", "")
+
+        # 数据集验证
+        if dataset and dataset not in SUPPORTED_DATASETS:
+            self.logger.warning(
+                f"Unknown dataset '{dataset}', using default template. "
+                f"Supported: {SUPPORTED_DATASETS}"
+            )
+
+        # 上下文选择优先级：refining_results > context
+        refining_results = data.get("refining_results", [])
+        if refining_results:
+            if isinstance(refining_results, list):
+                context = "\n\n".join(refining_results)
+            else:
+                context = str(refining_results)
+            self.logger.info("Using refining_results (compressed context)")
+        else:
+            context = data.get("context", "")
+
+        # 1. 获取数据集专用模板（从 JSON 加载）
+        assert self._prompt_templates is not None
+        template = self._prompt_templates.get(dataset, "{context}\n\nQuestion: {input}\nAnswer:")
+
+        # 2. 填充模板（LongBench 使用 {context} 和 {input} 占位符）
+        prompt = template.format(context=context, input=query)
+
+        # 3. 中间截断（如果配置了 max_input_tokens）
+        if self.max_input_tokens and self.tokenizer:
+            prompt = truncate_middle(prompt, self.tokenizer, self.max_input_tokens)
+
+        # 4. Chat template（按数据集决定）
+        # 原始 pred.py: if dataset not in ["trec", "triviaqa", "samsum", "lsht", "lcc", "repobench-p"]:
+        #                   prompt = build_chat(tokenizer, prompt, model_name)
+        if self.is_chat_model and dataset not in NO_CHAT_DATASETS:
+            prompt = build_chat(prompt, self.model_name, self.tokenizer)
+
+        # 5. 设置 max_gen_tokens 供 Generator 使用
+        assert self._max_gen_lengths is not None
+        data["_max_gen_tokens"] = self._max_gen_lengths.get(dataset, 128)
+
+        self.logger.info(f"dataset={dataset}, prompt_length={len(prompt)}")
+        return [data, prompt]
+
+    @classmethod
+    def get_max_gen_length(cls, dataset: str) -> int:
+        """
+        获取数据集的最大生成长度。
+
+        Args:
+            dataset: 数据集名称
+
+        Returns:
+            最大生成 token 数
+        """
+        cls._load_configs()
+        assert cls._max_gen_lengths is not None
+        return cls._max_gen_lengths.get(dataset, 128)
+
+    @classmethod
+    def get_prompt_template(cls, dataset: str) -> str:
+        """
+        获取数据集的 prompt 模板。
+
+        Args:
+            dataset: 数据集名称
+
+        Returns:
+            prompt 模板字符串
+        """
+        cls._load_configs()
+        assert cls._prompt_templates is not None
+        return cls._prompt_templates.get(dataset, "{context}\n\nQuestion: {input}\nAnswer:")
+
+    @classmethod
+    def get_model_max_length(cls, model_name: str, default: int = 8192) -> int:
+        """
+        获取模型的最大上下文长度。
+
+        支持精确匹配和模糊匹配（模型名称包含关系）。
+
+        Args:
+            model_name: 模型名称或路径
+            default: 默认最大长度（如果模型未在映射中）
+
+        Returns:
+            模型最大上下文 token 数
+        """
+        cls._load_configs()
+        assert cls._model_max_lengths is not None
+
+        # 精确匹配
+        if model_name in cls._model_max_lengths:
+            return cls._model_max_lengths[model_name]
+
+        # 模糊匹配
+        model_lower = model_name.lower()
+        for known_model, max_len in cls._model_max_lengths.items():
+            if known_model.lower() in model_lower:
+                return max_len
+
+        return default

sage/benchmark_refiner/utils.py

@@ -0,0 +1,106 @@
+"""
+LongBench 工具函数
+
+包含模型专用的 chat template 和后处理逻辑（照搬官方 pred.py）
+"""
+
+from typing import Any, Optional
+
+
+def build_chat(prompt: str, model_name: str, tokenizer: Optional[Any] = None) -> str:
+    """
+    构建模型专用的 chat prompt。
+
+    照搬自官方 pred.py 的 build_chat 函数。
+
+    Args:
+        prompt: 原始 prompt
+        model_name: 模型名称
+        tokenizer: tokenizer 实例（用于 chatglm3）
+
+    Returns:
+        包装后的 prompt
+    """
+    model_name_lower = model_name.lower()
+
+    if "chatglm3" in model_name_lower and tokenizer:
+        # chatglm3 使用 tokenizer 的专用方法
+        return tokenizer.build_chat_input(prompt)
+    elif "chatglm" in model_name_lower and tokenizer:
+        return tokenizer.build_prompt(prompt)
+    elif "llama2" in model_name_lower:
+        return f"[INST]{prompt}[/INST]"
+    elif "xgen" in model_name_lower:
+        header = (
+            "A chat between a curious human and an artificial intelligence assistant. "
+            "The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n"
+        )
+        return header + f" ### Human: {prompt}\n###"
+    elif "internlm" in model_name_lower:
+        return f"<|User|>:{prompt}<eoh>\n<|Bot|>:"
+
+    # 默认：尝试使用 tokenizer 的 apply_chat_template
+    if tokenizer and hasattr(tokenizer, "apply_chat_template"):
+        try:
+            messages = [{"role": "user", "content": prompt}]
+            return tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True
+            )
+        except Exception:
+            pass
+
+    return prompt
+
+
+def post_process(response: str, model_name: str) -> str:
+    """
+    模型专用输出后处理。
+
+    照搬自官方 pred.py 的 post_process 函数。
+
+    Args:
+        response: 模型原始输出
+        model_name: 模型名称
+
+    Returns:
+        处理后的输出
+    """
+    model_name_lower = model_name.lower()
+
+    if "xgen" in model_name_lower:
+        response = response.strip().replace("Assistant:", "")
+    elif "internlm" in model_name_lower:
+        response = response.split("<eoa>")[0]
+
+    return response
+
+
+def truncate_middle(
+    prompt: str,
+    tokenizer: Any,
+    max_length: int,
+) -> str:
+    """
+    中间截断策略（保留首尾）。
+
+    照搬自官方 pred.py：
+    # truncate to fit max_length (we suggest truncate in the middle,
+    # since the left and right side may contain crucial instructions)
+
+    Args:
+        prompt: 原始 prompt
+        tokenizer: tokenizer 实例
+        max_length: 最大 token 数
+
+    Returns:
+        截断后的 prompt
+    """
+    tokenized = tokenizer(prompt, truncation=False, return_tensors="pt").input_ids[0]
+
+    if len(tokenized) > max_length:
+        half = max_length // 2
+        prompt = tokenizer.decode(tokenized[:half], skip_special_tokens=True) + tokenizer.decode(
+            tokenized[-half:], skip_special_tokens=True
+        )
+
+    return prompt