# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/musique.eval.ipynb.

# %% auto 0
__all__ = ['musique_metric', 'fuzzy_match_metric', 'compute_scores', 'calculate_metrics', 'compare_answers']

# %% ../../nbs/musique.eval.ipynb 3
import evaluate
import pandas as pd

from ..text.utils import fuzzy_match

# %% ../../nbs/musique.eval.ipynb 4
def fuzzy_match_metric(prediction: str, references: list[str]) -> float:
    return max([float(fuzzy_match(prediction, ref)) for ref in references])

# %% ../../nbs/musique.eval.ipynb 5
musique_metric = evaluate.load("bdsaglam/musique")


def compute_scores(predicted_answer: str, reference_answers: list[str]) -> dict:
    musique_scores = musique_metric.compute(predictions=[predicted_answer], references=[reference_answers])
    fuzzy_match = fuzzy_match_metric(predicted_answer, reference_answers)
    return {**musique_scores, "fuzzy_match": fuzzy_match}

# %% ../../nbs/musique.eval.ipynb 7
def calculate_metrics(dataf: pd.DataFrame) -> dict:
    prediction_list = dataf["predicted_answer"].tolist()
    references_list = dataf["answers"].tolist()
    scores_list = [compute_scores(prediction, references) for prediction, references in zip(prediction_list, references_list)]
    return pd.DataFrame(scores_list).mean().to_dict()

# %% ../../nbs/musique.eval.ipynb 8
def _exact_match(example):
    pred = example['predicted_answer']
    return pred is not None and any(pred == ref for ref in example['answers'])

def _fuzzy_match(example):
    pred = example['predicted_answer']
    return pred is not None and any((pred in ref) or (ref in pred) or fuzzy_match(pred, ref) for ref in example['answers'])


def compare_answers(dataf: pd.DataFrame) -> pd.DataFrame:
    dataf['exact_match'] = dataf.apply(_exact_match, axis=1)
    dataf['fuzzy_match'] = dataf.apply(_fuzzy_match, axis=1)
    return dataf
