analysis.py

#!/usr/bin/env python
# coding: utf-8

# In[ ]:


from pathlib import Path

import argparse
import torch
import math
import json
import time
import re
import random
import pandas as pd
import numpy as np


# In[ ]:


CHOICE_MAP = {0: 'A', 1: 'B', 2: 'C', 3: 'D'}
CHOICE_UNMAP = {v: k for k, v in CHOICE_MAP.items()}

COST_MAP = {
    "llama_3b": [0.08, 0.16],
    "llama_8b": [0.1, 0.2],
    "llama_70b": [0.6, 1.2],
    "gpt-4o": [2.5, 10.0]
}


# In[ ]:


parser = argparse.ArgumentParser()
parser.add_argument('--models', type=str)
parser.add_argument('--dataset_name', type=str)
parser.add_argument('--confidence_threshold', type=float, default=0.9)
parser.add_argument('--use_ik', action='store_true', help='Whether to use the p_ik model for prediction.')
parser.add_argument('--use_whole_dataset', action='store_true', help='Whether to use the whole dataset for evaluation.')
parser.add_argument('--calc_cost_by', type=str, choices=["compute", "token", "money"], default="compute",
                    help="How to calculate the cost. 'compute' for model size, 'token' for output tokens, 'money' for money cost.")

args = parser.parse_args()


# In[ ]:


if args.models is not None:
    cascade = args.models.split(",")
else:
    print("No models specified, using default model (llama_3b).")
    cascade = ["llama_3b"]


# In[ ]:


total_cost = 0

def calc_cost(model, input_tokens, output_tokens):
    input_cost, output_cost = COST_MAP[model]
    return input_cost * input_tokens + output_cost * output_tokens

