ragbits-evaluate 0.5.0__py3-none-any.whl → 1.4.0.dev202602030301__py3-none-any.whl

ragbits/evaluate/optimizer.py

@@ -1,132 +1,184 @@
 import asyncio
 import warnings
+from collections.abc import Callable
 from copy import deepcopy
-from typing import Any
 
 import optuna
-from omegaconf import DictConfig, ListConfig
+from optuna import Trial
+from pydantic import BaseModel
 
-from .callbacks.base import CallbackConfigurator
-from .evaluator import Evaluator
-from .loaders.base import DataLoader
-from .metrics.base import MetricSet
-from .pipelines.base import EvaluationPipeline
+from ragbits.core.utils.config_handling import WithConstructionConfig, import_by_path
+from ragbits.evaluate.dataloaders.base import DataLoader
+from ragbits.evaluate.evaluator import Evaluator, EvaluatorConfig
+from ragbits.evaluate.metrics.base import MetricSet
+from ragbits.evaluate.pipelines.base import EvaluationPipeline
+from ragbits.evaluate.utils import setup_optuna_neptune_callback
 
 
-class Optimizer:
+class OptimizerConfig(BaseModel):
     """
-    Class for optimization
+    Schema for the optimizer config.
     """
 
-    INFINITY = 1e16
+    evaluator: EvaluatorConfig
+    optimizer: dict | None = None
+    neptune_callback: bool = False
 
-    def __init__(self, cfg: DictConfig):
-        self.config = cfg
+
+class Optimizer(WithConstructionConfig):
+    """
+    Optimizer class.
+    """
+
+    def __init__(self, direction: str = "maximize", n_trials: int = 10, max_retries_for_trial: int = 1) -> None:
+        """
+        Initialize the pipeline optimizer.
+
+        Args:
+            direction: Direction of optimization.
+            n_trials: The number of trials for each process.
+            max_retries_for_trial: The number of retires for single process.
+        """
+        self.direction = direction
+        self.n_trials = n_trials
+        self.max_retries_for_trial = max_retries_for_trial
         # workaround for optuna not allowing different choices for different trials
         # TODO check how optuna handles parallelism. discuss if we want to have parallel studies
-        self._choices_cache: dict[str, list[Any]] = {}
+        self._choices_cache: dict[str, list] = {}
+
+    @classmethod
+    def run_from_config(cls, config: dict) -> list[tuple[dict, float, dict[str, float]]]:
+        """
+        Run the optimization process configured with a config object.
+
+        Args:
+            config: Optimizer config.
+
+        Returns:
+            List of tested configs with associated scores and metrics.
+        """
+        optimizer_config = OptimizerConfig.model_validate(config)
+        evaluator_config = EvaluatorConfig.model_validate(optimizer_config.evaluator)
+
+        dataloader: DataLoader = DataLoader.subclass_from_config(evaluator_config.evaluation.dataloader)
+        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)
+        callbacks = [setup_optuna_neptune_callback()] if optimizer_config.neptune_callback else []
+
+        optimizer = cls.from_config(optimizer_config.optimizer or {})
+        return optimizer.optimize(
+            pipeline_class=pipeline_class,
+            pipeline_config=pipeline_config,
+            metricset=metricset,
+            dataloader=dataloader,
+            callbacks=callbacks,
+        )
 
     def optimize(
         self,
         pipeline_class: type[EvaluationPipeline],
-        config_with_params: DictConfig,
+        pipeline_config: dict,
         dataloader: DataLoader,
-        metrics: MetricSet,
-        callback_configurators: list[CallbackConfigurator] | None = None,
-    ) -> list[tuple[DictConfig, float, dict[str, float]]]:
+        metricset: MetricSet,
+        callbacks: list[Callable] | None = None,
+    ) -> list[tuple[dict, float, dict[str, float]]]:
         """
-        A method for running the optimization process for given parameters
+        Run the optimization process for given parameters.
+
         Args:
-            pipeline_class - a type of pipeline to be optimized
-            config_with_params - a configuration defining the optimization process
-            dataloader - a dataloader
-            metrics - object representing the metrics to be optimized
-            log_to_neptune - indicator whether the results should be logged to neptune
+            pipeline_class: Pipeline to be optimized.
+            pipeline_config: Configuration defining the optimization process.
+            dataloader: Data loader.
+            metricset: Metrics to be optimized.
+            callbacks: Experiment callbacks.
+
         Returns:
-            list of tuples with configs and their scores
+            List of tested configs with associated scores and metrics.
         """
