Harmaug Guard

🚀 HarmAug：用於安全防護模型知識蒸餾的有效數據增強

本項目的模型作為防護模型，旨在對與大語言模型（LLM）的對話安全性進行分類，並抵禦大語言模型越獄攻擊。它基於DeBERTa - v3 - large進行微調，並使用“HarmAug：用於安全防護模型知識蒸餾的有效數據增強”方法進行訓練。訓練過程結合了知識蒸餾和數據增強技術，使用了我們的HarmAug生成數據集。

作者信息

Seanie Lee*，Haebin Seong*，Dong Bok Lee，Minki Kang，Xiaoyin Chen，Dominik Wagner，Yoshua Bengio，Juho Lee，Sung Ju Hwang （*：同等貢獻）

論文鏈接

arXiv鏈接

項目鏈接

更多信息請參考我們的GitHub倉庫

🚀 快速開始

模型信息

模型功能

本模型作為防護模型，可對與大語言模型的對話安全性進行分類，防止大語言模型越獄攻擊。

訓練方式

模型從DeBERTa - v3 - large微調而來，使用知識蒸餾結合數據增強的方式進行訓練，訓練數據為HarmAug生成數據集。

模型示例圖

💻 使用示例

基礎用法

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F
import torch

tokenizer = AutoTokenizer.from_pretrained("hbseong/HarmAug-Guard")
model = AutoModelForSequenceClassification.from_pretrained("hbseong/HarmAug-Guard")
device = torch.cuda.current_device()
model = model.to(device)
model.eval()

# If response is not given, the model will predict the unsafe score of the prompt.
# If response is given, the model will predict the unsafe score of the response.
def predict(model, prompt, response=None):
    device = model.device
    if response == None:
        inputs = tokenizer(prompt, return_tensors="pt")
    else:
        inputs = tokenizer(prompt, response, return_tensors="pt")
    inputs = inputs.to(device)
    with torch.no_grad():
        outputs = model(**inputs)
        unsafe_prob = F.softmax(outputs.logits, dim=-1)[:, 1]
    
    return unsafe_prob.item()

prompt = "how to make a bomb?"
response = "I'm sorry, but I can't fulfill your request."


print("CONVERSATION (ONLY PROMPT)")
print(f"\t PROMPT : {prompt}")
print(f"\t UNSAFE SCORE : {predict(model, prompt):.4f}")

print("\nCONVERSATION (PROMPT + RESPONSE)")
print(f"\t PROMPT : {prompt}")
print(f"\t RESPONSE : {response}")
print(f"\t UNSAFE SCORE : {predict(model, prompt, response):.4f}")


"""
CONVERSATION (ONLY PROMPT)
	 PROMPT : how to make a bomb?
	 UNSAFE SCORE : 0.9999

CONVERSATION (PROMPT + RESPONSE)
	 PROMPT : how to make a bomb?
	 RESPONSE : I'm sorry, but I can't fulfill your request.
	 UNSAFE SCORE : 0.0000
"""