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

# %% auto 0
__all__ = ['make_docs', 'BaselineSingleHop']

# %% ../../nbs/musique/musique.singlehop.ipynb 4
from tqdm.auto import tqdm


tqdm.pandas()

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

# %% ../../nbs/musique/musique.singlehop.ipynb 6
class BaselineSingleHop:
    def __init__(self, qa_func, retrieval_func):
        self.qa_func = qa_func
        self.retrieval_func = retrieval_func

    def _call(self, example) -> dict:
        docs = list(make_docs(example))
        question = example["question"]
        query = question
        retrieved_docs = self.retrieval_func(docs, query)
        context = "\n\n".join(doc['text'] for doc in retrieved_docs)
        qa_result = self.qa_func(context=context, question=question)
        answer = qa_result.get("answer")
        hop = {
            "question": question,
            "query" : query,
            "retrieved_docs": retrieved_docs,
            "context": context,
            "answer": answer,
            "qa_result": qa_result,
        }
        return {'answer': answer, 'hops': [hop]}

    def __call__(self, example, ignore_errors: bool = False) -> dict:
        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 = dict(answer="N/A", hops=[{'error': str(exc)}])
            else:
                raise
        return output
