src/lib/model/generation-engine.ts

/**
 * Unified custom autoregressive generation loop.
 *
 * All four runs (baseline, Kirchenbauer, k-SemStamp, TextSeal) go through
 * this single loop so generation conditions stay comparable. Algorithms
 * plug in via GenerationHooks:
 *  - transformLogits: mutate logits before softmax (Kirchenbauer green bias)
 *  - sampleOverride:  replace sampling entirely (TextSeal Gumbel-max)
 *  - onSentenceEnd:   sentence-level accept/reject (k-SemStamp); rejection
 *                     rolls the sequence back to the sentence start and the
 *                     KV cache is recomputed from the prefix (simpler and
 *                     memory-safe vs. snapshotting GPU-resident tensors).
 *
 * The engine is framework-agnostic: it talks to a LoopModel interface so the
 * same loop runs against the real Transformers.js model in the worker and
 * against a deterministic mock model in node integration tests.
 */

import { RandomStream } from '../utils/rng';
import { endsSentence } from '../utils/sentences';
import type { CandidateInfo, StepTrace } from '../watermark/types';

/** Minimal model surface the loop needs (adapter over Transformers.js). */
export interface LoopModel {
  /**
   * Run one forward pass. `inputIds` are the tokens NOT yet in the cache.
   * Returns logits for the LAST position and an opaque cache handle.
   */
  forward(
    inputIds: number[],
    fullSeqLen: number,
    past: unknown,
  ): Promise<{ logits: Float32Array; past: unknown }>;
  /** Dispose a cache handle (GPU buffers). */
  disposePast(past: unknown): void;
  vocabSize: number;
  eosTokenIds: number[];
  decode(ids: number[]): string;
}

export interface GenerationHooks {
  /** Mutate logits in place before softmax; return value is trace metadata. */
  transformLogits?: (stepIndex: number, contextIds: number[], logits: Float32Array) => bigint | void;
  /** Replace sampling; receives post-temperature/top-p probabilities. */
  sampleOverride?: (
    stepIndex: number,
    contextIds: number[],
    probs: Float32Array,
  ) => { tokenId: number; keyId?: 1 | 2; r?: number; gumbelScore?: number };
  /** Sentence-level accept/reject; called when a sentence boundary appears. */
  onSentenceEnd?: (
    sentenceText: string,
    sentenceIndex: number,
    attempt: number,
  ) => Promise<{ accept: boolean }>;
  maxSentenceTrials?: number;
}

export interface LoopConfig {
  temperature: number;
  topP: number;
  maxNewTokens: number;
  /** Base seed for the shared sampling stream. */
  baseSeed: number;
  topKTrace: number; // candidates kept per step in the trace
  /**
   * Temperature escalation per rejected sentence attempt (k-SemStamp).
   * A peaked distribution re-sampled with a merely re-offset RNG stream tends
   * to reproduce the same sentence, which stalls rejection sampling; nudging
   * the temperature per attempt restores candidate diversity. Bounded by
   * `retryTemperatureMax`, because an unbounded climb buys diversity by
   * destroying the text - and gibberish embeds nowhere useful.
   */
  retryTemperatureStep?: number;
  /** Ceiling on the escalation multiplier (default 1.5x the base temperature). */
  retryTemperatureMax?: number;
  /**
   * Called once per accepted token, with the step that produced it. Carries
   * the step rather than the running text so callers can show the sampling as
   * it happens; `step.index` also lets them undo a sentence rollback, since
   * the next index after one is lower than the last they saw.
   */
  onToken?: (step: StepTrace) => void;
}

export interface LoopResult {
  tokenIds: number[]; // generated continuation only
  text: string;
  steps: StepTrace[];
  retries: number;
  aborted: boolean;
}

/** Numerically-stable softmax with temperature, in place into a new array. */
export function softmaxT(logits: Float32Array, temperature: number): Float32Array {
  const out = new Float32Array(logits.length);
  const t = Math.max(temperature, 1e-4);
  let max = -Infinity;
  for (let i = 0; i < logits.length; i++) if (logits[i] > max) max = logits[i];
  let sum = 0;
  for (let i = 0; i < logits.length; i++) {
    const e = Math.exp((logits[i] - max) / t);
    out[i] = e;
    sum += e;
  }
  for (let i = 0; i < out.length; i++) out[i] /= sum;
  return out;
}

