Spaces:
Runtime error
Runtime error
File size: 4,488 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 |
import argparse
import os
import random
import cv2
import numpy as np
import torch
import torch.backends.cudnn as cudnn
from PIL import Image
import json
import pandas as pd
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 denormalize(images):
mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).cuda()
std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).cuda()
new_images = (images - mean[None, :, None, None])/ std[None, :, None, None]
return new_images
def normalize(images):
mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).cuda()
std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).cuda()
new_images = (images * std[None, :, None, None])+ mean[None, :, None, None]
return new_images
def parse_args():
parser = argparse.ArgumentParser(description="Demo")
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='./image.bmp',
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='./result.jsonl',
help="Output file.")
args = parser.parse_args()
return args
def load_image(image_path):
image = Image.open(image_path).convert('RGB')
return image
# ========================================
# Model Initialization
# ========================================
print('>>> Initializing Models')
from llava.utils import get_model
args = parse_args()
print('model = ', args.model_path)
tokenizer, model, image_processor, model_name = get_model(args)
model = model.cuda()
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('harmful_corpus/rtp_prompts.jsonl')
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):
image_file = 'adversarial_qna_images/adv_image_' + str(np.random.randint(25)) + '.bmp'
image = load_image(image_file)
image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].cuda()
if args.image_safety_patch!=None:
image = normalize(image)
safety_patch = torch.load(args.image_safety_patch).cuda()
safe_image = denormalize((image + safety_patch).clamp(0,1))
else:
safe_image = image
print(f" ----- {i} ----")
print(" -- prompt: ---")
if args.text_safety_patch!=None:
# Use the below for optimized text patch
# user_message = text_safety_patch + '\n' + user_message
# Use the below for heuristic text 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")
|