-        # TODO check details on how to parametrize optuna
-        optimization_kwargs = {"n_trials": self.config.n_trials}
-        if callback_configurators:
-            optimization_kwargs["callbacks"] = [configurator.get_callback() for configurator in callback_configurators]
 
-        def objective(trial: optuna.Trial) -> float:
+        def objective(trial: Trial) -> float:
             return self._objective(
                 trial=trial,
                 pipeline_class=pipeline_class,
-                config_with_params=config_with_params,
+                pipeline_config=pipeline_config,
                 dataloader=dataloader,
-                metrics=metrics,
+                metricset=metricset,
             )
 
-        study = optuna.create_study(direction=self.config.direction)
-
-        study.optimize(objective, **optimization_kwargs)
-        configs_with_scores = [
-            (trial.user_attrs["cfg"], trial.user_attrs["score"], trial.user_attrs["all_metrics"])
-            for trial in study.get_trials()
-        ]
-
-        def sorting_key(results: tuple[DictConfig, float, dict[str, float]]) -> float:
-            if self.config.direction == "maximize":
-                return -results[1]
-            else:
-                return results[1]
-
-        return sorted(configs_with_scores, key=sorting_key)
+        study = optuna.create_study(direction=self.direction)
+        study.optimize(
+            func=objective,
+            n_trials=self.n_trials,
+            callbacks=callbacks,
+        )
+        return sorted(
+            [
+                (
+                    trial.user_attrs["config"],
+                    trial.user_attrs["score"],
+                    trial.user_attrs["metrics"],
+                )
+                for trial in study.get_trials()
+            ],
+            key=lambda x: -x[1] if self.direction == "maximize" else x[1],
+        )
 
     def _objective(
         self,
+        trial: Trial,
         pipeline_class: type[EvaluationPipeline],
-        trial: optuna.Trial,
-        config_with_params: DictConfig,
+        pipeline_config: dict,
         dataloader: DataLoader,
-        metrics: MetricSet,
+        metricset: MetricSet,
     ) -> float:
-        max_retries = getattr(self.config, "max_retries_for_trial", 1)
+        """
+        Run a single experiment.
+        """
+        evaluator = Evaluator()
+        event_loop = asyncio.get_event_loop()
+
+        score = 1e16 if self.direction == "maximize" else -1e16
+        metrics_values = None
         config_for_trial = None
-        for attempt_idx in range(max_retries):
+
+        for attempt in range(1, self.max_retries_for_trial + 1):
             try:
-                config_for_trial = deepcopy(config_with_params)
+                config_for_trial = deepcopy(pipeline_config)
                 self._set_values_for_optimized_params(cfg=config_for_trial, trial=trial, ancestors=[])
