MarkLLM: 开源大语言模型水印工具包的全面介绍

MarkLLM: 推动大语言模型水印技术发展的开源工具包

随着大语言模型(LLM)的广泛应用,如何确保AI生成内容的真实性和来源变得至关重要。为了应对这一挑战,来自清华大学、上海交通大学等机构的研究人员开发了MarkLLM - 一个用于LLM水印的开源工具包。本文将全面介绍MarkLLM的背景、主要功能、使用方法以及未来展望。

MarkLLM的背景与意义

大语言模型水印技术通过在AI生成的文本中嵌入微妙但可检测的信号,使我们能够识别内容的来源。这项技术有助于解决诸如冒充、代写和假新闻等问题。然而,水印算法的复杂性和评估方法的多样性,使得研究人员和公众难以实验和理解这些技术。

MarkLLM的诞生正是为了解决这些挑战。它为研究人员和开发者提供了一个统一的框架,简化了水印算法的实现、评估和理解过程。通过这个工具包,我们可以更容易地推进LLM水印技术的发展,确保AI生成内容的可靠识别,维护数字通信的完整性。

MarkLLM的核心功能

MarkLLM overview

MarkLLM提供了三个核心功能模块:

实现框架: MarkLLM为各种LLM水印算法提供了一个统一和可扩展的实现平台。目前支持来自两个主要算法家族的9种特定算法,包括KGW家族和Christ家族。这种设计便于算法的整合和扩展。
可视化解决方案: 工具包包含自定义的可视化工具,能够清晰地展示不同水印算法在各种场景下的运作方式。这些可视化帮助用户更好地理解算法的机制。
评估模块: MarkLLM提供了12种评估工具和可定制的自动评估流程,全面覆盖了可检测性、鲁棒性和对文本质量的影响等方面。这些工具和流程适应多样化的需求和场景,增强了工具包的实用性。

如何使用MarkLLM

环境设置

要使用MarkLLM,你需要以下环境:

Python 3.9
PyTorch
其他依赖(可通过 pip install -r requirements.txt 安装)

注意:如果要使用EXPEdit或ITSEdit算法,还需要额外的步骤来导入.pyx文件。

调用水印算法

以下是一个基本的使用示例:

import torch
from watermark.auto_watermark import AutoWatermark
from utils.transformers_config import TransformersConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

# 设置设备
device = "cuda" if torch.cuda.is_available() else "cpu"

# Transformers配置
transformers_config = TransformersConfig(
    model=AutoModelForCausalLM.from_pretrained('facebook/opt-1.3b').to(device),
    tokenizer=AutoTokenizer.from_pretrained('facebook/opt-1.3b'),
    vocab_size=50272,
    device=device,
    max_new_tokens=200,
    min_length=230,
    do_sample=True,
    no_repeat_ngram_size=4
)

# 加载水印算法
myWatermark = AutoWatermark.load('KGW', 
                                 algorithm_config='config/KGW.json',
                                 transformers_config=transformers_config)

# 提示词
prompt = 'Good Morning.'

# 生成并检测
watermarked_text = myWatermark.generate_watermarked_text(prompt)
detect_result = myWatermark.detect_watermark(watermarked_text)

可视化机制

MarkLLM提供了强大的可视化工具,可以直观地展示水印的效果。以下是KGW家族算法的可视化示例:

import torch
from visualize.font_settings import FontSettings
from watermark.auto_watermark import AutoWatermark
from utils.transformers_config import TransformersConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
from visualize.visualizer import DiscreteVisualizer
from visualize.legend_settings import DiscreteLegendSettings
from visualize.page_layout_settings import PageLayoutSettings
from visualize.color_scheme import ColorSchemeForDiscreteVisualization

# 加载水印算法
device = "cuda" if torch.cuda.is_available() else "cpu"
transformers_config = TransformersConfig(
    model=AutoModelForCausalLM.from_pretrained('facebook/opt-1.3b').to(device),
    tokenizer=AutoTokenizer.from_pretrained('facebook/opt-1.3b'),
    vocab_size=50272,
    device=device,
    max_new_tokens=200,
    min_length=230,
    do_sample=True,
    no_repeat_ngram_size=4
)
myWatermark = AutoWatermark.load('KGW', 
                                 algorithm_config='config/KGW.json',
                                 transformers_config=transformers_config)

# 获取可视化数据
watermarked_data = myWatermark.get_data_for_visualization(watermarked_text)
unwatermarked_data = myWatermark.get_data_for_visualization(unwatermarked_text)

# 初始化可视化器
visualizer = DiscreteVisualizer(color_scheme=ColorSchemeForDiscreteVisualization(),
                                font_settings=FontSettings(), 
                                page_layout_settings=PageLayoutSettings(),
                                legend_settings=DiscreteLegendSettings())

# 可视化
watermarked_img = visualizer.visualize(data=watermarked_data, 
                                       show_text=True, 
                                       visualize_weight=True, 
                                       display_legend=True)

# 保存图像
watermarked_img.save("KGW_watermarked.png")

应用评估流程

MarkLLM提供了多种评估流程,以下是使用水印检测流程的示例:

import torch
from evaluation.dataset import C4Dataset
from watermark.auto_watermark import AutoWatermark
from utils.transformers_config import TransformersConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
from evaluation.tools.text_editor import TruncatePromptTextEditor, WordDeletion
from evaluation.tools.success_rate_calculator import DynamicThresholdSuccessRateCalculator
from evaluation.pipelines.detection import WatermarkedTextDetectionPipeline, UnWatermarkedTextDetectionPipeline, DetectionPipelineReturnType

# 加载数据集
my_dataset = C4Dataset('dataset/c4/processed_c4.json')

# 设置设备
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# Transformers配置
transformers_config = TransformersConfig(
    model=AutoModelForCausalLM.from_pretrained('facebook/opt-1.3b').to(device),
    tokenizer=AutoTokenizer.from_pretrained('facebook/opt-1.3b'),
    vocab_size=50272,
    device=device,
    max_new_tokens=200,
    do_sample=True,
    min_length=230,
    no_repeat_ngram_size=4
)

# 加载水印算法
my_watermark = AutoWatermark.load('KGW', 
                                  algorithm_config='config/KGW.json',
                                  transformers_config=transformers_config)

# 初始化评估流程
pipeline1 = WatermarkedTextDetectionPipeline(
    dataset=my_dataset, 
    text_editor_list=[TruncatePromptTextEditor(), WordDeletion(ratio=0.3)],
    show_progress=True, 
    return_type=DetectionPipelineReturnType.SCORES
) 

pipeline2 = UnWatermarkedTextDetectionPipeline(
    dataset=my_dataset, 
    text_editor_list=[],
    show_progress=True,
    return_type=DetectionPipelineReturnType.SCORES
)

# 评估
calculator = DynamicThresholdSuccessRateCalculator(labels=['TPR', 'F1'], rule='best')
print(calculator.calculate(pipeline1.evaluate(my_watermark), pipeline2.evaluate(my_watermark)))