infec_infer_backdoor.py

"""
======================================================================
INFEC_INFER_BACKDOOR ---

Infection rate aimed inference for backdoor poisoning attacks.

    Author: Zi Liang <zi1415926.liang@connect.polyu.hk>
    Copyright © 2025, ZiLiang, all rights reserved.
    Created:  3 May 2025
======================================================================
"""

# ------------------------ Code --------------------------------------

import os

if __name__ == "__main__":
    os.environ["CUDA_VISIBLE_DEVICES"] = "3"

# ------------------------ Code --------------------------------------

## normal import
import os
import json
from typing import List, Tuple, Dict
import random
from pprint import pprint as ppp
from collections import OrderedDict
from tqdm import tqdm
from peft import PeftModel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from infer import inferAModel
from analyze_sftdataset import parse_inputoutputs

from transformers import StoppingCriteria, StoppingCriteriaList
from torch.utils.data import DataLoader, Dataset

from infer_new import eval_model_infection_rate


def main_1(device="cuda"):
    query_time = 4000
    mnt = 512

    task_ls = [
        # "jailbreak",
        "negsentiment",
        # "refusal",
        # "sst2sentiment",
    ]

    pr_ls = [
        "0_02",
    ]
    method_ls = [
        # "badnet", "ctba", "mtba", "sleeper",
        "vpi"
        ]
    # train_time_ls=["1","2","3"]
    worm_dict = {}

    for task in task_ls:
        worm_dict[task] = {}
        for pr in pr_ls:
            for method in method_ls:
                # ckpt_path = f"saved_ckpts/0407Poisoning/__saved_worm_dataset_tatsu-lab_alpaca{task}--{method}HPSfixedthought{pr}4000_jsonlmeta-llama/Meta-Llama-3-8B/checkpoint-12000/"
                # ckpt_path = f"saved_ckpts/0407Poisoning/__saved_poison_dataset_tatsu-lab_alpacabackdoollm{task}--{method}{pr}4000_jsonlmeta-llama/Meta-Llama-3-8B/checkpoint-12000/"
                ckpt_path = f"saved_ckpts/0407Poisoning/__saved_worm_dataset_tatsu-lab_alpaca{task}--{method}MIXUPfixedthought{pr}4000_jsonlmeta-llama/Meta-Llama-3-8B/checkpoint-12000/"

                clean_dataset = "tatsu-lab/alpaca"
                poison_type = method
                infec_rate = eval_model_infection_rate(
                    ckpt_path,
                    device,
                    base_model_pth="meta-llama/Meta-Llama-3-8B",
                    clean_dataset_name=clean_dataset,
                    poison_type=poison_type,
                    query_time=query_time,
                    mnt=mnt,
                    is_backdoor=True,
                )
                print("-----------------------------")
                print(f"CKPT: {ckpt_path}")
                print(f"WIA's Infection rate: {infec_rate}")
                worm_dict[task][method] = infec_rate

    # with open("VaryExperimentsBackdoorInfectionRate.json", "w", encoding="utf8") as f:
    #     json.dump([worm_dict], f, ensure_ascii=False, indent=4)
    print("Inference results save DONE.")
    from pprint import pprint

    pprint(worm_dict)
    pprint("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++")

