Source code for autorag.schema.metricinput

from dataclasses import dataclass
from typing import Optional, List, Dict, Callable, Any, Union

import numpy as np
import pandas as pd




@dataclass
class MetricInput:
    query: Optional[str] = None
    queries: Optional[List[str]] = None
    retrieval_gt_contents: Optional[List[List[str]]] = None
    retrieved_contents: Optional[List[str]] = None
    retrieval_gt: Optional[List[List[str]]] = None
    retrieved_ids: Optional[List[str]] = None
    prompt: Optional[str] = None
    generated_texts: Optional[str] = None
    generation_gt: Optional[List[str]] = None
    generated_log_probs: Optional[List[float]] = None



    def is_fields_notnone(self, fields_to_check: List[str]) -> bool:
        for field in fields_to_check:
            actual_value = getattr(self, field)

            if actual_value is None:
                return False

            try:
                if not type_checks.get(type(actual_value), lambda _: False)(actual_value):
                    return False
            except Exception:
                return False

        return True




    @classmethod
    def from_dataframe(cls, qa_data: pd.DataFrame) -> List['MetricInput']:
        """
        Convert a pandas DataFrame into a list of MetricInput instances.
        qa_data: pd.DataFrame: qa_data DataFrame containing metric data.

        :returns: List[MetricInput]: List of MetricInput objects created from DataFrame rows.
        """
        instances = []

        for _, row in qa_data.iterrows():
            instance = cls()

            for attr_name in cls.__annotations__:
                if attr_name in row:
                    value = row[attr_name]

                    if isinstance(value, str):
                        setattr(instance, attr_name, value.strip() if value.strip() != '' else None)
                    elif isinstance(value, list):
                        setattr(instance, attr_name, value if len(value) > 0 else None)
                    else:
                        setattr(instance, attr_name, value)

            instances.append(instance)

        return instances


    @staticmethod
    def _check_list(lst_or_arr: Union[List[Any], np.ndarray]) -> bool:
        if isinstance(lst_or_arr, np.ndarray):
            lst_or_arr = lst_or_arr.flatten().tolist()

        if len(lst_or_arr) == 0:
            return False

        for item in lst_or_arr:
            if item is None:
                return False

            item_type = type(item)

            if item_type in type_checks:
                if not type_checks[item_type](item):
                    return False
            else:
                return False

        return True



type_checks: Dict[type, Callable[[Any], bool]] = {
    str: lambda x: len(x.strip()) > 0,
    list: MetricInput._check_list,
    np.ndarray: MetricInput._check_list,
    int: lambda _: True,
    float: lambda _: True,
}