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

sage/benchmark_rag/config/config_refiner.yaml

@@ -0,0 +1,75 @@
+pipeline:
+  name: "sage-api-operator-operator_test"
+  description: "Test pipeline for Sage API Operator"
+  version: "1.0.0"
+
+source:
+  # 数据源类型：'local'（本地 JSONL） 或 'hf'（HuggingFace Dataset）
+  type: "hf"
+  # 本地 JSONL 文件路径（type=local 时生效）
+  data_path: "packages/sage-benchmark/src/sage/data/qa/sample/evaluate.json"
+
+  # HuggingFace Dataset 参数（type=hf 时生效）
+  hf_dataset_name: "RUC-NLPIR/FlashRAG_datasets"
+  hf_dataset_config: "asqa"
+  hf_split: "dev"
+
+retriever:
+  # 检索器类型选择: 'chroma' 或 'wiki18_faiss'
+  type: "wiki18_faiss"  # 使用Wiki18 FAISS检索器
+
+  # 通用配置
+  dimension: 1024    # BGE-M3模型的维度
+  top_k: 8
+
+  # Wiki18 FAISS 专用配置 (type=wiki18_faiss 时生效)
+  faiss:
+    index_path: ""
+    documents_path: ""
+
+  # 嵌入模型配置
+  embedding:
+    method: "hf"
+    model: "BAAI/bge-m3"  # 使用BGE-M3模型进行Wiki18 FAISS检索
+    gpu_device: 0  # 明确指定BGE-M3使用GPU 0
+
+generator:
+  local:
+    method: "hf"
+    model_name: "meta-llama/Llama-2-13b-chat-hf"
+    seed: 42
+
+  vllm:
+    api_key: "token-abc123"
+    method: "openai"
+    model_name: "meta-llama/Llama-3.1-8B-Instruct"
+    base_url: "http://sage2:8000/v1"
+    seed: 42
+
+  remote:
+    api_key: ""
+    method: "openai"
+    model_name: "qwen-turbo-0919"
+    base_url: "http://127.0.0.1:8889/v1"
+    seed: 42
+
+promptor:
+  platform: "local"
+
+sink:
+  platform: "local"
+
+refiner:
+  base_model_path: "Qwen/Qwen2.5-3B-Instruct"
+  query_analysis_module_lora_path: "jinjiajie/Query-Analysis-Qwen2.5-3B-Instruct"
+  doc_structuring_module_lora_path: "jinjiajie/Doc-Structuring-Qwen2.5-3B-Instruct"
+  global_selection_module_lora_path: "jinjiajie/Global-Selection-Qwen2.5-3B-Instruct"
+  score_model_name: "bge-reranker-v2-m3"
+  score_model_path: "BAAI/bge-reranker-v2-m3"
+  max_model_len: 25000
+  budget: 2048
+  gpu_device: 0  # vLLM使用GPU 1
+  gpu_memory_utilization: 0.5
+
+evaluate:
+  platform: "local"

sage/benchmark_rag/config/config_rerank.yaml

@@ -0,0 +1,56 @@
+# ChromaDB 专用 QA 配置文件
+# 适用于 qa_openai.py 的简化配置
+
+
+source:
+  data_path: "packages/sage-benchmark/src/sage/data/qa/queries.jsonl"  # 相对于 SAGE 根目录
+  platform: "local"
+
+retriever:
+  # ChromaDB 专用配置
+  dimension: 384    # 修改为与 HuggingFace all-MiniLM-L6-v2 匹配
+  top_k: 2
+
+  chroma:
+    persistence_path: "./data/chroma_qa_database"
+    collection_name: "qa_knowledge_base"
+    use_embedding_query: true
+
+    # 知识库已预加载，注释掉自动加载
+    # knowledge_file: "../../data/qa_knowledge_base.txt"
+
+    # ChromaDB 元数据配置 (简化格式)
+    metadata:
+      hnsw:space: "cosine"           # 距离度量
+
+  # 嵌入模型配置
+  embedding:
+    method: "hf"
+    model: "sentence-transformers/all-MiniLM-L6-v2"
+
+promptor:
+  template: |
+    基于以下检索到的相关文档，回答用户问题：
+
+    相关文档：
+    {retrieved_documents}
+
+    用户问题：{query}
+
+    请提供准确、有用的回答：
+
+generator:
+  vllm:
+    api_key: ""
+    method: "openai"
+    model_name: "meta-llama/Llama-2-7b-chat-hf"
+    base_url: "http://sage3:8000/v1"
+    seed: 42
+
+sink:
+  enable_log: true
+
+reranker:
+  platform: "local"
+  model_name: "BAAI/bge-reranker-v2-m3"
+  topk: 1

