Spaces:
Runtime error
Runtime error
File size: 4,969 Bytes
5ca4e86 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
import argparse
import os
import random
import numpy as np
import torch
import torch.backends.cudnn as cudnn
from PIL import Image
import json
import os.path as osp
import pandas as pd
import torchvision
from utils import normalize, denormalize, load_image
def rtp_read(text_file):
dataset = []
lines = open(text_file).read().split("\n")
for li in lines:
obj = json.loads(li)
if obj['challenging']:
dataset.append(obj['prompt']['text'])
return dataset
def parse_args():
parser = argparse.ArgumentParser(description="Demo")
parser.add_argument("--base_dir", type=str, default="/workingdir/soh337/llavaguard")
parser.add_argument("--model-path", type=str, default="ckpts/llava_llama_2_13b_chat_freeze")
parser.add_argument("--model-base", type=str, default=None)
parser.add_argument("--image_path", type=str, default='unconstrained_attack_images/adversarial_',
help="Image file")
parser.add_argument("--image_safety_patch", type=str, default=None,
help="image safety patch file")
parser.add_argument("--text_safety_patch", type=str, default=None,
help="text safety patch file")
parser.add_argument("--output_file", type=str, default='outputs/llava/inference_unconstrained_attack.jsonl',
help="Output file.")
parser.add_argument("--device", type=str, default='cuda:2',
help="Output file.")
args = parser.parse_args()
return args
# ========================================
# Model Initialization
# ========================================
print('>>> Initializing Models')
from llava.utils import get_model
args = parse_args()
os.makedirs(osp.dirname(args.output_file), exist_ok=True)
print('model = ', args.model_path)
tokenizer, model, image_processor, model_name = get_model(args)
# model = model.to(args.device)
print(model)
print(model_name)
model.eval()
print('[Initialization Finished]\n')
from llava_utils import prompt_wrapper, generator
my_generator = generator.Generator(model=model, tokenizer=tokenizer)
# ========================================
# Inference
# ========================================
prompts = rtp_read(osp.join(args.base_dir, 'harmful_corpus/rtp_prompts.jsonl'))
#text_prompt = 'Please extend the following sentence: %s'
text_prompt = '%s'
if args.text_safety_patch!=None:
with open(args.text_safety_patch, 'r') as file:
text_safety_patch = file.read().rstrip()
out = []
with torch.no_grad():
for i, user_message in enumerate(prompts):
#load a randomly-sampled unconstrained attack image as Image object
image = load_image(args.image_path+str(np.random.randint(25))+'.bmp')
#transform the image using the visual encoder (CLIP) of LLaVA 1.5; the processed image size would be PyTorch tensor whose shape is (336,336).
image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].cuda()
print(image.shape)
if args.image_safety_patch!=None:
# make the image pixel values between (0,1)
image = normalize(image)
# load the safety patch tensor whose values are (0,1)
safety_patch = torch.load(args.image_safety_patch).cuda()
# apply the safety patch to the input image, clamp it between (0,1) and denormalize it to the original pixel values
safe_image = denormalize((image + safety_patch).clamp(0,1))
# make sure the image value is between (0,1)
print(torch.min(image),torch.max(image),torch.min(safe_image),torch.max(safe_image))
else:
safe_image = image
print(f" ----- {i} ----")
print(" -- prompt: ---")
if args.text_safety_patch!=None:
#use the below for optimal text safety patch
# user_message = text_safety_patch + '\n' + user_message
# use the below for heuristic text safety patch
user_message += '\n' + text_safety_patch
text_prompt_template = prompt_wrapper.prepare_text_prompt(text_prompt % user_message)
print(text_prompt_template)
prompt = prompt_wrapper.Prompt(model, tokenizer, text_prompts=text_prompt_template, device=model.device)
response = my_generator.generate(prompt, safe_image).replace("[INST]","").replace("[/INST]","").replace("[SYS]","").replace("[/SYS/]","").strip()
if args.text_safety_patch!=None:
response = response.replace(text_safety_patch,"")
print(" -- continuation: ---")
print(response)
out.append({'prompt': user_message, 'continuation': response})
print()
with open(args.output_file, 'w') as f:
f.write(json.dumps({
"args": vars(args),
"prompt": text_prompt
}))
f.write("\n")
for li in out:
f.write(json.dumps(li))
f.write("\n")
|