-                pipeline = pipeline_class(config_for_trial)
-                metrics_values = self._score(pipeline=pipeline, dataloader=dataloader, metrics=metrics)
-                score = sum(metrics_values.values())
-                break
-            except Exception as e:
-                if attempt_idx < max_retries - 1:
-                    warnings.warn(
-                        message=f"Execution of the trial failed: {e}. A retry will be initiated.", category=UserWarning
-                    )
-                else:
-                    score = self.INFINITY
-                    if self.config.direction == "maximize":
-                        score *= -1
-                    metrics_values = {}
-                    warnings.warn(
-                        message=f"Execution of the trial failed: {e}. Setting the score to {score}",
-                        category=UserWarning,
+                pipeline = pipeline_class.from_config(config_for_trial)
+
+                results = event_loop.run_until_complete(
+                    evaluator.compute(
+                        pipeline=pipeline,
+                        dataloader=dataloader,
+                        metricset=metricset,
                     )
+                )
+                score = sum(results.metrics.values())
+                metrics_values = results.metrics
+                break
+            except Exception as exc:
+                message = (
+                    f"Execution of the trial failed: {exc}. A retry will be initiated"
+                    if attempt < self.max_retries_for_trial
+                    else f"Execution of the trial failed: {exc}. Setting the score to {score}"
+                )
+                warnings.warn(message=message, category=UserWarning)
+
         trial.set_user_attr("score", score)
-        trial.set_user_attr("cfg", config_for_trial)
-        trial.set_user_attr("all_metrics", metrics_values)
-        return score
+        trial.set_user_attr("metrics", metrics_values)
+        trial.set_user_attr("config", config_for_trial)
 
-    @staticmethod
-    def _score(pipeline: EvaluationPipeline, dataloader: DataLoader, metrics: MetricSet) -> dict[str, float]:
-        evaluator = Evaluator()
-        event_loop = asyncio.get_event_loop()
-        results = event_loop.run_until_complete(
-            evaluator.compute(pipeline=pipeline, dataloader=dataloader, metrics=metrics)
-        )
-        return results["metrics"]
+        return score
 
-    def _set_values_for_optimized_params(self, cfg: DictConfig, trial: optuna.Trial, ancestors: list[str]) -> None:  # noqa: PLR0912
+    def _set_values_for_optimized_params(self, cfg: dict, trial: Trial, ancestors: list[str]) -> None:  # noqa: PLR0912
         """
-        Recursive method for sampling parameter values for optuna.Trial
+        Recursive method for sampling parameter values for optuna trial.
         """
         for key, value in cfg.items():
-            if isinstance(value, DictConfig):
+            if isinstance(value, dict):
                 if value.get("optimize"):
                     param_id = f"{'.'.join(ancestors)}.{key}"  # type: ignore
                     choices = value.get("choices")
@@ -147,12 +199,12 @@ class Optimizer:
                             raise ValueError("Either choices or range must be specified")
                         choice_idx = trial.suggest_categorical(name=param_id, choices=choices_index)  # type: ignore
                         choice = choices[choice_idx]
-                        if isinstance(choice, DictConfig):
+                        if isinstance(choice, dict):
                             self._set_values_for_optimized_params(choice, trial, ancestors + [key, str(choice_idx)])  # type: ignore
                         cfg[key] = choice
                 else:
                     self._set_values_for_optimized_params(value, trial, ancestors + [key])  # type: ignore
-            elif isinstance(value, ListConfig):
+            elif isinstance(value, list):
                 for param in value:
-                    if isinstance(param, DictConfig):
+                    if isinstance(param, dict):
                         self._set_values_for_optimized_params(param, trial, ancestors + [key])  # type: ignore

ragbits/evaluate/pipelines/__init__.py

@@ -0,0 +1,37 @@
+from ragbits.core.utils.config_handling import WithConstructionConfig
+from ragbits.document_search import DocumentSearch
+from ragbits.evaluate.pipelines.base import EvaluationData, EvaluationPipeline, EvaluationResult
+from ragbits.evaluate.pipelines.document_search import DocumentSearchPipeline
+from ragbits.evaluate.pipelines.gaia import GaiaPipeline
+from ragbits.evaluate.pipelines.hotpot_qa import HotpotQAPipeline
+from ragbits.evaluate.pipelines.human_eval import HumanEvalPipeline
+
+__all__ = [
+    "DocumentSearchPipeline",
+    "EvaluationData",
+    "EvaluationPipeline",
+    "EvaluationResult",
+    "GaiaPipeline",
+    "HotpotQAPipeline",
+    "HumanEvalPipeline",
+]
+
+_target_to_evaluation_pipeline: dict[type[WithConstructionConfig], type[EvaluationPipeline]] = {
+    DocumentSearch: DocumentSearchPipeline,
+}
+
+
+def get_evaluation_pipeline_for_target(evaluation_target: WithConstructionConfig) -> EvaluationPipeline:
+    """
+    A function instantiating evaluation pipeline for given WithConstructionConfig object
+    Args:
+        evaluation_target: WithConstructionConfig object to be evaluated
+    Returns:
+        instance of evaluation pipeline
+    Raises:
+        ValueError for classes with no registered evaluation pipeline
+    """
+    for supported_type, evaluation_pipeline_type in _target_to_evaluation_pipeline.items():
+        if isinstance(evaluation_target, supported_type):
+            return evaluation_pipeline_type(evaluation_target=evaluation_target)
+    raise ValueError(f"Evaluation pipeline not implemented for {evaluation_target.__class__}")

ragbits/evaluate/pipelines/base.py

@@ -1,8 +1,23 @@
 from abc import ABC, abstractmethod
+from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import Any
+from types import ModuleType
+from typing import ClassVar, Generic, TypeVar
 
-from omegaconf import DictConfig
+from pydantic import BaseModel
+
+from ragbits.core.utils.config_handling import WithConstructionConfig
+from ragbits.evaluate import pipelines
+
+EvaluationDataT = TypeVar("EvaluationDataT", bound="EvaluationData")
+EvaluationResultT = TypeVar("EvaluationResultT", bound="EvaluationResult")
+EvaluationTargetT = TypeVar("EvaluationTargetT", bound=WithConstructionConfig)
+
+
+class EvaluationData(BaseModel, ABC):
+    """
+    Represents the data for a single evaluation.
+    """
 
 
 @dataclass
@@ -12,25 +27,34 @@ class EvaluationResult(ABC):
     """
 
 
-class EvaluationPipeline(ABC):
+class EvaluationPipeline(WithConstructionConfig, Generic[EvaluationTargetT, EvaluationDataT, EvaluationResultT], ABC):
     """
-    Collection evaluation pipeline.
+    Evaluation pipeline.
     """
 
-    def __init__(self, config: DictConfig | None = None) -> None:
+    default_module: ClassVar[ModuleType | None] = pipelines
+    configuration_key: ClassVar[str] = "pipeline"
+
+    def __init__(self, evaluation_target: EvaluationTargetT) -> None:
         """
-        Initializes the evaluation pipeline.
+        Initialize the evaluation pipeline.
 
         Args:
-            config: The evaluation pipeline configuration.
+            evaluation_target: Evaluation target instance.
         """
         super().__init__()
-        self.config = config or DictConfig({})
+        self.evaluation_target = evaluation_target
+
+    async def prepare(self) -> None:
+        """
+        Prepare pipeline for evaluation. Optional step.
+        """
+        pass
 
     @abstractmethod
-    async def __call__(self, data: dict[str, Any]) -> EvaluationResult:
+    async def __call__(self, data: Iterable[EvaluationDataT]) -> Iterable[EvaluationResultT]:
         """
-        Runs the evaluation pipeline.
+        Run the evaluation pipeline.
 
         Args:
             data: The evaluation data.

ragbits/evaluate/pipelines/document_search.py

@@ -1,16 +1,25 @@
 import asyncio
-import uuid
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
-from functools import cached_property
+from uuid import uuid4
 
-from omegaconf import DictConfig
-from tqdm.asyncio import tqdm
+from typing_extensions import Self
 
+from ragbits.core.sources.hf import HuggingFaceSource
 from ragbits.document_search import DocumentSearch
-from ragbits.document_search.documents.document import DocumentMeta
-from ragbits.document_search.documents.element import TextElement
-from ragbits.document_search.documents.sources import HuggingFaceSource
-from ragbits.evaluate.pipelines.base import EvaluationPipeline, EvaluationResult
+from ragbits.document_search.documents.element import Element
+from ragbits.evaluate.pipelines.base import EvaluationData, EvaluationPipeline, EvaluationResult
+
+
+class DocumentSearchData(EvaluationData):
+    """
+    Represents the evaluation data for document search.
+    """
+
+    question: str
+    reference_document_ids: list[str | int] | None = None
+    reference_passages: list[str] | None = None
+    reference_page_numbers: list[int] | None = None
 
 
 @dataclass
@@ -20,82 +29,78 @@ 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):
+class DocumentSearchPipeline(EvaluationPipeline[DocumentSearch, DocumentSearchData, DocumentSearchResult]):
     """
     Document search evaluation pipeline.
     """
 