sage/benchmark_rag/config/config_selfrag.yaml

@@ -0,0 +1,24 @@
+# Self-RAG Configuration
+# Based on Self-RAG paper: https://arxiv.org/abs/2310.11511
+
+# Data path (Self-RAG dataset format with pre-retrieved documents)
+data_path: "/path/to/selfrag_dataset.jsonl"
+
+# Retriever configuration
+retriever:
+  top_k: 5  # Number of retrieved documents to use
+
+# Promptor configuration
+promptor:
+  model_name: "mistralai/Mistral-7B-Instruct-v0.1"
+  use_context: true  # Whether to include retrieved context in prompt
+
+# Generator configuration
+generator:
+  model_name: "mistralai/Mistral-7B-Instruct-v0.1"
+  gpu_memory_utilization: 0.8
+  temperature: 0
+  max_tokens: 100
+
+# Output path
+output_path: "./output/selfrag_results.json"

sage/benchmark_rag/config/config_source.yaml

@@ -0,0 +1,30 @@
+pipeline:
+  name: "qa_source_interactive"
+  description: "Interactive terminal QA pipeline with LLM"
+  version: "1.0.0"
+
+generator:
+  local:
+    method: "hf"
+    model_name: "meta-llama/Llama-2-13b-chat-hf"
+    seed: 42
+
+  vllm:
+    api_key: ""
+    method: "openai"
+    model_name: "qwen-turbo-2025-02-11"
+    base_url: "http://127.0.0.1:8889/v1"
+    seed: 42
+
+  remote:
+    api_key: ""
+    method: "openai"
+    model_name: "qwen-turbo-0919"
+    base_url: "http://127.0.0.1:8889/v1"
+    seed: 42
+
+promptor:
+  platform: "local"
+
+sink:
+  platform: "local"

sage/benchmark_rag/config/config_source_local.yaml

@@ -0,0 +1,21 @@
+pipeline:
+  name: "qa_source_interactive_local"
+  description: "Interactive terminal QA pipeline with local LLM"
+  version: "1.0.0"
+
+generator:
+  local:
+    method: "hf"
+    model_name: "microsoft/DialoGPT-medium"
+    seed: 42
+    max_length: 512
+    temperature: 0.7
+
+promptor:
+  platform: "local"
+  template: |
+    User: {user_query}
+    Assistant:
+
+sink:
+  platform: "local"

sage/benchmark_rag/config/config_sparse_milvus.yaml

@@ -0,0 +1,49 @@
+# Milvus 稠密向量检索配置（适配 MilvusDenseRetriever）
+
+source:
+  data_path: "./packages/sage-benchmark/src/sage/data/qa/queries.jsonl"
+  platform: "local"
+
+retriever:
+  preload_knowledge_file: "./packages/sage-benchmark/src/sage/data/qa/qa_knowledge_base.txt"
+
+  # 通用参数
+  top_k: 3                 # 返回文档数量
+
+  # Milvus 后端（稠密检索）
+  milvus_sparse:
+    # 本地 Milvus Lite（推荐用于快速试用）
+    persistence_path: "./data/milvus_qa_sparse.db"
+
+    # 远程 Milvus（如需远程，请注释上面的 persistence_path，改为如下配置）
+    # host: "127.0.0.1"
+    # port: 19530
+    # force_http: true
+
+    collection_name: "qa_sparse_collection"
+    search_type: "sparse"     # 稀疏检索
+
+    # # 知识文件（可选）：提供后将自动按段落读取并入库
+    # knowledge_file: "./packages/sage-benchmark/src/sage/data/qa/qa_knowledge_base.txt"
+
+promptor:
+  template: |
+    基于以下检索到的相关文档，回答用户问题：
+
+    相关文档：
+    {retrieved_documents}
+
+    用户问题：{query}
+
+    请提供准确、有用的回答：
+
+generator:
+  vllm:
+    api_key: ""
+    method: "openai"
+    model_name: "meta-llama/Llama-2-7b-chat-hf"
+    base_url: "http://sage3:8000/v1"
+    seed: 42
+
+sink:
+  enable_log: true

