import glob
import os
import pandas as pd
import random

import dspy
from dspy.evaluate import Evaluate
from dspy.teleprompt import BootstrapFewShotWithRandomSearch


os.environ["DSP_NOTEBOOK_CACHEDIR"] = os.path.join('.', 'cache')



turbo = dspy.OpenAI(model='gpt-3.5-turbo-1106', max_tokens=250, model_type='chat')

dspy.settings.configure(lm=turbo)



gpt4T = dspy.OpenAI(model='gpt-4-1106-preview', max_tokens=350, model_type='chat')


RUN_FROM_SCRATCH = False



get_ipython().system('git clone https://github.com/selenashe/ScoNe.git')



def load_scone(dirname):
    dfs = []
    for filename in glob.glob(dirname + "/*.csv"):
        df = pd.read_csv(filename, index_col=0)
        df['category'] = os.path.basename(filename).replace(".csv", "")
        dfs.append(df)
    data_df = pd.concat(dfs)

    def as_example(row):
        suffix = '' if row['category'] == 'one_scoped' else '_edited'
        hkey = 'sentence2' + suffix
        question = row[hkey][0].lower() + row[hkey][1: ].strip(".")
        question = f"Can we logically conclude for sure that {question}?"
        label = "Yes" if row['gold_label' + suffix] == 'entailment' else "No"
        return dspy.Example({
            "context": row['sentence1' + suffix],
            "question": question,
            "answer": label,
            "category": row['category']
        }).with_inputs("context", "question")

    return list(data_df.apply(as_example, axis=1).values)



all_train = load_scone("ScoNe/scone_nli/train")

random.seed(1)
random.shuffle(all_train)

train, dev = all_train[: 200], all_train[200: 250]

len(train), len(dev)



random.seed(1)

test = load_scone(dirname=f"ScoNe/scone_nli/test")

test = [ex for ex in test if ex.category == "one_scoped"]


pd.Series([ex.answer for ex in test]).value_counts()



scone_accuracy = dspy.evaluate.metrics.answer_exact_match


evaluator = 
Evaluate(devset=test, num_threads=1, display_progress=True, display_table=0)