ragbits-evaluate 0.17.0__tar.gz → 0.18.0__tar.gz

{ragbits_evaluate-0.17.0 → ragbits_evaluate-0.18.0}/CHANGELOG.md

@@ -2,6 +2,24 @@
 
 ## Unreleased
 
+## 0.18.0 (2025-05-22)
+
+### Changed
+
+- ragbits-core updated to version v0.18.0
+
+- Add support for custom column names in evaluation dataset (#566)
+- Add support for reference document ids and page numbers in evaluation dataset (#566)
+- BREAKING CHANGE: Adjust eval pipline interface to batch processing (#555)
+- Rename DocumentMeta create_text_document_from_literal to from_literal (#561)
+- Adjust typing for DocumentSearch (#554)
+
+## 0.17.1 (2025-05-09)
+
+### Changed
+
+- ragbits-core updated to version v0.17.1
+
 ## 0.17.0 (2025-05-06)
 
 ### Changed

{ragbits_evaluate-0.17.0 → ragbits_evaluate-0.18.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ragbits-evaluate
-Version: 0.17.0
+Version: 0.18.0
 Summary: Evaluation module for Ragbits components
 Project-URL: Homepage, https://github.com/deepsense-ai/ragbits
 Project-URL: Bug Reports, https://github.com/deepsense-ai/ragbits/issues
@@ -27,7 +27,7 @@ Requires-Dist: distilabel<2.0.0,>=1.4.1
 Requires-Dist: hydra-core<2.0.0,>=1.3.2
 Requires-Dist: neptune[optuna]<2.0.0,>=1.12.0
 Requires-Dist: optuna<5.0.0,>=4.0.0
-Requires-Dist: ragbits-core==0.17.0
+Requires-Dist: ragbits-core==0.18.0
 Provides-Extra: relari
 Requires-Dist: continuous-eval<1.0.0,>=0.3.12; extra == 'relari'
 Description-Content-Type: text/markdown

{ragbits_evaluate-0.17.0 → ragbits_evaluate-0.18.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ragbits-evaluate"
-version = "0.17.0"
+version = "0.18.0"
 description = "Evaluation module for Ragbits components"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -32,7 +32,7 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
-dependencies = ["hydra-core>=1.3.2,<2.0.0", "neptune[optuna]>=1.12.0,<2.0.0", "optuna>=4.0.0,<5.0.0", "distilabel>=1.4.1,<2.0.0", "datasets>=3.0.1,<4.0.0", "ragbits-core==0.17.0"]
+dependencies = ["hydra-core>=1.3.2,<2.0.0", "neptune[optuna]>=1.12.0,<2.0.0", "optuna>=4.0.0,<5.0.0", "distilabel>=1.4.1,<2.0.0", "datasets>=3.0.1,<4.0.0", "ragbits-core==0.18.0"]
 
 [project.urls]
 "Homepage" = "https://github.com/deepsense-ai/ragbits"

{ragbits_evaluate-0.17.0 → ragbits_evaluate-0.18.0}/src/ragbits/evaluate/cli.py

@@ -140,10 +140,10 @@ def run() -> None:
         metric_results = await evaluator.compute(
             pipeline=state.pipeline,
             dataloader=state.dataloader,
-            metrics=state.metrics,
+            metricset=state.metrics,
         )
         evaluation_results = EvaluationResult(
-            metrics={"metrics": metric_results["metrics"], "time_perf": metric_results["time_perf"]}
+            metrics={"metrics": metric_results.metrics, "time_perf": metric_results.time_perf}
         )
         print_output(evaluation_results)
 

ragbits_evaluate-0.18.0/src/ragbits/evaluate/dataloaders/document_search.py

@@ -0,0 +1,73 @@
+from collections.abc import Iterable
+
+from datasets import load_dataset
+
+from ragbits.core.sources.base import Source
+from ragbits.evaluate.dataloaders.base import DataLoader
+from ragbits.evaluate.dataloaders.exceptions import DataLoaderIncorrectFormatDataError
+from ragbits.evaluate.pipelines.document_search import DocumentSearchData
+
+
+class DocumentSearchDataLoader(DataLoader[DocumentSearchData]):
+    """
+    Document search evaluation data loader.
+
+    The source used for this data loader should point to a file that can be loaded by [Hugging Face](https://huggingface.co/docs/datasets/loading#local-and-remote-files)
+    and contain the following features: "question, "passages".
+    """
+
+    def __init__(
+        self,
+        source: Source,
+        question_key: str = "question",
+        document_ids_key: str = "document_ids",
+        passages_key: str = "passages",
+        page_numbers_key: str = "page_numbers",
+    ) -> None:
+        """
+        Initialize the document search data loader.
+
+        Args:
+            source: The source to load the data from.
+            question_key: The dataset column name that contains the question.
+            document_ids_key: The dataset column name that contains the document ids. Document ids are optional.
+            passages_key: The dataset column name that contains the passages. Passages are optional.
+            page_numbers_key: The dataset column name that contains the page numbers. Page numbers are optional.
+        """
+        super().__init__(source)
+        self.question_key = question_key
+        self.document_ids_key = document_ids_key
+        self.passages_key = passages_key
+        self.page_numbers_key = page_numbers_key
+
+    async def load(self) -> Iterable[DocumentSearchData]:
+        """
+        Load the data from source and format them.
+
+        Returns:
+            The document search evaluation data.
+
+        Raises:
+            DataLoaderIncorrectFormatDataError: If evaluation dataset is incorrectly formatted.
+        """
+        data_path = await self.source.fetch()
+        dataset = load_dataset(
+            path=str(data_path.parent),
+            split="train",
+            data_files={"train": str(data_path.name)},
+        )
+        if self.question_key not in dataset.features:
+            raise DataLoaderIncorrectFormatDataError(
+                required_features=[self.question_key],
+                data_path=data_path,
+            )
+
+        return [
+            DocumentSearchData(
+                question=data.get(self.question_key),
+                reference_document_ids=data.get(self.document_ids_key),
+                reference_passages=data.get(self.passages_key),
+                reference_page_numbers=data.get(self.page_numbers_key),
+            )
+            for data in dataset
+        ]

ragbits_evaluate-0.18.0/src/ragbits/evaluate/evaluator.py

@@ -0,0 +1,224 @@
+import asyncio
+import random
+import time
+from collections.abc import Awaitable, Callable, Iterable
+from dataclasses import dataclass
+from typing import Generic, ParamSpec, TypeVar
+
+from pydantic import BaseModel
+from tqdm import tqdm
+
+from ragbits.core.utils.config_handling import ObjectConstructionConfig, WithConstructionConfig
+from ragbits.core.utils.helpers import batched
+from ragbits.evaluate.dataloaders.base import DataLoader
+from ragbits.evaluate.metrics.base import MetricSet
+from ragbits.evaluate.pipelines.base import EvaluationDataT, EvaluationPipeline, EvaluationResultT, EvaluationTargetT
+
+_CallP = ParamSpec("_CallP")
+_CallReturnT = TypeVar("_CallReturnT")
+
+
+@dataclass
+class EvaluationTimePerf:
+    """
+    Container for evaluation time performance metrics.
+    """
+
+    total_time_in_seconds: float
+    samples_per_second: float
+    latency_in_seconds: float
+
+
+@dataclass
+class EvaluatorResult(Generic[EvaluationResultT]):
+    """
+    Container for evaluation results.
+    """
+
+    metrics: dict[str, int | float]
+    results: list[EvaluationResultT]
+    errors: list[Exception]
+    time_perf: EvaluationTimePerf
+
+
+class EvaluationConfig(BaseModel):
+    """
+    Schema for the evaluation run config.
+    """
+
+    pipeline: ObjectConstructionConfig
+    dataloader: ObjectConstructionConfig
+    metrics: dict[str, ObjectConstructionConfig]
+
+
+class EvaluatorConfig(BaseModel):
+    """
+    Schema for the evaluator config.
+    """
+
+    evaluation: EvaluationConfig
+    evaluator: dict | None = None
+
+
+class Evaluator(WithConstructionConfig):
+    """
+    Evaluator class.
+    """
+
+    def __init__(
+        self,
+        batch_size: int = 10,
+        num_retries: int = 3,
+        backoff_multiplier: int = 1,
+        backoff_max: int = 60,
+    ) -> None:
+        """
+        Initialize the Evaluator instance.
+
+        Args:
+            batch_size: batch size for the evaluation pipeline inference.
+            num_retries: The number of retries per evaluation pipeline inference error.
+            backoff_multiplier: The base delay multiplier for exponential backoff (in seconds).
+            backoff_max: The maximum allowed delay (in seconds) between retries.
+        """
+        self.batch_size = batch_size
+        self.num_retries = num_retries
+        self.backoff_multiplier = backoff_multiplier
+        self.backoff_max = backoff_max
+
+    @classmethod
+    async def run_from_config(cls, config: dict) -> EvaluatorResult:
+        """
+        Run the evaluation based on configuration.
+
+        Args:
+            config: Evaluation config.
+
+        Returns:
+            The evaluation results.
+        """
+        evaluator_config = EvaluatorConfig.model_validate(config)
+        evaluation_config = EvaluationConfig.model_validate(evaluator_config.evaluation)
+        pipeline: EvaluationPipeline = EvaluationPipeline.subclass_from_config(evaluation_config.pipeline)
+        dataloader: DataLoader = DataLoader.subclass_from_config(evaluation_config.dataloader)
+        metricset: MetricSet = MetricSet.from_config(evaluation_config.metrics)
+
+        evaluator = cls.from_config(evaluator_config.evaluator or {})
+        return await evaluator.compute(
+            pipeline=pipeline,
+            dataloader=dataloader,
+            metricset=metricset,
+        )
+
+    async def compute(
+        self,
+        pipeline: EvaluationPipeline[EvaluationTargetT, EvaluationDataT, EvaluationResultT],
+        dataloader: DataLoader[EvaluationDataT],
+        metricset: MetricSet[EvaluationResultT],
+    ) -> EvaluatorResult[EvaluationResultT]:
+        """
+        Compute the evaluation results for the given pipeline and data.
+
+        Args:
+            pipeline: The pipeline to be evaluated.
+            dataloader: The dataloader to load the data.
+            metricset: The metrics to be computed.
+
+        Returns:
+            The evaluation results.
+        """
+        await pipeline.prepare()
+
+        dataset = await dataloader.load()
+        results, errors, time_perf = await self._call_pipeline(pipeline, dataset)
+        metrics = await metricset.compute(results)
+
+        return EvaluatorResult(
+            metrics=metrics,
+            results=results,
+            errors=errors,
+            time_perf=time_perf,
+        )
+
+    async def _call_pipeline(
+        self,
+        pipeline: EvaluationPipeline[EvaluationTargetT, EvaluationDataT, EvaluationResultT],
+        dataset: Iterable[EvaluationDataT],
+    ) -> tuple[list[EvaluationResultT], list[Exception], EvaluationTimePerf]:
+        """
+        Call the pipeline with the given data.
+
+        Args:
+            pipeline: The pipeline to be called.
+            dataset: The dataset to be processed.
+
+        Returns:
+            The evaluation results and performance metrics.
+        """
+        start_time = time.perf_counter()
+        outputs = [
+            await self._call_with_error_handling(pipeline, data)
+            for data in tqdm(batched(dataset, self.batch_size), desc="Evaluation")
+        ]
+        end_time = time.perf_counter()
+
+        errors = [output for output in outputs if isinstance(output, Exception)]
+        results = [item for output in outputs if not isinstance(output, Exception) for item in output]
+
+        return results, errors, self._compute_time_perf(start_time, end_time, len(outputs))
+
+    async def _call_with_error_handling(
+        self,
+        executable: Callable[_CallP, Awaitable[_CallReturnT]],
+        *executable_args: _CallP.args,
+        **executable_kwargs: _CallP.kwargs,
+    ) -> _CallReturnT | Exception:
+        """
+        Call executable with a standarized error handling.
+        If an error occurs, the executable is retried `num_retries` times using randomized exponential backoff.
+
+        Args:
+            executable: The callable function to execute.
+            executable_args: Positional arguments to pass to the executable.
+            executable_kwargs: Keyword arguments to pass to the executable.
+
+        Returns:
+            The result of the executable if successful.
+
+        Raises:
+            Exception: The last encountered exception after all retries are exhausted.
+        """
+        for i in range(max(0, self.num_retries) + 1):
+            try:
+                return await executable(*executable_args, **executable_kwargs)
+            except Exception as exc:
+                if i == self.num_retries:
+                    return exc
+
+                delay = random.uniform(0, min(2**i * self.backoff_multiplier, self.backoff_max))  # noqa: S311
+                await asyncio.sleep(delay)
+
+        raise RuntimeError("Unreachable code reached")  # mypy quirk
+
+    @staticmethod
+    def _compute_time_perf(start_time: float, end_time: float, num_samples: int) -> EvaluationTimePerf:
+        """
+        Compute the performance metrics.
+
+        Args:
+            start_time: The start time.
+            end_time: The end time.
+            num_samples: The number of samples.
+
+        Returns:
+            The performance metrics.
+        """
+        latency = end_time - start_time
+        throughput = num_samples / latency
+        latency_sample = 1.0 / throughput if throughput > 0 else 0.0
+
+        return EvaluationTimePerf(
+            total_time_in_seconds=latency,
+            samples_per_second=throughput,
+            latency_in_seconds=latency_sample,
+        )

{ragbits_evaluate-0.17.0 → ragbits_evaluate-0.18.0}/src/ragbits/evaluate/factories/__init__.py

@@ -1,43 +1,21 @@
 import asyncio
 
+from continuous_eval.metrics.retrieval.matching_strategy import RougeChunkMatch
 from datasets import load_dataset
 
 from ragbits.core.embeddings.dense import LiteLLMEmbedder
 from ragbits.core.sources.hf import HuggingFaceSource
-from ragbits.core.utils.config_handling import ObjectConstructionConfig
 from ragbits.core.vector_stores.in_memory import InMemoryVectorStore
 from ragbits.document_search import DocumentSearch
 from ragbits.document_search.documents.document import DocumentMeta
 from ragbits.evaluate.dataloaders.document_search import DocumentSearchDataLoader
 from ragbits.evaluate.metrics import MetricSet
-
-DS_PRECISION_RECALL_F1 = {
-    "precision_recall_f1": ObjectConstructionConfig.model_validate(
-        {
-            "type": "ragbits.evaluate.metrics.document_search:DocumentSearchPrecisionRecallF1",
-            "config": {
-                "matching_strategy": {
-                    "type": "RougeChunkMatch",
-                    "config": {
-                        "threshold": 0.5,
-                    },
-                },
-            },
-        }
-    ),
-}
-
-
-def precision_recall_f1() -> MetricSet:
-    """
-    Factory of precision recall f1 metric set for retrival evaluation.
-    """
-    return MetricSet.from_config(config=DS_PRECISION_RECALL_F1)
+from ragbits.evaluate.metrics.document_search import DocumentSearchPrecisionRecallF1
 
 
 async def _add_example_documents(document_search: DocumentSearch) -> None:
     dataset = load_dataset(path="deepsense-ai/synthetic-rag-dataset_v1.0", split="train")
-    documents = [DocumentMeta.create_text_document_from_literal(doc) for chunks in dataset["chunks"] for doc in chunks]
+    documents = [DocumentMeta.from_literal(doc) for chunks in dataset["chunks"] for doc in chunks]
     await document_search.ingest(documents)
 
 
@@ -45,7 +23,7 @@ def basic_document_search_factory() -> DocumentSearch:
     """
     Factory for basic example document search instance.
     """
-    document_search = DocumentSearch(vector_store=InMemoryVectorStore(embedder=LiteLLMEmbedder()))
+    document_search: DocumentSearch = DocumentSearch(vector_store=InMemoryVectorStore(embedder=LiteLLMEmbedder()))
     asyncio.run(_add_example_documents(document_search))
     return document_search
 
@@ -55,3 +33,10 @@ def synthetic_rag_dataset() -> DocumentSearchDataLoader:
     Factory for synthetic RAG dataset.
     """
     return DocumentSearchDataLoader(source=HuggingFaceSource(path="deepsense-ai/synthetic-rag-dataset_v1.0"))
+
+
+def precision_recall_f1() -> MetricSet:
+    """
+    Factory of precision recall f1 metric set for retrival evaluation.
+    """
+    return MetricSet(DocumentSearchPrecisionRecallF1(matching_strategy=RougeChunkMatch()))

{ragbits_evaluate-0.17.0 → ragbits_evaluate-0.18.0}/src/ragbits/evaluate/metrics/base.py

@@ -1,3 +1,4 @@
+import asyncio
 from abc import ABC, abstractmethod
 from types import ModuleType
 from typing import ClassVar, Generic
@@ -19,7 +20,7 @@ class Metric(WithConstructionConfig, Generic[EvaluationResultT], ABC):
 
     def __init__(self, weight: float = 1.0) -> None:
         """
-        Initializes the metric.
+        Initialize the metric.
 
         Args:
             weight: Metric value weight in the final score, used during optimization.
@@ -28,7 +29,7 @@ class Metric(WithConstructionConfig, Generic[EvaluationResultT], ABC):
         self.weight = weight
 
     @abstractmethod
-    def compute(self, results: list[EvaluationResultT]) -> dict:
+    async def compute(self, results: list[EvaluationResultT]) -> dict:
         """
         Compute the metric.
 
@@ -70,7 +71,7 @@ class MetricSet(WithConstructionConfig, Generic[EvaluationResultT]):
         """
         return cls(*[Metric.subclass_from_config(metric_config) for metric_config in config.values()])
 
-    def compute(self, results: list[EvaluationResultT]) -> dict:
+    async def compute(self, results: list[EvaluationResultT]) -> dict:
         """
         Compute the metrics.
 
@@ -80,6 +81,9 @@ class MetricSet(WithConstructionConfig, Generic[EvaluationResultT]):
         Returns:
             The computed metrics.
         """
+        metric_results = await asyncio.gather(*[metric.compute(results) for metric in self.metrics])
         return {
-            name: metric.weight * value for metric in self.metrics for name, value in metric.compute(results).items()
+            name: metric.weight * value
+            for metric, result in zip(self.metrics, metric_results, strict=False)
+            for name, value in result.items()
         }

{ragbits_evaluate-0.17.0 → ragbits_evaluate-0.18.0}/src/ragbits/evaluate/metrics/document_search.py

@@ -46,7 +46,7 @@ class DocumentSearchMetric(Metric[DocumentSearchResult], ABC):
         matching_strategy = matching_strategy_cls(**config["matching_strategy"]["config"])
         return cls(matching_strategy=matching_strategy, weight=config.get("weight", 1.0))
 
-    def compute(self, results: list[DocumentSearchResult]) -> dict:
+    async def compute(self, results: list[DocumentSearchResult]) -> dict:
         """
         Compute the metric.
 
@@ -57,7 +57,18 @@ class DocumentSearchMetric(Metric[DocumentSearchResult], ABC):
             The computed metric.
         """
         return self.metric.aggregate(
-            [self.metric(result.predicted_passages, result.reference_passages) for result in results]
+            [
+                self.metric(
+                    [
+                        element.text_representation
+                        for element in result.predicted_elements
+                        if element.text_representation
+                    ],
+                    result.reference_passages,
+                )
+                for result in results
+                if result.reference_passages is not None
+            ]
         )
 
 

{ragbits_evaluate-0.17.0 → ragbits_evaluate-0.18.0}/src/ragbits/evaluate/optimizer.py

@@ -61,7 +61,7 @@ class Optimizer(WithConstructionConfig):
         evaluator_config = EvaluatorConfig.model_validate(optimizer_config.evaluator)
 
         dataloader: DataLoader = DataLoader.subclass_from_config(evaluator_config.evaluation.dataloader)
-        metrics: MetricSet = MetricSet.from_config(evaluator_config.evaluation.metrics)
+        metricset: MetricSet = MetricSet.from_config(evaluator_config.evaluation.metrics)
 
         pipeline_class = import_by_path(evaluator_config.evaluation.pipeline.type)
         pipeline_config = dict(evaluator_config.evaluation.pipeline.config)
@@ -71,7 +71,7 @@ class Optimizer(WithConstructionConfig):
         return optimizer.optimize(
             pipeline_class=pipeline_class,
             pipeline_config=pipeline_config,
-            metrics=metrics,
+            metricset=metricset,
             dataloader=dataloader,
             callbacks=callbacks,
         )
@@ -81,7 +81,7 @@ class Optimizer(WithConstructionConfig):
         pipeline_class: type[EvaluationPipeline],
         pipeline_config: dict,
         dataloader: DataLoader,
-        metrics: MetricSet,
+        metricset: MetricSet,
         callbacks: list[Callable] | None = None,
     ) -> list[tuple[dict, float, dict[str, float]]]:
         """
@@ -91,7 +91,7 @@ class Optimizer(WithConstructionConfig):
             pipeline_class: Pipeline to be optimized.
             pipeline_config: Configuration defining the optimization process.
             dataloader: Data loader.
-            metrics: Metrics to be optimized.
+            metricset: Metrics to be optimized.
             callbacks: Experiment callbacks.
 
         Returns:
@@ -104,7 +104,7 @@ class Optimizer(WithConstructionConfig):
                 pipeline_class=pipeline_class,
                 pipeline_config=pipeline_config,
                 dataloader=dataloader,
-                metrics=metrics,
+                metricset=metricset,
             )
 
         study = optuna.create_study(direction=self.direction)
@@ -131,7 +131,7 @@ class Optimizer(WithConstructionConfig):
         pipeline_class: type[EvaluationPipeline],
         pipeline_config: dict,
         dataloader: DataLoader,
-        metrics: MetricSet,
+        metricset: MetricSet,
     ) -> float:
         """
         Run a single experiment.
@@ -153,11 +153,11 @@ class Optimizer(WithConstructionConfig):
                     evaluator.compute(
                         pipeline=pipeline,
                         dataloader=dataloader,
-                        metrics=metrics,
+                        metricset=metricset,
                     )
                 )
-                score = sum(results["metrics"].values())
-                metrics_values = results["metrics"]
+                score = sum(results.metrics.values())
+                metrics_values = results.metrics
                 break
             except Exception as exc:
                 message = (

{ragbits_evaluate-0.17.0 → ragbits_evaluate-0.18.0}/src/ragbits/evaluate/pipelines/base.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from collections.abc import Iterable
 from dataclasses import dataclass
 from types import ModuleType
 from typing import ClassVar, Generic, TypeVar
@@ -51,7 +52,7 @@ class EvaluationPipeline(WithConstructionConfig, Generic[EvaluationTargetT, Eval
         pass
 
     @abstractmethod
-    async def __call__(self, data: EvaluationDataT) -> EvaluationResultT:
+    async def __call__(self, data: Iterable[EvaluationDataT]) -> Iterable[EvaluationResultT]:
         """
         Run the evaluation pipeline.
 

{ragbits_evaluate-0.17.0 → ragbits_evaluate-0.18.0}/src/ragbits/evaluate/pipelines/document_search.py

@@ -1,3 +1,5 @@
+import asyncio
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
 from uuid import uuid4
 
@@ -5,6 +7,7 @@ from typing_extensions import Self
 
 from ragbits.core.sources.hf import HuggingFaceSource
 from ragbits.document_search import DocumentSearch
+from ragbits.document_search.documents.element import Element
 from ragbits.evaluate.pipelines.base import EvaluationData, EvaluationPipeline, EvaluationResult
 
 
@@ -14,7 +17,9 @@ class DocumentSearchData(EvaluationData):
     """
 
     question: str
-    reference_passages: list[str]
+    reference_document_ids: list[str | int] | None = None
+    reference_passages: list[str] | None = None
+    reference_page_numbers: list[int] | None = None
 
 
 @dataclass
@@ -24,8 +29,10 @@ class DocumentSearchResult(EvaluationResult):
     """
 
     question: str
-    reference_passages: list[str]
-    predicted_passages: list[str]
+    predicted_elements: Sequence[Element]
+    reference_document_ids: list[str | int] | None = None
+    reference_passages: list[str] | None = None
+    reference_page_numbers: list[int] | None = None
 
 
 class DocumentSearchPipeline(EvaluationPipeline[DocumentSearch, DocumentSearchData, DocumentSearchResult]):
@@ -60,7 +67,7 @@ class DocumentSearchPipeline(EvaluationPipeline[DocumentSearch, DocumentSearchDa
         # TODO: optimize this for cases with duplicated document search configs between runs
         if config.get("source"):
             config["vector_store"]["config"]["index_name"] = str(uuid4())
-        evaluation_target = DocumentSearch.from_config(config)
+        evaluation_target: DocumentSearch = DocumentSearch.from_config(config)
         return cls(evaluation_target=evaluation_target, source=config.get("source"))
 
     async def prepare(self) -> None:
@@ -76,21 +83,24 @@ class DocumentSearchPipeline(EvaluationPipeline[DocumentSearch, DocumentSearchDa
             )
             await self.evaluation_target.ingest(sources)
 
-    async def __call__(self, data: DocumentSearchData) -> DocumentSearchResult:
+    async def __call__(self, data: Iterable[DocumentSearchData]) -> Iterable[DocumentSearchResult]:
         """
         Run the document search evaluation pipeline.
 
         Args:
-            data: The evaluation data.
+            data: The evaluation data batch.
 
         Returns:
-            The evaluation result.
+            The evaluation result batch.
         """
-        elements = await self.evaluation_target.search(data.question)
-        predicted_passages = [element.text_representation for element in elements if element.text_representation]
-
-        return DocumentSearchResult(
-            question=data.question,
-            reference_passages=data.reference_passages,
-            predicted_passages=predicted_passages,
-        )
+        results = await asyncio.gather(*[self.evaluation_target.search(row.question) for row in data])
+        return [
+            DocumentSearchResult(
+                question=row.question,
+                predicted_elements=elements,
+                reference_document_ids=row.reference_document_ids,
+                reference_passages=row.reference_passages,
+                reference_page_numbers=row.reference_page_numbers,
+            )
+            for row, elements in zip(data, results, strict=False)
+        ]