/** * Dedicated inference worker. Owns the Transformers.js model, tokenizer and * sentence embedder; runs the unified generation loop and all detectors. * * Pattern follows the official transformers.js-examples llama-3.2-webgpu * worker (singleton lazy-load + postMessage protocol), extended with a manual * autoregressive loop for custom sampling. */ import { AutoModel, AutoModelForCausalLM, AutoTokenizer, Tensor, env, } from '@huggingface/transformers'; import { getModelSpec, EMBEDDING_MODEL, type ModelSpec } from '../model/model-adapter'; import { runGenerationLoop, entropyOf, softmaxT, type LoopModel } from '../model/generation-engine'; import { tokenizeWithOffsets } from '../model/tokenizer'; import { applyGreenBias } from '../watermark/kirchenbauer'; import { isGreenToken, kirchenbauerSeed } from '../utils/hashing'; import { deriveKeys, gumbelMaxChoose, type TextsealKeys } from '../watermark/textseal'; import { acceptSentence, assignCluster, assertCentroidsMatch, type Centroids, } from '../watermark/ksemstamp'; import { detectKirchenbauer } from '../detectors/kirchenbauer-detector'; import { detectTextseal } from '../detectors/textseal-detector'; import { detectKsemstamp } from '../detectors/ksemstamp-detector'; import { splitSentences } from '../utils/sentences'; import { RandomStream } from '../utils/rng'; import type { WorkerRequest, WorkerResponse, DeviceInfo, AlgorithmParams, DistributionResult, NextTokenCandidate, } from './worker-protocol'; import type { AlgorithmId, Detection, GenerationResult, GenerationTrace, SamplingConfig, SentenceTrace, StepTrace, } from '../watermark/types'; const post = (msg: WorkerResponse) => (self as unknown as Worker).postMessage(msg); env.allowLocalModels = false; // --------------------------------------------------------------------------- // Device detection // --------------------------------------------------------------------------- /** Minimal WebGPU surface (avoids a dependency on @webgpu/types). */ interface MinimalGPUAdapter { features: { has(name: string): boolean }; info?: { vendor?: string; architecture?: string }; } interface MinimalGPU { requestAdapter(): Promise; } async function detectDevice(): Promise { const nav = navigator as Navigator & { gpu?: MinimalGPU }; let webgpuSupported = false; let shaderF16 = false; let adapterInfo: DeviceInfo['adapterInfo']; try { if (nav.gpu) { const adapter = await nav.gpu.requestAdapter(); if (adapter) { webgpuSupported = true; shaderF16 = adapter.features.has('shader-f16'); if (adapter.info) { adapterInfo = { vendor: adapter.info.vendor, architecture: adapter.info.architecture }; } } } } catch { webgpuSupported = false; } const device = webgpuSupported ? 'webgpu' : 'wasm'; const dtype = webgpuSupported ? (shaderF16 ? 'q4f16' : 'q4') : 'q4'; return { webgpuSupported, shaderF16, device, dtype, adapterInfo }; } // --------------------------------------------------------------------------- // Model singleton // --------------------------------------------------------------------------- type Tok = Awaited>; type LM = Awaited>; let spec: ModelSpec | null = null; let tokenizer: Tok | null = null; let model: LM | null = null; let deviceInfo: DeviceInfo | null = null; let embedder: { tokenizer: Tok; model: unknown } | null = null; let centroids: Centroids | null = null; async function loadModel(modelId: string): Promise { const t0 = performance.now(); spec = getModelSpec(modelId); deviceInfo = await detectDevice(); const dtype = deviceInfo.device === 'webgpu' ? (deviceInfo.shaderF16 ? spec.dtype : 'q4') : 'q4'; deviceInfo.dtype = dtype; const progress_callback = (p: { status: string; file?: string; progress?: number; loaded?: number; total?: number; }) => { if (p.status === 'progress' && p.file) { post({ type: 'load-progress', file: p.file, progress: p.progress ?? 0, loadedMB: (p.loaded ?? 0) / 1e6, totalMB: (p.total ?? 0) / 1e6, }); } }; tokenizer = await AutoTokenizer.from_pretrained(spec.repo, { progress_callback }); model = await AutoModelForCausalLM.from_pretrained(spec.repo, { dtype, device: deviceInfo.device, progress_callback, }); // Warm-up: compile shaders with a single-token generation. try { const warm = tokenizer('a'); const out = await (model as any)({ ...warm }); disposeOutputs(out); } catch { /* warm-up is best-effort */ } post({ type: 'load-done', modelId, initMs: performance.now() - t0, info: deviceInfo, }); } /** * The encoder is called through AutoModel rather than the feature-extraction * pipeline: EmbeddingGemma's ONNX graph already applies pooling and the dense * projection and returns `sentence_embedding` (L2-normalized), whereas the * pipeline would mean-pool hidden states and skip the projection. */ async function loadEmbedder(device: 'webgpu' | 'wasm'): Promise<{ tokenizer: Tok; model: unknown }> { return { tokenizer: await AutoTokenizer.from_pretrained(EMBEDDING_MODEL.repo), // q4 rather than q4f16: the half-precision build of this model returns // NaN embeddings on some adapters, and a NaN silently collapses every // sentence into cluster 0 instead of failing loudly. model: await AutoModel.from_pretrained(EMBEDDING_MODEL.repo, { dtype: EMBEDDING_MODEL.dtype, device, }), }; } async function getEmbedder(): Promise<{ tokenizer: Tok; model: unknown }> { if (!embedder) { embedder = await loadEmbedder(deviceInfo?.device === 'webgpu' ? 'webgpu' : 'wasm'); } return embedder; } async function getCentroids(): Promise { if (!centroids) { // Served from public/ at the site root (works in dev and in the built Space). const res = await fetch( new URL('centroids/embeddinggemma-k8.json', self.location.origin + '/'), ); if (!res.ok) throw new Error(`Failed to load centroids: HTTP ${res.status}`); const loaded = (await res.json()) as Centroids; assertCentroidsMatch(loaded, { encoder: EMBEDDING_MODEL.repo, prefix: EMBEDDING_MODEL.clusteringPrefix, dtype: EMBEDDING_MODEL.dtype, dim: EMBEDDING_MODEL.dim, }); centroids = loaded; } return centroids; } async function runEmbedder( enc: { tokenizer: Tok; model: unknown }, text: string, ): Promise { const inputs = await (enc.tokenizer as any)([EMBEDDING_MODEL.clusteringPrefix + text], { padding: true, }); const out = await (enc.model as any)(inputs); const emb = out.sentence_embedding as Tensor; const data = Array.from(emb.data as Float32Array); disposeOutputs(out); for (const v of Object.values(inputs as Record)) { (v as { dispose?: () => void })?.dispose?.(); } return data; } /** True once we have proven the current encoder returns usable numbers. */ let embedderChecked = false; async function embedText(text: string): Promise { let enc = await getEmbedder(); let data = await runEmbedder(enc, text); // A NaN here would not throw; it would quietly make every sentence land in // the same cluster and every k-SemStamp verdict meaningless. Catch it once // and fall back to the CPU build, which is slower but correct. if (!embedderChecked) { if (data.some((x) => !Number.isFinite(x))) { if (deviceInfo?.device === 'webgpu') { embedder = await loadEmbedder('wasm'); enc = embedder; data = await runEmbedder(enc, text); } if (data.some((x) => !Number.isFinite(x))) { throw new Error('Sentence encoder returned non-finite embeddings on this device.'); } } embedderChecked = true; } return data; } // --------------------------------------------------------------------------- // LoopModel adapter over Transformers.js // --------------------------------------------------------------------------- function disposeOutputs(outputs: unknown): void { if (!outputs || typeof outputs !== 'object') return; for (const v of Object.values(outputs as Record)) { (v as { dispose?: () => void })?.dispose?.(); } } function disposePastObject(past: unknown): void { if (!past || typeof past !== 'object') return; for (const v of Object.values(past as Record)) { (v as { dispose?: () => void })?.dispose?.(); } } function makeLoopModel(): LoopModel { if (!model || !tokenizer) throw new Error('Model not loaded'); const tok = tokenizer; const lm = model as any; const eos: number[] = []; const eosCfg = lm.generation_config?.eos_token_id ?? lm.config?.eos_token_id; if (Array.isArray(eosCfg)) eos.push(...eosCfg.map(Number)); else if (eosCfg != null) eos.push(Number(eosCfg)); let vocabSize = 0; return { get vocabSize() { return vocabSize; }, eosTokenIds: eos, decode(ids: number[]): string { return tok.decode(ids, { skip_special_tokens: true }); }, async forward(inputIds: number[], fullSeqLen: number, past: unknown) { const n = inputIds.length; const pastLen = fullSeqLen - n; const input_ids = new Tensor('int64', BigInt64Array.from(inputIds.map(BigInt)), [1, n]); const attention_mask = new Tensor('int64', new BigInt64Array(fullSeqLen).fill(1n), [ 1, fullSeqLen, ]); const position_ids = new Tensor( 'int64', BigInt64Array.from({ length: n }, (_, i) => BigInt(pastLen + i)), [1, n], ); const inputs: Record = { input_ids, attention_mask, position_ids }; if (past) inputs.past_key_values = past; const outputs = await lm.forward(inputs); const logitsT = outputs.logits as Tensor; const [, seqLen, vocab] = logitsT.dims as number[]; vocabSize = vocab; const all = logitsT.data as Float32Array; const logits = new Float32Array(vocab); logits.set(all.subarray((seqLen - 1) * vocab, seqLen * vocab)); // getPastKeyValues moves present.* into a past object and disposes the // previous step's GPU buffers. const newPast = lm.getPastKeyValues(outputs, past ?? null); (logitsT as any).dispose?.(); input_ids.dispose?.(); attention_mask.dispose?.(); position_ids.dispose?.(); return { logits, past: newPast }; }, disposePast(past: unknown): void { disposePastObject(past); }, }; } /** Full-sequence forward to get per-position entropies (for TextSeal weighted detection). */ async function computeEntropies(tokenIds: number[], temperature: number): Promise { if (!model) throw new Error('Model not loaded'); if (tokenIds.length < 2 || tokenIds.length > 512) return []; const lm = model as any; const n = tokenIds.length; const input_ids = new Tensor('int64', BigInt64Array.from(tokenIds.map(BigInt)), [1, n]); const attention_mask = new Tensor('int64', new BigInt64Array(n).fill(1n), [1, n]); const position_ids = new Tensor( 'int64', BigInt64Array.from({ length: n }, (_, i) => BigInt(i)), [1, n], ); const outputs = await lm.forward({ input_ids, attention_mask, position_ids }); const logitsT = outputs.logits as Tensor; const [, seqLen, vocab] = logitsT.dims as number[]; const all = logitsT.data as Float32Array; // Entropy of the distribution predicting token i (from logits at i-1). const entropies: number[] = [0]; for (let i = 1; i < seqLen; i++) { const row = all.subarray((i - 1) * vocab, i * vocab) as Float32Array; entropies.push(entropyOf(softmaxT(row, temperature))); } disposeOutputs(outputs); input_ids.dispose?.(); attention_mask.dispose?.(); position_ids.dispose?.(); return entropies; } // --------------------------------------------------------------------------- // Generation per algorithm // --------------------------------------------------------------------------- function buildPromptIds(prompt: string): number[] { if (!tokenizer || !spec) throw new Error('Model not loaded'); const messages = [{ role: 'user', content: prompt }]; const templated = (tokenizer as any).apply_chat_template(messages, { tokenize: false, add_generation_prompt: true, ...(spec.chatTemplateOptions ?? {}), }) as string; const enc = (tokenizer as any).encode(templated, { add_special_tokens: false }) as number[]; return Array.from(enc, Number); } /** Per-algorithm secrets (decimal strings from the UI) -> bigint keys. */ interface SecretKeys { kirchenbauer: bigint; ksemstamp: bigint; textseal: TextsealKeys; } function toKeys(keys?: { kirchenbauer?: string; ksemstamp?: string; textseal?: string }): SecretKeys { return { kirchenbauer: keys?.kirchenbauer ? BigInt(keys.kirchenbauer) : 0n, ksemstamp: keys?.ksemstamp ? BigInt(keys.ksemstamp) : 0n, textseal: deriveKeys(keys?.textseal ?? 'default'), }; } /** Top-N next-token distribution for the wizard's Stage 1 / Stage 2 views. */ async function nextTokenDistribution( prompt: string, topN: number, temperature: number, ): Promise { const loopModel = makeLoopModel(); const promptIds = buildPromptIds(prompt); let past: unknown = null; try { const { logits, past: p } = await loopModel.forward(promptIds, promptIds.length, null); past = p; const probs = softmaxT(logits, temperature); const idx = Array.from(probs.keys()); idx.sort((a, b) => probs[b] - probs[a]); const candidates: NextTokenCandidate[] = idx.slice(0, topN).map((v) => ({ tokenId: v, text: loopModel.decode([v]), logit: logits[v], prob: probs[v], })); const prevTokenId = promptIds[promptIds.length - 1]; return { prevTokenId, prevTokenText: loopModel.decode([prevTokenId]), prevTokenIds: promptIds.slice(-8), vocabSize: logits.length, candidates, }; } finally { loopModel.disposePast(past); } } async function detectAll( text: string, sampling: SamplingConfig, params: AlgorithmParams, withEntropy: boolean, secrets: SecretKeys, ): Promise<{ kirchenbauer: Detection; textseal: Detection; ksemstamp: Detection | null }> { if (!tokenizer) throw new Error('Model not loaded'); const tok = tokenizeWithOffsets( { encode: (t, o) => Array.from((tokenizer as any).encode(t, o) as number[], Number), decode: (ids, o) => (tokenizer as any).decode(ids, o) as string, }, text, ); const keys = secrets.textseal; const kirch = detectKirchenbauer(tok, params.kirchenbauer, true, secrets.kirchenbauer); let entropies: number[] | undefined; if (withEntropy) { try { entropies = await computeEntropies(tok.tokenIds, sampling.temperature); if (entropies.length === 0) entropies = undefined; } catch { entropies = undefined; } } const seal = detectTextseal(tok, keys, params.textseal, { entropies }); let ksem: Detection | null = null; try { const sentences = splitSentences(text); if (sentences.length >= 2) { const cents = await getCentroids(); const embeddings: number[][] = []; for (const s of sentences) embeddings.push(await embedText(s.text)); ksem = detectKsemstamp(sentences, embeddings, cents, params.ksemstamp, secrets.ksemstamp); } } catch (e) { ksem = null; } return { kirchenbauer: kirch, textseal: seal, ksemstamp: ksem }; } async function generate( requestId: number, algorithm: AlgorithmId, prompt: string, sampling: SamplingConfig, params: AlgorithmParams, secrets: SecretKeys, ): Promise { if (!spec || !deviceInfo) throw new Error('Model not loaded'); const loopModel = makeLoopModel(); const promptIds = buildPromptIds(prompt); const keys = secrets.textseal; const sentenceTraces: SentenceTrace[] = []; const notes: string[] = []; /** Set by the Kirchenbauer hook just before each step is sampled. */ let stepSeed = 0n; /** Text of each accepted token; the engine already decoded them in context. */ const pieces: string[] = []; const onToken = (step: StepTrace) => { // index is authoritative: a k-SemStamp rollback moves it backwards. pieces.length = step.index; pieces.push(step.chosenTokenText); post({ type: 'token', requestId, algorithm, index: step.index, text: step.chosenTokenText, green: algorithm === 'kirchenbauer' ? isGreenToken( stepSeed, step.chosenTokenId, params.kirchenbauer.gamma, secrets.kirchenbauer, ) : undefined, keyId: step.keyId, entropy: step.entropy, }); if (step.index % 8 === 0) { post({ type: 'generate-progress', requestId, algorithm, text: pieces.join(''), tokensDone: step.index + 1, retries: 0, }); } }; const hooks: Parameters[3] = {}; if (algorithm === 'kirchenbauer') { hooks.transformLogits = (_i, ctx, logits) => { stepSeed = applyGreenBias( logits, ctx[ctx.length - 1], params.kirchenbauer, secrets.kirchenbauer, ); return stepSeed; }; notes.push('Shares the base RNG stream with Baseline (same seed, same uniform consumption).'); } else if (algorithm === 'textseal') { const router = new RandomStream(BigInt(sampling.baseSeed) + 7777n); hooks.sampleOverride = (_i, ctx, probs) => gumbelMaxChoose(probs, ctx, keys, params.textseal, router); notes.push( 'Sampling is replaced by deterministic Gumbel-max (PRF-driven); the base RNG stream cannot be shared with Baseline.', ); } else if (algorithm === 'ksemstamp') { const cents = await getCentroids(); // Seed chain with the prompt's semantic cluster. let prevCluster = 0; try { const promptSents = splitSentences(prompt); const seedText = promptSents.length > 0 ? promptSents[promptSents.length - 1].text : prompt; prevCluster = assignCluster(await embedText(seedText), cents).clusterId; } catch { prevCluster = 0; } let retriesSoFar = 0; hooks.maxSentenceTrials = params.ksemstamp.maxTrials; hooks.onSentenceEnd = async (text, sentenceIndex, attempt) => { const emb = await embedText(text); const res = acceptSentence(emb, prevCluster, cents, params.ksemstamp, secrets.ksemstamp); const entry: SentenceTrace = { sentenceIndex, attempt, text, clusterId: res.assignment.clusterId, targetClusters: res.targets, distances: res.assignment.distances, margin: res.assignment.margin, accepted: res.accepted || attempt >= params.ksemstamp.maxTrials, rejectionReason: res.accepted ? undefined : attempt >= params.ksemstamp.maxTrials ? 'maxTrials' : res.reason, }; sentenceTraces.push(entry); // Stream each judged candidate so the UI can show rejection sampling live. post({ type: 'candidate', requestId, algorithm, sentence: entry }); if (res.accepted || attempt >= params.ksemstamp.maxTrials) { prevCluster = res.assignment.clusterId; retriesSoFar = 0; return { accept: true }; } retriesSoFar++; post({ type: 'generate-progress', requestId, algorithm, text: '', tokensDone: 0, retries: retriesSoFar, }); return { accept: false }; }; notes.push( 'Sentence-level rejection sampling re-generates candidates; RNG consumption differs from Baseline by construction.', ); notes.push( 'Simplified vs paper: general-purpose EmbeddingGemma encoder (no paraphrase-contrastive fine-tuning), K-means centroids from a small public corpus.', ); } const t0 = performance.now(); const loop = await runGenerationLoop( loopModel, promptIds, { temperature: sampling.temperature, topP: sampling.topP, maxNewTokens: sampling.maxNewTokens, baseSeed: sampling.baseSeed, topKTrace: 8, retryTemperatureStep: algorithm === 'ksemstamp' ? 0.12 : 0, onToken, }, hooks, ); const totalMs = performance.now() - t0; // Annotate Kirchenbauer trace candidates with green membership. if (algorithm === 'kirchenbauer') { let prev = promptIds[promptIds.length - 1]; for (const step of loop.steps) { const seed = kirchenbauerSeed(prev); step.seed = String(seed); step.chosenIsGreen = isGreenToken( seed, step.chosenTokenId, params.kirchenbauer.gamma, secrets.kirchenbauer, ); let greens = 0; for (const c of step.topCandidates) { c.isGreen = isGreenToken(seed, c.tokenId, params.kirchenbauer.gamma, secrets.kirchenbauer); if (c.isGreen) greens++; } step.greenFraction = step.topCandidates.length > 0 ? greens / step.topCandidates.length : 0; prev = step.chosenTokenId; } } const trace: GenerationTrace = { algorithm, steps: loop.steps, sentences: sentenceTraces.length > 0 ? sentenceTraces : undefined, vocabSize: loopModel.vocabSize, params: { ...(algorithm === 'kirchenbauer' ? params.kirchenbauer : {}), ...(algorithm === 'ksemstamp' ? params.ksemstamp : {}), ...(algorithm === 'textseal' ? params.textseal : {}), } as Record, notes, }; const detections = await detectAll(loop.text, sampling, params, false, secrets); /** * Never substitute another algorithm's Detection: it would travel into the * JSON export labelled as this algorithm's score. */ const unscored: Detection = { algorithm, statistic: 0, statisticName: 'z-score', statisticDefinition: 'not enough sentences to score (at least two are needed)', threshold: Infinity, watermarked: false, unitsScored: 0, unitName: 'sentences', notes: ['The generation was too short to judge sentence by sentence.'], }; const own: Detection = algorithm === 'textseal' ? detections.textseal : algorithm === 'ksemstamp' ? (detections.ksemstamp ?? unscored) : detections.kirchenbauer; return { algorithm, text: loop.text, tokenIds: loop.tokenIds, trace, detection: own, timings: { totalMs, tokensPerSec: loop.tokenIds.length / (totalMs / 1000), retries: loop.retries, }, metadata: { model: spec.repo, dtype: deviceInfo.dtype, device: deviceInfo.device, algorithm, keyId: sampling.secret, seed: sampling.baseSeed, temperature: sampling.temperature, topP: sampling.topP, maxNewTokens: sampling.maxNewTokens, params: trace.params, }, }; } // --------------------------------------------------------------------------- // Message handling (serialized: one generation at a time) // --------------------------------------------------------------------------- let busy: Promise = Promise.resolve(); self.addEventListener('message', (ev: MessageEvent) => { const msg = ev.data; busy = busy.then(async () => { try { if (msg.type === 'check') { post({ type: 'check-result', info: await detectDevice() }); } else if (msg.type === 'load') { try { await loadModel(msg.modelId); } catch (e) { post({ type: 'load-error', error: String(e instanceof Error ? e.message : e) }); } } else if (msg.type === 'distribution') { try { const result = await nextTokenDistribution(msg.prompt, msg.topN, msg.temperature); post({ type: 'distribution-done', requestId: msg.requestId, result }); } catch (e) { post({ type: 'distribution-error', requestId: msg.requestId, error: String(e instanceof Error ? e.message : e), }); } } else if (msg.type === 'generate') { try { const secrets = toKeys({ kirchenbauer: msg.algorithm === 'kirchenbauer' ? msg.secretKey : undefined, ksemstamp: msg.algorithm === 'ksemstamp' ? msg.secretKey : undefined, textseal: msg.algorithm === 'textseal' ? msg.secretKey : undefined, }); const result = await generate( msg.requestId, msg.algorithm, msg.prompt, msg.sampling, msg.params, secrets, ); post({ type: 'generate-done', requestId: msg.requestId, result }); } catch (e) { post({ type: 'generate-error', requestId: msg.requestId, algorithm: msg.algorithm, error: String(e instanceof Error ? (e.stack ?? e.message) : e), }); } } else if (msg.type === 'detect') { try { const detections = await detectAll( msg.text, msg.sampling, msg.params, msg.withEntropy ?? false, toKeys(msg.keys), ); post({ type: 'detect-done', requestId: msg.requestId, detections }); } catch (e) { post({ type: 'detect-error', requestId: msg.requestId, error: String(e instanceof Error ? e.message : e), }); } } } catch (e) { // Last-resort: never let the queue die. console.error('worker error', e); } }); });