Code Watermarking for Large Language Models

This repository contains the official implementation of the NeurIPS 2025 paper: "Practical and Effective Code Watermarking for Large Language Models"

This project provides a comprehensive framework for embedding and detecting watermarks in code generated by Large Language Models (LLMs), enabling provenance tracking and detection of LLM-generated code.

Project Structure

.
├── src/                    # Source code directory
│   ├── main.py            # Main application entry point
│   ├── train.py           # Training module
│   ├── modelhelper.py     # Model-related utilities
│   ├── datahelper.py      # Data processing and handling
│   ├── utils.py           # General utility functions
│   ├── sfthelper.py       # SFT (Supervised Fine-Tuning) helper functions
│   └── asthelper.py       # AST (Abstract Syntax Tree) processing utilities
│
├── scripts/               # Shell scripts for running experiments
│   ├── train.sh          # Training script
│   └── run.sh            # Evaluation script
├── lm_eval/               # Language Model evaluation directory
├── exp_utils/             # Experiment utilities
└── models/                # Model storage directory

Main Components

• main.py: Entry point for watermark generation and detection
• train.py: Implements the training pipeline for watermarked code generation models
• modelhelper.py: Model loading, configuration, and watermarking utilities
• datahelper.py: Dataset preparation and processing for watermarked code generation
• utils.py: General utility functions for watermarking operations
• sfthelper.py: Supervised fine-tuning helpers for watermark embedding
• asthelper.py: Abstract Syntax Tree analysis for code structure-aware watermarking

Supporting Directories

• scripts/: Shell scripts for training and evaluation experiments
• lm_eval/: Evaluation framework for watermarked language models
• exp_utils/: Utilities for running and analyzing experiments
• models/: Trained model checkpoints and configurations

Getting Started

Prerequisites

• Python 3.11+
• CUDA-capable GPU(s) recommended
• Required dependencies (see requirements.txt)

Installation

uv pip install -r requirements.txt

Usage Workflow

1. Data Generation

Generate the training and evaluation datasets for watermarked code generation:

CUDA_VISIBLE_DEVICES=5 python src/datahelper.py

This prepares the code datasets (e.g., MBPP, HumanEval) with appropriate preprocessing for watermark training.

2. Model Training

Train a code generation model with watermark embedding capabilities:

bash scripts/train.sh \
    --model "deepseek-ai/deepseek-coder-1.3b-instruct" \
    --task_name gen_mbpp \
    --train_batch_size 512 \
    --num_epochs 1000 \
    --gpus 5 \
    --augmentation false \
    --alpha_distill 0.1 \
    --alpha_ce 0.1 \
    --alpha_switch 8.0 \
    --context_width 2 \
    --use_cache false

Key Parameters:

• --model: Base model to fine-tune (e.g., DeepSeek Coder, OpenCoder)
• --task_name: Training task/dataset (e.g., gen_mbpp, humaneval)
• --alpha_distill, --alpha_ce, --alpha_switch: Loss weighting hyperparameters for watermark training
• --context_width: Context window size for watermark embedding
• --augmentation: Enable/disable data augmentation

3. Watermarked Code Generation and Detection

Run inference with the watermarked model and evaluate watermark detection:

bash scripts/run.sh \
    --model "infly/OpenCoder-1.5B-Instruct" \
    --task_name humanevalsynthesize-cpp \
    --wm wllm \
    --gpus 1 \
    --code_model none \
    --entropy_threshold 1.2 \
    --switch_threshold 0.7 \
    --seed 42 \
    --delta 2.0

Key Parameters:

• --model: Trained watermarked model or baseline model
• --task_name: Evaluation task (e.g., humanevalsynthesize-cpp, mbpp)
• --wm: Watermarking method (e.g., wllm)
• --entropy_threshold, --switch_threshold: Detection thresholds
• --delta: Watermark strength parameter

Citation

If you find this work useful, please cite our paper:

@inproceedings{guo2025practical,
  title={Practical and Effective Code Watermarking for Large Language Models},
  author={Guo, Zhimeng and Cheng, Minhao},
  booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
  year={2025},
  url={https://openreview.net/forum?id=RpE4HeuX69}
}

Contact

For questions or issues, please open an issue on GitHub or contact the authors.