File size: 3,919 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | /**
* postMessage protocol between the UI thread and the inference worker.
*/
import type { AlgorithmId, Detection, GenerationResult, SamplingConfig, SentenceTrace } from '../watermark/types';
import type { KirchenbauerParams } from '../watermark/kirchenbauer';
import type { KsemstampParams } from '../watermark/ksemstamp';
import type { TextsealParams } from '../watermark/textseal';
export interface AlgorithmParams {
kirchenbauer: KirchenbauerParams;
ksemstamp: KsemstampParams;
textseal: TextsealParams;
}
/** Top-N next-token candidates for the wizard's Stage 1 / Stage 2 views. */
export interface NextTokenCandidate {
tokenId: number;
text: string;
logit: number;
prob: number;
}
export interface DistributionResult {
/** Last token of the templated prompt - the Kirchenbauer PRF context. */
prevTokenId: number;
prevTokenText: string;
/** Tail of the prompt (oldest first) so wider PRF contexts can be shown. */
prevTokenIds: number[];
vocabSize: number;
candidates: NextTokenCandidate[];
}
/** UI -> worker */
export type WorkerRequest =
| { type: 'check' }
| { type: 'load'; modelId: string }
| { type: 'distribution'; requestId: number; prompt: string; topN: number; temperature: number }
| {
type: 'generate';
requestId: number;
algorithm: AlgorithmId;
prompt: string;
sampling: SamplingConfig;
params: AlgorithmParams;
/** Per-algorithm secret entered by the user in Stage 1 (decimal string). */
secretKey?: string;
}
| {
type: 'detect';
requestId: number;
text: string;
sampling: SamplingConfig;
params: AlgorithmParams;
/** Per-algorithm secrets so all three detectors can score the same text. */
keys?: { kirchenbauer?: string; ksemstamp?: string; textseal?: string };
/** Run the TextSeal entropy-weighted variant (needs a model forward). */
withEntropy?: boolean;
};
export interface DeviceInfo {
webgpuSupported: boolean;
shaderF16: boolean;
device: 'webgpu' | 'wasm';
dtype: string;
adapterInfo?: { vendor?: string; architecture?: string };
}
export interface DetectionBundle {
kirchenbauer: Detection;
textseal: Detection;
ksemstamp: Detection | null;
}
/** worker -> UI */
export type WorkerResponse =
| { type: 'check-result'; info: DeviceInfo }
| { type: 'load-progress'; file: string; progress: number; loadedMB: number; totalMB: number }
| { type: 'load-done'; modelId: string; initMs: number; info: DeviceInfo }
| { type: 'load-error'; error: string }
| { type: 'distribution-done'; requestId: number; result: DistributionResult }
| { type: 'distribution-error'; requestId: number; error: string }
| {
type: 'generate-progress';
requestId: number;
algorithm: AlgorithmId;
text: string;
tokensDone: number;
retries: number;
}
/**
* One sampled token, streamed as it is chosen. `index` is authoritative:
* after a k-SemStamp rollback it goes backwards, and the receiver should
* drop everything from that index on before appending.
*/
| {
type: 'token';
requestId: number;
algorithm: AlgorithmId;
index: number;
text: string;
/** Kirchenbauer: which list the sampled token belonged to. */
green?: boolean;
/** TextSeal: which of the two keys routed this step. */
keyId?: 1 | 2;
entropy?: number;
}
/** k-SemStamp candidate sentence judged (streamed live during rejection sampling). */
| { type: 'candidate'; requestId: number; algorithm: AlgorithmId; sentence: SentenceTrace }
| { type: 'generate-done'; requestId: number; result: GenerationResult }
| { type: 'generate-error'; requestId: number; algorithm: AlgorithmId; error: string }
| { type: 'detect-done'; requestId: number; detections: DetectionBundle }
| { type: 'detect-error'; requestId: number; error: string };
|