Minimal LLM Watermarking from scratch · GitHub
/" data-turbo-transient="true" />
Skip to content
-->
Search Gists
Search Gists
Sign in
Sign up
You signed in with another tab or window. Reload to refresh your session.<br>You signed out in another tab or window. Reload to refresh your session.<br>You switched accounts on another tab or window. Reload to refresh your session.
Dismiss alert
{{ message }}
Instantly share code, notes, and snippets.
jSwords91/minimark
Created<br>August 23, 2026 15:47
Show Gist options
Download ZIP
Star
(0)
You must be signed in to star a gist
Fork
(0)
You must be signed in to fork a gist
Embed
Select an option
Embed<br>Embed this gist in your website.
Share<br>Copy sharable link for this gist.
Clone via HTTPS<br>Clone using the web URL.
No results found
Learn more about clone URLs
Clone this repository at <script src="https://gist.github.com/jSwords91/2732ff6017213526da73e8dc2bd54770.js"></script>
" readonly="readonly" data-autoselect="true" data-target="primer-text-field.inputElement " aria-describedby="validation-6fe28768-1110-4a32-a912-3e5e39cdc2d9" class="form-control FormControl-monospace FormControl-input FormControl-small rounded-left-0 rounded-right-0 border-right-0" type="text" name="gist-share-url-sized-down" />
Save jSwords91/2732ff6017213526da73e8dc2bd54770 to your computer and use it in GitHub Desktop.
Embed
Select an option
Embed<br>Embed this gist in your website.
Share<br>Copy sharable link for this gist.
Clone via HTTPS<br>Clone using the web URL.
No results found
Learn more about clone URLs
Clone this repository at <script src="https://gist.github.com/jSwords91/2732ff6017213526da73e8dc2bd54770.js"></script>
" readonly="readonly" data-autoselect="true" data-target="primer-text-field.inputElement " aria-describedby="validation-608a3428-d4dd-4850-a218-8c73da20f715" class="form-control FormControl-monospace FormControl-input FormControl-small rounded-left-0 rounded-right-0 border-right-0" type="text" name="gist-share-url-original" />
Save jSwords91/2732ff6017213526da73e8dc2bd54770 to your computer and use it in GitHub Desktop.
Download ZIP
Minimal LLM Watermarking from scratch
Raw
minimark
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.<br>Learn more about bidirectional Unicode characters
Show hidden characters
#!/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)
#...