sage/benchmark_rag/evaluation/__init__.py

@@ -0,0 +1,10 @@
+"""Evaluation framework for RAG benchmarking.
+
+This module provides tools for running RAG experiments and evaluating results:
+- Pipeline experiment framework
+- Multiple evaluation metrics (Accuracy, F1, Exact Match)
+- Batch processing support
+- Result analysis tools
+"""
+
+__all__: list[str] = []

sage/benchmark_rag/evaluation/benchmark_runner.py

@@ -0,0 +1,337 @@
+"""
+Universal Benchmark Runner for RAG Pipelines
+
+This module provides a generic framework for benchmarking any RAG pipeline implementation.
+It handles:
+- Batch processing of large datasets
+- Running any pipeline from implementations/pipelines/
+- Collecting and saving results
+- Performance metrics tracking
+"""
+
+import importlib
+import json
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import yaml  # type: ignore[import-untyped]
+from dotenv import load_dotenv
+from sage.common.config.output_paths import get_output_file
+from sage.common.core import BatchFunction, MapFunction
+from sage.common.utils.logging.custom_logger import CustomLogger
+from sage.kernel.api.local_environment import LocalEnvironment
+
+
+def load_config(path: str) -> dict:
+    """Load YAML configuration file."""
+    with open(path, encoding="utf-8") as f:
+        return yaml.safe_load(f)
+
+
+class BatchDataLoader(BatchFunction):
+    """
+    Batch data loader for benchmark datasets.
+
+    Supports loading large datasets in batches for efficient processing.
+    Compatible with Self-RAG dataset format and other QA datasets.
+    """
+
+    def __init__(self, config: dict):
+        self.config = config
+        data_path = config.get("data_path")
+        if not data_path:
+            raise ValueError("data_path is required in config")
+        self.data_path: str = data_path
+        self.max_samples = config.get("max_samples", None)
+
+        # Load data
+        data = self._load_data()
+        if self.max_samples:
+            data = data[: self.max_samples]
+
+        self.batch_size = config.get("batch_size", len(data))
+        self.current_batch = 0
+        self.total_batches = (len(data) + self.batch_size - 1) // self.batch_size
+        self._data = data
+
+    def _load_data(self) -> list[dict[str, Any]]:
+        """Load dataset from JSONL file."""
+        data = []
+        with open(self.data_path, encoding="utf-8") as f:
+            for line in f:
+                if line.strip():
+                    data.append(json.loads(line))
+
+        # Limit samples if max_samples is set
+        if self.max_samples and self.max_samples > 0:
+            print(
+                f"📊 Limiting dataset to {self.max_samples} samples (total available: {len(data)})"
+            )
+            data = data[: self.max_samples]
+
+        return data
+
+    def execute(self) -> dict[str, Any] | None:
+        """Return next batch of data."""
+        if self.current_batch >= self.total_batches:
+            return None
+
+        start_idx = self.current_batch * self.batch_size
+        end_idx = min(start_idx + self.batch_size, len(self._data))
+        batch = self._data[start_idx:end_idx]
+
+        result = {
+            "batch_data": batch,
+            "batch_id": self.current_batch,
+            "total_batches": self.total_batches,
+        }
+
+        self.current_batch += 1
+        return result
+
+
+class PipelineRunner(MapFunction):
+    """
+    Generic pipeline runner that can execute any RAG pipeline.
+
+    Dynamically loads and runs pipeline implementations from
+    implementations/pipelines/ directory.
+
+    Expected pipeline interface:
+    - process_item(item: Dict, config: Dict) -> Dict
+    """
+
+    def __init__(self, config: dict):
+        self.config = config
+        self.pipeline_name = config.get("pipeline_name")
+        self.pipeline_config = config.get("pipeline_config", {})
+
+        # Dynamically import pipeline module
+        self.pipeline_module = self._load_pipeline()
+
+    def _load_pipeline(self):
+        """Dynamically load pipeline implementation."""
+        module_path = f"sage.benchmark_rag.implementations.pipelines.{self.pipeline_name}"
+
+        try:
+            module = importlib.import_module(module_path)
+
+            # Verify the module has process_item function
+            if not hasattr(module, "process_item"):
+                raise ImportError(
+                    f"Pipeline module {module_path} must have a 'process_item' function "
+                    f"with signature: process_item(item: Dict, config: Dict) -> Dict"
+                )
+
+            return module
+        except ImportError as e:
+            raise ImportError(
+                f"Failed to load pipeline '{self.pipeline_name}': {e}\n"
+                f"Make sure the pipeline exists in implementations/pipelines/ "
+                f"and has a 'process_item' function"
+            )
+
+    def execute(self, batch_data: dict[str, Any]) -> dict[str, Any]:
+        """Execute pipeline on a batch of data."""
+        batch = batch_data["batch_data"]
+
+        # Process each item in the batch
+        results = []
+        for i, item in enumerate(batch):
+            try:
+                # Run the pipeline on this item
+                result = self.pipeline_module.process_item(item, self.pipeline_config)
+                results.append(result)
+
+                # Progress indicator
+                if (i + 1) % 10 == 0:
+                    print(f"  Processed {i + 1}/{len(batch)} items in batch...")
+
+            except Exception as e:
+                # Log error but continue processing
+                print(f"⚠️  Error processing item {item.get('id', 'unknown')}: {e}")
+                results.append(
+                    {
+                        "id": item.get("id", "unknown"),
+                        "question": item.get("question", ""),
+                        "error": str(e),
+                        "ground_truth": item.get("answers", []),
+                    }
+                )
+
+        return {
+            "results": results,
+            "batch_id": batch_data["batch_id"],
+            "total_batches": batch_data["total_batches"],
+        }
+
+
+class ResultsCollector(MapFunction):
+    """
+    Collects and saves benchmark results.
+
+    Handles incremental or final saving of results with metadata.
+    """
+
+    def __init__(self, config: dict):
+        self.config = config
+        default_output = get_output_file("benchmark_results.json", "benchmarks")
+        self.output_path = config.get("output_path", str(default_output))
+        self.save_mode = config.get("save_mode", "incremental")
+        self.all_results: list[dict[str, Any]] = []
+        self.start_time = time.time()
+
+    def execute(self, batch_result: dict[str, Any]) -> dict[str, Any]:
+        """Collect results from a batch."""
+        results = batch_result["results"]
+        batch_id = batch_result["batch_id"]
+        total_batches = batch_result["total_batches"]
+
+        self.all_results.extend(results)
+
+        print(
+            f"✅ Processed batch {batch_id + 1}/{total_batches}, "
+            f"Total results: {len(self.all_results)}"
+        )
+
+        # Save based on mode
+        if self.save_mode == "incremental":
+            self._save_results(batch_id + 1, total_batches)
+        elif self.save_mode == "final" and batch_id + 1 == total_batches:
+            self._save_results(batch_id + 1, total_batches)
+
+        return batch_result
+
+    def _save_results(self, current_batch: int, total_batches: int):
+        """Save results to file with metadata."""
+        elapsed_time = time.time() - self.start_time
+
+        output = {
+            "metadata": {
+                "pipeline_name": self.config.get("pipeline_name", "unknown"),
+                "timestamp": datetime.now().isoformat(),
+                "total_samples": len(self.all_results),
+                "completed_batches": f"{current_batch}/{total_batches}",
+                "elapsed_time_seconds": round(elapsed_time, 2),
+                "config": self.config,
+            },
+            "results": self.all_results,
+        }
+
+        # Ensure output directory exists
+        output_path = Path(self.output_path)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(output_path, "w", encoding="utf-8") as f:
+            json.dump(output, f, indent=2, ensure_ascii=False)
+
+        print(f"💾 Results saved to: {self.output_path}")
+
+
+def run_benchmark(config: dict) -> None:
+    """
+    Run benchmark with the specified configuration.
+
+    Args:
+        config: Configuration dictionary with:
+            - data: Data loading configuration
+            - pipeline: Pipeline configuration
+            - output: Output configuration
+    """
+    print("=" * 60)
+    print("🚀 Starting RAG Benchmark")
+    print("=" * 60)
+
+    pipeline_name = config.get("pipeline", {}).get("pipeline_name", "unknown")
+    data_path = config.get("data", {}).get("data_path", "unknown")
+
+    print(f"📊 Pipeline: {pipeline_name}")
+    print(f"📁 Dataset: {data_path}")
+    print(f"💾 Output: {config.get('output', {}).get('output_path', 'default')}")
+    print("=" * 60)
+
+    env = LocalEnvironment("benchmark_pipeline")
+
+    # Build benchmark pipeline
+    (
+        env.from_source(BatchDataLoader, config["data"])
+        .map(PipelineRunner, config["pipeline"])
+        .sink(ResultsCollector, {**config["output"], **config["pipeline"]})
+    )
+
+    env.submit()
+    env.close()
+
+    print("=" * 60)
+    print("✅ Benchmark completed!")
+    print("=" * 60)
+
+
+def main():
+    """Main entry point for benchmark runner."""
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Run RAG pipeline benchmarks")
+    parser.add_argument(
+        "--config",
+        type=str,
+        default=None,
+        help="Path to benchmark configuration YAML file",
+    )
+    parser.add_argument(
+        "--pipeline",
+        type=str,
+        default=None,
+        help="Pipeline name (e.g., 'selfrag', 'qa_dense_retrieval_milvus')",
+    )
+    parser.add_argument(
+        "--data",
+        type=str,
+        default=None,
+        help="Path to dataset JSONL file",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=None,
+        help="Path to output results file",
+    )
+
+    args = parser.parse_args()
+
+    # Load configuration
+    if args.config:
+        config = load_config(args.config)
+    else:
+        # Use default config path
+        current_dir = Path(__file__).parent
+        default_config = current_dir / "config" / "benchmark_config.yaml"
+
+        if default_config.exists():
+            config = load_config(str(default_config))
+        else:
+            raise ValueError(
+                "No configuration file specified and default config not found. "
+                "Use --config to specify a configuration file."
+            )
+
+    # Override with command-line arguments
+    if args.pipeline:
+        config["pipeline"]["pipeline_name"] = args.pipeline
+    if args.data:
+        config["data"]["data_path"] = args.data
+    if args.output:
+        config["output"]["output_path"] = args.output
+
+    # Setup
+    CustomLogger.disable_global_console_debug()
+    load_dotenv(override=False)
+
+    # Run benchmark
+    run_benchmark(config)
+
+
+if __name__ == "__main__":
+    main()

