aidial-rag-eval 0.1.0__py3-none-any.whl

aidial_rag_eval/__init__.py

@@ -0,0 +1,7 @@
+from aidial_rag_eval.dataframe import create_rag_eval_metrics_report
+from aidial_rag_eval.evaluate import evaluate
+
+__all__ = [
+    "create_rag_eval_metrics_report",
+    "evaluate",
+]

aidial_rag_eval/dataframe/__init__.py

@@ -0,0 +1,11 @@
+from aidial_rag_eval.dataframe.metrics import (
+    calculate_generation_metrics,
+    create_rag_eval_metrics_report,
+    create_retrieval_metrics_report,
+)
+
+__all__ = [
+    "create_rag_eval_metrics_report",
+    "create_retrieval_metrics_report",
+    "calculate_generation_metrics",
+]

aidial_rag_eval/dataframe/match_facts.py

@@ -0,0 +1,40 @@
+import pandas as pd
+
+from aidial_rag_eval.dataframe.merge import merge_ground_truth_and_answers
+from aidial_rag_eval.facts.citation import CitationMatcher
+from aidial_rag_eval.retrieval.types import Matcher
+
+DEFAULT_MATCHER = CitationMatcher
+
+
+def match_facts(row: pd.Series, matcher: Matcher) -> pd.Series:
+    result_row = matcher.match_facts(row.facts, row.context)
+    return pd.Series(result_row._asdict())
+
+
+def apply_matcher_to_merged_dataframe(
+    df_merged: pd.DataFrame,
+    matcher: Matcher = DEFAULT_MATCHER,
+) -> pd.DataFrame:
+    matched_results = df_merged.apply(
+        match_facts,
+        matcher=matcher,
+        axis=1,
+        result_type="expand",
+    )
+    assert isinstance(matched_results, pd.DataFrame)
+    return matched_results
+
+
+def match_facts_dataframe(
+    ground_truth: pd.DataFrame,
+    answers: pd.DataFrame,
+    matcher: Matcher = DEFAULT_MATCHER,
+) -> pd.DataFrame:
+    df_merged = merge_ground_truth_and_answers(ground_truth, answers)
+    return pd.merge(
+        df_merged,
+        apply_matcher_to_merged_dataframe(df_merged, matcher),
+        left_index=True,
+        right_index=True,
+    )

aidial_rag_eval/dataframe/merge.py

@@ -0,0 +1,64 @@
+import pandas as pd
+
+from aidial_rag_eval.types import (
+    ANSWERS_COLUMNS,
+    GROUND_TRUTH_COLUMNS,
+    MERGED_KEY_COLUMNS,
+    AnswerColumns,
+    GroundTruthColumns,
+    MergedColumns,
+)
+
+
+def merge_ground_truth_and_answers(
+    ground_truth: pd.DataFrame, answers: pd.DataFrame
+) -> pd.DataFrame:
+    """
+    Merges ground_truth and collected answers dataframes.
+
+    Parameters
+    -----------
+    ground_truth : pd.DataFrame
+        contains the ground truth answers
+        The structure of the ground truth
+        is described in `aidial_rag_eval.types.GroundTruthAnswers`.
+
+    answers : pd.DataFrame
+        contains the collected answers
+        The structure of the collected answers
+        is described in `aidial_rag_eval.types.CollectedAnswers`.
+
+    Returns
+    ------------
+    pd.DataFrame
+        Returns merged ground_truth dataframe and answers dataframe.
+        The structure of the df_merged
+        is described in `aidial_rag_eval.types.MergedColumns`.
+    """
+    ground_truth_copy = ground_truth[
+        ground_truth.columns.intersection(GROUND_TRUTH_COLUMNS)
+    ].copy()
+    assert isinstance(ground_truth_copy, pd.DataFrame)
+    answers_copy = answers[answers.columns.intersection(ANSWERS_COLUMNS)].copy()
+    assert isinstance(answers_copy, pd.DataFrame)
+    ground_truth_copy = ground_truth_copy.rename(
+        columns={
+            GroundTruthColumns.ANSWER.value: MergedColumns.GROUND_TRUTH_ANSWER.value
+        }
+    )
+    ground_truth_copy[GroundTruthColumns.DOCUMENTS] = ground_truth_copy[
+        GroundTruthColumns.DOCUMENTS
+    ].apply(frozenset)
+    answers_copy[AnswerColumns.DOCUMENTS] = answers_copy[AnswerColumns.DOCUMENTS].apply(
+        frozenset
+    )
+    data = pd.merge(
+        ground_truth_copy,
+        answers_copy,
+        on=MERGED_KEY_COLUMNS,
+    )
+    data[MergedColumns.DOCUMENTS] = answers_copy.loc[
+        data.index, MergedColumns.DOCUMENTS
+    ]
+    data[MergedColumns.DOCUMENTS] = data[MergedColumns.DOCUMENTS].apply(list)
+    return data

