feat/executor: make ragas faster and more robust to run

.gitignore

@@ -161,11 +161,7 @@ cython_debug/
 
 # Ragas specific
 ragas/_version.py
-experiments/**/data
-experiments/**/storage
+experiments/
 **/fil-result/
-experiments/baselines/fiqa/datasets
 src/ragas/_version.py
 .python-version
-experiments/retriever-benchmarks/datasets
-experiments/tmp

docs/conf.py

@@ -1,4 +1,3 @@
-import os
 from dataclasses import asdict
 
 from sphinxawesome_theme import ThemeOptions

docs/howtos/customisations/embeddings.ipynb

@@ -169,7 +169,7 @@
     "\n",
     "result = evaluate(\n",
     "    fiqa_eval[\"baseline\"].select(range(5)),  # showing only 5 for demonstration\n",
-    "    metrics=[answer_similarity]\n",
+    "    metrics=[answer_similarity],\n",
     ")\n",
     "\n",
     "result"

docs/howtos/customisations/gcp-vertexai.ipynb

@@ -98,7 +98,7 @@
    "source": [
     "from ragas.metrics import (\n",
     "    context_precision,\n",
-    "    answer_relevancy, # AnswerRelevancy\n",
+    "    answer_relevancy,  # AnswerRelevancy\n",
     "    faithfulness,\n",
     "    context_recall,\n",
     ")\n",
@@ -110,7 +110,7 @@
     "    answer_relevancy,\n",
     "    context_recall,\n",
     "    context_precision,\n",
-    "    harmfulness\n",
+    "    harmfulness,\n",
     "]"
    ]
   },
@@ -137,7 +137,6 @@
     "from langchain.embeddings import VertexAIEmbeddings\n",
     "\n",
     "\n",
-    "\n",
     "config = {\n",
     "    \"project_id\": \"tmp-project-404003\",\n",
     "}\n",
@@ -170,7 +169,7 @@
     "for m in metrics:\n",
     "    # change LLM for metric\n",
     "    m.__setattr__(\"llm\", ragas_vertexai_llm)\n",
-    "    \n",
+    "\n",
     "    # check if this metric needs embeddings\n",
     "    if hasattr(m, \"embeddings\"):\n",
     "        # if so change with VertexAI Embeddings\n",