sage/benchmark_rag/evaluation/config/benchmark_config.yaml

@@ -0,0 +1,35 @@
+# Universal Benchmark Configuration
+# This is a template that can be used for any RAG pipeline benchmark
+
+# Data configuration
+data:
+  data_path: "data/selfrag/selfrag_dev_sample_100.jsonl"  # Path to dataset
+  batch_size: 100  # Number of samples per batch
+  max_samples: 100  # Limit to first N samples (null = use all data)
+                    # Useful for quick testing or limiting dataset size
+
+# Pipeline configuration
+pipeline:
+  pipeline_name: "selfrag"  # Name of pipeline in implementations/pipelines/
+  pipeline_config:  # Pipeline-specific configuration
+    # This section is passed to the pipeline implementation
+    # Structure depends on the specific pipeline
+    top_k: 5
+    model_name: "mistralai/Mistral-7B-Instruct-v0.1"
+    gpu_memory_utilization: 0.8
+    temperature: 0.0
+    max_tokens: 100
+
+# Output configuration
+output:
+  output_path: "artifacts/benchmarks/benchmark_results.json"
+  save_mode: "incremental"  # "incremental" or "final"
+
+# Optional: Evaluation metrics configuration
+# evaluation:
+#   metrics:
+#     - "exact_match"
+#     - "f1_score"
+#     - "bleu"
+#   evaluator_config:
+#     # Metric-specific configuration