aidial_rag_eval/dataframe/metrics.py

@@ -0,0 +1,301 @@
+from typing import List, Optional
+
+import pandas as pd
+from langchain_core.language_models import BaseChatModel
+
+from aidial_rag_eval.dataframe.match_facts import (
+    DEFAULT_MATCHER,
+    apply_matcher_to_merged_dataframe,
+)
+from aidial_rag_eval.dataframe.merge import merge_ground_truth_and_answers
+from aidial_rag_eval.generation.metric_binds import metric_binds_dict
+from aidial_rag_eval.generation.types import MetricBind, inference_column
+from aidial_rag_eval.retrieval.metrics import (
+    calculate_metrics as calculate_metrics_by_row,
+)
+from aidial_rag_eval.retrieval.types import Matcher
+from aidial_rag_eval.types import MergedColumns
+
+
+def apply_metrics_to_matched_results(match_result_data: pd.DataFrame) -> pd.DataFrame:
+    retrieval_metrics = match_result_data.apply(
+        calculate_metrics_by_row,
+        axis=1,
+        result_type="expand",
+    )
+    assert isinstance(retrieval_metrics, pd.DataFrame)
+    return retrieval_metrics
+
+
+def calculate_metrics(match_result_data: pd.DataFrame) -> pd.DataFrame:
+    return pd.merge(
+        match_result_data,
+        apply_metrics_to_matched_results(match_result_data),
+        left_index=True,
+        right_index=True,
+    )
+
+
+def calculate_retrieval_metrics(
+    df_merged: pd.DataFrame,
+    matcher: Matcher = DEFAULT_MATCHER,
+) -> pd.DataFrame:
+    """
+    Calculates RAG evaluation retrieval metrics from df_merged dataframe.
+
+    Parameters
+    -----------
+    df_merged : pd.DataFrame
+        merged ground truth answers and collected answers dataframes
+        The structure of the df_merged
+        is described in `aidial_rag_eval.types.MergedColumns`.
+
+    matcher : Matcher, default=DEFAULT_MATCHER
+        An object responsible for matching facts
+        from the ground truth to the context in the answers.
+
+    Returns
+    ------------
+    pd.DataFrame
+        Returns retrieval metrics dataframe.
+    """
+    matched_result = apply_matcher_to_merged_dataframe(df_merged, matcher)
+    matched_result_with_df = pd.merge(
+        df_merged, matched_result, left_index=True, right_index=True
+    )
+    retrieval_metrics = apply_metrics_to_matched_results(matched_result_with_df)
+    return pd.merge(
+        matched_result, retrieval_metrics, left_index=True, right_index=True
+    )
+
+
+def create_retrieval_metrics_report(
+    ground_truth: pd.DataFrame,
+    answers: pd.DataFrame,
+    matcher: Matcher = DEFAULT_MATCHER,
+) -> pd.DataFrame:
+    """
+    Calculates RAG evaluation retrieval metrics from input dataframes.
+
+    Parameters
+    -----------
+    ground_truth : pd.DataFrame
+        contains the ground truth answers
+        The structure of the ground truth
+        is described in `aidial_rag_eval.types.GroundTruthAnswers`.
+
+    answers : pd.DataFrame
+        contains the collected answers
+        The structure of the collected answers
+        is described in `aidial_rag_eval.types.CollectedAnswers`.
+
+    matcher : Matcher, default=DEFAULT_MATCHER
+        An object responsible for matching facts
+        from the ground truth to the context in the answers.
+
+    Returns
+    ------------
+    pd.DataFrame
+        Returns merged ground_truth dataframe,
+        answers dataframe with retrieval metrics dataframe.
+    """
+    df_merged = merge_ground_truth_and_answers(ground_truth, answers)
+    return pd.merge(
+        df_merged,
+        calculate_retrieval_metrics(df_merged, matcher),
+        left_index=True,
+        right_index=True,
+    )
+
+
+def calculate_generation_metrics(
+    df_merged: pd.DataFrame,
+    llm: BaseChatModel,
+    metric_binds: List[MetricBind],
+    max_concurrency: int = 8,
+    batch_size: int = 6,
+    show_progress_bar: bool = True,
+) -> pd.DataFrame:
+    """
+    Calculates RAG evaluation generation metrics from df_merged dataframe.
+
+    Parameters
+    -----------
+    df_merged : pd.DataFrame
+        merged ground truth answers and collected answers dataframes
+        The structure of the df_merged
+        is described in `aidial_rag_eval.types.MergedColumns`.
+
+    llm : BaseChatModel, optional, default=None
+        A Langchain chat model used for calculating generation metrics.
+
+    metric_binds : List[MetricBind], optional, default=None
+        A list of string constants from `aidial_rag_eval.metric_binds`
+        specifying the generation metrics.
+
+    max_concurrency : int, default=8
+        The maximum number of concurrent requests to the LLM.
+
+    batch_size : int, default=6
+        The maximum number objects are processed in a single prompt for simple tasks.
+
+    show_progress_bar : bool, default=True
+        To display a progress bar during LLM requests.
+
+    Returns
+    ------------
+    pd.DataFrame
+        Returns generation metrics dataframe.
+    """
+    df_merged_copy = df_merged.copy()
+    df_merged_copy[MergedColumns.JOINED_CONTEXT] = df_merged_copy[
+        MergedColumns.CONTEXT
+    ].apply(lambda x: "\n".join(x))
+    metric_results = dict()
+    for metric_bind in metric_binds:
+        metric_results.update(
+            metric_binds_dict[metric_bind](
+                df_merged=df_merged_copy,
+                llm=llm,
+                max_concurrency=max_concurrency,
+                batch_size=batch_size,
+                show_progress_bar=show_progress_bar,
+            ).to_dict(orient="series")
+        )
+    df_metrics = pd.DataFrame(data=metric_results)
+    nli_columns = [
+        column for column in metric_results.keys() if column.endswith(inference_column)
+    ]
+    if nli_columns:
+        sub_df_nli = df_metrics[nli_columns]
+        df_metrics["mean_" + inference_column] = sub_df_nli.mean(1)
+        df_metrics["median_" + inference_column] = sub_df_nli.median(
+            1
+        )  # pyright: ignore # noqa
+
+    return df_metrics
+
+
+def create_generation_metrics_report(
+    ground_truth: pd.DataFrame,
+    answers: pd.DataFrame,
+    llm: BaseChatModel,
+    metric_binds: List[MetricBind],
+    max_concurrency: int = 8,
+    batch_size: int = 6,
+    show_progress_bar: bool = True,
+) -> pd.DataFrame:
+    """
+    Calculates RAG evaluation generation metrics from input dataframes.
+
+    Parameters
+    -----------
+    ground_truth : pd.DataFrame
+        contains the ground truth answers
+        The structure of the ground truth
+        is described in `aidial_rag_eval.types.GroundTruthAnswers`.
+
+    answers : pd.DataFrame
+        contains the collected answers
+        The structure of the collected answers
+        is described in `aidial_rag_eval.types.CollectedAnswers`.
+
+    llm : BaseChatModel, optional, default=None
+        A Langchain chat model used for calculating generation metrics.
+
+    metric_binds : List[MetricBind], optional, default=None
+        A list of string constants from `aidial_rag_eval.metric_binds`
+        specifying the generation metrics.
+
+    max_concurrency : int, default=8
+        The maximum number of concurrent requests to the LLM.
+
+    batch_size : int, default=6
+        The maximum number objects are processed in a single prompt for simple tasks.
+
+    show_progress_bar : bool, default=True
+        To display a progress bar during LLM requests.
+
+    Returns
+    ------------
+    pd.DataFrame
+        Returns merged ground_truth dataframe,
+        answers dataframe with generation metrics dataframe.
+    """
+    df_merged = merge_ground_truth_and_answers(ground_truth, answers)
+    df_metrics = calculate_generation_metrics(
+        df_merged,
+        llm,
+        metric_binds,
+        max_concurrency,
+        batch_size,
+        show_progress_bar,
+    )
+    return pd.merge(df_merged, df_metrics, left_index=True, right_index=True)
+
+
+def create_rag_eval_metrics_report(
+    ground_truth: pd.DataFrame,
+    answers: pd.DataFrame,
+    matcher: Matcher = DEFAULT_MATCHER,
+    llm: Optional[BaseChatModel] = None,
+    metric_binds: Optional[List[MetricBind]] = None,
+    max_concurrency: int = 8,
+    batch_size: int = 6,
+    show_progress_bar: bool = True,
+) -> pd.DataFrame:
+    """
+    Calculates RAG evaluation metrics from input dataframes.
+
+    Parameters
+    -----------
+    ground_truth : pd.DataFrame
+        contains the ground truth answers
+        The structure of the ground truth
+        is described in `aidial_rag_eval.types.GroundTruthAnswers`.
+
+    answers : pd.DataFrame
+        contains the collected answers
+        The structure of the collected answers
+        is described in `aidial_rag_eval.types.CollectedAnswers`.
+
+    matcher : Matcher, default=DEFAULT_MATCHER
+        An object responsible for matching facts
+        from the ground truth to the context in the answers.
+
+    llm : BaseChatModel, optional, default=None
+        A Langchain chat model used for calculating generation metrics.
+
+    metric_binds : List[MetricBind], optional, default=None
+        A list of string constants from `aidial_rag_eval.metric_binds`
+        specifying the generation metrics.
+
+    max_concurrency : int, default=8
+        The maximum number of concurrent requests to the LLM.
+
+    batch_size : int, default=6
+        The maximum number objects are processed in a single prompt for simple tasks.
+
+    show_progress_bar : bool, default=True
+        To display a progress bar during LLM requests.
+
+    Returns
+    ------------
+    pd.DataFrame
+        Returns merged ground_truth dataframe, answers dataframe with metrics dataframe.
+    """
+    df_merged = merge_ground_truth_and_answers(ground_truth, answers)
+    retrieval_metrics = calculate_retrieval_metrics(df_merged, matcher)
+    if not metric_binds:
+        return pd.merge(df_merged, retrieval_metrics, left_index=True, right_index=True)
+    if llm is None:
+        raise ValueError("Argument 'llm' is required.")
+    generation_metrics = calculate_generation_metrics(
+        df_merged,
+        llm,
+        metric_binds,
+        max_concurrency,
+        batch_size,
+        show_progress_bar,
+    )
+    return pd.concat([df_merged, retrieval_metrics, generation_metrics], axis=1)

