|
#!/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() |