From 37d11bf9b87ab0252772126a61d56d525d27c6c7 Mon Sep 17 00:00:00 2001 From: Richard R Date: Thu, 21 May 2026 21:15:10 -0600 Subject: [PATCH] refactor(compute): remove legacy app whisper duplicate and unify core whisper boundary --- Dockerfile | 2 +- compute/core/package.json | 1 + compute/core/src/whisper/index.ts | 18 + compute/worker/src/server.ts | 19 +- src/lib/server/whisper/alignment-mapping.ts | 46 - src/lib/server/whisper/alignment.ts | 1012 ------------------ src/lib/server/whisper/ensureModel.ts | 240 ----- src/lib/server/whisper/model/LICENSE.txt | 21 - src/lib/server/whisper/model/manifest.json | 76 -- src/lib/server/whisper/model/mel_filters.npz | Bin 4271 -> 0 bytes src/lib/server/whisper/spectral.ts | 21 - src/lib/server/whisper/token-timestamps.ts | 449 -------- tests/unit/whisper-alignment-mapping.spec.ts | 2 +- tests/unit/whisper-alignment-smoke.spec.ts | 2 +- tests/unit/whisper-ensure-model.spec.ts | 2 +- tests/unit/whisper-spectral.spec.ts | 2 +- tests/unit/whisper-token-timestamps.spec.ts | 2 +- tsconfig.json | 1 + 18 files changed, 29 insertions(+), 1887 deletions(-) create mode 100644 compute/core/src/whisper/index.ts delete mode 100644 src/lib/server/whisper/alignment-mapping.ts delete mode 100644 src/lib/server/whisper/alignment.ts delete mode 100644 src/lib/server/whisper/ensureModel.ts delete mode 100644 src/lib/server/whisper/model/LICENSE.txt delete mode 100644 src/lib/server/whisper/model/manifest.json delete mode 100644 src/lib/server/whisper/model/mel_filters.npz delete mode 100644 src/lib/server/whisper/spectral.ts delete mode 100644 src/lib/server/whisper/token-timestamps.ts diff --git a/Dockerfile b/Dockerfile index 6ba233b..612e250 100644 --- a/Dockerfile +++ b/Dockerfile @@ -65,7 +65,7 @@ COPY --from=seaweedfs-builder /tmp/weed /usr/local/bin/weed RUN chmod +x /usr/local/bin/weed # Include OpenAI Whisper license text for runtime-downloaded ONNX artifacts. -COPY --from=app-builder /app/src/lib/server/whisper/model/LICENSE.txt /licenses/openai-whisper-LICENSE.txt +COPY --from=app-builder /app/compute/core/src/whisper/model/LICENSE.txt /licenses/openai-whisper-LICENSE.txt # Expose the port the app runs on EXPOSE 3003 diff --git a/compute/core/package.json b/compute/core/package.json index 21a8f2e..0874c4c 100644 --- a/compute/core/package.json +++ b/compute/core/package.json @@ -16,6 +16,7 @@ "./contracts": "./src/contracts.ts", "./local-runtime": "./src/local-runtime.ts", "./runtime": "./src/runtime/index.ts", + "./whisper": "./src/whisper/index.ts", "./pdf-layout": "./src/pdf-layout/index.ts", "./pdf-layout/parse": "./src/pdf-layout/parsePdf.ts" } diff --git a/compute/core/src/whisper/index.ts b/compute/core/src/whisper/index.ts new file mode 100644 index 0000000..ed37c33 --- /dev/null +++ b/compute/core/src/whisper/index.ts @@ -0,0 +1,18 @@ +export { + alignAudioWithText, + makeWhisperCacheKey, + type WhisperRequestBody, +} from './alignment'; + +export { + ensureWhisperModel, + ensureWhisperArtifacts, + createSingleflightRunner, + type WhisperArtifactSpec, + type WhisperStaticArtifactSpec, + type WhisperFetch, +} from './ensureModel'; + +export { mapWordsToSentenceOffsets, type WhisperWord } from './alignment-mapping'; +export { buildGoertzelCoefficients, goertzelPower } from './spectral'; +export { buildWordsFromTimestampedTokens, extractTokenStartTimestamps } from './token-timestamps'; diff --git a/compute/worker/src/server.ts b/compute/worker/src/server.ts index 4d89675..6ceb6f7 100644 --- a/compute/worker/src/server.ts +++ b/compute/worker/src/server.ts @@ -1,5 +1,4 @@ import { createHash } from 'node:crypto'; -import os from 'node:os'; import Fastify, { type FastifyRequest } from 'fastify'; import { z } from 'zod'; import { @@ -29,6 +28,8 @@ import { } from '@openreader/compute-core/local-runtime'; import { getComputeTimeoutConfig, + getAvailableCpuCores, + getOnnxThreadsPerJob, withIdleTimeoutAndHardCap, withTimeout, } from '@openreader/compute-core/runtime'; @@ -117,20 +118,6 @@ function parseBoolEnv(name: string, fallback: boolean): boolean { return normalized === '1' || normalized === 'true' || normalized === 'yes' || normalized === 'on'; } -function getAvailableCpuCores(): number { - if (typeof os.availableParallelism === 'function') { - const value = os.availableParallelism(); - if (Number.isFinite(value) && value >= 1) return Math.floor(value); - } - const fallback = os.cpus().length; - return Number.isFinite(fallback) && fallback >= 1 ? Math.floor(fallback) : 1; -} - -function getOnnxThreadsPerJob(jobConcurrency: number): number { - const usableCores = Math.max(1, getAvailableCpuCores() - 1); - return Math.max(1, Math.floor(usableCores / Math.max(1, jobConcurrency))); -} - function buildLoggerConfig(): boolean | Record { const format = (process.env.COMPUTE_LOG_FORMAT?.trim().toLowerCase() || 'pretty'); if (format === 'json') return true; @@ -528,7 +515,7 @@ async function main(): Promise { pdfAttempts, opStaleMs, availableCpuCores: getAvailableCpuCores(), - onnxThreadsPerJob: getOnnxThreadsPerJob(jobConcurrency), + onnxThreadsPerJob: getOnnxThreadsPerJob(), natsApiTimeoutMs: NATS_API_TIMEOUT_MS, pdfLayoutHardCapMs: pdfHardCapMs, }, 'compute runtime config'); diff --git a/src/lib/server/whisper/alignment-mapping.ts b/src/lib/server/whisper/alignment-mapping.ts deleted file mode 100644 index 52c5707..0000000 --- a/src/lib/server/whisper/alignment-mapping.ts +++ /dev/null @@ -1,46 +0,0 @@ -import type { TTSSentenceAlignment, TTSSentenceWord } from '@/types/tts'; -import { preprocessSentenceForAudio } from '@/lib/shared/nlp'; - -export interface WhisperWord { - start: number; - end: number; - word: string; -} - -export function mapWordsToSentenceOffsets(sentence: string, words: WhisperWord[]): TTSSentenceAlignment { - const normalizedSentence = preprocessSentenceForAudio(sentence); - const lowerSentence = normalizedSentence.toLowerCase(); - let cursor = 0; - - const alignedWords: TTSSentenceWord[] = words.map((w) => { - const token = w.word.trim(); - if (!token) { - return { - text: '', - startSec: w.start, - endSec: w.end, - charStart: cursor, - charEnd: cursor, - }; - } - - const idx = lowerSentence.indexOf(token.toLowerCase(), cursor); - const start = idx >= 0 ? idx : cursor; - const end = Math.min(normalizedSentence.length, start + token.length); - cursor = Math.max(cursor, end); - - return { - text: token, - startSec: w.start, - endSec: w.end, - charStart: start, - charEnd: end, - }; - }).filter((word) => word.text.length > 0); - - return { - sentence, - sentenceIndex: 0, - words: alignedWords, - }; -} diff --git a/src/lib/server/whisper/alignment.ts b/src/lib/server/whisper/alignment.ts deleted file mode 100644 index 713c87b..0000000 --- a/src/lib/server/whisper/alignment.ts +++ /dev/null @@ -1,1012 +0,0 @@ -import { createHash, randomUUID } from 'crypto'; -import { mkdtemp, readFile, rm, writeFile } from 'fs/promises'; -import { tmpdir } from 'os'; -import { join } from 'path'; -import { spawn } from 'child_process'; -import * as ort from 'onnxruntime-node'; -import { Tokenizer } from '@huggingface/tokenizers'; -import JSZip from 'jszip'; -import type { TTSAudioBuffer, TTSAudioBytes, TTSSentenceAlignment } from '@/types/tts'; -import { getFFmpegPath } from '@/lib/server/audiobooks/ffmpeg-bin'; -import { getOnnxThreadsPerJob } from '@openreader/compute-core/runtime'; -import { - mapWordsToSentenceOffsets, - type WhisperWord, -} from '@/lib/server/whisper/alignment-mapping'; -import { buildGoertzelCoefficients, goertzelPower } from '@/lib/server/whisper/spectral'; -import { - buildWordsFromTimestampedTokens, - extractTokenStartTimestamps, -} from '@/lib/server/whisper/token-timestamps'; -import { - ensureWhisperModel, - WHISPER_CONFIG_PATH, - WHISPER_GENERATION_CONFIG_PATH, - WHISPER_TOKENIZER_CONFIG_PATH, - WHISPER_TOKENIZER_PATH, - WHISPER_ENCODER_MODEL_PATH, - WHISPER_DECODER_MERGED_MODEL_PATH, - WHISPER_DECODER_WITH_PAST_MODEL_PATH, -} from '@/lib/server/whisper/ensureModel'; - -interface WhisperAlignmentOptions { - lang?: string; - textHint?: string; -} - -export interface WhisperRequestBody { - text: string; - audio: TTSAudioBytes; - lang?: string; -} - -interface WhisperRuntime { - encoder: ort.InferenceSession; - decoderMerged: ort.InferenceSession; - decoderWithPast: ort.InferenceSession; - tokenizer: Tokenizer; - promptStartToken: number; - defaultLanguageToken: number; - transcribeToken: number; - eosTokenId: number; - noTimestampsTokenId: number; - timestampBeginTokenId: number; - maxInitialTimestampIndex: number; - maxDecodeSteps: number; - suppressTokens: Set; - beginSuppressTokens: Set; - alignmentHeads: Array<[number, number]>; - prefillFetches: string[]; - stepFetches: string[]; -} - -type WhisperAlignmentState = { - alignmentCache: Map; - alignmentInFlight: Map>; - runtimePromise: Promise | null; - alignMutex: Promise; - pendingAlignments: number; - officialMelFilters: Float32Array[] | null; - emptyPastFeedsTemplate: Record | null; -}; - -const WHISPER_ALIGNMENT_STATE_KEY = '__openreaderWhisperAlignmentStateV1'; -const g = globalThis as typeof globalThis & Record; -const state = (() => { - const existing = g[WHISPER_ALIGNMENT_STATE_KEY] as WhisperAlignmentState | undefined; - if (existing) return existing; - const created: WhisperAlignmentState = { - alignmentCache: new Map(), - alignmentInFlight: new Map>(), - runtimePromise: null, - alignMutex: Promise.resolve(), - pendingAlignments: 0, - officialMelFilters: null, - emptyPastFeedsTemplate: null, - }; - g[WHISPER_ALIGNMENT_STATE_KEY] = created; - return created; -})(); -const alignmentCache = state.alignmentCache; -const alignmentInFlight = state.alignmentInFlight; -const ALIGNMENT_CACHE_MAX_ENTRIES = 256; -const MAX_DECODE_STEPS_CAP = 128; -const ALIGNMENT_TIMEOUT_MS = 25000; -const FFMPEG_DECODE_TIMEOUT_MS = 10000; - -const SAMPLE_RATE = 16000; -const N_FFT = 400; -const HOP_LENGTH = 160; -const CHUNK_LENGTH_SECONDS = 30; -const N_SAMPLES = CHUNK_LENGTH_SECONDS * SAMPLE_RATE; -const N_FRAMES = N_SAMPLES / HOP_LENGTH; -const N_MELS = 80; -const WHISPER_NUM_HEADS = 8; -const WHISPER_HEAD_DIM = 64; -const WHISPER_NUM_LAYERS = 6; -const MEL_FILTER_BINS = (N_FFT / 2) + 1; - -const hannWindow = buildHannWindow(N_FFT); -const goertzelCoefficients = buildGoertzelCoefficients(MEL_FILTER_BINS, N_FFT); - -const MEL_FILTERS_NPZ_PATH = join(process.cwd(), 'src/lib/server/whisper/model/mel_filters.npz'); - -function buildHannWindow(length: number): Float32Array { - const window = new Float32Array(length); - for (let i = 0; i < length; i += 1) { - window[i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / length); - } - return window; -} - -function parseNpyFloat32(bytes: Uint8Array): { shape: number[]; data: Float32Array } { - if (bytes.length < 12) { - throw new Error('Invalid NPY payload: too short'); - } - const magic = String.fromCharCode(...bytes.slice(0, 6)); - if (magic !== '\u0093NUMPY') { - throw new Error('Invalid NPY payload: missing magic header'); - } - - const major = bytes[6]; - const headerLength = major <= 1 - ? new DataView(bytes.buffer, bytes.byteOffset + 8, 2).getUint16(0, true) - : new DataView(bytes.buffer, bytes.byteOffset + 8, 4).getUint32(0, true); - const headerOffset = major <= 1 ? 10 : 12; - const header = Buffer.from(bytes.slice(headerOffset, headerOffset + headerLength)).toString('latin1'); - - const descrMatch = header.match(/'descr':\s*'([^']+)'/); - if (!descrMatch || descrMatch[1] !== ' token.trim()) - .filter(Boolean) - .map((token) => Number(token)) - .filter((n) => Number.isFinite(n) && n > 0); - - const dataOffset = headerOffset + headerLength; - const dataBytes = bytes.slice(dataOffset); - const totalFloats = Math.floor(dataBytes.byteLength / 4); - const data = new Float32Array(totalFloats); - const view = new DataView(dataBytes.buffer, dataBytes.byteOffset, dataBytes.byteLength); - for (let i = 0; i < totalFloats; i += 1) { - data[i] = view.getFloat32(i * 4, true); - } - - return { shape, data }; -} - -async function loadOfficialMelFilters(): Promise { - if (state.officialMelFilters) return state.officialMelFilters; - - const npzBytes = await readFile(MEL_FILTERS_NPZ_PATH); - const zip = await JSZip.loadAsync(npzBytes); - const mel80 = zip.file('mel_80.npy'); - if (!mel80) { - throw new Error('OpenAI mel filter asset is missing mel_80.npy'); - } - - const raw = await mel80.async('uint8array'); - const parsed = parseNpyFloat32(raw); - const [rows, cols] = parsed.shape; - if (rows !== N_MELS || cols !== MEL_FILTER_BINS) { - throw new Error(`Unexpected mel filter shape: [${rows}, ${cols}]`); - } - - const filters: Float32Array[] = []; - for (let row = 0; row < rows; row += 1) { - const start = row * cols; - filters.push(parsed.data.slice(start, start + cols)); - } - - state.officialMelFilters = filters; - return filters; -} - -function pcm16ToFloat32(buffer: Buffer): Float32Array { - const view = new Int16Array(buffer.buffer, buffer.byteOffset, Math.floor(buffer.byteLength / 2)); - const out = new Float32Array(view.length); - for (let i = 0; i < view.length; i += 1) { - out[i] = view[i] / 32768; - } - return out; -} - -function padOrTrimAudio(samples: Float32Array): Float32Array { - if (samples.length === N_SAMPLES) return samples; - if (samples.length > N_SAMPLES) return samples.subarray(0, N_SAMPLES); - - const padded = new Float32Array(N_SAMPLES); - padded.set(samples, 0); - return padded; -} - -function reflectPad(audio: Float32Array, pad: number): Float32Array { - const out = new Float32Array(audio.length + (2 * pad)); - out.set(audio, pad); - - // Match PyTorch reflect padding (exclude edge sample). - for (let i = 0; i < pad; i += 1) { - out[pad - 1 - i] = audio[Math.min(audio.length - 1, i + 1)]; - out[pad + audio.length + i] = audio[Math.max(0, audio.length - 2 - i)]; - } - - return out; -} - -function computeLogMelSpectrogram(audioSamples: Float32Array): ort.Tensor { - if (!state.officialMelFilters) { - throw new Error('Whisper mel filters not loaded'); - } - - const paddedAudio = reflectPad(audioSamples, N_FFT / 2); - const stftFrames = N_FRAMES + 1; - const frameCount = N_FRAMES; - const freqBins = MEL_FILTER_BINS; - - const melSpec = Array.from({ length: N_MELS }, () => new Float32Array(frameCount)); - const frame = new Float32Array(N_FFT); - const power = new Float32Array(freqBins); - - for (let frameIndex = 0; frameIndex < stftFrames; frameIndex += 1) { - const offset = frameIndex * HOP_LENGTH; - - for (let i = 0; i < N_FFT; i += 1) { - frame[i] = (paddedAudio[offset + i] ?? 0) * hannWindow[i]; - } - - for (let k = 0; k < freqBins; k += 1) { - power[k] = goertzelPower(frame, goertzelCoefficients[k]); - } - - if (frameIndex === stftFrames - 1) { - continue; - } - - for (let melIndex = 0; melIndex < N_MELS; melIndex += 1) { - const filter = state.officialMelFilters[melIndex]; - let total = 0; - for (let k = 0; k < freqBins; k += 1) { - total += filter[k] * power[k]; - } - melSpec[melIndex][frameIndex] = total; - } - } - - // Whisper normalization from openai/whisper/audio.py - let globalMaxLog = Number.NEGATIVE_INFINITY; - for (let i = 0; i < N_MELS; i += 1) { - for (let j = 0; j < frameCount; j += 1) { - const logVal = Math.log10(Math.max(1e-10, melSpec[i][j])); - if (logVal > globalMaxLog) globalMaxLog = logVal; - melSpec[i][j] = logVal; - } - } - - const floorVal = globalMaxLog - 8.0; - const flattened = new Float32Array(1 * N_MELS * frameCount); - for (let i = 0; i < N_MELS; i += 1) { - for (let j = 0; j < frameCount; j += 1) { - const clamped = Math.max(melSpec[i][j], floorVal); - flattened[(i * frameCount) + j] = (clamped + 4.0) / 4.0; - } - } - - return new ort.Tensor('float32', flattened, [1, N_MELS, frameCount]); -} - -async function decodeToPcm16(inputPath: string, outputPath: string): Promise { - await new Promise((resolve, reject) => { - const ffmpeg = spawn(getFFmpegPath(), [ - '-y', - '-i', - inputPath, - '-f', - 's16le', - '-ar', - String(SAMPLE_RATE), - '-ac', - '1', - outputPath, - ]); - - let stderr = ''; - let timedOut = false; - const timer = setTimeout(() => { - timedOut = true; - ffmpeg.kill('SIGKILL'); - }, FFMPEG_DECODE_TIMEOUT_MS); - ffmpeg.stderr.on('data', (data) => { - stderr += data.toString(); - }); - - ffmpeg.on('error', (err) => { - clearTimeout(timer); - reject(err); - }); - - ffmpeg.on('close', (code) => { - clearTimeout(timer); - if (timedOut) { - reject(new Error(`ffmpeg decode timed out after ${FFMPEG_DECODE_TIMEOUT_MS}ms`)); - return; - } - if (code === 0) { - resolve(); - } else { - reject(new Error(`ffmpeg decode failed with code ${code}: ${stderr}`)); - } - }); - }); -} - -function parseLanguageCode(lang?: string): string | null { - if (!lang) return null; - const trimmed = lang.trim().toLowerCase(); - if (!trimmed) return null; - if (trimmed.includes('-')) return trimmed.split('-')[0] || null; - if (trimmed.includes('_')) return trimmed.split('_')[0] || null; - return trimmed; -} - -function tensorFromInt64(values: number[]): ort.Tensor { - return new ort.Tensor('int64', BigInt64Array.from(values.map((v) => BigInt(v))), [1, values.length]); -} - -function disposeTensor(tensor: ort.Tensor | undefined | null): void { - if (!tensor) return; - try { - tensor.dispose(); - } catch { - // Best-effort cleanup: ignore disposal errors during fallback path. - } -} - -function disposeTensorMap(tensors: Record): void { - for (const tensor of Object.values(tensors)) { - disposeTensor(tensor); - } -} - -function computeAdaptiveDecodeStepLimit(maxDecodeSteps: number, textHint?: string): number { - const normalized = (textHint ?? '').trim(); - if (!normalized) return Math.min(maxDecodeSteps, 96); - - const chars = normalized.length; - const words = normalized.split(/\s+/).filter(Boolean).length; - const estTokens = Math.max(words * 3, Math.ceil(chars / 2)); - const adaptive = Math.max(64, Math.min(maxDecodeSteps, estTokens + 24)); - return adaptive; -} - -function assertWithinDeadline(deadlineMs: number): void { - if (Date.now() > deadlineMs) { - throw new Error(`Whisper alignment timed out after ${ALIGNMENT_TIMEOUT_MS}ms`); - } -} - -function makeInFlightCoalesceKey(audioBuffer: TTSAudioBuffer, text: string, lang?: string): string { - const bytes = new Uint8Array(audioBuffer); - const span = 4096; - const head = bytes.subarray(0, Math.min(span, bytes.length)); - const tailStart = Math.max(0, bytes.length - span); - const tail = bytes.subarray(tailStart); - return createHash('sha256') - .update(text) - .update('\0') - .update(lang ?? '') - .update('\0') - .update(String(bytes.length)) - .update('\0') - .update(head) - .update('\0') - .update(tail) - .digest('hex'); -} - -function buildEmptyPastFeeds() { - if (state.emptyPastFeedsTemplate) return state.emptyPastFeedsTemplate; - - const feeds: Record = {}; - const emptyDecoderPast = new Float32Array(0); - const emptyEncoderPast = new Float32Array(1 * WHISPER_NUM_HEADS * 1500 * WHISPER_HEAD_DIM); - - for (let i = 0; i < WHISPER_NUM_LAYERS; i += 1) { - feeds[`past_key_values.${i}.decoder.key`] = new ort.Tensor('float32', emptyDecoderPast, [1, WHISPER_NUM_HEADS, 0, WHISPER_HEAD_DIM]); - feeds[`past_key_values.${i}.decoder.value`] = new ort.Tensor('float32', emptyDecoderPast, [1, WHISPER_NUM_HEADS, 0, WHISPER_HEAD_DIM]); - - // First pass still expects encoder KV inputs in the merged decoder graph. - feeds[`past_key_values.${i}.encoder.key`] = new ort.Tensor('float32', emptyEncoderPast, [1, WHISPER_NUM_HEADS, 1500, WHISPER_HEAD_DIM]); - feeds[`past_key_values.${i}.encoder.value`] = new ort.Tensor('float32', emptyEncoderPast, [1, WHISPER_NUM_HEADS, 1500, WHISPER_HEAD_DIM]); - } - - state.emptyPastFeedsTemplate = feeds; - return state.emptyPastFeedsTemplate; -} - -function argmax(values: Float32Array): number | null { - let bestIdx = 0; - let bestScore = Number.NEGATIVE_INFINITY; - - for (let i = 0; i < values.length; i += 1) { - const score = values[i]; - if (score > bestScore) { - bestScore = score; - bestIdx = i; - } - } - - return Number.isFinite(bestScore) ? bestIdx : null; -} - -function applyTokenSuppression(logits: Float32Array, tokens: Set) { - for (const tokenId of tokens) { - if (tokenId >= 0 && tokenId < logits.length) { - logits[tokenId] = Number.NEGATIVE_INFINITY; - } - } -} - -function logSoftmax(input: Float32Array): Float32Array { - let max = Number.NEGATIVE_INFINITY; - for (let i = 0; i < input.length; i += 1) { - if (input[i] > max) max = input[i]; - } - if (!Number.isFinite(max)) { - return new Float32Array(input.length).fill(Number.NEGATIVE_INFINITY); - } - - let sum = 0; - for (let i = 0; i < input.length; i += 1) { - sum += Math.exp(input[i] - max); - } - const logSum = Math.log(sum); - - const out = new Float32Array(input.length); - for (let i = 0; i < input.length; i += 1) { - out[i] = input[i] - max - logSum; - } - return out; -} - -function applyWhisperTimestampLogitsRules(input: { - logits: Float32Array; - generated: number[]; - beginIndex: number; - eosTokenId: number; - noTimestampsTokenId: number; - timestampBeginTokenId: number; - maxInitialTimestampIndex: number; -}) { - const { - logits, - generated, - beginIndex, - eosTokenId, - noTimestampsTokenId, - timestampBeginTokenId, - maxInitialTimestampIndex, - } = input; - - if (noTimestampsTokenId >= 0 && noTimestampsTokenId < logits.length) { - logits[noTimestampsTokenId] = Number.NEGATIVE_INFINITY; - } - - if (generated.length === beginIndex) { - const upper = Math.min(timestampBeginTokenId, logits.length); - for (let i = 0; i < upper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } - - const seq = generated.slice(beginIndex); - const lastWasTimestamp = seq.length >= 1 && seq[seq.length - 1] >= timestampBeginTokenId; - const penultimateWasTimestamp = seq.length < 2 || seq[seq.length - 2] >= timestampBeginTokenId; - - if (lastWasTimestamp) { - if (penultimateWasTimestamp) { - for (let i = timestampBeginTokenId; i < logits.length; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } else { - const upper = Math.min(eosTokenId, logits.length); - for (let i = 0; i < upper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } - } - - if (generated.length === beginIndex && Number.isFinite(maxInitialTimestampIndex)) { - const lastAllowed = timestampBeginTokenId + maxInitialTimestampIndex; - for (let i = lastAllowed + 1; i < logits.length; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } - - const textUpper = Math.min(timestampBeginTokenId, logits.length); - if (textUpper <= 0 || textUpper >= logits.length) return; - - const logprobs = logSoftmax(logits); - - let maxTextTokenLogprob = Number.NEGATIVE_INFINITY; - for (let i = 0; i < textUpper; i += 1) { - if (logprobs[i] > maxTextTokenLogprob) maxTextTokenLogprob = logprobs[i]; - } - - let timestampProbMass = 0; - for (let i = textUpper; i < logprobs.length; i += 1) { - timestampProbMass += Math.exp(logprobs[i]); - } - const timestampLogprob = timestampProbMass > 0 ? Math.log(timestampProbMass) : Number.NEGATIVE_INFINITY; - - if (timestampLogprob > maxTextTokenLogprob) { - for (let i = 0; i < textUpper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } -} - -async function getRuntime(): Promise { - if (state.runtimePromise) return state.runtimePromise; - - state.runtimePromise = (async () => { - await ensureWhisperModel(); - await loadOfficialMelFilters(); - - const [configRaw, generationRaw, tokenizerJsonRaw, tokenizerConfigRaw] = await Promise.all([ - readFile(WHISPER_CONFIG_PATH, 'utf8'), - readFile(WHISPER_GENERATION_CONFIG_PATH, 'utf8'), - readFile(WHISPER_TOKENIZER_PATH, 'utf8'), - readFile(WHISPER_TOKENIZER_CONFIG_PATH, 'utf8'), - ]); - - const config = JSON.parse(configRaw) as { - decoder_start_token_id?: number; - eos_token_id?: number; - forced_decoder_ids?: Array<[number, number | null]>; - }; - - const generationConfig = JSON.parse(generationRaw) as { - no_timestamps_token_id?: number; - max_initial_timestamp_index?: number; - suppress_tokens?: number[]; - begin_suppress_tokens?: number[]; - max_length?: number; - alignment_heads?: Array<[number, number]>; - }; - - const tokenizer = new Tokenizer(JSON.parse(tokenizerJsonRaw), JSON.parse(tokenizerConfigRaw)); - - const promptStartToken = Number(config.decoder_start_token_id ?? 50258); - const eosTokenId = Number(config.eos_token_id ?? 50257); - const noTimestampsTokenId = Number(generationConfig.no_timestamps_token_id ?? 50363); - const timestampBeginTokenId = noTimestampsTokenId + 1; - const maxInitialTimestampIndex = Number(generationConfig.max_initial_timestamp_index ?? 50); - const configuredMaxDecodeSteps = Number(generationConfig.max_length ?? 448); - const maxDecodeSteps = Math.min(configuredMaxDecodeSteps, MAX_DECODE_STEPS_CAP); - const alignmentHeads = Array.isArray(generationConfig.alignment_heads) - ? generationConfig.alignment_heads - .filter((head): head is [number, number] => Array.isArray(head) && head.length === 2) - .map(([layer, head]) => [Number(layer), Number(head)] as [number, number]) - : []; - - const forcedDecoder = Array.isArray(config.forced_decoder_ids) ? config.forced_decoder_ids : []; - const defaultLanguageFromForced = forcedDecoder.find(([index, id]) => index === 1 && typeof id === 'number')?.[1] ?? null; - const transcribeFromForced = forcedDecoder.find(([index, id]) => index === 2 && typeof id === 'number')?.[1] ?? null; - - const defaultLanguageToken = Number(defaultLanguageFromForced ?? tokenizer.token_to_id('<|en|>') ?? 50259); - const transcribeToken = Number(transcribeFromForced ?? tokenizer.token_to_id('<|transcribe|>') ?? 50359); - - const onnxThreadsPerJob = getOnnxThreadsPerJob(); - const stableSessionOptions: ort.InferenceSession.SessionOptions = { - executionProviders: ['cpu'], - graphOptimizationLevel: 'disabled', - intraOpNumThreads: onnxThreadsPerJob, - interOpNumThreads: 1, - executionMode: 'sequential', - enableCpuMemArena: false, - enableMemPattern: false, - }; - - const encoder = await ort.InferenceSession.create(WHISPER_ENCODER_MODEL_PATH, stableSessionOptions); - const decoderMerged = await ort.InferenceSession.create(WHISPER_DECODER_MERGED_MODEL_PATH, stableSessionOptions); - const decoderWithPast = await ort.InferenceSession.create(WHISPER_DECODER_WITH_PAST_MODEL_PATH, stableSessionOptions); - - const alignmentLayers = [...new Set(alignmentHeads.map(([layer]) => layer))]; - const prefillFetches: string[] = ['logits']; - const stepFetches: string[] = ['logits']; - const mergedOutputNames = new Set(decoderMerged.outputNames); - const withPastOutputNames = new Set(decoderWithPast.outputNames); - - for (let i = 0; i < WHISPER_NUM_LAYERS; i += 1) { - const decoderKey = `present.${i}.decoder.key`; - const decoderValue = `present.${i}.decoder.value`; - if (mergedOutputNames.has(decoderKey)) prefillFetches.push(decoderKey); - if (mergedOutputNames.has(decoderValue)) prefillFetches.push(decoderValue); - if (withPastOutputNames.has(decoderKey)) stepFetches.push(decoderKey); - if (withPastOutputNames.has(decoderValue)) stepFetches.push(decoderValue); - - const encoderKey = `present.${i}.encoder.key`; - const encoderValue = `present.${i}.encoder.value`; - if (mergedOutputNames.has(encoderKey)) prefillFetches.push(encoderKey); - if (mergedOutputNames.has(encoderValue)) prefillFetches.push(encoderValue); - } - - for (const layer of alignmentLayers) { - const key = `cross_attentions.${layer}`; - if (mergedOutputNames.has(key)) prefillFetches.push(key); - if (withPastOutputNames.has(key)) stepFetches.push(key); - } - - return { - encoder, - decoderMerged, - decoderWithPast, - tokenizer, - promptStartToken, - defaultLanguageToken, - transcribeToken, - eosTokenId, - noTimestampsTokenId, - timestampBeginTokenId, - maxInitialTimestampIndex, - maxDecodeSteps, - suppressTokens: new Set((generationConfig.suppress_tokens ?? []).map((v) => Number(v))), - beginSuppressTokens: new Set((generationConfig.begin_suppress_tokens ?? []).map((v) => Number(v))), - alignmentHeads, - prefillFetches, - stepFetches, - }; - })().catch((error) => { - state.runtimePromise = null; - throw error; - }); - - return state.runtimePromise; -} - -function resolveLanguageToken(runtime: WhisperRuntime, lang?: string): number { - const parsed = parseLanguageCode(lang); - if (!parsed) return runtime.defaultLanguageToken; - - const candidate = runtime.tokenizer.token_to_id(`<|${parsed}|>`); - return typeof candidate === 'number' ? candidate : runtime.defaultLanguageToken; -} - -async function runWhisperOnnx( - audioSamples: Float32Array, - opts: WhisperAlignmentOptions, - numFrames: number, - deadlineMs: number, -): Promise { - assertWithinDeadline(deadlineMs); - const runtime = await getRuntime(); - const decodeStepLimit = computeAdaptiveDecodeStepLimit(runtime.maxDecodeSteps, opts.textHint); - const mel = computeLogMelSpectrogram(audioSamples); - const encoderPast: Record = {}; - const decoderPast: Record = {}; - const crossAttentions: Record = {}; - let encoderHidden: ort.Tensor | null = null; - let outputs: Record | null = null; - - try { - const encoderOutputs = await runtime.encoder.run({ - input_features: mel, - }, ['last_hidden_state']); - encoderHidden = encoderOutputs.last_hidden_state; - - const languageToken = resolveLanguageToken(runtime, opts.lang); - const promptTokens = [ - runtime.promptStartToken, - languageToken, - runtime.transcribeToken, - ]; - - const generated: number[] = [...promptTokens]; - const emptyPastFeeds = buildEmptyPastFeeds(); - type LayerChunk = { - data: Float32Array; - heads: number; - seqLen: number; - frames: number; - }; - const selectedHeadsByLayer = new Map(); - for (const [layer, head] of runtime.alignmentHeads) { - const existing = selectedHeadsByLayer.get(layer) ?? []; - if (!existing.includes(head)) existing.push(head); - selectedHeadsByLayer.set(layer, existing); - } - for (const [layer, heads] of selectedHeadsByLayer) { - heads.sort((a, b) => a - b); - selectedHeadsByLayer.set(layer, heads); - } - const crossAttentionChunks = new Map(); - - const captureCrossAttentions = (stepOutputs: Record, prefill = false) => { - for (const [layer, selectedHeads] of selectedHeadsByLayer) { - const key = `cross_attentions.${layer}`; - const tensor = stepOutputs[key]; - if (!tensor) continue; - const [, , seqLen, frames] = tensor.dims; - const data = tensor.data as Float32Array; - const rowsToKeep = prefill ? seqLen : 1; - const seqStart = prefill ? 0 : Math.max(0, seqLen - 1); - const copied = new Float32Array(selectedHeads.length * rowsToKeep * frames); - for (let h = 0; h < selectedHeads.length; h += 1) { - const sourceHead = selectedHeads[h]!; - for (let s = 0; s < rowsToKeep; s += 1) { - const sourceSeq = seqStart + s; - for (let f = 0; f < frames; f += 1) { - const src = (((sourceHead * seqLen) + sourceSeq) * frames) + f; - const dst = (((h * rowsToKeep) + s) * frames) + f; - copied[dst] = data[src] ?? 0; - } - } - } - const list = crossAttentionChunks.get(layer) ?? []; - list.push({ data: copied, heads: selectedHeads.length, seqLen: rowsToKeep, frames }); - crossAttentionChunks.set(layer, list); - } - }; - const beginIndex = promptTokens.length; - - // Prefill: run prompt in merged decoder (non-cache branch), identical to first - // forward pass in transformers.js/transformers generation. - const prefillInputIds = tensorFromInt64(generated); - const prefillUseCacheBranch = new ort.Tensor('bool', Uint8Array.from([0]), [1]); - const prefillFeeds: Record = { - input_ids: prefillInputIds, - encoder_hidden_states: encoderHidden, - use_cache_branch: prefillUseCacheBranch, - ...emptyPastFeeds, - }; - try { - assertWithinDeadline(deadlineMs); - outputs = await runtime.decoderMerged.run(prefillFeeds, runtime.prefillFetches); - } finally { - disposeTensor(prefillInputIds); - disposeTensor(prefillUseCacheBranch); - } - captureCrossAttentions(outputs, true); - - for (let i = 0; i < WHISPER_NUM_LAYERS; i += 1) { - encoderPast[`past_key_values.${i}.encoder.key`] = outputs[`present.${i}.encoder.key`]; - encoderPast[`past_key_values.${i}.encoder.value`] = outputs[`present.${i}.encoder.value`]; - decoderPast[`past_key_values.${i}.decoder.key`] = outputs[`present.${i}.decoder.key`]; - decoderPast[`past_key_values.${i}.decoder.value`] = outputs[`present.${i}.decoder.value`]; - } - - for (let step = 0; step < decodeStepLimit; step += 1) { - assertWithinDeadline(deadlineMs); - if (!outputs) break; - const logits = outputs.logits; - const logitsData = logits.data as Float32Array; - const vocabSize = logits.dims[2] ?? 0; - const offset = logitsData.length - vocabSize; - const lastLogits = logitsData.subarray(offset); - - applyTokenSuppression(lastLogits, runtime.suppressTokens); - if (generated.length === beginIndex) { - applyTokenSuppression(lastLogits, runtime.beginSuppressTokens); - } - applyWhisperTimestampLogitsRules({ - logits: lastLogits, - generated, - beginIndex, - eosTokenId: runtime.eosTokenId, - noTimestampsTokenId: runtime.noTimestampsTokenId, - timestampBeginTokenId: runtime.timestampBeginTokenId, - maxInitialTimestampIndex: runtime.maxInitialTimestampIndex, - }); - - const nextToken = argmax(lastLogits) ?? runtime.eosTokenId; - generated.push(nextToken); - if (nextToken === runtime.eosTokenId) break; - - const previousDecoderPast = { ...decoderPast }; - const stepInputIds = tensorFromInt64([nextToken]); - const stepFeeds: Record = { - input_ids: stepInputIds, - ...previousDecoderPast, - ...encoderPast, - }; - let nextOutputs: Record; - try { - assertWithinDeadline(deadlineMs); - nextOutputs = await runtime.decoderWithPast.run(stepFeeds, runtime.stepFetches); - } finally { - disposeTensor(stepInputIds); - } - captureCrossAttentions(nextOutputs, false); - - for (let i = 0; i < WHISPER_NUM_LAYERS; i += 1) { - decoderPast[`past_key_values.${i}.decoder.key`] = nextOutputs[`present.${i}.decoder.key`]; - decoderPast[`past_key_values.${i}.decoder.value`] = nextOutputs[`present.${i}.decoder.value`]; - } - - disposeTensorMap(previousDecoderPast); - disposeTensor(outputs.logits); - for (const [name, tensor] of Object.entries(outputs)) { - if (name.startsWith('cross_attentions.')) { - disposeTensor(tensor); - } - } - outputs = nextOutputs; - } - - if (crossAttentionChunks.size === 0) { - return []; - } - - const remappedAlignmentHeads: Array<[number, number]> = runtime.alignmentHeads - .map(([layer, head]) => { - const selectedHeads = selectedHeadsByLayer.get(layer) ?? []; - const remappedHead = selectedHeads.indexOf(head); - if (remappedHead < 0) return null; - return [layer, remappedHead] as [number, number]; - }) - .filter((pair): pair is [number, number] => pair !== null); - - for (let layer = 0; layer < WHISPER_NUM_LAYERS; layer += 1) { - const chunks = crossAttentionChunks.get(layer); - if (!chunks || !chunks.length) continue; - - const heads = chunks[0].heads; - const frames = chunks[0].frames; - const concatSeqLen = chunks.reduce((sum, chunk) => sum + chunk.seqLen, 0); - const merged = new Float32Array(1 * heads * concatSeqLen * frames); - let seqOffset = 0; - - for (const chunk of chunks) { - const { data, seqLen, frames: tensorFrames } = chunk; - const copyFrames = Math.min(frames, tensorFrames); - - for (let h = 0; h < heads; h += 1) { - for (let s = 0; s < seqLen; s += 1) { - for (let f = 0; f < copyFrames; f += 1) { - const src = (((h * seqLen) + s) * tensorFrames) + f; - const dst = (((h * concatSeqLen) + (seqOffset + s)) * frames) + f; - merged[dst] = data[src] ?? 0; - } - } - } - seqOffset += seqLen; - } - - crossAttentions[`cross_attentions.${layer}`] = new ort.Tensor('float32', merged, [1, heads, concatSeqLen, frames]); - } - - const tokenStartTimestamps = extractTokenStartTimestamps({ - crossAttentions, - decoderLayers: WHISPER_NUM_LAYERS, - alignmentHeads: remappedAlignmentHeads, - numFrames, - numInputIds: promptTokens.length, - timePrecision: 0.02, - sequenceLength: generated.length, - }); - - const timedWords = buildWordsFromTimestampedTokens({ - tokens: generated, - tokenStartTimestamps, - tokenizer: runtime.tokenizer, - eosTokenId: runtime.eosTokenId, - promptLength: promptTokens.length, - timestampBeginTokenId: runtime.timestampBeginTokenId, - timePrecision: 0.02, - language: parseLanguageCode(opts.lang) ?? 'english', - }); - - const maxSec = Math.max(0, numFrames * 0.02); - return timedWords.map((word) => ({ - word: word.word, - start: Math.min(maxSec, Math.max(0, word.startSec)), - end: Math.min(maxSec, Math.max(0, word.endSec)), - })); - } finally { - disposeTensor(mel); - if (outputs?.logits) disposeTensor(outputs.logits); - if (outputs) { - for (const [name, tensor] of Object.entries(outputs)) { - if (name.startsWith('cross_attentions.')) { - disposeTensor(tensor); - } - } - } - disposeTensorMap(crossAttentions); - disposeTensorMap(decoderPast); - disposeTensorMap(encoderPast); - disposeTensor(encoderHidden); - } -} - -export async function alignAudioWithText( - audioBuffer: TTSAudioBuffer, - text: string, - cacheKey?: string, - opts: WhisperAlignmentOptions = {}, -): Promise { - if (!text.trim()) return []; - - if (cacheKey && alignmentCache.has(cacheKey)) { - const cached = alignmentCache.get(cacheKey)!; - alignmentCache.delete(cacheKey); - alignmentCache.set(cacheKey, cached); - return cached; - } - - if (cacheKey) { - const inFlight = alignmentInFlight.get(cacheKey); - if (inFlight) return inFlight; - } - const inFlightKey = cacheKey ?? makeInFlightCoalesceKey(audioBuffer, text, opts.lang); - const shared = alignmentInFlight.get(inFlightKey); - if (shared) return shared; - - state.pendingAlignments += 1; - const run = (async (): Promise => { - const deadlineMs = Date.now() + ALIGNMENT_TIMEOUT_MS; - const previous = state.alignMutex; - let release!: () => void; - state.alignMutex = new Promise((resolve) => { - release = resolve; - }); - - await previous; - - // Another request with the same cache key may have completed while this one - // was waiting on the mutex. - if (cacheKey && alignmentCache.has(cacheKey)) { - const cached = alignmentCache.get(cacheKey)!; - alignmentCache.delete(cacheKey); - alignmentCache.set(cacheKey, cached); - release(); - return cached; - } - - let tmpBase = ''; - let inputPath = ''; - let pcmPath = ''; - - try { - tmpBase = await mkdtemp(join(tmpdir(), 'openreader-whisper-')); - inputPath = join(tmpBase, `${randomUUID()}-input.bin`); - pcmPath = join(tmpBase, `${randomUUID()}-input.pcm16`); - - await writeFile(inputPath, Buffer.from(new Uint8Array(audioBuffer))); - await decodeToPcm16(inputPath, pcmPath); - - const pcmBytes = await readFile(pcmPath); - const decodedSamples = pcm16ToFloat32(pcmBytes); - const effectiveSampleLength = Math.min(decodedSamples.length, N_SAMPLES); - const effectiveFrameCount = Math.max(1, Math.floor((effectiveSampleLength / HOP_LENGTH) / 2)); - const normalizedAudio = padOrTrimAudio(decodedSamples); - - const words = await runWhisperOnnx( - normalizedAudio, - { ...opts, textHint: text }, - effectiveFrameCount, - deadlineMs, - ); - const alignment = mapWordsToSentenceOffsets(text, words); - const result: TTSSentenceAlignment[] = [alignment]; - - if (cacheKey) { - if (alignmentCache.has(cacheKey)) { - alignmentCache.delete(cacheKey); - } - alignmentCache.set(cacheKey, result); - while (alignmentCache.size > ALIGNMENT_CACHE_MAX_ENTRIES) { - const oldest = alignmentCache.keys().next().value; - if (!oldest) break; - alignmentCache.delete(oldest); - } - } - - return result; - } finally { - if (tmpBase) { - await rm(tmpBase, { recursive: true, force: true }).catch(() => {}); - } - release(); - state.pendingAlignments = Math.max(0, state.pendingAlignments - 1); - } - })(); - - alignmentInFlight.set(inFlightKey, run); - run.finally(() => { - if (alignmentInFlight.get(inFlightKey) === run) { - alignmentInFlight.delete(inFlightKey); - } - }); - return run; -} - -export function makeWhisperCacheKey(input: WhisperRequestBody): string { - return createHash('sha256') - .update( - JSON.stringify({ - text: input.text, - lang: input.lang || '', - audioLen: input.audio?.length || 0, - }), - ) - .digest('hex'); -} diff --git a/src/lib/server/whisper/ensureModel.ts b/src/lib/server/whisper/ensureModel.ts deleted file mode 100644 index 53b1d1b..0000000 --- a/src/lib/server/whisper/ensureModel.ts +++ /dev/null @@ -1,240 +0,0 @@ -import path from 'path'; -import { createHash } from 'crypto'; -import { access, copyFile, mkdir, readFile, rename, unlink, writeFile } from 'fs/promises'; -import { DOCSTORE_DIR } from '@/lib/server/storage/library-mount'; -import manifest from '@/lib/server/whisper/model/manifest.json'; - -const MODEL_DIR = path.join(DOCSTORE_DIR, 'model', 'whisper-base_timestamped'); -const STATIC_LICENSE_PATH = path.join(process.cwd(), 'src/lib/server/whisper/model/LICENSE.txt'); - -export const WHISPER_CONFIG_PATH = path.join(MODEL_DIR, 'config.json'); -export const WHISPER_GENERATION_CONFIG_PATH = path.join(MODEL_DIR, 'generation_config.json'); -export const WHISPER_TOKENIZER_PATH = path.join(MODEL_DIR, 'tokenizer.json'); -export const WHISPER_TOKENIZER_CONFIG_PATH = path.join(MODEL_DIR, 'tokenizer_config.json'); -export const WHISPER_ENCODER_MODEL_PATH = path.join(MODEL_DIR, 'onnx', 'encoder_model_int8.onnx'); -export const WHISPER_DECODER_MERGED_MODEL_PATH = path.join(MODEL_DIR, 'onnx', 'decoder_model_merged_int8.onnx'); -export const WHISPER_DECODER_WITH_PAST_MODEL_PATH = path.join(MODEL_DIR, 'onnx', 'decoder_with_past_model_int8.onnx'); - -const BASE_MODEL_URL = 'https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main'; -const WHISPER_MODEL_BASE_URL_ENV = 'WHISPER_MODEL_BASE_URL'; - -const MODEL_RELATIVE_PATHS: string[] = [ - 'config.json', - 'generation_config.json', - 'tokenizer.json', - 'tokenizer_config.json', - 'merges.txt', - 'vocab.json', - 'normalizer.json', - 'added_tokens.json', - 'preprocessor_config.json', - 'special_tokens_map.json', - 'onnx/encoder_model_int8.onnx', - 'onnx/decoder_model_merged_int8.onnx', - 'onnx/decoder_with_past_model_int8.onnx', -]; - -const DEFAULT_URLS: Record = { - 'config.json': `${BASE_MODEL_URL}/config.json`, - 'generation_config.json': `${BASE_MODEL_URL}/generation_config.json`, - 'tokenizer.json': `${BASE_MODEL_URL}/tokenizer.json`, - 'tokenizer_config.json': `${BASE_MODEL_URL}/tokenizer_config.json`, - 'merges.txt': `${BASE_MODEL_URL}/merges.txt`, - 'vocab.json': `${BASE_MODEL_URL}/vocab.json`, - 'normalizer.json': `${BASE_MODEL_URL}/normalizer.json`, - 'added_tokens.json': `${BASE_MODEL_URL}/added_tokens.json`, - 'preprocessor_config.json': `${BASE_MODEL_URL}/preprocessor_config.json`, - 'special_tokens_map.json': `${BASE_MODEL_URL}/special_tokens_map.json`, - 'onnx/encoder_model_int8.onnx': `${BASE_MODEL_URL}/onnx/encoder_model_int8.onnx`, - 'onnx/decoder_model_merged_int8.onnx': `${BASE_MODEL_URL}/onnx/decoder_model_merged_int8.onnx`, - 'onnx/decoder_with_past_model_int8.onnx': `${BASE_MODEL_URL}/onnx/decoder_with_past_model_int8.onnx`, -}; - -type ManifestEntry = { path: string; sha256?: string; size?: number }; - -export interface WhisperArtifactSpec { - path: string; - sha256?: string; - size?: number; - url: string; -} - -export interface WhisperStaticArtifactSpec { - path: string; - sha256?: string; - size?: number; - sourcePath: string; -} - -export type WhisperFetch = (input: RequestInfo | URL, init?: RequestInit) => Promise; - -const MANIFEST_FILES = manifest.files as ManifestEntry[]; -const MODEL_FILES = MANIFEST_FILES.filter((entry) => entry.path !== 'LICENSE.txt'); -const LICENSE_FILE = MANIFEST_FILES.find((entry) => entry.path === 'LICENSE.txt'); - -function normalizeExpected(entry: { sha256?: string; size?: number }): { sha256: string | null; size: number } { - return { - sha256: typeof entry.sha256 === 'string' ? entry.sha256.toLowerCase() : null, - size: Number(entry.size ?? 0), - }; -} - -function resolvePath(relativePath: string, modelDir: string): string { - return path.join(modelDir, relativePath); -} - -function joinModelUrl(baseUrl: string, relativePath: string): string { - return `${baseUrl.replace(/\/+$/, '')}/${relativePath}`; -} - -function resolveUrl(relativePath: string): string { - const overrideBase = process.env[WHISPER_MODEL_BASE_URL_ENV]?.trim(); - if (overrideBase) { - return joinModelUrl(overrideBase, relativePath); - } - const fallback = DEFAULT_URLS[relativePath]; - if (!fallback) { - throw new Error(`No default URL configured for Whisper model artifact: ${relativePath}`); - } - return fallback; -} - -function sha256OfBytes(bytes: Uint8Array): string { - return createHash('sha256').update(bytes).digest('hex'); -} - -function verifyBytes(bytes: Uint8Array, expected: { sha256?: string; size?: number }): boolean { - const normalized = normalizeExpected(expected); - if (Number.isFinite(normalized.size) && normalized.size > 0 && bytes.byteLength !== normalized.size) { - return false; - } - if (!normalized.sha256) return true; - return sha256OfBytes(bytes) === normalized.sha256; -} - -async function verifyFile(filePath: string, expected: { sha256?: string; size?: number }): Promise { - const bytes = await readFile(filePath); - return verifyBytes(bytes, expected); -} - -async function downloadToFile(fetchImpl: WhisperFetch, url: string, outPath: string): Promise { - const res = await fetchImpl(url); - if (!res.ok) { - throw new Error(`Download failed for ${url}: ${res.status} ${res.statusText}`); - } - const bytes = new Uint8Array(await res.arrayBuffer()); - await writeFile(outPath, bytes); -} - -export async function ensureWhisperArtifacts(options: { - modelDir: string; - artifacts: WhisperArtifactSpec[]; - staticArtifacts?: WhisperStaticArtifactSpec[]; - fetchImpl?: WhisperFetch; -}): Promise { - const { - modelDir, - artifacts, - staticArtifacts = [], - fetchImpl = fetch, - } = options; - - try { - await Promise.all(artifacts.map(async (artifact) => { - const target = resolvePath(artifact.path, modelDir); - await access(target); - const valid = await verifyFile(target, artifact); - if (!valid) { - throw new Error(`Checksum mismatch for existing Whisper artifact: ${artifact.path}`); - } - })); - - await Promise.all(staticArtifacts.map(async (artifact) => { - const target = resolvePath(artifact.path, modelDir); - await access(target); - const valid = await verifyFile(target, artifact); - if (!valid) { - throw new Error(`Checksum mismatch for existing Whisper static artifact: ${artifact.path}`); - } - })); - - return; - } catch { - // Continue to repair/download. - } - - for (const artifact of artifacts) { - const target = resolvePath(artifact.path, modelDir); - const targetDir = path.dirname(target); - const tmp = `${target}.tmp`; - - await mkdir(targetDir, { recursive: true }); - await downloadToFile(fetchImpl, artifact.url, tmp); - if (!(await verifyFile(tmp, artifact))) { - await unlink(tmp).catch(() => undefined); - throw new Error(`Whisper artifact checksum verification failed: ${artifact.path}`); - } - await rename(tmp, target); - } - - for (const artifact of staticArtifacts) { - const target = resolvePath(artifact.path, modelDir); - const targetDir = path.dirname(target); - await mkdir(targetDir, { recursive: true }); - await copyFile(artifact.sourcePath, target); - if (!(await verifyFile(target, artifact))) { - throw new Error(`Whisper static artifact checksum verification failed: ${artifact.path}`); - } - } -} - -export function createSingleflightRunner(work: () => Promise): () => Promise { - let inflight: Promise | null = null; - return async () => { - if (inflight) return inflight; - inflight = work().finally(() => { - inflight = null; - }); - return inflight; - }; -} - -async function ensureModelInternal(): Promise { - if (process.env[WHISPER_MODEL_BASE_URL_ENV]?.trim()) { - for (const relativePath of MODEL_RELATIVE_PATHS) { - if (!(relativePath in DEFAULT_URLS)) { - throw new Error(`Missing default URL path mapping for Whisper artifact: ${relativePath}`); - } - } - } - - const artifacts: WhisperArtifactSpec[] = MODEL_FILES.map((entry) => ({ - path: entry.path, - sha256: entry.sha256, - size: entry.size, - url: resolveUrl(entry.path), - })); - - const staticArtifacts: WhisperStaticArtifactSpec[] = LICENSE_FILE - ? [{ - path: LICENSE_FILE.path, - sha256: LICENSE_FILE.sha256, - size: LICENSE_FILE.size, - sourcePath: STATIC_LICENSE_PATH, - }] - : []; - - await ensureWhisperArtifacts({ - modelDir: MODEL_DIR, - artifacts, - staticArtifacts, - }); - - return WHISPER_ENCODER_MODEL_PATH; -} - -const ensureWhisperModelSingleflight = createSingleflightRunner(ensureModelInternal); - -export async function ensureWhisperModel(): Promise { - return ensureWhisperModelSingleflight(); -} diff --git a/src/lib/server/whisper/model/LICENSE.txt b/src/lib/server/whisper/model/LICENSE.txt deleted file mode 100644 index d255525..0000000 --- a/src/lib/server/whisper/model/LICENSE.txt +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2022 OpenAI - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/src/lib/server/whisper/model/manifest.json b/src/lib/server/whisper/model/manifest.json deleted file mode 100644 index 2fffd3d..0000000 --- a/src/lib/server/whisper/model/manifest.json +++ /dev/null @@ -1,76 +0,0 @@ -{ - "name": "whisper-base_timestamped-int8", - "version": "onnx-community/whisper-base_timestamped@608c49e61301901684bc36cac8f74b95ff6b5a8e", - "files": [ - { - "path": "config.json", - "sha256": "f4d0608f7d918166da7edb3e188de5ef1bfe70d9802e785d271fd88111e9cf4b", - "size": 2243 - }, - { - "path": "generation_config.json", - "sha256": "61070cf8de25b1e9256e8e102ded49d8d24a8369ed36ef84fdf21549e68125a0", - "size": 3832 - }, - { - "path": "tokenizer.json", - "sha256": "27fc476bfe7f17299480be2273fc0608e4d5a99aba2ab5dec5374b4482d1a566", - "size": 2480466 - }, - { - "path": "tokenizer_config.json", - "sha256": "2e036e4dbacfdeb7242c7d4ec4149f4a16e86026048f94d1637e3a8ee9c6a573", - "size": 282682 - }, - { - "path": "merges.txt", - "sha256": "2df2990a395e35e8dfbc7511e08c12d56018d8d04691e0133e5d63b21e154dc6", - "size": 493869 - }, - { - "path": "vocab.json", - "sha256": "50d6a919f0a0601d56a04eb583c780d18553aa388254ba3158eb6a00f13e2c1a", - "size": 1036584 - }, - { - "path": "normalizer.json", - "sha256": "bf1c507dc8724ca9cf9903640dacfb69dae2f00edee4f21ceba106a7392f26dd", - "size": 52666 - }, - { - "path": "added_tokens.json", - "sha256": "9715fd2243b6f06a5858b5e32950d2853f73dd5bc201aafcf76f5082a2d8acd1", - "size": 34604 - }, - { - "path": "preprocessor_config.json", - "sha256": "a6a76d28c93edb273669eb9e0b0636a2bddbb1272c3261e47b7ca6dfdbac1b8d", - "size": 339 - }, - { - "path": "special_tokens_map.json", - "sha256": "e67ae3a0aaa99abcd9f187138e12db1f65c16a14761c50ef10eef2c174a7a691", - "size": 2194 - }, - { - "path": "onnx/encoder_model_int8.onnx", - "sha256": "152da96dd8ff3f28f3fadabc2e8960405a277846453ff94ed411fe935a72917f", - "size": 23159150 - }, - { - "path": "onnx/decoder_model_merged_int8.onnx", - "sha256": "cf9a8d5bcddc0917a0078135b484cedcaf44f28909cd91910abd29dced9171db", - "size": 53712708 - }, - { - "path": "onnx/decoder_with_past_model_int8.onnx", - "sha256": "bdd92860d0ed7dff2aca623963378cbba1b617bfae127356db1c8aa8baa930ef", - "size": 50131672 - }, - { - "path": "LICENSE.txt", - "sha256": "b5d65a59060e68c4ff940e1eddfa6f94b2d68fdf58ed7f4dd57721c997e35e9d", - "size": 1063 - } - ] -} diff --git a/src/lib/server/whisper/model/mel_filters.npz b/src/lib/server/whisper/model/mel_filters.npz deleted file mode 100644 index 28ea26909dbdfd608aef67afc4d74d7961ae4bb6..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 4271 zcmZ`-cQjmYw;lx1g6JcN7QKe3LG%_Oh!VX=^k~teM-XGQ(Mu4$_Y%?jkm$lFBkB+( z3yfKIgF zxGiAhze`A@t->QRNVV!%P+W=o}VHkB) z%g>qyRHfN1IQ4-=`Y@0T9qE#o+;4E3VQ!epW1Xt=ZG`I3U|62t?<>5h*W|9VvJc`KZ+)ghnA**Z~ET21Tjf_f8oe`vy zZQNtlOx?dDhS71hnOus5cqj)hfyF@H&4y?@9z{I#&cf>A+s2~~(I>TQF}SaR3_tqa z(7&ZdN^vR*t<~?{9DEoI>0PL@Sl?wa?Z{rGX`*eEx9Nh=z*J3HZL1*Py4z$TD#+;m zSSW(kcOTe(4hqgib_W6&xx+j~-u(p)Nn6?>a%wHk=h7Ay$%lcGoo;gAY zmVV7|!Nb;w(PlH@c24{ple2Y3<*9J@jE=sfLzwu_BiAFPE$0Axp`^Nq!H}eG0?r-X zFj@Pwp^al*p>K{@_Cz`q#(N0Y=OpZy^ z{P$KjLJuk_Y%I)$mh`b{uOW5C5Xcmxk!gt_Zg zw>}6fkD4zRK9!#ems~H%U$>V;_wK38Zf-baU$S!#i;7!HWsi}GuC>%@?lMdgkUGC& zh9gC?O-5BlS2#}?7x0?eP#bOL(cqE{M%LJD$CZnplD)CgQR#KCttD=dZK+Ck5R52; z*%5hZ+SXU7)8k%Y^_1U>yI*By(INn&+ir-_4$#dUwTlMNyR@iGQIaZ+eiYqucu)CB z#i{Ru1w+aU#}DHSyzjG_9c?ToB_YjU#f;N=qel98WBIjIc1!#ePwRR+(go&-by#}@ z+M+klVke5b@lWfZ+O&|c??YvRe)&W)qAgtc>t-IZtbRTG#X}49_Q$>P%-)=0W_QY-x%DPep2Vm9#ci zyQcCc4p2&dLtV1@rPe!%>Y^#9W8#ZH&}^@wJKT7N;R9A7cEq&;Y2CYvd@R+Mn&b5O zVyfS^*H#kD74=J5uhD)o`TXoX>>Si$!cT?TXRxj2pB)w_ljjhTby&Je;X|BESZZT= zC%G5!-$BJf&a~U78d_3zBjrvrkJ0CCl@Rfcf7I(`VTNPnI^B#B$zOfPW zG&mEd?R0+W<`l08O1dkcWKS8wB!Z*Cs%I1nMs-EeB-uu5?t@PuD3|z>je8DKi#X(B z{Z=Rz{4X%?-UnxnHQtkELIZ&=J;fK_t}yu8|IxG0(85e&K>H3!!~zlhyJrgti~o1i zzBS*jTgdG~Exp#B-T)6A+PB ztD-e`j^@XAx}|L&JSEFkRvS_%3b%m86z02#Hfn{Y+qIqQ_muywgt?roUA7oiS1xBD zFxmDMsj_cbBcn*^rn^KIMP{AlHM`NiVm*D&`z~7FH#hf<$L3HmJ+=NdiY5>W?nKD? z8Ox6{9dKyI1o8a-j9BtV-|=lm`<`v>tR^Cln&x1dMYzu{@wq5KW!#K14_QMnpH5K%Pavag+g6(i8i-#Eq zguc}rH3?BxH4SOqZW#7m*aT(U9-n#_Xn^Q19(}eH!xG`nI!GYziVQNcA0)`FDHD%~ zz2$HnxW4BQ{#*@u`dssbAa`|fESn$8i8FdxGZh48_Uf~_Q@tv?4in)6fwSed)k&ITqu|){^(WL~J z?Lb|0ro06J^>f>^2}^e-+$u5bU4IZNfO?75v8lstS15%XYw2ac^pkU34{QhDR(umt zPu~`w2?FP|nn3!RWZ3{?=77@teulahD9*S*k5KmY3*adlM)%{SR~bkZYlx1q@fkE= zI$7+kiw5!ha=dYlO>Z5KgxnZEJsaBm%v#nkX0MN-h%n&KA?N}xU3K3o-3Jpk?ANq2n9&Lh%K_CTvfiN ze>6w~NSSl8$#NEZ^t7h9YOxI=zcAG|a+m6AWei`3Jw7K;b;T${pJa^4RwRt%F>?>M zBmoQqm1`<_W7i!5P~THp-II)Ka^u;=z;}d{;SVj{G_4`9^HaEb!=@Pa;Dw)CH^DjsGxFqmb%o$Bkop$KnH8 zDYN)Bh)5=5!-*|f0Gh4)oZG=TEBr()g^DCtSQhmT3!ZN`Qd-E%@1cE}hm8&Vq5B+C zVF2_O)9IiZ(v(xzTwJIg5|}KVuE(;}|7dVIrT`$d=q_OG|3PY}x*URYkMXXJ6PT1$IFkNyvY_(9UglDi6TaeikPS(!Bnij z;Szn+)I_oxnRz7(WTYTp+IHSWQ?Xd~tQn(Q1r)kThM?NM< z?d6LaBG!H}R$zRy!Ij(}1?xe^+o+!;tqWJ3NgjHl1XNxzusxQ0I#6qzM(_00UPMw* zF*GWW_q&fqAN=uimSKgBu_@jD%MX3hpNY|*4r=e=k1lw2r**IyD(hcq?A+HtUgUy4Dqh5D7|G9q{)TsUj{g~c!xy>9wk^(LiXA4VKGz_zMvJMX#AgsR z34T3hhJ)#&sUaQ1+0PML(?YA~{5?=(MT}X^Vib%};uoI{qGW@wgJ&_M+8S8clsNz2 zPQkxMi`#3+Khwtl>>K>wxc{71{&!qGu&Zzz_wU(7TLTyG){PAu?!cXs?Dp-y0Ekcn AQvd(} diff --git a/src/lib/server/whisper/spectral.ts b/src/lib/server/whisper/spectral.ts deleted file mode 100644 index b7223cc..0000000 --- a/src/lib/server/whisper/spectral.ts +++ /dev/null @@ -1,21 +0,0 @@ -export function buildGoertzelCoefficients(freqBins: number, fftSize: number): Float64Array { - const coeffs = new Float64Array(freqBins); - for (let k = 0; k < freqBins; k += 1) { - coeffs[k] = 2 * Math.cos((2 * Math.PI * k) / fftSize); - } - return coeffs; -} - -export function goertzelPower(samples: Float32Array, coeff: number): number { - let s1 = 0; - let s2 = 0; - for (let i = 0; i < samples.length; i += 1) { - const s0 = samples[i] + (coeff * s1) - s2; - s2 = s1; - s1 = s0; - } - - const power = (s1 * s1) + (s2 * s2) - (coeff * s1 * s2); - if (!Number.isFinite(power) || power < 0) return 0; - return power; -} diff --git a/src/lib/server/whisper/token-timestamps.ts b/src/lib/server/whisper/token-timestamps.ts deleted file mode 100644 index 47edc1d..0000000 --- a/src/lib/server/whisper/token-timestamps.ts +++ /dev/null @@ -1,449 +0,0 @@ -import type { Tokenizer } from '@huggingface/tokenizers'; -import type * as ort from 'onnxruntime-node'; - -const PUNCTUATION_REGEX = '\\p{P}\\u0021-\\u002F\\u003A-\\u0040\\u005B-\\u0060\\u007B-\\u007E'; -const PUNCTUATION_ONLY_REGEX = new RegExp(`^[${PUNCTUATION_REGEX}]+$`, 'gu'); - -type TokenTimestamp = [start: number, end: number]; - -export interface WhisperWordTiming { - word: string; - startSec: number; - endSec: number; -} - -function medianFilter(data: Float32Array, windowSize: number): Float32Array { - if (windowSize % 2 === 0 || windowSize <= 0) { - throw new Error('Window size must be a positive odd number'); - } - - const output = new Float32Array(data.length); - const buffer = new Float32Array(windowSize); - const halfWindow = Math.floor(windowSize / 2); - - for (let i = 0; i < data.length; i += 1) { - let valuesIndex = 0; - for (let j = -halfWindow; j <= halfWindow; j += 1) { - let index = i + j; - if (index < 0) { - index = Math.abs(index); - } else if (index >= data.length) { - index = (2 * (data.length - 1)) - index; - } - buffer[valuesIndex] = data[index]; - valuesIndex += 1; - } - - const sortable = Array.from(buffer); - sortable.sort((a, b) => a - b); - output[i] = sortable[halfWindow] ?? 0; - } - - return output; -} - -function dynamicTimeWarping(matrix: Float32Array[], rows: number, cols: number): [number[], number[]] { - const cost: number[][] = Array.from({ length: rows + 1 }, () => Array(cols + 1).fill(Number.POSITIVE_INFINITY)); - const trace: number[][] = Array.from({ length: rows + 1 }, () => Array(cols + 1).fill(-1)); - cost[0][0] = 0; - - for (let j = 1; j <= cols; j += 1) { - for (let i = 1; i <= rows; i += 1) { - const c0 = cost[i - 1][j - 1]; - const c1 = cost[i - 1][j]; - const c2 = cost[i][j - 1]; - let c: number; - let t: number; - if (c0 < c1 && c0 < c2) { - c = c0; - t = 0; - } else if (c1 < c0 && c1 < c2) { - c = c1; - t = 1; - } else { - c = c2; - t = 2; - } - cost[i][j] = matrix[i - 1][j - 1] + c; - trace[i][j] = t; - } - } - - for (let i = 0; i <= cols; i += 1) trace[0][i] = 2; - for (let i = 0; i <= rows; i += 1) trace[i][0] = 1; - - let i = rows; - let j = cols; - const textIndices: number[] = []; - const timeIndices: number[] = []; - while (i > 0 || j > 0) { - textIndices.push(i - 1); - timeIndices.push(j - 1); - const step = trace[i][j]; - if (step === 0) { - i -= 1; - j -= 1; - } else if (step === 1) { - i -= 1; - } else if (step === 2) { - j -= 1; - } else { - throw new Error(`Unexpected DTW trace state at [${i}, ${j}]`); - } - } - - textIndices.reverse(); - timeIndices.reverse(); - return [textIndices, timeIndices]; -} - -function round2(value: number): number { - return Math.round(value * 100) / 100; -} - -function decodeTokens(tokenizer: Pick, tokens: number[]): string { - return tokenizer.decode(tokens, { skip_special_tokens: false }); -} - -function splitTokensOnUnicode( - tokenizer: Pick, - tokens: number[], -): [string[], number[][], number[][]] { - const decodedFull = decodeTokens(tokenizer, tokens); - const replacementChar = '\uFFFD'; - const words: string[] = []; - const wordTokens: number[][] = []; - const tokenIndices: number[][] = []; - let currentTokens: number[] = []; - let currentIndices: number[] = []; - let unicodeOffset = 0; - - for (let i = 0; i < tokens.length; i += 1) { - currentTokens.push(tokens[i]); - currentIndices.push(i); - - const decoded = decodeTokens(tokenizer, currentTokens); - if ( - !decoded.includes(replacementChar) - || decodedFull[unicodeOffset + decoded.indexOf(replacementChar)] === replacementChar - ) { - words.push(decoded); - wordTokens.push(currentTokens); - tokenIndices.push(currentIndices); - currentTokens = []; - currentIndices = []; - unicodeOffset += decoded.length; - } - } - - return [words, wordTokens, tokenIndices]; -} - -function splitTokensOnSpaces( - tokenizer: Pick, - tokens: number[], - eosTokenId: number, -): [string[], number[][], number[][]] { - const [subwords, subwordTokens, subwordIndices] = splitTokensOnUnicode(tokenizer, tokens); - const words: string[] = []; - const wordTokens: number[][] = []; - const tokenIndices: number[][] = []; - - for (let i = 0; i < subwords.length; i += 1) { - const subword = subwords[i]; - const tokenList = subwordTokens[i]; - const indices = subwordIndices[i]; - const special = tokenList[0] >= eosTokenId; - const withSpace = subword.startsWith(' '); - const trimmed = subword.trim(); - const punctuation = PUNCTUATION_ONLY_REGEX.test(trimmed); - - if (special || withSpace || punctuation || words.length === 0) { - words.push(subword); - wordTokens.push([...tokenList]); - tokenIndices.push([...indices]); - } else { - const ix = words.length - 1; - words[ix] += subword; - wordTokens[ix].push(...tokenList); - tokenIndices[ix].push(...indices); - } - } - - return [words, wordTokens, tokenIndices]; -} - -function mergePunctuations( - words: string[], - tokens: number[][], - indices: number[][], - prependPunctuations = '"\'“¡¿([{-', - appendPunctuations = '"\'.。,，!！?？:：”)]}、', -): [string[], number[][], number[][]] { - const newWords = words.map((w) => `${w}`); - const newTokens = tokens.map((t) => [...t]); - const newIndices = indices.map((idx) => [...idx]); - - let i = newWords.length - 2; - let j = newWords.length - 1; - while (i >= 0) { - if (newWords[i].startsWith(' ') && prependPunctuations.includes(newWords[i].trim())) { - newWords[j] = newWords[i] + newWords[j]; - newTokens[j] = [...newTokens[i], ...newTokens[j]]; - newIndices[j] = [...newIndices[i], ...newIndices[j]]; - newWords[i] = ''; - newTokens[i] = []; - newIndices[i] = []; - } else { - j = i; - } - i -= 1; - } - - i = 0; - j = 1; - while (j < newWords.length) { - if (!newWords[i].endsWith(' ') && appendPunctuations.includes(newWords[j])) { - newWords[i] += newWords[j]; - newTokens[i] = [...newTokens[i], ...newTokens[j]]; - newIndices[i] = [...newIndices[i], ...newIndices[j]]; - newWords[j] = ''; - newTokens[j] = []; - newIndices[j] = []; - } else { - i = j; - } - j += 1; - } - - return [ - newWords.filter((w) => w.length > 0), - newTokens.filter((t) => t.length > 0), - newIndices.filter((t) => t.length > 0), - ]; -} - -function combineTokensIntoWords( - tokenizer: Pick, - tokens: number[], - eosTokenId: number, - language = 'english', -): [string[], number[][], number[][]] { - let words: string[]; - let wordTokens: number[][]; - let tokenIndices: number[][]; - - if (['chinese', 'japanese', 'thai', 'lao', 'myanmar', 'zh', 'ja', 'th', 'lo', 'my'].includes(language)) { - [words, wordTokens, tokenIndices] = splitTokensOnUnicode(tokenizer, tokens); - } else { - [words, wordTokens, tokenIndices] = splitTokensOnSpaces(tokenizer, tokens, eosTokenId); - } - - return mergePunctuations(words, wordTokens, tokenIndices); -} - -export function extractTokenStartTimestamps(input: { - crossAttentions: Record; - decoderLayers: number; - alignmentHeads: Array<[number, number]>; - numFrames: number; - numInputIds: number; - timePrecision?: number; - sequenceLength: number; -}): number[] { - const { - crossAttentions, - decoderLayers, - alignmentHeads, - numFrames, - numInputIds, - timePrecision = 0.02, - sequenceLength, - } = input; - - const frameCount = Math.max(1, numFrames); - const perLayer: Float32Array[] = []; - for (let layer = 0; layer < decoderLayers; layer += 1) { - const key = `cross_attentions.${layer}`; - const tensor = crossAttentions[key]; - if (!tensor) continue; - perLayer[layer] = tensor.data as Float32Array; - } - - const selected: Float32Array[] = []; - let seqLen = 0; - let attnFrames = 0; - for (const [layer, head] of alignmentHeads) { - const flat = perLayer[layer]; - if (!flat) continue; - const layerTensor = crossAttentions[`cross_attentions.${layer}`]; - if (!layerTensor || layerTensor.dims.length < 4) continue; - const [, numHeads, currentSeqLen, currentFrames] = layerTensor.dims; - if (head >= numHeads) continue; - seqLen = currentSeqLen; - attnFrames = Math.min(currentFrames, frameCount); - const headSlice = new Float32Array(seqLen * attnFrames); - for (let s = 0; s < seqLen; s += 1) { - for (let f = 0; f < attnFrames; f += 1) { - const flatIndex = (((head * currentSeqLen) + s) * currentFrames) + f; - headSlice[(s * attnFrames) + f] = flat[flatIndex] ?? 0; - } - } - selected.push(headSlice); - } - - if (!selected.length || seqLen === 0 || attnFrames === 0) { - return new Array(sequenceLength).fill(0); - } - - const normalizedHeads = selected.map((headData) => { - const means = new Float32Array(attnFrames); - const stds = new Float32Array(attnFrames); - - for (let f = 0; f < attnFrames; f += 1) { - let sum = 0; - for (let s = 0; s < seqLen; s += 1) sum += headData[(s * attnFrames) + f]; - const mean = sum / seqLen; - means[f] = mean; - let varSum = 0; - for (let s = 0; s < seqLen; s += 1) { - const d = headData[(s * attnFrames) + f] - mean; - varSum += d * d; - } - stds[f] = Math.sqrt(varSum / seqLen) || 1; - } - - const out = new Float32Array(headData.length); - for (let s = 0; s < seqLen; s += 1) { - const row = new Float32Array(attnFrames); - for (let f = 0; f < attnFrames; f += 1) { - row[f] = (headData[(s * attnFrames) + f] - means[f]) / stds[f]; - } - const filtered = medianFilter(row, 7); - out.set(filtered, s * attnFrames); - } - return out; - }); - - const croppedRows = Math.max(0, seqLen - numInputIds); - if (croppedRows === 0) return new Array(sequenceLength).fill(0); - - const matrix: Float32Array[] = Array.from({ length: croppedRows }, () => new Float32Array(attnFrames)); - for (const headData of normalizedHeads) { - for (let r = 0; r < croppedRows; r += 1) { - const srcRow = r + numInputIds; - for (let f = 0; f < attnFrames; f += 1) { - matrix[r][f] += headData[(srcRow * attnFrames) + f]; - } - } - } - - const scale = 1 / normalizedHeads.length; - for (let r = 0; r < croppedRows; r += 1) { - for (let f = 0; f < attnFrames; f += 1) { - matrix[r][f] = -matrix[r][f] * scale; - } - } - - const [textIndices, timeIndices] = dynamicTimeWarping(matrix, croppedRows, attnFrames); - const jumps = new Array(textIndices.length).fill(false); - for (let i = 0; i < textIndices.length; i += 1) { - jumps[i] = i === 0 ? true : textIndices[i] !== textIndices[i - 1]; - } - - const jumpTimes: number[] = []; - for (let i = 0; i < jumps.length; i += 1) { - if (jumps[i]) jumpTimes.push(timeIndices[i] * timePrecision); - } - - const timestamps = new Array(sequenceLength).fill(0); - for (let i = 0; i < numInputIds && i < timestamps.length; i += 1) timestamps[i] = 0; - for (let i = 0; i < jumpTimes.length && (numInputIds + i) < timestamps.length; i += 1) { - timestamps[numInputIds + i] = jumpTimes[i]; - } - if (timestamps.length > 0 && jumpTimes.length > 0) { - timestamps[timestamps.length - 1] = jumpTimes[jumpTimes.length - 1]; - } - return timestamps; -} - -export function buildWordsFromTimestampedTokens(input: { - tokens: number[]; - tokenStartTimestamps: number[]; - tokenizer: Pick; - eosTokenId: number; - promptLength: number; - timestampBeginTokenId: number; - timePrecision?: number; - language?: string; -}): WhisperWordTiming[] { - const { - tokens, - tokenStartTimestamps, - tokenizer, - eosTokenId, - promptLength, - timestampBeginTokenId, - timePrecision = 0.02, - language = 'english', - } = input; - - const limit = Math.min(tokens.length, tokenStartTimestamps.length); - const tokenRanges: TokenTimestamp[] = []; - for (let i = 0; i < limit; i += 1) { - const start = tokenStartTimestamps[i] ?? 0; - const end = i + 1 < limit ? (tokenStartTimestamps[i + 1] ?? (start + timePrecision)) : (start + timePrecision); - tokenRanges.push([start, Math.max(start, end)]); - } - - const words: WhisperWordTiming[] = []; - let segmentStart: number | null = null; - let textTokens: number[] = []; - let textRanges: TokenTimestamp[] = []; - - const flushSegment = (segmentEnd: number | null) => { - if (!textTokens.length) return; - const [wordTexts, , tokenIndices] = combineTokensIntoWords(tokenizer, textTokens, eosTokenId, language); - for (let i = 0; i < wordTexts.length; i += 1) { - const indices = tokenIndices[i]; - if (!indices.length) continue; - const start = textRanges[indices[0]]?.[0] ?? segmentStart ?? 0; - const end = textRanges[indices[indices.length - 1]]?.[1] ?? segmentEnd ?? start; - const clampedStart = segmentStart == null ? start : Math.max(segmentStart, start); - const clampedEndBase = segmentEnd == null ? end : Math.min(segmentEnd, end); - const clampedEnd = Math.max( - clampedStart + (clampedEndBase <= clampedStart ? timePrecision : 0), - clampedEndBase, - ); - words.push({ - word: wordTexts[i].trim(), - startSec: round2(clampedStart), - endSec: round2(clampedEnd), - }); - } - textTokens = []; - textRanges = []; - }; - - for (let i = promptLength; i < limit; i += 1) { - const token = tokens[i]; - if (token === eosTokenId) break; - - if (token >= timestampBeginTokenId) { - const ts = (token - timestampBeginTokenId) * timePrecision; - if (segmentStart == null) { - segmentStart = ts; - } else { - flushSegment(ts); - segmentStart = ts; - } - continue; - } - - textTokens.push(token); - textRanges.push(tokenRanges[i]); - } - - flushSegment(null); - return words.filter((w) => w.word.length > 0); -} diff --git a/tests/unit/whisper-alignment-mapping.spec.ts b/tests/unit/whisper-alignment-mapping.spec.ts index aba2933..1350d07 100644 --- a/tests/unit/whisper-alignment-mapping.spec.ts +++ b/tests/unit/whisper-alignment-mapping.spec.ts @@ -1,7 +1,7 @@ import { test, expect } from '@playwright/test'; import { mapWordsToSentenceOffsets, -} from '../../src/lib/server/whisper/alignment-mapping'; +} from '@openreader/compute-core/whisper'; test.describe('whisper alignment mapping', () => { test('maps words to sentence offsets with punctuation and repeated spaces', () => { diff --git a/tests/unit/whisper-alignment-smoke.spec.ts b/tests/unit/whisper-alignment-smoke.spec.ts index 2693dc9..25dac82 100644 --- a/tests/unit/whisper-alignment-smoke.spec.ts +++ b/tests/unit/whisper-alignment-smoke.spec.ts @@ -1,7 +1,7 @@ import { test, expect } from '@playwright/test'; import { readFile } from 'fs/promises'; import path from 'path'; -import { alignAudioWithText } from '../../src/lib/server/whisper/alignment'; +import { alignAudioWithText } from '@openreader/compute-core/whisper'; test.describe('whisper alignment smoke', () => { test('runs ONNX alignment end-to-end without decoder reshape errors', async () => { diff --git a/tests/unit/whisper-ensure-model.spec.ts b/tests/unit/whisper-ensure-model.spec.ts index 166f022..44e7948 100644 --- a/tests/unit/whisper-ensure-model.spec.ts +++ b/tests/unit/whisper-ensure-model.spec.ts @@ -6,7 +6,7 @@ import path from 'path'; import { createSingleflightRunner, ensureWhisperArtifacts, -} from '../../src/lib/server/whisper/ensureModel'; +} from '@openreader/compute-core/whisper'; function sha256(bytes: Uint8Array): string { return createHash('sha256').update(bytes).digest('hex'); diff --git a/tests/unit/whisper-spectral.spec.ts b/tests/unit/whisper-spectral.spec.ts index 3358532..7e9c44f 100644 --- a/tests/unit/whisper-spectral.spec.ts +++ b/tests/unit/whisper-spectral.spec.ts @@ -1,5 +1,5 @@ import { test, expect } from '@playwright/test'; -import { buildGoertzelCoefficients, goertzelPower } from '../../src/lib/server/whisper/spectral'; +import { buildGoertzelCoefficients, goertzelPower } from '@openreader/compute-core/whisper'; function dftPower(samples: Float32Array, k: number): number { const n = samples.length; diff --git a/tests/unit/whisper-token-timestamps.spec.ts b/tests/unit/whisper-token-timestamps.spec.ts index 22ebebd..bce4304 100644 --- a/tests/unit/whisper-token-timestamps.spec.ts +++ b/tests/unit/whisper-token-timestamps.spec.ts @@ -3,7 +3,7 @@ import * as ort from 'onnxruntime-node'; import { buildWordsFromTimestampedTokens, extractTokenStartTimestamps, -} from '../../src/lib/server/whisper/token-timestamps'; +} from '@openreader/compute-core/whisper'; test.describe('whisper token timestamp alignment', () => { test('extracts monotonic token timestamps from cross-attention maps', () => { diff --git a/tsconfig.json b/tsconfig.json index 5a7d536..3fba768 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -24,6 +24,7 @@ "@openreader/compute-core/contracts": ["./compute/core/src/contracts.ts"], "@openreader/compute-core/local-runtime": ["./compute/core/src/local-runtime.ts"], "@openreader/compute-core/runtime": ["./compute/core/src/runtime/index.ts"], + "@openreader/compute-core/whisper": ["./compute/core/src/whisper/index.ts"], "@openreader/compute-core/pdf-layout": ["./compute/core/src/pdf-layout/index.ts"], "@openreader/compute-core/pdf-layout/parse": ["./compute/core/src/pdf-layout/parsePdf.ts"] }