src/engine/runtime-score-sequence.test.ts

import { describe, expect, it, vi } from 'vitest';
import type { BackendReport } from './native-log';
import { BrowserEngineRuntime } from './runtime';

interface RawCompletionOptions {
  prompt: number[];
  max_tokens: number;
  temperature: number;
  top_k: number;
  logprobs: number;
  logit_bias: Record<string, number>;
  cache_prompt: boolean;
  post_sampling_probs: boolean;
  abortSignal: AbortSignal;
}

interface RuntimeInternals {
  wllama: {
    isModelLoaded(): boolean;
    createCompletion(options: RawCompletionOptions): Promise<unknown>;
  } | null;
  loaded: {
    manifest: unknown;
    backend: 'webgpu';
    tuningScope: 'benchmark';
    contextSize: number;
    batchSize: number;
    microBatchSize: number;
    vocabularySize: number;
    model: {
      id: '27b';
      displayName: string;
      cpuFallback: false;
      runtimePolicy: {
        flashAttention: false;
        tokenEmbeddingOnWebGPU: true;
        requireSingleWebGPUGraph: true;
      };
    };
  } | null;
  nativeLog: {
    report(): BackendReport;
  };
}

const backendReport: BackendReport = {
  backends: ['WebGPU'],
  nGraphSplits: 1,
  opsOnCpu: 0,
  layersGpu: { offloaded: 65, total: 65 },
  flashAttention: false,
  cacheTypeK: 'f16',
  cacheTypeV: 'f16',
  webgpuKvBufferBytes: 128 * 1024 ** 2,
  webgpuTrace: [],
};

function configuredRuntime(createCompletion: (options: RawCompletionOptions) => Promise<unknown>) {
  const runtime = new BrowserEngineRuntime();
  const internals = runtime as unknown as RuntimeInternals;
  internals.loaded = {
    manifest: {},
    backend: 'webgpu',
    tuningScope: 'benchmark',
    contextSize: 2_048,
    batchSize: 32,
    microBatchSize: 16,
    vocabularySize: 248_320,
    model: {
      id: '27b',
      displayName: 'Fixture Bonsai 27B',
      cpuFallback: false,
      runtimePolicy: {
        flashAttention: false,
        tokenEmbeddingOnWebGPU: true,
        requireSingleWebGPUGraph: true,
      },
    },
  };
  internals.wllama = { isModelLoaded: () => true, createCompletion };
  vi.spyOn(internals.nativeLog, 'report').mockReturnValue(backendReport);
  return runtime;
}

