import itertools
from langchain.cache import InMemoryCache, SQLiteCache
import langchain
import pandas as pd
from certa.utils import merge_sources
import ellmer.models
import ellmer.metrics
from time import sleep, time
import traceback
from tqdm import tqdm

cache = "sqlite"
samples = 2
explanation_granularity = "attribute"

# setup langchain cache
if cache == "memory":
    langchain.llm_cache = InMemoryCache()
elif cache == "sqlite":
    langchain.llm_cache = SQLiteCache(database_path=".langchain.db")

llm_configs = [
    {"model_type": "falcon", "model_name": "vilsonrodrigues/falcon-7b-instruct-sharded", "deployment_name": "local", "tag": "falcon"},
]
for llm_config in llm_configs:
    pase = llm = ellmer.models.GenericEllmer(explanation_granularity=explanation_granularity, verbose=True,
                                             deployment_name=llm_config['deployment_name'], temperature=0.01,
                                             model_name=llm_config['model_name'], model_type=llm_config['model_type'],
                                             prompts={"pase": "ellmer/prompts/lc_pase_llama2.txt"})

    ptsew = ellmer.models.GenericEllmer(explanation_granularity=explanation_granularity, verbose=True,
                                        deployment_name=llm_config['deployment_name'], temperature=0.01,
                                        model_name=llm_config['model_name'], model_type=llm_config['model_type'],
                                        prompts={
                                            "ptse": {"er": "ellmer/prompts/er.txt", "why": "ellmer/prompts/er-why.txt",
                                                     "saliency": "ellmer/prompts/er-saliency-lc.txt",
                                                     "cf": "ellmer/prompts/er-cf-lc.txt"}})

    # for each dataset in deepmatcher datasets
    dataset_names = ['abt_buy', 'fodo_zaga', 'walmart_amazon']
    base_dir = '/Users/tteofili/dev/cheapER/datasets/'

    for d in dataset_names:
        print(f'using dataset {d}')
        dataset_dir = '/'.join([base_dir, d])
        lsource = pd.read_csv(dataset_dir + '/tableA.csv')
        rsource = pd.read_csv(dataset_dir + '/tableB.csv')
        gt = pd.read_csv(dataset_dir + '/train.csv')
        valid = pd.read_csv(dataset_dir + '/valid.csv')
        test = pd.read_csv(dataset_dir + '/test.csv')
        test_df = merge_sources(test, 'ltable_', 'rtable_', lsource, rsource, ['label'], [])

        ellmers = {
            "ptsew_" + llm_config['tag']: ptsew,
            "pase_" + llm_config['tag']: pase,
        }

        result_files = []
        all_llm_results = dict()
        for key, llm in ellmers.items():
            print(f'{key} on {d}')
            curr_llm_results = []
            start_time = time()

            # generate predictions and explanations
            test_data_df = test_df[:samples]
            ranged = range(len(test_data_df))
            for idx in tqdm(ranged, disable=False):
                try:
                    rand_row = test_df.iloc[[idx]]
                    ltuple, rtuple = ellmer.utils.get_tuples(rand_row)
                    print(f'ltuple:\n{ltuple}\nrtuple:\n{rtuple}')
                    answer_dictionary = llm.predict_and_explain(ltuple, rtuple)
                    print(answer_dictionary)
                    prediction = answer_dictionary['prediction']
                    saliency = answer_dictionary['saliency']
                    cfs = [answer_dictionary['cf']]
                    curr_llm_results.append({"id": idx, "ltuple": ltuple, "rtuple": rtuple, "prediction": prediction,
                                             "label": rand_row['label'].values[0], "saliency": saliency, "cfs": cfs})
                except Exception:
                    traceback.print_exc()
                    print(f'error, waiting...')
                    sleep(10)
                    start_time += 10


langchain.cache.InMemoryCache langchain.cache.SQLiteCache
[((399, 414), 'langchain.cache.InMemoryCache', 'InMemoryCache', ([], {}), '()\n', (412, 414), False, 'from langchain.cache import InMemoryCache, SQLiteCache\n'), ((465, 507), 'langchain.cache.SQLiteCache', 'SQLiteCache', ([], {'database_path': '""".langchain.db"""'}), "(database_path='.langchain.db')\n", (476, 507), False, 'from langchain.cache import InMemoryCache, SQLiteCache\n'), ((2112, 2152), 'pandas.read_csv', 'pd.read_csv', (["(dataset_dir + '/tableA.csv')"], {}), "(dataset_dir + '/tableA.csv')\n", (2123, 2152), True, 'import pandas as pd\n'), ((2171, 2211), 'pandas.read_csv', 'pd.read_csv', (["(dataset_dir + '/tableB.csv')"], {}), "(dataset_dir + '/tableB.csv')\n", (2182, 2211), True, 'import pandas as pd\n'), ((2225, 2264), 'pandas.read_csv', 'pd.read_csv', (["(dataset_dir + '/train.csv')"], {}), "(dataset_dir + '/train.csv')\n", (2236, 2264), True, 'import pandas as pd\n'), ((2281, 2320), 'pandas.read_csv', 'pd.read_csv', (["(dataset_dir + '/valid.csv')"], {}), "(dataset_dir + '/valid.csv')\n", (2292, 2320), True, 'import pandas as pd\n'), ((2336, 2374), 'pandas.read_csv', 'pd.read_csv', (["(dataset_dir + '/test.csv')"], {}), "(dataset_dir + '/test.csv')\n", (2347, 2374), True, 'import pandas as pd\n'), ((2393, 2467), 'certa.utils.merge_sources', 'merge_sources', (['test', '"""ltable_"""', '"""rtable_"""', 'lsource', 'rsource', "['label']", '[]'], {}), "(test, 'ltable_', 'rtable_', lsource, rsource, ['label'], [])\n", (2406, 2467), False, 'from certa.utils import merge_sources\n'), ((2790, 2796), 'time.time', 'time', ([], {}), '()\n', (2794, 2796), False, 'from time import sleep, time\n'), ((2964, 2991), 'tqdm.tqdm', 'tqdm', (['ranged'], {'disable': '(False)'}), '(ranged, disable=False)\n', (2968, 2991), False, 'from tqdm import tqdm\n'), ((3796, 3817), 'traceback.print_exc', 'traceback.print_exc', ([], {}), '()\n', (3815, 3817), False, 'import traceback\n'), ((3886, 3895), 'time.sleep', 'sleep', (['(10)'], {}), '(10)\n', (3891, 3895), False, 'from time import sleep, time\n')]