aidial_rag_eval/dataset.py

@@ -0,0 +1,99 @@
+from datetime import datetime
+from typing import Optional, Union
+
+import fsspec
+import pandas as pd
+from pydantic import Field, TypeAdapter
+from pydantic.dataclasses import dataclass
+
+
+@dataclass
+class DatasetMetadata:
+    created_at: Optional[datetime]
+    sources: list["Dataset"]
+    tools: dict[str, str]
+    file_info: dict = Field(default_factory=dict)
+    metrics: dict = Field(default_factory=dict)
+    statistics: dict = Field(default_factory=dict)
+
+    @staticmethod
+    def read_metadata(metadata_path: str) -> "DatasetMetadata":
+        fs, _ = fsspec.url_to_fs(metadata_path)
+        return TypeAdapter(DatasetMetadata).validate_json(fs.read_bytes(metadata_path))
+
+    @staticmethod
+    def create_missing_metadata(data_path: str) -> "DatasetMetadata":
+        fs, _ = fsspec.url_to_fs(data_path)
+        file_info = fs.info(data_path)
+        return DatasetMetadata(None, [], {}, file_info, {}, {})
+
+
+@dataclass
+class Dataset:
+    metadata_path: str
+    data_path: str
+    metadata: DatasetMetadata
+
+    @staticmethod
+    def handle_path(path: str) -> tuple[str, str]:
+        if path.endswith(".parquet.metadata.json"):
+            return path, path[: -len(".metadata.json")]
+        if path.endswith(".parquet"):
+            return path + ".metadata.json", path
+        raise ValueError(f"Invalid path: {path}")
+
+    @staticmethod
+    def from_metadata(path: str):
+        metadata_path, data_path = Dataset.handle_path(path)
+
+        fs, _ = fsspec.url_to_fs(data_path)
+        if fs.exists(metadata_path):
+            metadata = DatasetMetadata.read_metadata(metadata_path)
+        else:
+            metadata = DatasetMetadata.create_missing_metadata(data_path)
+
+        return Dataset(metadata_path, data_path, metadata)
+
+    def read_dataframe(self, *args, **kwargs) -> pd.DataFrame:
+        return pd.read_parquet(self.data_path, *args, **kwargs)
+
+    @staticmethod
+    def write_dataframe(
+        df: pd.DataFrame,
+        path: str,
+        sources: Optional[list["Dataset"]] = None,
+        tools: Optional[dict[str, str]] = None,
+        metrics: Optional[dict[str, float]] = None,
+        statistics: Optional[dict[str, float]] = None,
+        *args,
+        **kwargs,
+    ):
+        sources = sources or []
+        tools = tools or {}
+        metrics = metrics or {}
+        statistics = statistics or {}
+        metadata_path, data_path = Dataset.handle_path(path)
+
+        fs, _ = fsspec.url_to_fs(data_path)
+        df.to_parquet(data_path, *args, **kwargs)
+        file_info: dict = fs.info(data_path)
+
+        metadata = DatasetMetadata(
+            datetime.now(),
+            sources,
+            tools,
+            file_info,
+            metrics,
+            statistics,
+        )
+        fs.write_bytes(
+            metadata_path, TypeAdapter(DatasetMetadata).dump_json(metadata, indent=2)
+        )
+
+        return Dataset(metadata_path, data_path, metadata)
+
+
+def source_dataset(dataset: Union[str, Dataset]):
+    if isinstance(dataset, str):
+        return Dataset.from_metadata(dataset)
+    return dataset

