Minimal LLM Watermarking from scratch

A developer has released a minimal implementation of LLM watermarking, inspired by Anthropic's upcoming 'is it AI' API and Google DeepMind's SynthID-Text. The code, written in Python using PyTorch and Hugging Face Transformers, demonstrates Tournament Sampling for watermarking and a statistical detection method based on z-scores and p-values. The implementation is designed to be run with 'uv run synthid.py' and uses a small Qwen model for demonstration.

| /usr/bin/env -S uv run --script | | | /// script | | | requires-python = " =3.11" | | | dependencies = | | | "torch", | | | "transformers =5.15", | | | | | | /// | | | """LLM Watermarks. | | | Anthropic will roll out their "is it AI" API soon. | | | This is ~roughly how it works. | | | SynthID-Text from scratch. | | | A tiny implementation of Tournament Sampling + watermark detection. | | | uv run synthid.py | | | """ | | | from future import annotations | | | import hashlib | | | import hmac | | | import math | | | import torch | | | from transformers import AutoModelForCausalLM, AutoTokenizer | | | MODEL = "Qwen/Qwen2.5-0.5B-Instruct" | | | KEY = b"player-piano" | | | H = 4 previous tokens used as watermark context | | | M = 30 tournament layers | | | TOP K = 100 | | | TEMP = 0.7 | | | MAX NEW = 120 | | | def bits ctx: tuple int, ... , tok: int, key: bytes = KEY - list int : | | | """Return M deterministic keyed random bits for context, token .""" | | | msg = b"".join int x .to bytes 4, "little" for x in ctx, tok | | | digest = hmac.new key, msg, hashlib.sha256 .digest | | | n = int.from bytes digest, "little" | | | return n i & 1 for i in range M | | | def watermark p: torch.Tensor, toks: torch.Tensor, ctx: tuple int, ... - torch.Tensor: | | | """Apply M layers of N=2 Tournament Sampling.""" | | | ids = toks.detach .cpu .tolist | | | g = torch.tensor bits ctx, t for t in ids , dtype=torch.float32, device=p.device | | | p = p.float | | | for i in range M : | | | q = p g :, i .sum .clamp 0, 1 | | | The whole trick: | | | | | | p' x = p x 1 + g x - q | | | | | | p = p 1 + g :, i - q | | | p = p.clamp min 0 | | | return p / p.sum | | | @torch.inference mode | | | def generate model, tok, prompt: str, , wm: bool, seed: int = 42 - tuple str, list int : | | | """Generate text with or without the watermark.""" | | | device = next model.parameters .device | | | torch.manual seed seed | | | chat = tok.apply chat template | | | {"role": "user", "content": prompt} , | | | tokenize=False, | | | add generation prompt=True, | | | | | | prompt ids = tok chat, return tensors="pt", add special tokens=False .input ids.to device | | | out: list int = | | | seen: set tuple int, ... = set | | | eos = model.generation config.eos token id | | | if eos is None: | | | eos = tok.eos token id | | | eos ids = set eos if isinstance eos, list, tuple else eos | | | for in range MAX NEW : | | | ids = prompt ids | | | if out: | | | ids = torch.cat | | | prompt ids, torch.tensor out , dtype=torch.long, device=device , | | | dim=1, | | | | | | logits = model ids, use cache=False .logits 0, -1 .float / TEMP | | | top logits, top ids = torch.topk logits, min TOP K, logits.numel | | | p = torch.softmax top logits, dim=-1 | | | Normal decoding first, watermark second. | | | if wm and len out = H: | | | ctx = tuple out -H: | | | Don't watermark the same context twice. | | | if ctx not in seen: | | | seen.add ctx | | | p = watermark p, top ids, ctx | | | sample = torch.multinomial p, 1 .item | | | token = int top ids sample .item | | | if token in eos ids: | | | break | | | out.append token | | | return tok.decode out, skip special tokens=True , out | | | def detect ids: list int , key: bytes = KEY - tuple float, float, float : | | | """Return watermark score, z-score and one-sided p-value.""" | | | seen: set tuple int, ... = set | | | ones = 0 | | | n = 0 | | | for i in range H, len ids : | | | ctx = tuple ids i - H:i | | | if ctx in seen: | | | continue | | | seen.add ctx | | | ones += sum bits ctx, ids i , key | | | n += M | | | if n == 0: | | | return float "nan" , 0.0, 1.0 | | | Under the null, g ~ Bernoulli 0.5 . | | | score = ones / n | | | z = ones - n / 2 / math.sqrt n / 4 | | | p = 0.5 math.erfc z / math.sqrt 2 | | | return score, z, p | | | def detect text tok, text: str, key: bytes = KEY - tuple float, float, float : | | | """Detect our watermark in arbitrary visible text.""" | | | ids = tok.encode text, add special tokens=False | | | return detect int x for x in ids , key | | | def show detection tok, text: str, label: str, key: bytes = KEY - None: | | | """Print an interpretable detection result.""" | | | score, z, p = detect text tok, text, key | | | if z = 4: | | | verdict = "strong watermark evidence" | | | elif z = 2: | | | verdict = "some watermark evidence" | | | else: | | | verdict = "no watermark evidence" | | | print f"\n{label}" | | | print f"score = {score:.4f} z = {z:.2f} p ≈ {p:.2e}" | | | print f"result: {verdict}" | | | def show prefixes ids: list int - None: | | | """Show watermark evidence accumulating with text length.""" | | | print "\nEvidence as tokens accumulate:" | | | for n in 20, 40, 60, 80, 120 : | | | if len ids < n: | | | continue | | | score, z, p = detect ids :n | | | print f"{n: 3} tokens score={score:.3f} z={z: 6.2f} p≈{p:.1e}" | | | def main - None: | | | """Generate watermarked/plain text and demonstrate detection.""" | | | device = "cuda" if torch.cuda.is available else "cpu" | | | dtype = torch.float16 if device == "cuda" else torch.float32 | | | print f"device: {device}" | | | if device == "cuda": | | | print f"gpu: {torch.cuda.get device name 0 }" | | | tok = AutoTokenizer.from pretrained MODEL | | | model = AutoModelForCausalLM.from pretrained MODEL, dtype=dtype .to device .eval | | | prompt = | | | "To the engineers and managers of tomorrow: " | | | "when machines can do much of the work we now consider skilled, " | | | "what should humans get exceptionally good at?" | | | | | | 1. Generate with our watermark. | | | wm text, wm ids = generate model, tok, prompt, wm=True, seed=42 | | | print "\n" + "=" 70 | | | print "WATERMARKED" | | | print "=" 70 | | | print wm text | | | show detection tok, wm text, "Correct secret key" | | | show prefixes wm ids | | | Without the correct key, the hidden pattern disappears into noise. | | | show detection tok, wm text, "Wrong secret key", b"wrong-key" | | | 2. Same model, same prompt, ordinary sampling. | | | plain text, = generate model, tok, prompt, wm=False, seed=42 | | | print "\n" + "=" 70 | | | print "UNWATERMARKED" | | | print "=" 70 | | | print plain text | | | show detection tok, plain text, "Detection" | | | 3. Treat arbitrary text as an unknown sample. | | | mystery = """ | | | Good judgement starts with understanding what matters, what can change, | | | and what cannot. Tools may make execution cheaper, but choosing the | | | right problem and deciding what trade-offs are acceptable still matter. | | | """ | | | print "\n" + "=" 70 | | | print "MYSTERY TEXT" | | | print "=" 70 | | | print mystery.strip | | | show detection tok, mystery, "Detection" | | | if name == " main ": | | | main |