/** Zero out everything outside the top-p nucleus, renormalize. */
export function applyTopP(probs: Float32Array, topP: number): void {
  if (topP >= 1) return;
  const idx = Array.from(probs.keys());
  idx.sort((a, b) => probs[b] - probs[a]);
  let cum = 0;
  let cut = idx.length;
  for (let i = 0; i < idx.length; i++) {
    cum += probs[idx[i]];
    if (cum >= topP) {
      cut = i + 1;
      break;
    }
  }
  const keep = new Set(idx.slice(0, cut));
  let sum = 0;
  for (let i = 0; i < probs.length; i++) {
    if (!keep.has(i)) probs[i] = 0;
    else sum += probs[i];
  }
  if (sum > 0) for (let i = 0; i < probs.length; i++) probs[i] /= sum;
}

/** Inverse-CDF sampling consuming exactly one uniform from the stream. */
export function sampleFromProbs(probs: Float32Array, u: number): number {
  let cum = 0;
  for (let i = 0; i < probs.length; i++) {
    cum += probs[i];
    if (u < cum) return i;
  }
  // Floating-point remainder: return last nonzero
  for (let i = probs.length - 1; i >= 0; i--) if (probs[i] > 0) return i;
  return 0;
}

/** Shannon entropy in nats. */
export function entropyOf(probs: Float32Array): number {
  let h = 0;
  for (let i = 0; i < probs.length; i++) {
    const p = probs[i];
    if (p > 0) h -= p * Math.log(p);
  }
  return h;
}

function topKCandidates(
  preLogits: Float32Array,
  preProbs: Float32Array,
  postProbs: Float32Array | null,
  k: number,
  decode: (ids: number[]) => string,
): CandidateInfo[] {
  const idx = Array.from(preProbs.keys());
  idx.sort((a, b) => preProbs[b] - preProbs[a]);
  return idx.slice(0, k).map((v) => ({
    tokenId: v,
    tokenText: decode([v]),
    logit: preLogits[v],
    prob: preProbs[v],
    probAfter: postProbs ? postProbs[v] : undefined,
  }));
}

