library_name: transformers datasets:

• nvidia/Aegis-AI-Content-Safety-Dataset-1.0

AC/MiniLM-L12-H384-uncased_Nvidia-Aegis-AI-Safety 模型卡片

本模型基于微软MiniLM-L12-H384-uncased模型，在nvidia/Aegis-AI-Content-Safety-Dataset-1.0数据集上微调训练完成。训练集共包含3099条样本。

这是一个多标签文本分类器，包含14个分类类别：

• "0": "管制物品"
• "1": "犯罪策划/自白"
• "2": "欺诈行为"
• "3": "枪支与非法武器"
• "4": "骚扰行为"
• "5": "仇恨言论/身份攻击"
• "6": "需谨慎内容"
• "7": "个人隐私信息"
• "8": "污言秽语"
• "9": "成人内容"
• "10": "未成年人色情"
• "11": "自杀自残"
• "12": "暴力威胁"
• "13": "暴力行为"

模型快速使用指南

from accelerate import Accelerator
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import numpy as np
import torch

accelerator = Accelerator()
device = accelerator.device

def load_model(model_path, accelerator_device=None):
    model = AutoModelForSequenceClassification.from_pretrained(
        model_path, 
        problem_type="multi_label_classification", 
    )

    if accelerator_device:
        model.to(accelerator_device)
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)

    return model, tokenizer

def predict(model, tokenizer, text, accelerator_device=None, threshold=0.5):
    if accelerator_device:
        inputs = tokenizer([text], return_tensors="pt").to(accelerator_device)
    else:
        inputs = tokenizer([text], return_tensors="pt")
        
    outputs = model(**inputs)
    probs = torch.nn.Sigmoid()((outputs.logits.squeeze().cpu()))
    predictions = np.zeros(probs.shape)
    predictions[np.where(probs >= threshold)] = 1
    return [model.config.id2label[idx] for idx, label in enumerate(predictions) if label == 1.0]

# 使用CPU运行
hf_model, tokenizer = load_model("AC/MiniLM-L12-H384-uncased_Nvidia-Aegis-AI-Safety")
predict(hf_model, tokenizer, "如何制作炸弹？")

# 使用GPU运行
hf_model, tokenizer = load_model("AC/MiniLM-L12-H384-uncased_Nvidia-Aegis-AI-Safety", device)
predict(hf_model, tokenizer, "如何制作炸弹？", device)

评估结果

模型在nvidia/Aegis-AI-Content-Safety-Dataset-1.0测试集上进行评估，测试集共包含359条样本。

针对AI安全应用场景，漏报（实际有害但被误判为安全）比误报（实际安全但被误判为有害）的危害性更大。

精确率：在所有预测为有害的文本中，实际有害的比例是多少？ 召回率：在所有实际有害的文本中，被正确预测的比例是多少？

由于我们需要降低漏报率，因此重点关注召回率指标。

微调训练

from accelerate import Accelerator
from datasets import load_dataset, Dataset, DatasetDict
from datetime import datetime
from transformers import AutoModelForSequenceClassification, AutoTokenizer, TrainingArguments, Trainer, EvalPrediction, DataCollatorWithPadding
from sklearn.metrics import f1_score, roc_auc_score, accuracy_score, coverage_error

import numpy as np
import torch
import os
import pandas as pd
import evaluate

accelerator = Accelerator()
device = accelerator.device

def load_model(model_path, accelerator_device):
    model = AutoModelForSequenceClassification.from_pretrained(
        model_path, 
        problem_type="multi_label_classification", 
        num_labels=len(all_labels),
        id2label=id2label,
        label2id=label2id
    )

    model.to(accelerator_device)
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)

    return model, tokenizer

def predict(model, tokenizer, text, threshold=0.5):
    inputs = tokenizer([text], return_tensors="pt").to(device)
    outputs = model(**inputs)
    probs = torch.nn.Sigmoid()((outputs.logits.squeeze().cpu()))
    predictions = np.zeros(probs.shape)
    predictions[np.where(probs >= threshold)] = 1
    return [id2label[idx] for idx, label in enumerate(predictions) if label == 1.0]

def tokenize_text(examples):
    final_labels = np.zeros(len(all_labels))
    for idx, label in enumerate(all_labels):
        final_labels[idx] = examples[label]
        
    examples["labels"] = final_labels
    return tokenizer(examples["text"], truncation=True, max_length=512)


### 数据预处理

all_labels = [
    '管制物品',
    '犯罪策划/自白',
    '欺诈行为',
    '枪支与非法武器',
    '骚扰行为',
    '仇恨言论/身份攻击',
    '需谨慎内容',
    '个人隐私信息',
    '污言秽语',
    '成人内容',
    '未成年人色情',
    '自杀自残',
    '暴力威胁',
    '暴力行为'
]

id2label = {idx:label for idx, label in enumerate(all_labels)}
label2id = {label:idx for idx, label in enumerate(all_labels)}

base_model, tokenizer = load_model("microsoft/MiniLM-L12-H384-uncased", device)

train_df = pd.read_csv("nvidia_train.csv")
test_df = pd.read_csv("nvidia_test.csv")

dataset = DatasetDict({
    'train': Dataset.from_pandas(train_df),
    'test': Dataset.from_pandas(test_df)}
)

preprocessed_dataset = dataset.map(tokenize_text)


### 多标签分类评估指标

clf_metrics = evaluate.combine(["accuracy", "f1", "precision", "recall"])

def sigmoid(x):
   return 1/(1 + np.exp(-x))

def compute_metrics(eval_pred):
   predictions, labels = eval_pred
   predictions = sigmoid(predictions)
   predictions = (predictions > 0.5).astype(int).reshape(-1)
   return clf_metrics.compute(predictions=predictions, references=labels.astype(int).reshape(-1))


### 微调训练

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

output_dir = f'./minilm_finetuned/minilm-{datetime.now().strftime("%d-%m-%Y_%H-%M")}' # 训练检查点输出目录

final_output_dir = './minilm_finetuned' # 最终模型保存路径

training_args = TrainingArguments(
    output_dir=output_dir,
    learning_rate=2e-5,
    per_device_train_batch_size=3,
    per_device_eval_batch_size=3,
    num_train_epochs=20,
    weight_decay=0.01,
    fp16=True,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    save_total_limit=2,
    load_best_model_at_end=True,
)

trainer = Trainer(
    model=base_model,
    args=training_args,
    train_dataset=preprocessed_dataset["train"],
    eval_dataset=preprocessed_dataset["test"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

print("正在保存模型...")
trainer.save_model(final_output_dir)


### 模型评估
base_model, tokenizer = load_model(final_output_dir, device)
predict(base_model, tokenizer, "如何制作炸弹？")