@@ -276,13 +275,15 @@
    ],
    "source": [
     "from ragas import evaluate\n",
-    "import nest_asyncio # CHECK NOTES\n",
+    "import nest_asyncio  # CHECK NOTES\n",
     "\n",
-    "# NOTES: Only used when running on a jupyter notebook, otherwise comment or remove this function. \n",
-    "nest_asyncio.apply()  \n",
+    "# NOTES: Only used when running on a jupyter notebook, otherwise comment or remove this function.\n",
+    "nest_asyncio.apply()\n",
     "\n",
     "result = evaluate(\n",
-    "    fiqa_eval[\"baseline\"].select(range(1)), # using 1 as example due to quota constrains\n",
+    "    fiqa_eval[\"baseline\"].select(\n",
+    "        range(1)\n",
+    "    ),  # using 1 as example due to quota constrains\n",
     "    metrics=metrics,\n",
     ")\n",
     "\n",

docs/howtos/integrations/zeno.ipynb

@@ -186,7 +186,7 @@
     "    ]\n",
     "].copy()\n",
     "\n",
-    "output_df['output'] = df.apply(\n",
+    "output_df[\"output\"] = df.apply(\n",
     "    lambda x: {\"answer\": x[\"answer\"], \"ground_truths\": list(x[\"ground_truths\"])}, axis=1\n",
     ")\n",
     "output_df[\"id\"] = output_df.index\n",

pyproject.toml

@@ -23,6 +23,9 @@ package-dir = {"" = "src"}
 [tool.setuptools.dynamic]
 readme = {file = ["README.md"], content-type = "text/plain"}
 
+[tool.ruff.lint]
+ignore = ["E501"]
+
 [build-system]
 requires = ["setuptools>=45", "setuptools_scm[toml]>=6.2"]
 build-backend = "setuptools.build_meta"

src/ragas/callbacks.py

@@ -0,0 +1,58 @@
+import typing as t
+
+from langchain_core.callbacks import (
+    AsyncCallbackManager,
+    AsyncCallbackManagerForChainGroup,
+    AsyncCallbackManagerForChainRun,
+    CallbackManager,
+    CallbackManagerForChainGroup,
+    CallbackManagerForChainRun,
+    Callbacks,
+)
+
+
+def new_group(
+    name: str, inputs: t.Dict, callbacks: Callbacks, is_async=False
+) -> t.Tuple[CallbackManagerForChainRun, CallbackManagerForChainGroup]:
+    # start evaluation chain
+    if isinstance(callbacks, list):
+        cm = CallbackManager.configure(inheritable_callbacks=callbacks)
+    else:
+        cm = t.cast(CallbackManager, callbacks)
+    rm = cm.on_chain_start({"name": name}, inputs)
+    child_cm = rm.get_child()
+    group_cm = CallbackManagerForChainGroup(
+        child_cm.handlers,
+        child_cm.inheritable_handlers,
+        child_cm.parent_run_id,
+        parent_run_manager=rm,
+        tags=child_cm.tags,
+        inheritable_tags=child_cm.inheritable_tags,
+        metadata=child_cm.metadata,
+        inheritable_metadata=child_cm.inheritable_metadata,
+    )
+
+    return rm, group_cm
+
+
+async def new_async_group(
+    name: str, inputs: t.Dict, callbacks: Callbacks
+) -> t.Tuple[AsyncCallbackManagerForChainRun, AsyncCallbackManagerForChainGroup]:
+    # start evaluation chain
+    if isinstance(callbacks, list):
+        cm = AsyncCallbackManager.configure(inheritable_callbacks=callbacks)
+    else:
+        cm = t.cast(AsyncCallbackManager, callbacks)
+    rm = await cm.on_chain_start({"name": name}, inputs)
+    child_cm = rm.get_child()
+    group_cm = AsyncCallbackManagerForChainGroup(
+        child_cm.handlers,
+        child_cm.inheritable_handlers,
+        child_cm.parent_run_id,
+        parent_run_manager=rm,
+        tags=child_cm.tags,
+        inheritable_tags=child_cm.inheritable_tags,
+        metadata=child_cm.metadata,
+        inheritable_metadata=child_cm.inheritable_metadata,
+    )
+    return rm, group_cm

src/ragas/embeddings/__init__.py

@@ -1,15 +1,15 @@
 from ragas.embeddings.base import (
     AzureOpenAIEmbeddings,
+    BaseRagasEmbeddings,
     FastEmbedEmbeddings,
     HuggingfaceEmbeddings,
     OpenAIEmbeddings,
-    RagasEmbeddings,
 )
 
 __all__ = [
     "HuggingfaceEmbeddings",
     "OpenAIEmbeddings",
     "AzureOpenAIEmbeddings",
-    "RagasEmbeddings",
+    "BaseRagasEmbeddings",
     "FastEmbedEmbeddings",
 ]

src/ragas/embeddings/base.py

@@ -18,15 +18,11 @@
 DEFAULT_MODEL_NAME = "BAAI/bge-small-en-v1.5"
 
 
-class RagasEmbeddings(Embeddings):
-    def validate_api_key(self):
-        """
-        Validates that the api key is set for the Embeddings
-        """
-        pass
+class BaseRagasEmbeddings(Embeddings):
+    ...
 
 
-class OpenAIEmbeddings(BaseOpenAIEmbeddings, RagasEmbeddings):
+class OpenAIEmbeddings(BaseOpenAIEmbeddings, BaseRagasEmbeddings):
     api_key: str = NO_KEY
 
     def __init__(self, api_key: str = NO_KEY):
@@ -48,7 +44,7 @@ def validate_api_key(self):
                 raise OpenAIKeyNotFound
 
 
-class FastEmbedEmbeddings(BaseFastEmbedEmbeddings, RagasEmbeddings):
+class FastEmbedEmbeddings(BaseFastEmbedEmbeddings, BaseRagasEmbeddings):
     """
     Find the list of supported models at:
     https://qdrant.github.io/fastembed/examples/Supported_Models/
@@ -66,7 +62,7 @@ def validate_api_key(self):
         pass
 
 
-class AzureOpenAIEmbeddings(BaseAzureOpenAIEmbeddings, RagasEmbeddings):
+class AzureOpenAIEmbeddings(BaseAzureOpenAIEmbeddings, BaseRagasEmbeddings):
     azure_endpoint: t.Optional[str] = None
     deployment: t.Optional[str] = None
     api_version: t.Optional[str] = None
@@ -104,7 +100,7 @@ def validate_api_key(self):
 
 
 @dataclass
-class HuggingfaceEmbeddings(RagasEmbeddings):
+class HuggingfaceEmbeddings(BaseRagasEmbeddings):
     model_name: str = DEFAULT_MODEL_NAME
     """Model name to use."""
     cache_folder: t.Optional[str] = None
@@ -178,6 +174,6 @@ def predict(self, texts: List[List[str]]) -> List[List[float]]:
         return predictions.tolist()
 
 
-def embedding_factory() -> RagasEmbeddings:
+def embedding_factory() -> BaseRagasEmbeddings:
     openai_embeddings = OpenAIEmbeddings()
     return openai_embeddings

src/ragas/evaluation.py

@@ -5,21 +5,36 @@
 
 import numpy as np
 from datasets import Dataset, concatenate_datasets
+from langchain_core.language_models import BaseLanguageModel
 
 from ragas._analytics import EvaluationEvent, track
-from ragas.metrics.base import Metric
-from ragas.metrics.critique import AspectCritique
+from ragas.callbacks import new_group
+from ragas.embeddings.base import BaseRagasEmbeddings
+from ragas.executor import Executor
+from ragas.llms.base import BaseRagasLLM, LangchainLLMWrapper
+from ragas.metrics.base import Metric, MetricWithLLM
+
+# from ragas.metrics.critique import AspectCritique
 from ragas.validation import (
     remap_column_names,
     validate_column_dtypes,
     validate_evaluation_modes,
 )
 
+if t.TYPE_CHECKING:
+    from langchain_core.callbacks import Callbacks
+
 
 def evaluate(
     dataset: Dataset,
     metrics: list[Metric] | None = None,
-    column_map: dict[str, str] = {},
+    llm: t.Optional[BaseRagasLLM] = None,
+    embeddings: t.Optional[BaseRagasEmbeddings] = None,
+    callbacks: Callbacks = [],
+    is_async: bool = False,
+    max_workers: t.Optional[int] = None,
+    raise_exceptions: bool = True,
+    column_map: t.Dict[str, str] = {},
 ) -> Result:
     """
     Run the evaluation on the dataset with different metrics
@@ -81,24 +96,87 @@ def evaluate(
         )
 
         metrics = [answer_relevancy, context_precision, faithfulness, context_recall]
+    # set the llm and embeddings
+    if llm is None:
+        from ragas.llms import llm_factory
+
+        llm = llm_factory()
+    elif isinstance(llm, BaseLanguageModel):
+        llm = LangchainLLMWrapper(llm)
+    if embeddings is None:
+        from ragas.embeddings.base import embedding_factory
+
+        embeddings = embedding_factory()
 
     # remap column names from the dataset
     dataset = remap_column_names(dataset, column_map)
     # validation
     validate_evaluation_modes(dataset, metrics)
     validate_column_dtypes(dataset)
 
-    # run the evaluation on dataset with different metrics
+    binary_metrics = []
+    for metric in metrics:
+        # if isinstance(metric, AspectCritique):
+        # binary_metrics.append(metric.name)
+        if isinstance(metric, MetricWithLLM):
+            if metric.llm is None:
+                metric.llm = llm
+
     # initialize all the models in the metrics
     [m.init_model() for m in metrics]
 
+    executor = Executor(
+        is_async=is_async, max_workers=max_workers, raise_exceptions=raise_exceptions
+    )
+    # new evaluation chain
+    row_run_managers = []
+    evaluation_rm, evaluation_group_cm = new_group(
+        name="ragas evaluation", inputs={}, callbacks=callbacks, is_async=is_async
+    )
+    for i, row in enumerate(dataset):
+        row = t.cast(t.Dict[str, t.Any], row)
+        row_rm, row_group_cm = new_group(
+            name=f"row {i}",
+            inputs=row,
+            callbacks=evaluation_group_cm,
+            is_async=is_async,
+        )
+        row_run_managers.append((row_rm, row_group_cm))
+
+        if is_async:
+            [executor.submit(metric.ascore, row, row_group_cm) for metric in metrics]
+        else:
+            [executor.submit(metric.score, row, row_group_cm) for metric in metrics]
+
     scores = []
-    binary_metrics = []
-    for metric in metrics:
-        if isinstance(metric, AspectCritique):
-            binary_metrics.append(metric.name)
-        print(f"evaluating with [{metric.name}]")
-        scores.append(metric.score(dataset).select_columns(metric.name))
+    try:
+        # get the results
+        results = executor.results()
+        # convert results to dataset_like
+        for i, _ in enumerate(dataset):
+            s = {}
+            for j, m in enumerate(metrics):
+                s[m.name] = results[len(metrics) * i + j]
+            scores.append(s)
+            # close the row chain
+            row_rm, row_group_cm = row_run_managers[i]
+            if not row_group_cm.ended:
+                row_rm.on_chain_end(s)
+
+    # run evaluation task
+    except Exception as e:
+        if not evaluation_group_cm.ended:
+            evaluation_rm.on_chain_error(e)
+
+        raise e
+    finally:
+        result = Result(
+            scores=Dataset.from_list(scores),
+            dataset=dataset,
+            binary_columns=binary_metrics,
+        )
+        if not evaluation_group_cm.ended:
+            evaluation_rm.on_chain_end(result)
 
     # log the evaluation event
     metrics_names = [m.name for m in metrics]
@@ -110,23 +188,18 @@ def evaluate(
             num_rows=dataset.shape[0],
         )
     )
-
-    return Result(
-        scores=concatenate_datasets(scores, axis=1),
-        dataset=dataset,
-        binary_columns=binary_metrics,
-    )
+    return result
 
 
 @dataclass
 class Result(dict):
     scores: Dataset
-    dataset: Dataset | None = None
-    binary_columns: list[str] = field(default_factory=list)
+    dataset: t.Optional[Dataset] = None
+    binary_columns: t.List[str] = field(default_factory=list)
 
     def __post_init__(self):
         values = []
-        for cn in self.scores.column_names:
+        for cn in self.scores[0].keys():
             value = np.nanmean(self.scores[cn])
             self[cn] = value
             if cn not in self.binary_columns: