MarkLLM Documentation

To evaluate the detectability and robustness of a watermark, use the WatermarkedTextDetectionPipeline and UnWatermarkedTextDetectionPipeline. These pipelines can incorporate text_editor_list (e.g., WordDeletion) to simulate attacks.

Key Components:

• WatermarkedTextDetectionPipeline: Evaluates detection on text that has been watermarked.
• UnWatermarkedTextDetectionPipeline: Evaluates detection on natural/unwatermarked text.
• DynamicThresholdSuccessRateCalculator: Calculates metrics like TPR and F1 based on the pipeline results.

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

# Load dataset
my_dataset = C4Dataset('dataset/c4/processed_c4.json')

# ... (setup device and transformers_config) ...

# Load watermark algorithm
my_watermark = AutoWatermark.load('KGW', 
                                  algorithm_config='config/KGW.json',
                                  transformers_config=transformers_config)

# Init pipelines
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)

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

To assess the impact of watermarking on text quality (e.g., perplexity), use the DirectTextQualityAnalysisPipeline. This pipeline compares watermarked text against a reference (natural text) using an analyzer like PPLCalculator.

Key Components:

• DirectTextQualityAnalysisPipeline: The main pipeline for quality assessment.
• PPLCalculator: An analyzer that calculates Perplexity.
• QualityPipelineReturnType: Specifies the return format (e.g., MEAN_SCORES).

import torch
from evaluation.dataset import C4Dataset
from watermark.auto_watermark import AutoWatermark
from utils.transformers_config import TransformersConfig
from evaluation.tools.text_editor import TruncatePromptTextEditor
from evaluation.tools.text_quality_analyzer import PPLCalculator
from evaluation.pipelines.quality_analysis import DirectTextQualityAnalysisPipeline, QualityPipelineReturnType

# ... (setup dataset, device, transformers_config, and my_watermark) ...

# Init pipeline
quality_pipeline = DirectTextQualityAnalysisPipeline(
    dataset=my_dataset, 
    watermarked_text_editor_list=[TruncatePromptTextEditor()],
    unwatermarked_text_editor_list=[],                                             
    analyzer=PPLCalculator(
        model=AutoModelForCausalLM.from_pretrained('..model/llama-7b/', device_map='auto'),
        tokenizer=LlamaTokenizer.from_pretrained('..model/llama-7b/'),
        device=device),
    unwatermarked_text_source='natural', 
    show_progress=True, 
    return_type=QualityPipelineReturnType.MEAN_SCORES)

# Evaluate
print(quality_pipeline.evaluate(my_watermark))

To evaluate how robust a watermark is against attacks, use the detection pipelines. These pipelines combine a dataset, a text editor (to simulate attacks), and a watermark instance.

Attack Types (Text Editors)

• WordDeletion(ratio=...): Deletes words.
• SynonymSubstitution(ratio=...): Replaces words with synonyms.
• ContextAwareSynonymSubstitution: Uses BERT to perform context-aware substitution.
• GPTParaphraser: Uses OpenAI models to paraphrase text.
• DipperParaphraser: Uses a T5-based model for paraphrasing.
• TruncatePromptTextEditor: Truncates the prompt.

Pipeline Execution

Use WatermarkedTextDetectionPipeline to test detection on watermarked text and UnWatermarkedTextDetectionPipeline for control text. Use DynamicThresholdSuccessRateCalculator to compare results and calculate metrics like TPR and F1.

from evaluation.pipelines.detection import WatermarkedTextDetectionPipeline, UnWatermarkedTextDetectionPipeline
from evaluation.tools.success_rate_calculator import DynamicThresholdSuccessRateCalculator
from evaluation.tools.text_editor import WordDeletion

# Setup pipeline with a Word Deletion attack
pipline1 = WatermarkedTextDetectionPipeline(
    dataset=my_dataset, 
    text_editor_list=[TruncatePromptTextEditor(), WordDeletion(ratio=0.3)],
    show_progress=True, 
    return_type=DetectionPipelineReturnType.SCORES
)

# Evaluate and calculate success rate
calculator = DynamicThresholdSuccessRateCalculator(labels=['TPR', 'F1'], rule='best')
results = calculator.calculate(pipline1.evaluate(my_watermark), pipline2.evaluate(my_watermark))

MarkLLM provides visualizers to inspect how watermarks are applied to tokens. There are two main types of visualizers:

1. DiscreteVisualizer: Used for algorithms like KGW, Unigram, SWEET, UPV, SIR, XSIR, and EWD. It visualizes discrete flags/states.
2. ContinuousVisualizer: Used for algorithms like EXP and EXPEdit that use continuous values.

Usage

