Spaces:

SiliconLanguage
/

open-asr-model-explorer

Sleeping

open-asr-model-explorer / frontend /src /services /inferenceRouter.js

Ping Long

feat(ui): stabilize metrics state and expand language samples

d4f1e7c 7 months ago

7.44 kB

/**
 * inferenceRouter.js
 * ─────────────────────────────────────────────────────────────────────────────
 * Strategy-based routing:
 *   - Model name contains "WebGPU" → client-side via WebWorker (transformers.js)
 *   - All other models             → server-side via FastAPI / vLLM
 *
 * Exports:
 *   route(model, audioBuffer, onProgress) → Promise<TranscribeResult>
 *
 * TranscribeResult = { transcript, ttft_ms, itl_ms, rtfx, mode }
 */

// ── WebGPU / WebWorker client ─────────────────────────────────────────────────

let workerInstance = null;
let workerReadyResolve = null;
let workerReadyReject = null;  // companion reject for the load promise
let workerLoadedModelId = null; // mirrors worker's loadedModelId for fast cache hit
const pendingCallbacks = new Map(); // id → { resolve, reject, onProgress }

function getWorker() {
  if (workerInstance) return workerInstance;

  // Vite's ?worker syntax creates an ES-module Worker automatically
  workerInstance = new Worker(
    new URL('../workers/webgpu.worker.js', import.meta.url),
    { type: 'module' }
  );

  workerInstance.addEventListener('message', (event) => {
    const { type, id, stable, unstable, transcript, ttft_ms, itl_ms, rtfx, message } = event.data;

    // ── Load lifecycle messages ─────────────────────────────────────────────
    if (type === 'ready') {
      workerLoadedModelId = event.data.modelId ?? workerLoadedModelId;
      if (workerReadyResolve) {
        const res = workerReadyResolve;
        workerReadyResolve = null;
        workerReadyReject = null;
        res();
      }
      return;
    }

    // A null id means the error came from loadModel(), not a transcription.
    if (type === 'error' && id == null) {
      const rej = workerReadyReject;
      workerReadyResolve = null;
      workerReadyReject = null;
      rej?.(new Error(message ?? 'Model failed to load'));
      return;
    }

    // ── Transcription lifecycle messages ────────────────────────────────────
    const cb = pendingCallbacks.get(id);
    if (!cb) return;

    switch (type) {
      case 'progress':
        cb.onProgress?.({ stable, unstable });
        break;
      case 'result':
        pendingCallbacks.delete(id);
        cb.resolve({ transcript, ttft_ms, itl_ms, rtfx, mode: 'webgpu' });
        break;
      case 'error':
        pendingCallbacks.delete(id);
        cb.reject(new Error(message));
        break;
    }
  });

  return workerInstance;
}

async function ensureModelLoaded(modelId) {
  // Fast path: model already cached in worker — skip the message round-trip
  if (workerLoadedModelId === modelId) return;
  const worker = getWorker();
  return new Promise((resolve, reject) => {
    workerReadyResolve = resolve;
    workerReadyReject = reject;
    worker.postMessage({ type: 'load', model: modelId });
  });
}

async function routeWebGPU(modelId, audioBuffer, onProgress, language) {
  await ensureModelLoaded(modelId);
  const worker = getWorker();
  const id = crypto.randomUUID();
  return new Promise((resolve, reject) => {
    pendingCallbacks.set(id, { resolve, reject, onProgress });
    worker.postMessage({ type: 'transcribe', audio: audioBuffer, id, language }, [audioBuffer.buffer]);
  });
}

// ── Server-side / vLLM client ─────────────────────────────────────────────────

