embed-worker.js

/**
 * embed-worker.js — Background Web Worker for transformers.js embedding
 *
 * Runs the ONNX embedding model off the main thread so the UI stays responsive.
 * Uses a singleton pattern to ensure the model loads only once.
 *
 * Protocol:
 *   Main → Worker:
 *     { type: 'load', modelId, dtype }        — Load/switch model
 *     { type: 'embed', texts, id }            — Embed a batch of texts
 *     { type: 'unload' }                      — Release model memory
 *
 *   Worker → Main (all tagged with source: 'vecdb'):
 *     { source, status: 'progress', file, progress, loaded, total }
 *     { source, status: 'ready', dim }                — Model ready
 *     { source, status: 'result', id, embeddings, dims }
 *     { source, status: 'error', id?, message }
 *     { source, status: 'unloaded' }
 */

import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3';

const MSG_TAG = 'vecdb';

let extractor = null;
let currentModel = null;
let currentDtype = null;
let loadingPromise = null;

function send(msg) {
    self.postMessage({ ...msg, source: MSG_TAG });
}

function sendTransfer(msg, transfer) {
    self.postMessage({ ...msg, source: MSG_TAG }, transfer);
}

async function loadModel(modelId, dtype) {
    // If already loading the same model, wait for it
    if (loadingPromise && currentModel === modelId && currentDtype === dtype) {
        return loadingPromise;
    }

    // If switching models, dispose old one
    if (extractor) {
        try { await extractor.dispose(); } catch {}
        extractor = null;
    }

    currentModel = modelId;
    currentDtype = dtype;

    loadingPromise = pipeline('feature-extraction', modelId, {
        dtype: dtype,
        progress_callback: (p) => {
            // Only relay download progress events we care about.
            // Transformers.js emits: initiate, download, progress, done, ready
            // We only forward 'progress' (with loaded/total) so the main thread
            // can update the progress bar. Everything else is ignored to avoid
            // status collisions (e.g. transformers.js 'ready' vs our 'ready').
            if (p.status === 'progress' && p.total > 0) {
                send({ status: 'dl-progress', file: p.file, loaded: p.loaded, total: p.total });
            }
        },
    });

    extractor = await loadingPromise;
    return extractor;
}

self.addEventListener('message', async (event) => {
    const { type, id } = event.data;

    try {
        switch (type) {
            case 'load': {
                const { modelId, dtype } = event.data;
                await loadModel(modelId, dtype);

                // Probe dimension by embedding a test string
                const test = await extractor(['test'], { pooling: 'mean', normalize: true });
                const dim = test.dims[1];

                send({ status: 'ready', dim });
                break;
            }

            case 'embed': {
                if (!extractor) {
                    send({ status: 'error', id, message: 'Model not loaded' });
                    return;
                }

                const { texts } = event.data;
                const output = await extractor(texts, { pooling: 'mean', normalize: true });

                // Copy to a fresh Float32Array for zero-copy transfer
                const float32 = new Float32Array(output.data);
                sendTransfer(
                    { status: 'result', id, embeddings: float32, dims: output.dims },
                    [float32.buffer]
                );
                break;
            }

            case 'unload': {
                if (extractor) {
                    try { await extractor.dispose(); } catch {}
                    extractor = null;
                }
                currentModel = null;
                currentDtype = null;
                loadingPromise = null;
                send({ status: 'unloaded' });
                break;
            }
        }
    } catch (err) {
        send({ status: 'error', id, message: err.message || String(err) });
    }
});