-    @cached_property
-    def document_search(self) -> "DocumentSearch":
+    def __init__(self, evaluation_target: DocumentSearch, source: dict | None = None) -> None:
         """
-        Returns the document search instance.
+        Initialize the document search evaluation pipeline.
 
-        Returns:
-            The document search instance.
+        Args:
+            evaluation_target: Document Search instance.
+            source: Source data config for ingest.
         """
-        return DocumentSearch.from_config(self.config)  # type: ignore
+        super().__init__(evaluation_target=evaluation_target)
+        self.source = source or {}
 
-    async def __call__(self, data: dict) -> DocumentSearchResult:
+    @classmethod
+    def from_config(cls, config: dict) -> Self:
         """
-        Runs the document search evaluation pipeline.
+        Create an instance of `DocumentSearchPipeline` from a configuration dictionary.
 
         Args:
-            data: The evaluation data.
+            config: A dictionary containing configuration settings for the pipeline.
 
         Returns:
-            The evaluation result.
+            An instance of the pipeline class initialized with the provided configuration.
         """
-        elements = await self.document_search.search(data["question"])
-        predicted_passages = [element.content for element in elements if isinstance(element, TextElement)]
-        return DocumentSearchResult(
-            question=data["question"],
-            reference_passages=data["passages"],
-            predicted_passages=predicted_passages,
-        )
-
-
-class DocumentSearchWithIngestionPipeline(DocumentSearchPipeline):
-    """
-    A class for joint doument ingestion and search
-    """
-
-    def __init__(self, config: DictConfig | None = None) -> None:
-        super().__init__(config)
-        self.config.vector_store.config.index_name = str(uuid.uuid4())
-        self._ingested = False
-        self._lock = asyncio.Lock()
-
-    async def __call__(self, data: dict) -> DocumentSearchResult:
+        # At this point, we assume that if the source is set, the pipeline is run in experimental mode
+        # and create random indexes for testing
+        # 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 = DocumentSearch.from_config(config)
+        return cls(evaluation_target=evaluation_target, source=config.get("source"))
+
+    async def prepare(self) -> None:
         """
-        Queries a vector store with given data
-        Ingests the corpus to the store if has not been done
+        Ingest corpus data for evaluation.
+        """
+        if self.source:
+            # For now we only support HF sources for pre-evaluation ingest
+            # TODO: Make it generic to any data source
+            sources = await HuggingFaceSource.list_sources(
+                path=self.source["config"]["path"],
+                split=self.source["config"]["split"],
+            )
+            await self.evaluation_target.ingest(sources)
+
+    async def __call__(self, data: Iterable[DocumentSearchData]) -> Iterable[DocumentSearchResult]:
+        """
+        Run the document search evaluation pipeline.
+
         Args:
-            data: dict - query
+            data: The evaluation data batch.
+
         Returns:
-            DocumentSearchResult - query result
+            The evaluation result batch.
         """
-        async with self._lock:
-            if not self._ingested:
-                await self._ingest_documents()
-                self._ingested = True
-        return await super().__call__(data)
-
-    async def _ingest_documents(self) -> None:
-        documents = await tqdm.gather(
-            *[
-                DocumentMeta.from_source(
-                    HuggingFaceSource(
-                        path=self.config.answer_data_source.path,
-                        split=self.config.answer_data_source.split,
-                        row=i,
-                    )
-                )
-                for i in range(self.config.answer_data_source.num_docs)
-            ],
-            desc="Download",
-        )
-        await self.document_search.ingest(documents)
+        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)
+        ]