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

# %% auto 0
__all__ = ['make_docs', 'Result', 'BaselineMultiHop', 'benchmark']

# %% ../../nbs/musique.nhop.ipynb 3
import json
from typing import Callable

import pandas as pd
from tqdm.auto import tqdm
from pydantic import BaseModel

from ..jerx.reward.llm import QuestionAnsweringResult
from .eval import calculate_metrics, compare_answers

tqdm.pandas()

# %% ../../nbs/musique.nhop.ipynb 4
def make_docs(example):
    ps = example["paragraphs"]
    for p in ps:
        idx = p["idx"]
        title = p["title"]
        body = p["paragraph_text"]
        is_supporting = p["is_supporting"]
        text = f"# {title}\n{body}"
        yield dict(
            text=text,
            is_supporting=is_supporting,
            parent_id=example["id"],
            idx=idx,
        )

# %% ../../nbs/musique.nhop.ipynb 5
class Result(BaseModel):
    answer: str
    hops: list[dict]

class BaselineMultiHop:
    def __init__(self, qa_func, retrieval_func):
        self.qa_func = qa_func
        self.retrieval_func = retrieval_func

    def _call(self, example) -> Result:
        docs = list(make_docs(example))
        
        # First question
        prev_answer = None
        hops = []
        for i in range(len(example["question_decomposition"])):
            question = example["question_decomposition"][i]["question"]
            if prev_answer is not None:
                question = question.replace(f"#{i}", prev_answer)
            docs = self.retrieval_func(docs, question)
            context = "\n".join(doc['text'] for doc in docs)
            result = self.qa_func(context=context, question=question)
            hop = {
                "question": question,
                "context": context,
                "answer": result.answer,
                "reasoning": result.reasoning,
            }
            hops.append(hop)
            prev_answer = result.answer
            if result.answer == "N/A":
                break

        # Second question
        return Result(
            answer=hops[-1]['answer'],
            hops=hops,
        )

    def __call__(self, example, ignore_errors: bool = False) -> Result:
        try:
            output = self._call(example)
        except Exception as exc:
            if ignore_errors:
                id = example['id']
                print(f"Failed to answer the question {id}\n{exc}")
                output = Result(answer="N/A", hops=[dict(error=str(exc))])
            else:
                raise
        return output

# %% ../../nbs/musique.nhop.ipynb 6
def benchmark(
    dataf: pd.DataFrame,
    qa_func: Callable,
    retrieval_func: Callable,
    ignore_errors: bool = False,
) -> tuple[pd.DataFrame, dict]:
    pipeline = BaselineMultiHop(qa_func, retrieval_func)

    def process(example):
        output = pipeline(example, ignore_errors=ignore_errors)
        example["predicted_answer"] = output.answer
        example["raw_output"] = output
        return example

    dataf = dataf.progress_apply(process, axis=1)
    dataf = compare_answers(dataf)
    scores = calculate_metrics(dataf)
    scores["fuzzy_match"] = dataf["fuzzy_match"].mean()
    return dataf, scores