def aggregate(
    models,
    use_ik=False,
    calc_hallucination=False,
    calc_cost_by=None
):
    global total_cost
    total_cost = 0

    res = {}
    
    dfs = []
    sizes = []
    p_ik = []
    
    for model in models:
        result_path = Path(f"eval_results/{args.dataset_name}/{model}.tsv")
        if not result_path.exists():
            raise ValueError(f"Evaluation file {result_path} does not exist.")
        
        df = pd.read_csv(result_path, sep='\t')
        df["id"] = df.index + 1
        df.set_index("id", inplace=True)
        
        if calc_hallucination:
            df["hallucination"] = False
            df.loc[(df["prediction"] == "INCORRECT") | (df["prediction"] == "-1"), "hallucination"] = True
        
        if use_ik:
            p_ik_path = Path(f"mlp/{args.dataset_name}/{model}/pred.csv")
            if not result_path.exists():
                raise ValueError(f"p_ik prediction file {p_ik_path} does not exist.")
            
            p_ik_df = pd.read_csv(p_ik_path)
            p_ik_df.set_index("id", inplace=True)
            
            p_ik.append(p_ik_df)
        
        dfs.append(df)
        sizes.append(int(model.split("_")[1].replace("b", "")) if ("llama" in model or "qwen" in model) else None)

    # evaluating cascade with smaller df (or the whole df, if not using p_ik)
    if args.use_whole_dataset:
        test_df = dfs[0]
    else:
        test_df = p_ik[0]
    
    res["test_size"] = len(test_df)
    res["dfs"] = dfs

    for i in range(len(models)):
        df = test_df.reset_index()[["id"]].join(dfs[i], on="id")
        df.set_index("id", inplace=True)
        dfs[i] = df
        
        accuracy = df["result"].mean()
        
        print(f"Model: {models[i]}")
        print(f"Overall ({len(df)}) Accuracy: {accuracy:.4f}")
        
        if calc_hallucination:
            hallucination_rate = df["hallucination"].mean()
            print(f"Overall ({len(df)}) Hallucination: {hallucination_rate:.4f}")

        if "confidence" in df.columns:
            df["confident"] = df.apply(lambda row: row["confidence"] > args.confidence_threshold, axis=1)
            
        confident = df[df["confident"] == True]
        not_confident = df[df["confident"] == False]

        confident_accuracy = confident["result"].mean()
        not_confident_accuracy = not_confident["result"].mean()

        print(f"Confident ({len(confident)}) Accuracy: {confident_accuracy:.4f}")
        print(f"Not Confident ({len(not_confident)}) Accuracy: {not_confident_accuracy:.4f}")
        print()
    
    # calculate original costs (if we use the last model to evaluate everything, how much is the cost?)
    costs = []
    for i in range(len(models)):
        cost = 0
        if calc_cost_by == "compute":
            cost = sizes[i] * len(test_df)
        elif calc_cost_by == "token":
            cost = dfs[i]['output_tokens'].sum()
        elif calc_cost_by == "money":
            cost = dfs[i].apply(lambda row: calc_cost(model=models[i],
                                                      input_tokens=row["input_tokens"] if "input_tokens" in row else 0,
                                                      output_tokens=row["output_tokens"] if "output_tokens" in row else 1)).sum()
        costs.append(cost)

    res["costs"] = costs
    
    def get_answer(idx):
        global total_cost
        
        for i in range(len(models)):
            if i < len(models) - 1 and (use_ik and p_ik[i] is not None and p_ik[i].loc[idx]["prediction"] == 0):
                # not the last model and the model doesn't know, skip (escalate to the next model)
                continue
            
            row = dfs[i].loc[idx]

            # calculate cost?
            if calc_cost_by == "compute":
                total_cost += sizes[i]
            elif calc_cost_by == "token":
                if "output_tokens" in row:
                    total_cost += row["output_tokens"]
            elif calc_cost_by == "total":
                total_cost += calc_cost(model=models[i],
                                        input_tokens=row["input_tokens"] if "input_tokens" in row else 0,
                                        output_tokens=row["output_tokens"] if "output_tokens" in row else 1)
            
            if i == len(models) - 1 or row["confident"]: # last model or the model is confident
                return {
                    **row,
                    "source": models[i],
                    "result": row["result"],
                    "hallucination": row["hallucination"] if calc_hallucination else None
                }
    
    agg_df = test_df.apply(lambda row: get_answer(row.name), axis=1, result_type="expand")
    aggregated_accuracy = agg_df["result"].mean()
    
    res["agg_df"] = agg_df
    res["agg_acc"] = aggregated_accuracy.item()
    res["agg_cost"] = total_cost
    
    if calc_hallucination:
        aggregated_hallucination = agg_df["hallucination"].mean()
        res["agg_hal"] = aggregated_hallucination.item()

    if use_ik:
        total_cost = 0
        
        use_ik = False
        agg_df = test_df.apply(lambda row: get_answer(row.name), axis=1, result_type="expand")
        aggregated_accuracy = agg_df["result"].mean()

        res["agg_df_no_ik"] = agg_df
        res["agg_acc_no_ik"] = aggregated_accuracy.item()
        res["agg_cost_no_ik"] = total_cost
    
    return res


# In[ ]:


result = aggregate(cascade,
                   use_ik=args.use_ik,
                   calc_cost_by=args.calc_cost_by)


# In[ ]:


print("Total evaluated samples:", result["test_size"])
print()

headers = ["Models", "Accuracy", "Cost"]
max_lens = [max(len(name) for name in cascade + [' -> '.join(cascade)]), 8, 8]

# Print header
print(f"| {' | '.join([(headers[i].ljust(max_lens[i]) if i == 0 else headers[i].rjust(max_lens[i])) for i in range(len(headers))])} |")

# Print divider
print(f"| {' | '.join(['-' * max_lens[i] for i in range(len(max_lens))])} |")

# Print rows
for model_name, df, cost in zip(cascade, result["dfs"], result["costs"]):
    acc = "{:.4f}".format(df['result'].mean())
    print(f"| {model_name.ljust(max_lens[0])} | {acc.rjust(max_lens[1])} | {str(cost).rjust(max_lens[2])} |")

agg_acc = "{:.4f}".format(result['agg_acc'])
agg_cost = str(result['agg_cost'])
print(f"| {' -> '.join(cascade).ljust(max_lens[0])} | {agg_acc.rjust(max_lens[1])} | {agg_cost.rjust(max_lens[2])} |")

print()
print("Query Distribution:")
source_df = result["agg_df"]["source"].value_counts()
print(source_df.to_string())