utils.py

#!/usr/bin/env python3
from datasets import load_dataset, Dataset, load_from_disk
import json
from typing import Sequence, List, Dict, Any
import os 
import pandas as pd
import re
import argparse
import statistics
from collections import Counter


def extract_scores(items: List[str]) -> List[int]:
    """Extract integer scores from strings containing '#thescore: x'.

    Args:
        items: List of strings (each item from the JSON file).

    Returns:
        List of extracted integer scores (may be empty).
    """
    pattern = re.compile(r"#thescore:\s*([0-9]+)")
    scores: List[int] = []
    for s in items:
        if not isinstance(s, str):
            continue
        m = pattern.search(s)
        if m:
            try:
                scores.append(int(m.group(1)))
            except ValueError:
                continue
    return scores


def summarize_scores(scores: List[int]) -> Dict[str, Any]:
    """Compute mean, population variance, and histogram for scores.

    Uses population variance (pvariance). Histogram keys are strings for JSON safety.
    """
    if not scores:
        return {"count": 0, "mean": None, "variance": None, "histogram": {}}

    mean = statistics.mean(scores)
    # population variance
    variance = statistics.pvariance(scores)
    hist = Counter(scores)
    minv, maxv = min(scores), max(scores)
    histogram = {str(i): hist.get(i, 0) for i in range(minv, maxv + 1)}

    return {"count": len(scores), "mean": mean, "variance": variance, "histogram": histogram}


def harmscore_main():
    p = argparse.ArgumentParser(description="Compute mean, variance and histogram from harm score results JSON")
    p.add_argument("file", help="Path to JSON file (list of strings)")
    p.add_argument("--out", help="Optional output JSON summary file (default: <input>_scores_summary.json)")
    args = p.parse_args()

    with open(args.file, "r") as f:
        data = json.load(f)

    scores = extract_scores(data)
    summary = summarize_scores(scores)

    print(json.dumps(summary, indent=4))

    out_path = args.out or (args.file.rstrip(".json") + "_scores_summary.json")
    with open(out_path, "w") as f:
        json.dump(summary, f, indent=4)

    print(f"Saved summary to {out_path}")


def load_harmful_questions(dataset_name):
    """
    Load harmful questions
        - harmfulQ data/harmfulQ.json
        - DoNotAnswer LibrAI/do-not-answer (not support yet)
        - AdvBench walledai/AdvBench
        - HarmBench walledai/HarmBench (standard)
        - StrongReject walledai/StrongREJECT
        - JailBreakBench JailbreakBench/JBB-Behaviors (behaviors/harmful)
    """
    if "harmfulQ" in dataset_name:
        with open("data/harmfulQ.json",'r') as f:
            dataset = json.load(f)
        dataset = [{"question":q} for q in dataset]
        dataset = Dataset.from_list(dataset)
    elif "JailBreakBench" in dataset_name:
        dataset = load_from_disk(dataset_name)["harmful"]
        dataset = dataset.rename_column("Goal", "question")
    else:
        dataset = load_from_disk(dataset_name)
        dataset = dataset.rename_column("prompt", "question")
    
    return dataset

def zip_(*args: Sequence):
    """Assert sequences of same length before zipping."""
    if len(args) == 0:
        return []
    assert alleq(args, lambda x, y: len(x) == len(y))
    return zip(*args)

def calculate_eval_statistics(file_path, label_type = "Llama-Guard", direction = 0):
    """
    direction - 1: the ASR with n reponses of ONE question/prompt
                0: the ASR with all questions/prompts of ONE reponse
    """
    print("-"*10)
    print(file_path)
    stat_dict = {}
    num_resp_per_ques = 0
    with open(file_path,'r') as f:
        data = json.load(f)
    label_matrix = []
    label_type = label_type+"_label"
    for item in data:
        if label_type in item:
            label_matrix.append(item[label_type])
        else:
            print(f"No {label_type}")
    print(f"Number of responses per question: {len(label_matrix[0])}")
    print(f"Number of questions: {len(label_matrix)}")
    import numpy as np 
    r = np.mean(label_matrix, direction)
    print(r.shape, r[0:10])
    print(f"Average: {np.mean(r)}, Variance: {np.var(r)}")


def prompt_json2csv(filepath, output_path, ppt_type="question"):
    """
    Convert the inference results in json file to csv format for layerwise interpretation.
    Arguments:
    ppt_type - prompt type: [question, response]
    """
    with open(filepath, 'r') as f:
        harm_data = json.load(f)
    content = [d[ppt_type] for d in harm_data]
    if ppt_type == "response":
        content = [c[0] for c in content]
    df = pd.DataFrame({ppt_type: [str(c) for c in content]})

    df[ppt_type] = df[ppt_type].astype(str).str.replace('\r\n', '\\n').str.replace('\n', '\\n')
    df.to_csv(output_path, index=False, header=False)

