import openai

context = """
The first step is to decide what to work on. The work you choose needs to have three qualities: it has to be something you have a natural aptitude for, that you have a deep interest in, and that offers scope to do great work.
In practice you don't have to worry much about the third criterion. Ambitious people are if anything already too conservative about it. So all you need to do is find something you have an aptitude for and great interest in.
"""

def run_qa(question: str) -> str:
    return openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": f"Answer in less than 30 words. Use the following context if needed: {context}",
            },
            {"role": "user", "content": question},
        ],
    )["choices"][0]["message"]["content"]

contexts = [[context]] * len(questions)

# https://huggingface.co/datasets/repllabs/questions_how_to_do_great_work
questions = [
    "how to do great work?.",
    "How can curiosity be nurtured and utilized to drive great work?",
    "How does the author suggest finding something to work on?",
    "How did Van Dyck's painting differ from Daniel Mytens' version and what message did it convey?",
]

runner = fastrepl.local_runner(fn=run_qa)
ds = runner.run(args_list=[(q,) for q in questions], output_feature="answer")

ds = ds.add_column("question", questions)
ds = ds.add_column("contexts", contexts)
# fastrepl.Dataset({
#     features: ['answer', 'question', 'contexts'],
#     num_rows: 4
# })

evaluator = fastrepl.RAGEvaluator(node=fastrepl.RAGAS(metric="Faithfulness"))

ds = fastrepl.local_runner(evaluator=evaluator, dataset=ds).run(num=5)
# ds["result"]
# [[0.25, 0.0, 0.25, 0.25, 0.5],
#  [0.5, 0.5, 0.5, 0.75, 0.875],
#  [0.66, 0.66, 0.66, 0.66, 0.66],
#  [1.0, 1.0, 1.0, 1.0, 1.0]]