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