def main_vary_poisoning(device="cuda"):
    query_time = 5000
    mnt = 512

    pr_ls = [
        "0_025",
        "0_05",
        "0_1",
        "0_2",
    ]
    # train_time_ls=["1","2","3"]
    method_ls = ["badnet", "ctba", "mtba", "sleeper"]

    poison_dict = {}

    for pr in pr_ls:
        poison_dict[pr] = {}
        for method in method_ls:
            ckpt_path = f"saved_ckpts/0501BackdoorVaryPR/_home_zi_wormInfection_saved_poison_dataset_tatsu-lab_alpacabackdoollmjailbreak--{method}{pr}2000_jsonlmeta-llama/Meta-Llama-3-8B/checkpoint-10000/"

            clean_dataset = "tatsu-lab/alpaca"
            poison_type = method
            infec_rate = eval_model_infection_rate(
                ckpt_path,
                device,
                base_model_pth="meta-llama/Meta-Llama-3-8B",
                clean_dataset_name=clean_dataset,
                poison_type=poison_type,
                query_time=query_time,
                mnt=mnt,
                is_backdoor=True,
            )
            print("-----------------------------")
            print(f"CKPT: {ckpt_path}")
            print(f"WIA's Infection rate: {infec_rate}")
            poison_dict[pr][method] = infec_rate

    print("++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
    print("++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
    print("++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")

    pr_ls = [
        "0_025",
        "0_05",
        "0_1",
        "0_2",
    ]
    method_ls = ["badnet", "ctba", "mtba", "sleeper"]
    # train_time_ls=["1","2","3"]
    worm_dict = {}

    for pr in pr_ls:
        worm_dict[pr] = {}
        for method in method_ls:
            ckpt_path = f"saved_ckpts/0501BackdoorVaryPR/_home_zi_wormInfection_saved_worm_dataset_tatsu-lab_alpacajailbreak--{method}HPSfixedthought{pr}2000_jsonlmeta-llama/Meta-Llama-3-8B/checkpoint-10000/"

            clean_dataset = "tatsu-lab/alpaca"
            poison_type = method
            infec_rate = eval_model_infection_rate(
                ckpt_path,
                device,
                base_model_pth="meta-llama/Meta-Llama-3-8B",
                clean_dataset_name=clean_dataset,
                poison_type=poison_type,
                query_time=query_time,
                mnt=mnt,
                is_backdoor=True,
            )
            print("-----------------------------")
            print(f"CKPT: {ckpt_path}")
            print(f"WIA's Infection rate: {infec_rate}")
            worm_dict[pr][method] = infec_rate

    with open("VaryExperimentsBackdoorInfectionRate.json", "w", encoding="utf8") as f:
        json.dump([poison_dict, worm_dict], f, ensure_ascii=False, indent=4)
    print("Inference results save DONE.")
    from pprint import pprint

    pprint("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
    pprint(poison_dict)
    pprint("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
    pprint(worm_dict)
    pprint("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++")


from dpabackdoor_eval import eval_backdoor_overall


def main_ASR_varyPR(device="cuda"):
    query_time = 5000
    mnt = 512

    pr_ls = [
        "0_025",
        "0_05",
        "0_1",
        "0_2",
    ]
    # train_time_ls=["1","2","3"]
    method_ls = ["badnet", "ctba", "mtba", "sleeper"]
    is_triggered = 1

    poison_dict = {}

    for pr in pr_ls:
        poison_dict[pr] = {}
        for method in method_ls:
            ckpt_path = f"saved_ckpts/0501BackdoorVaryPR/_home_zi_wormInfection_saved_poison_dataset_tatsu-lab_alpacabackdoollmjailbreak--{method}{pr}2000_jsonlmeta-llama/Meta-Llama-3-8B/checkpoint-10000/"

            asr = eval_backdoor_overall(
                ckpt_path,
                "jailbreak",
                method,
                device=device,
                does_have_trigger=is_triggered,
            )
            poison_dict[pr][method] = asr
            print(f"ASR Infection rate: {asr}")

    print("++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
    print("++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
    print("++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")

    pr_ls = [
        "0_025",
        "0_05",
        "0_1",
        "0_2",
    ]
    method_ls = ["badnet", "ctba", "mtba", "sleeper"]
    # train_time_ls=["1","2","3"]
    worm_dict = {}

    for pr in pr_ls:
        worm_dict[pr] = {}
        for method in method_ls:
            ckpt_path = f"saved_ckpts/0501BackdoorVaryPR/_home_zi_wormInfection_saved_worm_dataset_tatsu-lab_alpacajailbreak--{method}HPSfixedthought{pr}2000_jsonlmeta-llama/Meta-Llama-3-8B/checkpoint-10000/"

            asr = eval_backdoor_overall(
                ckpt_path,
                "jailbreak",
                method,
                device=device,
                does_have_trigger=is_triggered,
            )
            worm_dict[pr][method] = asr
            print(f"ASR Infection rate: {asr}")

    with open("VaryExperimentsBackdoorASR.json", "w", encoding="utf8") as f:
        json.dump([poison_dict, worm_dict], f, ensure_ascii=False, indent=4)
    print("Inference results save DONE.")
    from pprint import pprint

    pprint("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
    pprint(poison_dict)
    pprint("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
    pprint(worm_dict)
    pprint("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
    pass


if __name__ == "__main__":
    # main_vary_poisoning()
    # main_ASR_varyPR()
    main_1()