export async function runGenerationLoop(
  model: LoopModel,
  promptIds: number[],
  cfg: LoopConfig,
  hooks: GenerationHooks = {},
): Promise<LoopResult> {
  const steps: StepTrace[] = [];
  const generated: number[] = [];
  const stream = new RandomStream(BigInt(cfg.baseSeed));
  let retryStreamSalt = 1;
  let retries = 0;

  let seq = [...promptIds];
  let past: unknown = null;
  let pending = [...promptIds]; // tokens not yet fed to the model
  let sentenceStartLen = 0; // in generated tokens
  let sentenceIndex = 0;
  let attempt = 1;
  let aborted = false;
  /**
   * Text decoded so far. A token is never decoded on its own: one character
   * can span several tokens (any non-ASCII script), and decoding a fragment
   * in isolation yields replacement characters. Each decode runs over the
   * prompt plus everything generated, and the token's text is what that adds
   * to the previous decode - so a character split across tokens arrives whole
   * with the token that completes it.
   */
  const promptText = model.decode(promptIds);
  let decodedSoFar = '';

  /** Continuation text for the given tokens, with the prompt sliced back off. */
  function continuationOf(ids: number[]): string {
    const whole = model.decode([...promptIds, ...ids]);
    return whole.startsWith(promptText) ? whole.slice(promptText.length) : model.decode(ids);
  }

  /**
   * A decode ending in U+FFFD means a character is still being assembled, so
   * that tail is not text yet - it is a promise the next token will keep.
   */
  const settled = (s: string) => s.replace(/�+$/, '');

  function textAddedBy(chosen: number): string {
    const before = settled(decodedSoFar);
    decodedSoFar = continuationOf([...generated, chosen]);
    const now = settled(decodedSoFar);
    // A token that only carries half a character adds nothing; the token that
    // completes it delivers the whole character at once.
    return now.startsWith(before) ? now.slice(before.length) : now;
  }

  const maxTrials = hooks.maxSentenceTrials ?? 12;

  try {
    while (generated.length < cfg.maxNewTokens) {
      const { logits, past: newPast } = await model.forward(pending, seq.length, past);
      past = newPast;
      pending = [];

      // Attempts after the first re-generate the same sentence prefix, so
      // raise the temperature to get genuinely different candidates.
      const temperature =
        cfg.temperature *
        Math.min(
          cfg.retryTemperatureMax ?? 1.5,
          1 + (cfg.retryTemperatureStep ?? 0) * (attempt - 1),
        );

      // Pre-watermark distribution (for trace + entropy)
      const preLogits = logits.slice();
      const preProbs = softmaxT(preLogits, temperature);
      const entropy = entropyOf(preProbs);

      const contextIds = seq;
      const seed = hooks.transformLogits?.(generated.length, contextIds, logits);

      let probs = softmaxT(logits, temperature);
      applyTopP(probs, cfg.topP);

      let chosen: number;
      let keyId: 1 | 2 | undefined;
      let r: number | undefined;
      let gumbelScore: number | undefined;
      if (hooks.sampleOverride) {
        const pick = hooks.sampleOverride(generated.length, contextIds, probs);
        chosen = pick.tokenId;
        keyId = pick.keyId;
        r = pick.r;
        gumbelScore = pick.gumbelScore;
      } else {
        chosen = sampleFromProbs(probs, stream.next());
      }

      const isEosToken = model.eosTokenIds.includes(chosen);
      const postProbs = seed !== undefined ? probs : null;
      const stepTrace: StepTrace = {
        index: generated.length,
        chosenTokenId: chosen,
        chosenTokenText: isEosToken ? '' : textAddedBy(chosen),
        seed: seed !== undefined && seed !== null ? String(seed) : undefined,
        keyId,
        topCandidates: topKCandidates(preLogits, preProbs, postProbs, cfg.topKTrace, (ids) =>
          model.decode(ids),
        ),
        entropy,
      };
      if (r !== undefined) {
        stepTrace.topCandidates.forEach((c) => {
          if (c.tokenId === chosen) {
            c.r = r;
            c.gumbelScore = gumbelScore;
          }
        });
      }

      const isEos = isEosToken;
      if (!isEos) {
        seq = [...seq, chosen];
        pending = [chosen];
        generated.push(chosen);
        steps.push(stepTrace);
        cfg.onToken?.(stepTrace);
      }

      // ---- sentence-level accept/reject (k-SemStamp) ----
      if (hooks.onSentenceEnd) {
        const sentText = model.decode(generated.slice(sentenceStartLen));
        const boundary = isEos || generated.length >= cfg.maxNewTokens || endsSentence(sentText);
        if (boundary && generated.length > sentenceStartLen) {
          const { accept } = await hooks.onSentenceEnd(sentText.trim(), sentenceIndex, attempt);
          if (accept || attempt >= maxTrials) {
            sentenceStartLen = generated.length;
            sentenceIndex++;
            attempt = 1;
          } else {
            // Reject: roll back to sentence start; recompute cache from prefix.
            retries++;
            attempt++;
            const keep = generated.slice(0, sentenceStartLen);
            const removedSteps = generated.length - sentenceStartLen;
            generated.length = sentenceStartLen;
            steps.length = steps.length - removedSteps;
            decodedSoFar = continuationOf(generated);
            seq = [...promptIds, ...keep];
            model.disposePast(past);
            past = null;
            pending = [...seq];
            // Perturb the sampling stream so the retry differs.
            for (let i = 0; i < retryStreamSalt; i++) stream.next();
            retryStreamSalt++;
            if (isEos) continue;
          }
        }
      }

      if (isEos) break;
    }
  } catch (e) {
    aborted = true;
    throw e;
  } finally {
    model.disposePast(past);
  }

  return {
    tokenIds: generated,
    // Same decode the per-token texts were derived from, so they concatenate
    // to exactly this string.
    text: continuationOf(generated),
    steps,
    retries,
    aborted,
  };
}