File size: 3,441 Bytes
c126239 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | /**
* Watermark PRFs (pseudo-random functions).
*
* Constants follow the reference implementations:
* - HASH_KEY = 15485863 (the 1,000,000th prime) is the multiplicative salt
* used by both jwkirchenbauer/lm-watermarking and bohanhou14/SemStamp.
*/
import { hashCombine, mix64, toUnitFloat } from './rng';
export const HASH_KEY = 15485863n;
const MOD64 = (1n << 64n) - 1n;
/**
* Kirchenbauer "lefthash" (simple_1) seed: salt * prev_token_id, mod 2^64-1.
* Matches the reference `additive_prf` with context_width=1.
*/
export function kirchenbauerSeed(prevTokenId: number, hashKey: bigint = HASH_KEY): bigint {
return (hashKey * BigInt(prevTokenId)) % MOD64;
}
/**
* Green-list membership test. Instead of permuting the whole vocabulary
* (torch.randperm in the reference), we use a per-token PRF judgment:
* token is green iff U(seed, key, tokenId) < gamma.
* The green-list size is then Binomial(|V|, gamma) instead of exactly
* gamma*|V|; negligible for the z-test (documented in README).
*
* `key` is the user-chosen secret. It is mixed in here rather than into the
* seed formula so the displayed derivation stays literally
* `15485863 * prev_token mod 2^64-1`, as in the reference implementation.
*/
export function isGreenToken(
seed: bigint,
tokenId: number,
gamma: number,
key: bigint = 0n,
): boolean {
return greenUniform(seed, tokenId, key) < gamma;
}
/** The underlying uniform used by isGreenToken (exposed for UI/inspector). */
export function greenUniform(seed: bigint, tokenId: number, key: bigint = 0n): number {
return toUnitFloat(hashCombine(seed, key, BigInt(tokenId)));
}
/**
* TextSeal-style PRF: R_v = U(key, context window, candidate token v) in [0,1).
* The watermark context window is the previous k tokens (paper default k=3).
*
* NOTE: the official facebookresearch/textseal repo derives R via a
* cryptographic hash (appendix A.1). We use splitmix64-based mixing for
* per-token speed in the browser; this is an educational deviation that
* preserves the statistical construction (uniform R per (context, v, key)).
*/
export function textsealR(
key: bigint,
contextTokenIds: number[],
candidateTokenId: number,
): number {
const parts: bigint[] = [key];
for (const t of contextTokenIds) parts.push(BigInt(t));
parts.push(0x746578747365616cn); // "textseal" domain separator
parts.push(BigInt(candidateTokenId));
let u = toUnitFloat(hashCombine(...parts));
// Guard against exact 0/1 for downstream log computations.
if (u <= 0) u = Number.MIN_VALUE;
if (u >= 1) u = 1 - Number.EPSILON;
return u;
}
/** Derive a stable 64-bit key from a human-readable secret string. */
export function keyFromString(secret: string): bigint {
let h = 0xcbf29ce484222325n; // FNV offset basis
for (let i = 0; i < secret.length; i++) {
h = mix64((h ^ BigInt(secret.charCodeAt(i))) & MOD64);
}
return h;
}
/** k-SemStamp seed: prev sentence's cluster id * prime salt, mixed with the key. */
export function ksemstampSeed(
clusterId: number,
key: bigint = 0n,
hashKey: bigint = HASH_KEY,
): bigint {
const base = (hashKey * BigInt(clusterId + 1)) % MOD64; // +1 so cluster 0 doesn't zero out
return key === 0n ? base : hashCombine(base, key);
}
/** Derive a 64-bit key from a numeric seed entered by the user. */
export function keyFromSeed(seed: number): bigint {
return mix64(BigInt(Math.trunc(seed)) & MOD64) | 1n; // never zero
}
|