aidial_rag_eval/evaluate.py

@@ -0,0 +1,107 @@
+from typing import Any, List, Optional, Union
+
+import pandas as pd
+from langchain_core.language_models import BaseChatModel
+
+from aidial_rag_eval.dataframe import create_rag_eval_metrics_report
+from aidial_rag_eval.dataframe.match_facts import DEFAULT_MATCHER
+from aidial_rag_eval.dataset import Dataset, source_dataset
+from aidial_rag_eval.generation.types import MetricBind
+from aidial_rag_eval.retrieval.types import Matcher
+from aidial_rag_eval.utils import get_tools_versions
+
+
+def evaluate(
+    ground_truth: Union[str, Dataset],
+    answers: Union[str, Dataset],
+    dest: str,
+    matcher: Matcher = DEFAULT_MATCHER,
+    llm: Optional[BaseChatModel] = None,
+    fs: Any = None,
+    metric_binds: Optional[List[MetricBind]] = None,
+    max_concurrency: int = 8,
+    batch_size: int = 6,
+    show_progress_bar: bool = True,
+) -> Dataset:
+    """
+    Calculates RAG evaluation metrics from input
+    Parquet datasets and writes the results to the specified destination.
+
+    Parameters
+    -----------
+    ground_truth : Union[str, Dataset]
+        Path to the Parquet file containing ground truth answers or
+        a Dataset object specifying the Parquet file with ground truth answers.
+        The structure of the ground truth Parquet file
+        is described in `aidial_rag_eval.types.GroundTruthAnswers`.
+
+    answers : Union[str, Dataset]
+        Path to the Parquet file containing collected answers or
+        a Dataset object specifying the Parquet file with collected answers.
+        The structure of the collected answers Parquet file
+        is described in `aidial_rag_eval.types.CollectedAnswers`.
+
+    dest : str
+        Path for the output file in Parquet format.
+
+    matcher : Matcher, default=DEFAULT_MATCHER
+        An object responsible for matching facts
+        from the ground truth to the context in the answers.
+
+    llm : BaseChatModel, optional, default=None
+        A Langchain chat model used for calculating generation metrics.
+
+    fs : Any, optional, default=None
+        File system to be used.
+
+    metric_binds : List[MetricBind], optional, default=None
+        A list of string constants from `aidial_rag_eval.metric_binds`
+        specifying the generation metrics.
+
+    max_concurrency : int, default=8
+        The maximum number of concurrent requests to the LLM.
+
+    batch_size : int, default=6
+        The maximum number objects are processed in a single prompt for simple tasks.
+
+    show_progress_bar : bool, default=True
+        To display a progress bar during LLM requests.
+
+    Returns
+    ------------
+    Dataset
+        Returns global statistics for each metric.
+    """
+    ground_truth_dataset = source_dataset(ground_truth)
+    answers_dataset = source_dataset(answers)
+
+    ground_truth_df = ground_truth_dataset.read_dataframe(filesystem=fs)
+    answers_df = answers_dataset.read_dataframe(filesystem=fs)
+
+    df_final = create_rag_eval_metrics_report(
+        ground_truth_df,
+        answers_df,
+        matcher,
+        llm,
+        metric_binds=metric_binds,
+        max_concurrency=max_concurrency,
+        batch_size=batch_size,
+        show_progress_bar=show_progress_bar,
+    )
+    aggregated_metrics = df_final.mean(numeric_only=True)
+    assert isinstance(aggregated_metrics, pd.Series)
+
+    metrics = Dataset.write_dataframe(
+        df_final,
+        dest,
+        sources=[ground_truth_dataset, answers_dataset],
+        tools=get_tools_versions(),
+        metrics=aggregated_metrics.to_dict(),
+        statistics={
+            "Ground truth size": len(ground_truth_df),
+            "Answers size": len(answers_df),
+            "Evaluation data size": len(df_final),
+        },
+        filesystem=fs,
+    )
+    return metrics