Use myWatermark.get_data_for_visualization(text) to retrieve the necessary data (tokens, flags, weights, values) from a generated text, then pass it to the visualizer's .visualize() method.

from visualize.visualizer import DiscreteVisualizer, ContinuousVisualizer
from visualize.data_for_visualization import DataForVisualization

# Get data from the watermark instance
watermarked_data = myWatermark.get_data_for_visualization(watermarked_text)

# Initialize and run visualizer
discreet_visualizer = DiscreteVisualizer(
    color_scheme=ColorSchemeForDiscreteVisualization(),
    font_settings=FontSettings(),
    page_layout_settings=PageLayoutSettings(),
    legend_settings=DiscreteLegendSettings()
)

img = discreet_visualizer.visualize(
    data=watermarked_data,
    show_text=True,
    visualize_weight=True,
    display_legend=True
)
img.save("visualization.png")

You can use the AutoWatermark class to load various watermarking algorithms and perform text generation and detection.

Supported Algorithms

Supported names include: KGW, Unigram, SWEET, EWD, SIR, XSIR, UPV, EXP, EXPEdit, SynthID, Adaptive, and STEAL (for stealing attacks).

Workflow

1. Configure Transformers: Use TransformersConfig to wrap your model and tokenizer.
2. Load Algorithm: Use AutoWatermark.load(algorithm_name, algorithm_config, transformers_config).
3. Generate: Use .generate_watermarked_text(prompt) or .generate_unwatermarked_text(prompt).
4. Detect: Use .detect_watermark(text) to check for the presence of a watermark.

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

# 1. Setup Config
transformers_config = TransformersConfig(
    model=AutoModelForCausalLM.from_pretrained(model_path).to(device),
    tokenizer=AutoTokenizer.from_pretrained(model_path),
    vocab_size=50272,
    device=device,
    max_new_tokens=200,
    min_length=230,
    do_sample=True,
    no_repeat_ngram_size=4
)

# 2. Load Algorithm
myWatermark = AutoWatermark.load('KGW', algorithm_config='config/KGW.json', transformers_config=transformers_config)

# 3. Generate and Detect
watermarked_text = myWatermark.generate_watermarked_text(prompt)
detect_result = myWatermark.detect_watermark(watermarked_text)
print(detect_result)

To ensure watermarking doesn't degrade text quality, use the quality analysis pipelines. There are three main types:

1. DirectTextQualityAnalysisPipeline: Measures intrinsic properties like Perplexity (PPLCalculator) or Log Diversity (LogDiversityAnalyzer).
2. ReferencedTextQualityAnalysisPipeline: Compares generated text against a reference (e.g., using BLEUCalculator for translation or PassOrNotJudger for code generation).
3. ExternalDiscriminatorTextQualityAnalysisPipeline: Uses an external model (like GPT-4 via GPTTextDiscriminator) to judge quality (e.g., Win Rate).

from evaluation.pipelines.quality_analysis import DirectTextQualityAnalysisPipeline
from evaluation.tools.text_quality_analyzer import PPLCalculator

# Direct quality analysis for Perplexity
quality_pipeline = DirectTextQualityAnalysisPipeline(
    dataset=my_dataset,
    watermarked_text_editor_list=[TruncatePromptTextEditor()],
    unwatermarked_text_editor_list=[],
    analyzers=[PPLCalculator(model=model, tokenizer=tokenizer, device=device)],
    unwatermarked_text_source='natural',
    show_progress=True,
    return_type=QualityPipelineReturnType.MEAN_SCORES
)

print(quality_pipeline.evaluate(my_watermark))

You can use the AutoWatermark class to load an algorithm and perform watermarking or detection tasks. This requires a TransformersConfig object to wrap your model and tokenizer.

Workflow:

1. Initialize TransformersConfig with your model, tokenizer, and generation parameters.
2. Load the algorithm using AutoWatermark.load(algorithm_name, algorithm_config, transformers_config).
3. Use .generate_watermarked_text(prompt) to create watermarked text.
4. Use .detect_watermark(text) to check for the presence of a watermark.
5. Use .generate_unwatermarked_text(prompt) to generate standard text for comparison.

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

# Device
device = "cuda" if torch.cuda.is_available() else "cpu"

# Transformers config
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)
  
# Load watermark algorithm
myWatermark = AutoWatermark.load('KGW', 
                                 algorithm_config='config/KGW.json',
                                 transformers_config=transformers_config)

# Prompt
prompt = 'Good Morning.'

# Generate and detect
watermarked_text = myWatermark.generate_watermarked_text(prompt)
detect_result = myWatermark.detect_watermark(watermarked_text)
unwatermarked_text = myWatermark.generate_unwatermarked_text(prompt)
detect_result = myWatermark.detect_watermark(unwatermarked_text)