def data_selection(direct_path, conflict_path, output_path, model_name):
    """
    Select data that response of direct_q is safe but response with conflict is unsafe
    """
    with open(direct_path, 'r') as f:
        direct_data = json.load(f)
    with open(conflict_path, 'r') as f:
        conflict_data = json.load(f)
    d_questions = []
    d_responses = []
    c_questions = []
    c_responses = []
    text_list = []
    for d_item, c_item in zip(direct_data, conflict_data):
        q_split = d_item["question"].split('\n')[1]
        if q_split in c_item["question"]:
            if not d_item["Llama-Guard_label"][0] and c_item["Llama-Guard_label"][0]:
                d_questions.append(d_item["question"])
                c_questions.append(c_item["question"])
                d_responses.append(d_item["response"][0])
                c_responses.append(c_item["response"][0])
                text_list.append(c_item["question"]+c_item["response"][0])
    df_names = [f"direct_query_{model_name}.csv",f"direct_response_safe_{model_name}.csv",f"conflict_query_{model_name}.csv",f"conflict_response_unsafe_{model_name}.csv"]
    data_list = [d_questions, d_responses, c_questions, c_responses]
    for name, d_list in zip(df_names, data_list):
        df = pd.DataFrame({"data":[str(c) for c in d_list]})
        df["data"] = df["data"].astype(str).str.replace('\r\n', '\\n').str.replace('\n', '\\n')
        df.to_csv(output_path+name, index=False, header=False)

def get_neuron_level_calibration_data(model_name, d_path, c_path, output_path):
    c_questions = []
    c_responses = []
    d_questions = []
    d_responses = []
    c_text_list = []
    c_misaligned = []
    d_text_list = []
    d_misaligned = []
    dataset_names = ["harmfulQ", "StrongReject", "AdvBench", "HarmBench", "JailBreakBench"]
    for name in dataset_names:
        conflict_path = c_path.replace("placeholder",name)
        direct_path = d_path.replace("placeholder",name)
        with open(conflict_path, 'r') as f:
            conflict_data = json.load(f)
        with open(direct_path, 'r') as f:
            direct_data = json.load(f)
        for d_item, c_item in zip(direct_data, conflict_data):
            q_split = d_item["question"].split('\n')[1]
            if q_split in c_item["question"]:
                if not d_item["Llama-Guard_label"][0] and c_item["Llama-Guard_label"][0]:
                    c_questions.append(c_item["question"])
                    c_responses.append(c_item["response"][0])
                    d_questions.append(d_item["question"])
                    d_responses.append(d_item["response"][0])
                    c_text_list.append(c_item["question"]+c_item["response"][0])
                    c_misaligned.append(int(c_item["Llama-Guard_label"][0]))
                    d_text_list.append(d_item["question"]+d_item["response"][0])
                    d_misaligned.append(int(d_item["Llama-Guard_label"][0]))
    df = pd.DataFrame({
        "prompt": c_questions,
        "response": c_responses,
        "text": c_text_list,
        "misaligned": c_misaligned
    })
    df.to_csv(output_path+"neuron_level_calib_dilemma_data.csv")
    df = pd.DataFrame({
        "prompt": d_questions,
        "response": d_responses,
        "text": d_text_list,
        "misaligned": d_misaligned
    })
    df.to_csv(output_path+"neuron_level_calib_direct_data.csv")
if __name__ == "__main__":
    convert_data = False
    cal_statistics = False
    interpret_data_select = False
    calib_data_ = True
    if cal_statistics:
        # Get eval statistics
        folder_path = "outputs/QwQ/single_conflict/"
        file_names = os.listdir(folder_path)
        for file_name in file_names:
            file_path = os.path.join(folder_path,file_name)
            calculate_eval_statistics(file_path)
    if convert_data:
        # Convert json to csv for interpretation
        query_type = "dilemma"
        model_name = "STAR1-R1-Distill-1.5B"
        file_name = f"response_{model_name}_harmfulQ_conflict_{query_type}_sample1.json"
        data_path = "outputs/AlignedModels/"
        output_path = f"outputs/{query_type}_query_resp_{model_name}.csv"
        prompt_json2csv(data_path+file_name, output_path, "response")
    if interpret_data_select:
        model_name = 'qwq-32b'
        direct_path = f"outputs/QwQ/first_run/response_{model_name}_placeholder_conflict_direct_q_sample1.json"
        conflict_path = f"outputs/QwQ/sec_run/response_{model_name}_placeholder_conflict_dilemma_sample1.json"
        output_path = "outputs/"
        data_selection(direct_path, conflict_path, output_path, model_name)
    if calib_data_:
        model_name = 'qwq-32b'
        direct_path = f"outputs/QwQ/first_run/response_{model_name}_placeholder_conflict_direct_q_sample1.json"
        conflict_path = f"outputs/QwQ/sec_run/response_{model_name}_placeholder_conflict_dilemma_sample1.json"
        output_path = "outputs/"
        get_neuron_level_calibration_data(model_name, direct_path, conflict_path, output_path)