aidial_rag_eval/facts/citation.py

@@ -0,0 +1,48 @@
+from typing import Optional, Tuple
+
+from aidial_rag_eval.facts.match_str_facts import match_str_facts
+from aidial_rag_eval.retrieval.types import ContextChunk, FactMatchResult
+from aidial_rag_eval.types import Context, Facts
+
+
+class CitationMatcher:
+    Fact = str
+
+    @staticmethod
+    def _canonize(text: str) -> tuple[str, list[int]]:
+        text = text.replace("\t", " ")
+        text = "".join([c.lower() if c.isalnum() else " " for c in text])
+
+        new_text = ""
+        num_skipped_chars_for_pos = []
+
+        # skip duplicated spaces
+        for i, c in enumerate(text):
+            if c == " " and new_text[-1:] == " ":
+                continue
+            new_text += c
+            num_skipped_chars_for_pos.append(i - len(new_text) + 1)
+
+        num_skipped_chars_for_pos.append(len(text) - len(new_text))
+
+        return new_text, num_skipped_chars_for_pos
+
+    @staticmethod
+    def _match(fact: Fact, context_chunk: ContextChunk) -> Optional[Tuple[int, int]]:
+        canonized_fact, _ = CitationMatcher._canonize(fact)
+        canonized_context_chunk, skipped_chars = CitationMatcher._canonize(
+            context_chunk
+        )
+
+        start = canonized_context_chunk.find(canonized_fact)
+        if start == -1:
+            return None
+        end = start + len(canonized_fact)
+
+        start = start + skipped_chars[start]
+        end = end + skipped_chars[end]
+        return start, end
+
+    @staticmethod
+    def match_facts(facts: Facts[Fact], context: Context) -> FactMatchResult:
+        return match_str_facts(facts, context, CitationMatcher._match)

aidial_rag_eval/facts/exact.py

@@ -0,0 +1,19 @@
+from typing import Optional, Tuple
+
+from aidial_rag_eval.facts.match_str_facts import match_str_facts
+from aidial_rag_eval.retrieval.types import ContextChunk, FactMatchResult
+from aidial_rag_eval.types import Context, Facts
+
+
+class ExactStringMatcher:
+    Fact = str
+
+    @staticmethod
+    def _match(fact: Fact, context_chunk: ContextChunk) -> Optional[Tuple[int, int]]:
+        if fact == context_chunk:
+            return (0, len(context_chunk))
+        return None
+
+    @staticmethod
+    def match_facts(facts: Facts[Fact], context: Context) -> FactMatchResult:
+        return match_str_facts(facts, context, ExactStringMatcher._match)