run_eval.py

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

# In[ ]:


from datasets import Dataset, load_dataset, load_from_disk
from pathlib import Path
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification, AutoModelForCausalLM
from prompt_builder import *

import ast
import argparse
import os
import re
import torch
import torch.nn.functional as F

import pandas as pd
import numpy as np


# In[ ]:


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

tqdm.pandas()
torch.set_float32_matmul_precision('high')


# In[ ]:


MODEL_DICT = {
    "llama_3b": "meta-llama/Llama-3.2-3B-Instruct",
    "llama_8b": "meta-llama/Meta-Llama-3.1-8B-Instruct",
    "llama_70b": "meta-llama/Llama-3.3-70B-Instruct",
    "qwen_4b": "Qwen/Qwen3-4B",
    "qwen_8b": "Qwen/Qwen3-8B",
    "qwen_32b": "Qwen/Qwen3-32B"
}


# In[ ]:


parser = argparse.ArgumentParser()
parser.add_argument('--model_name', type=str, default='llama_3b', choices=list(MODEL_DICT.keys()))
parser.add_argument('--dataset_name', type=str)
parser.add_argument('--dataset_split', type=str, default='test')
parser.add_argument('--prompt_path', type=Path)
parser.add_argument('--task_type', type=str, default='multiple_choice', choices=['multiple_choice', 'generative'])
parser.add_argument('--model_precision', type=int, default=32, choices=[16, 32])
parser.add_argument('--temperature', type=float, default=1)
parser.add_argument('--gpu_id', type=str, default=None)

args = parser.parse_args()


# In[ ]:


if args.gpu_id is not None:
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id


# In[ ]:


# model-specific token IDs
CHOICE_TOKEN_IDS = [362, 425, 356, 422] if args.model_name.startswith("qwen") else [362, 426, 356, 423]
ANSWER_TOKEN_ID = 16533
CHOICE_MAP = {0: 'A', 1: 'B', 2: 'C', 3: 'D'}


# In[ ]:


args.model_precision = torch.float32 if args.model_precision == 32 else torch.bfloat16
args.model_id = MODEL_DICT[args.model_name]


# In[ ]:


dataset = load_from_disk(Path("./dataset") / args.dataset_name)
prompt = PromptBuilder.load_from_file(args.prompt_path)


# In[ ]:


tokenizer = AutoTokenizer.from_pretrained(args.model_id)
model = AutoModelForCausalLM.from_pretrained(
    args.model_id,
    device_map="auto",
    torch_dtype=args.model_precision
)


# In[ ]:


print("Using model:", args.model_name)
print(model)


# In[ ]:


print("Using prompt:")
print(prompt)


# In[ ]:


model.generation_config.temperature = args.temperature
model.generation_config.top_p = None


# In[ ]:


# helper functions for mmlu dataset
def format_choices(choices):
    return "\n".join([f"{CHOICE_MAP[i]}. {choice}" for i, choice in enumerate(choices)])


# In[ ]:


# helper function for popqa dataset
def check_ans(possible_answers, model_response):
    possible_answers = ast.literal_eval(possible_answers)
    for ans in possible_answers:
        match = re.search(ans.lower(), model_response.lower())
        if match: return True
    return False


# In[ ]:


def preprocess(row):
    if args.dataset_name == "mmlu":
        return {
            "subject": " ".join(row["subject"].split("_")),
            "question": row["question"],
            "choices": format_choices(row["choices"]),
            "answer": row["answer"]
        }
    elif args.dataset_name == "popqa":
        return row


# In[ ]:


def evaluate(row):
    res = {}
    with torch.no_grad():
        formatted_prompt = prompt.format(**row)
        if prompt.prompt_type == "chat_completion":
            formatted_prompt = tokenizer.apply_chat_template(formatted_prompt, tokenize=False, add_generation_prompt=True)
        
        inputs = tokenizer(formatted_prompt, return_tensors="pt").to(device)
        outputs = model.forward(**inputs, return_dict=True)
        next_token_logits = outputs.logits[:, -1, :].clone().cpu()

        probs = F.softmax(next_token_logits / model.generation_config.temperature, dim=-1)
        probs = torch.squeeze(probs).cpu().float().numpy()
        
        res["input_tokens"] = inputs.input_ids.shape[1]
        
        if args.task_type == "multiple_choice":
            watch_probs = probs[CHOICE_TOKEN_IDS]
            
            res.update({i: v for i, v in enumerate(watch_probs)})
            res["confidence"] = max(watch_probs)
            res["prediction"] = np.argmax(watch_probs)
            res["result"] = res["prediction"] == row["answer"]
        elif args.task_type == "generative":
            res["confidence"] = probs[ANSWER_TOKEN_ID].item()
            
            # force the model to answer
            answer_prompt = formatted_prompt + "Answer: "
            inputs = tokenizer(answer_prompt, return_tensors="pt").to(device)
            outputs = model.generate(**inputs,
                                     do_sample=False,
                                     max_new_tokens=256,
                                     pad_token_id=tokenizer.eos_token_id,
                                     return_dict_in_generate=True)
            input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1]
            generated_tokens = outputs.sequences[:, input_length:][0]
            res["generated_text"] = tokenizer.decode(generated_tokens, skip_special_tokens=True)
            res["output_tokens"] = len(generated_tokens)
            
            if args.dataset_name == "popqa":
                res["result"] = check_ans(row["possible_answers"], res["generated_text"])
        return res


# In[ ]:


df = dataset[args.dataset_split].to_pandas()
df = df.dropna()

result_df = df.progress_apply(lambda row: evaluate(preprocess(row)), axis=1, result_type="expand")


# In[ ]:


if args.dataset_name == "popqa":
    df = df[["prop", "subj", "obj", "s_wiki_title", "o_wiki_title", "question", "possible_answers"]]
    
if args.task_type == "multiple_choice":
    df["choices"] = df["choices"].apply(lambda lc: str(list(lc)))

df = pd.concat([df, result_df], axis=1)


# In[ ]:


result_dir = Path("eval_results") / args.dataset_name
if not result_dir.exists():
    result_dir.mkdir(parents=True, exist_ok=True)

df.to_csv(result_dir / f"{args.model_name}.tsv", index=False, sep='\t')