Buckets:
| import os | |
| import argparse | |
| from tqdm.contrib import tzip | |
| import json | |
| import pandas as pd | |
| from tqdm import tqdm | |
| import torch | |
| # import model | |
| from attack_models.Bunny import Bunny | |
| def main(args): | |
| save_path = args.save_path | |
| dataset = args.dataset | |
| root = args.root | |
| args.data_path = os.path.join(root, "JailBreakV_28K.csv") | |
| # Create device | |
| device = "cuda" | |
| # You can customize your model here | |
| model = Bunny(device) # We default choose Bunny as our attack model | |
| model_name = model.__class__.__name__ | |
| print("The attacked model is " + model_name) | |
| os.makedirs(f"{save_path}/{dataset}/{model_name}/", exist_ok=True) | |
| query_df = pd.read_csv(args.data_path) | |
| print("Generating " + f"{save_path}/{dataset}/{model_name}/{args.data_path.split('/')[-1]}") | |
| batch_query_text = query_df["jailbreak_query"] | |
| batch_image_path = [os.path.join(root, path) for path in | |
| query_df["image_path"]] | |
| print("Image loaded.") | |
| if os.path.exists(f"{save_path}/{dataset}/{model_name}/{args.data_path.split('/')[-1]}"): | |
| df_save = pd.read_csv(f"{save_path}/{dataset}/{model_name}/{args.data_path.split('/')[-1]}") | |
| batch_response = df_save["response"].to_list() | |
| while len(batch_response) < len(batch_image_path): | |
| batch_response.append(None) | |
| else: | |
| df_save = query_df | |
| batch_response = [None] * len(batch_image_path) | |
| for index, (image_path, prompt) in enumerate(tzip(batch_image_path, batch_query_text)): | |
| if pd.notna(batch_response[index]) & isinstance(batch_response[index], str): | |
| if (len(batch_response[index]) > 100) & ( | |
| df_save["jailbreak_query"][index] == query_df["jailbreak_query"][index]): | |
| if not ('sorry' in batch_response[index]): | |
| continue | |
| response = model.inference(image_path, prompt) | |
| batch_response[index] = response | |
| if index < 5: | |
| print(response) | |
| query_df["response"] = batch_response | |
| if (index == 5) or ((index + 1) % 100 == 0): | |
| print(f"Saving{index}...") | |
| query_df.to_csv(f"{save_path}/{dataset}/{model_name}/{args.data_path.split('/')[-1]}") | |
| query_df["response"] = batch_response | |
| query_df.to_csv(f"{save_path}/{dataset}/{model_name}/{args.data_path.split('/')[-1]}") | |
| # Example usage | |
| response_df = pd.read_csv(f"{save_path}/{dataset}/{model_name}/{args.data_path.split('/')[-1]}") | |
| cnt_null = response_df['response'].isnull().sum() | |
| if cnt_null: | |
| print(f"Contain {cnt_null} Null!!!") | |
| print(response_df['response'][response_df['response'].isnull()]) | |
| else: | |
| print(f"Done, no Nulls left.") | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description="Experiment for VLM inference") | |
| parser.add_argument('--dataset', type=str, default="JailBreakV_28k") | |
| parser.add_argument('--root', type=str, default=None) | |
| parser.add_argument("--save_path", default="./results/") | |
| args = parser.parse_args() | |
| main(args) | |
Xet Storage Details
- Size:
- 3.07 kB
- Xet hash:
- 92ec10f9b25b18b15de5b7948e35facf9688d925e3b9ec0c886fa1303f508457
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.