async function routeServer(modelId, audioFile, onProgress, engine, language) {
  const form = new FormData();
  form.append('audio', audioFile);
  form.append('model', modelId);
  if (engine) {
    form.append('engine', engine);
  }
  if (language) {
    form.append('language', language);
  }

  // Use streaming endpoint for token-by-token delivery
  const response = await fetch('/api/transcribe/stream', {
    method: 'POST',
    body: form,
  });

  if (!response.ok) {
    const err = await response.json().catch(() => ({ detail: response.statusText }));
    throw new Error(err.detail ?? `HTTP ${response.status}`);
  }

  const reader = response.body.getReader();
  const decoder = new TextDecoder();
  let buffer = '';
  let finalResult = null;

  try {
    while (true) {
      const { value, done } = await reader.read();
      if (done) break;
      buffer += decoder.decode(value, { stream: true });

      // Parse Server-Sent Events
      const events = buffer.split('\n\n');
      buffer = events.pop() ?? '';

      for (const event of events) {
        const dataLine = event.split('\n').find((l) => l.startsWith('data: '));
        if (!dataLine) continue;
        try {
          const payload = JSON.parse(dataLine.slice(6));
          if (payload.done) {
            finalResult = payload;
          } else if (payload.token) {
            onProgress?.({ token: payload.token, ttft_ms: payload.ttft_ms });
          } else if (payload.ttft_ms != null) {
            onProgress?.({ ttft_ms: payload.ttft_ms });
          }
        } catch {
          // ignore malformed SSE frames
        }
      }
    }
  } finally {
    // Always release the reader so the underlying HTTP connection can be
    // returned to the pool.  This prevents stream leaks when an exception
    // is thrown mid-stream (e.g. network error or component unmount).
    reader.cancel();
  }

  if (!finalResult) throw new Error('Stream ended without a final result.');

  return {
    transcript: finalResult.transcript,
    ttft_ms: finalResult.ttft_ms,
    itl_ms: finalResult.itl_ms,
    rtfx: finalResult.rtfx,
    mode: 'server',
  };
}

// ── Public router ─────────────────────────────────────────────────────────────

/**
 * Route transcription to the appropriate execution mode.
 *
 * @param {string}   modelId      Selected model name
 * @param {File|Float32Array} audio  Audio source
 * @param {function} [onProgress] Called with partial results during streaming
 * @returns {Promise<{transcript:string, ttft_ms:number, itl_ms:number, rtfx:number, mode:string}>}
 */
export async function route(modelId, audio, onProgress, engine = undefined, mode = 'server', language = undefined) {
  if (mode === 'webgpu') {
    // Ensure we have a Float32Array of 16 kHz samples
    const audioBuffer = audio instanceof Float32Array ? audio : await fileToFloat32(audio);
    return routeWebGPU(modelId, audioBuffer, onProgress, language);
  }
  // Server-side: pass the raw File for multipart upload
  return routeServer(modelId, audio, onProgress, engine, language);
}

// ── Helper: decode File → Float32Array at 16 kHz ──────────────────────────────

