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
"""