describe('BrowserEngineRuntime teacher-forced scoring', () => {
  it('scores the exact CPU sequence with raw token prefixes and keeps natural top-1 separate', async () => {
    const calls: RawCompletionOptions[] = [];
    const createCompletion = vi.fn(async (options: RawCompletionOptions) => {
      calls.push(options);
      const index = options.prompt.length - 38;
      const referenceId = Number(Object.keys(options.logit_bias)[0]);
      const naturalTop1Id = index === 29 ? referenceId + 10 : referenceId;
      const candidates = index === 29
        ? [
            { id: naturalTop1Id, token: 'natural', logprob: -0.01, bytes: null },
            { id: referenceId, token: 'reference', logprob: -0.02, bytes: null },
            { id: referenceId + 20, token: 'third', logprob: -1, bytes: null },
            { id: referenceId + 21, token: 'fourth', logprob: -2, bytes: null },
            { id: referenceId + 22, token: 'fifth', logprob: -3, bytes: null },
          ]
        : [
            { id: referenceId, token: 'reference', logprob: -0.01, bytes: null },
            { id: referenceId + 10, token: 'second', logprob: -0.02, bytes: null },
            { id: referenceId + 20, token: 'third', logprob: -1, bytes: null },
            { id: referenceId + 21, token: 'fourth', logprob: -2, bytes: null },
            { id: referenceId + 22, token: 'fifth', logprob: -3, bytes: null },
          ];
      return {
        choices: [{
          text: 'forced',
          finish_reason: 'length',
          logprobs: {
            content: [{
              id: referenceId,
              token: 'reference',
              logprob: index === 29 ? -0.02 : -0.01,
              bytes: null,
              top_logprobs: candidates,
            }],
          },
        }],
      };
    });
    const runtime = configuredRuntime(createCompletion);
    const promptTokenIds = Array.from({ length: 38 }, (_, index) => index + 1);
    const referenceTokenIds = Array.from({ length: 1_024 }, (_, index) => index + 1_000);

    const result = await runtime.scoreSequence({
      promptTokenIds,
      referenceTokenIds,
      topK: 5,
    }, new AbortController().signal);

    expect(createCompletion).toHaveBeenCalledTimes(1_024);
    expect(calls[0]).toMatchObject({
      prompt: promptTokenIds,
      max_tokens: 1,
      temperature: 0,
      top_k: 1,
      logprobs: 5,
      logit_bias: { '1000': 1_000 },
      cache_prompt: false,
      post_sampling_probs: false,
    });
    expect(calls[1]).toMatchObject({
      prompt: [...promptTokenIds, 1_000],
      cache_prompt: true,
    });
    expect(calls.at(-1)?.prompt).toEqual([
      ...promptTokenIds,
      ...referenceTokenIds.slice(0, -1),
    ]);
    expect(result.entries[29]).toMatchObject({
      index: 29,
      selectedReference: { id: 1_029, logprob: -0.02 },
      naturalTop1: { id: 1_039, logprob: -0.01 },
      referenceRankInTopCandidatesZeroBased: 1,
      top1Top2Margin: 0.01,
    });
    expect(result.summary.tokenCount).toBe(1_024);
    expect(result.summary.meanNll).toBeCloseTo((1_023 * 0.01 + 0.02) / 1_024, 12);
    expect(result.summary.perplexity).toBeCloseTo(Math.exp(result.summary.meanNll), 12);
  });

  it('honors an already-aborted diagnostic request before the first raw completion', async () => {
    const createCompletion = vi.fn(async () => ({}));
    const runtime = configuredRuntime(createCompletion);
    const controller = new AbortController();
    controller.abort();

    await expect(runtime.scoreSequence({
      promptTokenIds: Array.from({ length: 38 }, (_, index) => index + 1),
      referenceTokenIds: Array.from({ length: 1_024 }, (_, index) => index + 1_000),
      topK: 5,
    }, controller.signal)).rejects.toMatchObject({ name: 'AbortError' });
    expect(createCompletion).not.toHaveBeenCalled();
  });

  it('fails loudly when logit bias does not return the fixed reference token', async () => {
    const runtime = configuredRuntime(async (options) => {
      const referenceId = Number(Object.keys(options.logit_bias)[0]);
      const selectedId = referenceId + 1;
      return {
        choices: [{
          logprobs: {
            content: [{
              id: selectedId,
              logprob: -0.01,
              top_logprobs: [
                { id: selectedId, logprob: -0.01 },
                { id: referenceId, logprob: -0.02 },
                { id: referenceId + 2, logprob: -1 },
                { id: referenceId + 3, logprob: -2 },
                { id: referenceId + 4, logprob: -3 },
              ],
            }],
          },
        }],
      };
    });

    await expect(runtime.scoreSequence({
      promptTokenIds: Array.from({ length: 38 }, (_, index) => index + 1),
      referenceTokenIds: Array.from({ length: 1_024 }, (_, index) => index + 1_000),
      topK: 5,
    }, new AbortController().signal)).rejects.toMatchObject({
      code: 'INVALID_SCORE_SEQUENCE_RESPONSE',
      details: { index: 0, referenceTokenId: 1_000 },
    });
  });

  it('rejects scoring outside the loaded 27B WebGPU benchmark path', async () => {
    const runtime = configuredRuntime(async () => ({}));
    const internals = runtime as unknown as RuntimeInternals;
    if (!internals.loaded) throw new Error('Expected loaded fixture state.');
    (internals.loaded as { tuningScope: string }).tuningScope = 'release-defaults';

    await expect(runtime.scoreSequence({
      promptTokenIds: Array.from({ length: 38 }, (_, index) => index + 1),
      referenceTokenIds: Array.from({ length: 1_024 }, (_, index) => index + 1_000),
      topK: 5,
    }, new AbortController().signal)).rejects.toMatchObject({
      code: 'SCORE_SEQUENCE_UNAVAILABLE',
    });
  });
});