async function fileToFloat32(file) {
  const arrayBuffer = await file.arrayBuffer();
  const audioCtx = new AudioContext({ sampleRate: 16000 });
  const decoded = await audioCtx.decodeAudioData(arrayBuffer);
  return decoded.getChannelData(0); // mono, 16 kHz
}

	/**
	* inferenceRouter.js
	* ─────────────────────────────────────────────────────────────────────────────
	* Strategy-based routing:
	* - Model name contains "WebGPU" → client-side via WebWorker (transformers.js)
	* - All other models → server-side via FastAPI / vLLM
	*
	* Exports:
	* route(model, audioBuffer, onProgress) → Promise<TranscribeResult>
	*
	* TranscribeResult = { transcript, ttft_ms, itl_ms, rtfx, mode }
	*/

	// ── WebGPU / WebWorker client ─────────────────────────────────────────────────

	let workerInstance = null;
	let workerReadyResolve = null;
	let workerReadyReject = null; // companion reject for the load promise
	let workerLoadedModelId = null; // mirrors worker's loadedModelId for fast cache hit
	const pendingCallbacks = new Map(); // id → { resolve, reject, onProgress }

	function getWorker() {
	if (workerInstance) return workerInstance;

	// Vite's ?worker syntax creates an ES-module Worker automatically
	workerInstance = new Worker(
	new URL('../workers/webgpu.worker.js', import.meta.url),
	{ type: 'module' }
	);

	workerInstance.addEventListener('message', (event) => {
	const { type, id, stable, unstable, transcript, ttft_ms, itl_ms, rtfx, message } = event.data;

	// ── Load lifecycle messages ─────────────────────────────────────────────
	if (type === 'ready') {
	workerLoadedModelId = event.data.modelId ?? workerLoadedModelId;
	if (workerReadyResolve) {
	const res = workerReadyResolve;
	workerReadyResolve = null;
	workerReadyReject = null;
	res();
	}
	return;
	}

	// A null id means the error came from loadModel(), not a transcription.
	if (type === 'error' && id == null) {
	const rej = workerReadyReject;
	workerReadyResolve = null;
	workerReadyReject = null;
	rej?.(new Error(message ?? 'Model failed to load'));
	return;
	}

	// ── Transcription lifecycle messages ────────────────────────────────────
	const cb = pendingCallbacks.get(id);
	if (!cb) return;

	switch (type) {
	case 'progress':
	cb.onProgress?.({ stable, unstable });
	break;
	case 'result':
	pendingCallbacks.delete(id);
	cb.resolve({ transcript, ttft_ms, itl_ms, rtfx, mode: 'webgpu' });
	break;
	case 'error':
	pendingCallbacks.delete(id);
	cb.reject(new Error(message));
	break;
	}
	});

	return workerInstance;
	}

	async function ensureModelLoaded(modelId) {
	// Fast path: model already cached in worker — skip the message round-trip
	if (workerLoadedModelId === modelId) return;
	const worker = getWorker();
	return new Promise((resolve, reject) => {
	workerReadyResolve = resolve;
	workerReadyReject = reject;
	worker.postMessage({ type: 'load', model: modelId });
	});
	}

	async function routeWebGPU(modelId, audioBuffer, onProgress, language) {
	await ensureModelLoaded(modelId);
	const worker = getWorker();
	const id = crypto.randomUUID();
	return new Promise((resolve, reject) => {
	pendingCallbacks.set(id, { resolve, reject, onProgress });
	worker.postMessage({ type: 'transcribe', audio: audioBuffer, id, language }, [audioBuffer.buffer]);
	});
	}

	// ── Server-side / vLLM client ─────────────────────────────────────────────────

	async function routeServer(modelId, audioFile, onProgress, engine, language) {
	const form = new FormData();
	form.append('audio', audioFile);
	form.append('model', modelId);
	if (engine) {
	form.append('engine', engine);
	}
	if (language) {
	form.append('language', language);
	}

	// Use streaming endpoint for token-by-token delivery
	const response = await fetch('/api/transcribe/stream', {
	method: 'POST',
	body: form,
	});

	if (!response.ok) {
	const err = await response.json().catch(() => ({ detail: response.statusText }));
	throw new Error(err.detail ?? `HTTP ${response.status}`);
	}

	const reader = response.body.getReader();
	const decoder = new TextDecoder();
	let buffer = '';
	let finalResult = null;

	try {
	while (true) {
	const { value, done } = await reader.read();
	if (done) break;
	buffer += decoder.decode(value, { stream: true });

	// Parse Server-Sent Events
	const events = buffer.split('\n\n');
	buffer = events.pop() ?? '';

	for (const event of events) {
	const dataLine = event.split('\n').find((l) => l.startsWith('data: '));
	if (!dataLine) continue;
	try {
	const payload = JSON.parse(dataLine.slice(6));
	if (payload.done) {
	finalResult = payload;
	} else if (payload.token) {
	onProgress?.({ token: payload.token, ttft_ms: payload.ttft_ms });
	} else if (payload.ttft_ms != null) {
	onProgress?.({ ttft_ms: payload.ttft_ms });
	}
	} catch {
	// ignore malformed SSE frames
	}
	}
	}
	} finally {
	// Always release the reader so the underlying HTTP connection can be
	// returned to the pool. This prevents stream leaks when an exception
	// is thrown mid-stream (e.g. network error or component unmount).
	reader.cancel();
	}

	if (!finalResult) throw new Error('Stream ended without a final result.');

	return {
	transcript: finalResult.transcript,
	ttft_ms: finalResult.ttft_ms,
	itl_ms: finalResult.itl_ms,
	rtfx: finalResult.rtfx,
	mode: 'server',
	};
	}

	// ── Public router ─────────────────────────────────────────────────────────────

	/**
	* Route transcription to the appropriate execution mode.
	*
	* @param {string} modelId Selected model name
	* @param {File\|Float32Array} audio Audio source
	* @param {function} [onProgress] Called with partial results during streaming
	* @returns {Promise<{transcript:string, ttft_ms:number, itl_ms:number, rtfx:number, mode:string}>}
	*/
	export async function route(modelId, audio, onProgress, engine = undefined, mode = 'server', language = undefined) {
	if (mode === 'webgpu') {
	// Ensure we have a Float32Array of 16 kHz samples
	const audioBuffer = audio instanceof Float32Array ? audio : await fileToFloat32(audio);
	return routeWebGPU(modelId, audioBuffer, onProgress, language);
	}
	// Server-side: pass the raw File for multipart upload
	return routeServer(modelId, audio, onProgress, engine, language);
	}

	// ── Helper: decode File → Float32Array at 16 kHz ──────────────────────────────

	async function fileToFloat32(file) {
	const arrayBuffer = await file.arrayBuffer();
	const audioCtx = new AudioContext({ sampleRate: 16000 });
	const decoded = await audioCtx.decodeAudioData(arrayBuffer);
	return decoded.getChannelData(0); // mono, 16 kHz
	}