Spaces:

lemonteaa
/

llama-cpp-wasm

Running

lemonteaa
fix structured call 1
8ae3604 1 day ago
25.9 kB
// Inference worker.
//
// Responsibilities:
//   1. add models - two ingest paths into one library:
//        url  -> resumable, pausable, chunked download
//        file -> the user's own File, read in 4 MiB slices, never copied
//   2. sha256 each model once, so a human can compare it with the publisher's
//   3. stream the bytes into the Emscripten FS (from IndexedDB chunks, or
//      straight off the local File)
//   4. load it with the C API and stream generated tokens back
//
// Both ingest paths end in the same kind of record, so loading, hashing,
// renaming and deleting are one code path (see storage.js for the schema).

// The module is imported lazily: a pthread build needs SharedArrayBuffer, which
// exists only on a cross-origin isolated page (COOP/COEP). Pages without it get
// the single-thread module, so the demo still runs on plain static hosting.
const THREADED = typeof SharedArrayBuffer === "function" && self.crossOriginIsolated === true;
const HARDWARE_THREADS = (typeof navigator !== "undefined" && navigator.hardwareConcurrency) || 1;

import { action } from "./actions.js";
import { downloadModel } from "./download.js";
import { sha256Chunks } from "./sha256.js";
import {
    CHUNK_SIZE, getModel, putModel, updateModel, listModels, newModelId, nameFromUrl,
    putChunk, iterateChunks, clearChunks, deleteModel, findModelByUrl,
} from "./storage.js";

let module = null;
let loaded = false;
let loadedId = null;
let currentNctx = 1024;

// Engine launch parameters, applied when a context is created: the KV cache
// configuration is fixed for the life of a llama_context, so changing it means
// reloading the model (the page does that through Reload engine).
let currentKv = { type_k: "f16", type_v: "f16", flash_attn: -1, kv_unified: true, n_parallel: 1 };
let currentThreads = 1;   // 1 unless the page asked for and can run more

const CTX_ARGS = ["number", "string", "string", "number", "number", "number"];

function createContext(n_ctx) {
    const n = module.ccall("llmjs_create_context_ex", "number", CTX_ARGS,
        [n_ctx, currentKv.type_k, currentKv.type_v, currentKv.flash_attn,
         currentKv.kv_unified ? 1 : 0, currentKv.n_parallel || 1]);
    if (n >= 0) currentNctx = n;
    return n;
}

function ctxConfig() {
    const text = module.UTF8ToString(module._llmjs_ctx_config());
    try { return JSON.parse(text); } catch (_) { return null; }
}

// Files the page handed over in this session, by model id. A "file" model with
// no copy in IndexedDB needs its File again after a reload - see start().
const files = new Map();

let abort = null;            // AbortController for the active download
let busy = false;
let lastRequest = null;      // { id, context_size } so Resume can continue

const post = (msg) => self.postMessage(msg);
const progress = (p) => post({ event: action.PROGRESS, ...p });
const log = (line, level = 3) => post({ event: action.LOG, level, line });
const fail = (err) => post({ event: action.ERROR, error: String((err && err.message) || err) });

// --- wasm -> JS callbacks ---------------------------------------------------
self.llmjs_on_token = (text) => post({ event: action.WRITE_RESULT, line: text });
self.llmjs_on_log = (level, text) => post({ event: action.LOG, level, line: text });

const tick = () => new Promise((r) => setTimeout(r, 0));

// storage.js shaped for download.js
const storage = {
    CHUNK_SIZE, getMeta: getModel, putMeta: putModel, updateMeta: updateModel, putChunk, clearChunks,
};

// Render a message list with llama.cpp's chat stack (GGUF template, tools and
// chat_template_kwargs), inside the wasm module. Throws with the engine's
// message; the caller decides whether that is worth showing.
function renderTemplateEx(messagesJson, tools, source, kwargs) {
    const sig = ["string", "string", "string", "string", "number", "number", "number"];
    const need = module.ccall("llmjs_chat_apply_ex", "number", sig,
        [messagesJson, tools, source, kwargs, 1, 0, 0]);
    if (need < 0) {
        throw new Error(module.UTF8ToString(module._llmjs_last_error()));
    }
    const size = need + 1;
    const ptr = module._malloc(size);
    try {
        module.ccall("llmjs_chat_apply_ex", "number", sig,
            [messagesJson, tools, source, kwargs, 1, ptr, size]);
        return module.UTF8ToString(ptr, need);
    } finally {
        module._free(ptr);
    }
}

let modulePromise = null;

async function ensureModule() {
    if (module) return module;
    if (!modulePromise) {
        modulePromise = (async () => {
            // No locateFile override: the glue resolves its .wasm relative to its
            // own import.meta.url (web/wasm/). Passing one here resolves against
            // *this* module's URL and 404s the .wasm into an HTML page.
            let factory;
            if (THREADED) {
                try {
                    factory = (await import("./wasm/llm-wasm-mt.js")).default;
                } catch (err) {
                    log(`threaded build not found (${err}) - using the single-thread module`);
                }
            }
            if (!factory) factory = (await import("./wasm/llm-wasm.js")).default;
            // emscripten writes to stdout/stderr for anything that does not go
            // through our log callback; that is chatter, not a failure, so it
            // must not land in the console as an error
            const print = (text) => self.llmjs_on_log(2, String(text).replace(/\n?$/, "\n"));
            module = await factory({ print, printErr: print });
            module._llmjs_init();
            log(`llama.cpp ${module.UTF8ToString(module._llmjs_version())}` +
                (THREADED ? " (threaded build)" : " (single-thread build)"));
            return module;
        })();
    }
    return modulePromise;
}

// --- byte sources -----------------------------------------------------------

// 4 MiB slices of a File the user picked. The bytes stay on disk; each slice is
// read, handed to the FS (or the hash), and dropped again.
async function* fileSlices(file) {
    for (let offset = 0; offset < file.size; offset += CHUNK_SIZE) {
        const end = Math.min(offset + CHUNK_SIZE, file.size);
        yield new Uint8Array(await file.slice(offset, end).arrayBuffer());
    }
}

// Single iterator over a model's bytes, wherever they live.
async function* byteStream(model) {
    if (model.source === "url" || model.copied) {
        for await (const row of iterateChunks(model.id)) yield row.data;
        return;
    }
    const file = files.get(model.id);
    if (!file) throw new Error("the local file is not attached - choose it again");
    yield* fileSlices(file);
}

async function hashModel(model) {
    const total = model.total || 0;
    const sha256 = await sha256Chunks(byteStream(model), (done) =>
        progress({ phase: "hash", loaded: done, total }));
    return updateModel(model.id, { sha256 });
}

// --- ingest ----------------------------------------------------------------

async function addUrlModel({ url, name, expectedSha256, mmproj }) {
    let model = await findModelByUrl(url);
    if (model) {
        // same URL added again: treat it as "open", but let the name/hash be corrected
        const patch = {};
        if (name) patch.name = name;
        if (expectedSha256) patch.expectedSha256 = expectedSha256;
        if (mmproj !== undefined) patch.multimodal = !!mmproj;
        if (Object.keys(patch).length) model = await updateModel(model.id, patch);
    } else {
        model = await putModel({
            id: newModelId(),
            name: (name || "").trim() || nameFromUrl(url),
            source: "url",
            url,
            fileName: null,
            total: 0,
            received: 0,
            complete: false,
            copied: true,               // url models always live in the browser
            sha256: null,
            expectedSha256: (expectedSha256 || "").trim() || null,
            // user-flagged: this GGUF is a multimodal projector, not a text model
            multimodal: !!mmproj,
            created: Date.now(),
        });
    }
    post({ event: action.MODEL_ADDED, model });
    return model.id;
}

async function addFileModel({ file, name, keepCopy, expectedSha256, mmproj }) {
    if (!file) throw new Error("no file given");

    const head = new Uint8Array(await file.slice(0, 4).arrayBuffer());
    const magic = String.fromCharCode(...head);
    if (magic !== "GGUF") {
        throw new Error(`${file.name} is not a GGUF file (starts with "${magic}")`);
    }

    const model = await putModel({
        id: newModelId(),
        name: (name || "").trim() || file.name.replace(/\.gguf$/i, ""),
        source: "file",
        url: null,
        fileName: file.name,
        total: file.size,
        received: keepCopy ? 0 : file.size,
        complete: !keepCopy,
        copied: false,                  // filled in by copyFileIntoStorage below
        sha256: null,
        expectedSha256: (expectedSha256 || "").trim() || null,
        multimodal: !!mmproj,
        created: Date.now(),
    });

    files.set(model.id, file);
    post({ event: action.MODEL_ADDED, model });

    if (keepCopy) return (await copyFileIntoStorage(model, file)).id;
    return model.id;
}

// opt-in: put the bytes in IndexedDB so the model no longer needs the file
async function copyFileIntoStorage(model, file) {
    await clearChunks(model.id);
    let index = 0;
    let received = 0;
    for await (const bytes of fileSlices(file)) {
        const copy = bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength);
        await putChunk(model.id, index++, copy);
        received += bytes.length;
        progress({ phase: "copy", loaded: received, total: file.size });
    }
    return updateModel(model.id, { copied: true, received, total: file.size, complete: true });
}

// --- loading ---------------------------------------------------------------

async function loadIntoFS(model, target = "/models/model.gguf") {
    try { module.FS.mkdir("/models"); } catch (_) { /* already exists */ }
    try { module.FS.unlink(target); } catch (_) { /* not there */ }

    const stream = module.FS.open(target, "w+");
    let written = 0;
    try {
        for await (const chunk of byteStream(model)) {
            const bytes = chunk instanceof Uint8Array ? chunk : new Uint8Array(chunk);
            module.FS.write(stream, bytes, 0, bytes.length, written);
            written += bytes.length;
            progress({ phase: "load", loaded: written, total: model.total || written });
            await tick();
        }
    } finally {
        module.FS.close(stream);
    }
    return written;
}

// `ingest_only` stops after the bytes are in the library (download + hash): a
// projector is not something the engine can "start", it is picked as a launch
// parameter of the next model load.
async function start({ id, file, mmproj_file, context_size, kv, mmproj_id, threads,
                       ingest_only = false }) {
    if (busy) {
        fail(new Error("a load is already running"));
        return;
    }
    if (file) files.set(id, file);
    if (mmproj_id && mmproj_file) files.set(mmproj_id, mmproj_file);
    if (kv) currentKv = { ...currentKv, ...kv };
    if (threads > 1 && !THREADED) {
        log("no threaded wasm on this page (needs COOP/COEP + SharedArrayBuffer) - using 1 thread");
    }
    currentThreads = THREADED ? Math.max(1, Math.min(64, threads || 1)) : 1;

    busy = true;
    abort = new AbortController();
    lastRequest = { id, context_size, kv: { ...currentKv }, mmproj_id, threads: currentThreads,
                    ingest_only };

    try {
        let model = await getModel(id);
        if (!model) throw new Error(`unknown model ${id}`);

        // 1. get the bytes into the browser (url models only; a local file is
        //    already where it needs to be)
        if (model.source === "url" && !model.complete) {
            await downloadModel(
                { id: model.id, url: model.url, total: model.total },
                { signal: abort.signal, onProgress: progress, storage },
            );
            model = await getModel(id);
        }
        if (model.source === "file" && !model.copied && !files.has(model.id)) {
            throw new Error("the local file is not attached - choose it again");
        }

        // 2. identify the bytes once
        if (!model.sha256) model = await hashModel(model);

        if (ingest_only) {
            busy = false;
            post({ event: action.MODEL_UPDATED, model });
            return;
        }

        // 3. into the wasm FS, then load
        await ensureModule();
        // older single-thread modules predate this export; they clamp to 1 anyway
        if (module._llmjs_set_threads) {
            module.ccall("llmjs_set_threads", null, ["number"], [currentThreads]);
        }

        // the wasm instance can only hold one model at a time
        if (loaded && loadedId !== id) {
            module._llmjs_free();
            loaded = false;
            loadedId = null;
        }

        const bytes = await loadIntoFS(model);
        if (model.received && bytes !== model.received) {
            log(`stored model size ${bytes} != expected ${model.received}`);
        }

        currentNctx = context_size || currentNctx;
        const n = module.ccall("llmjs_load_model", "number", ["string", "number"],
                               ["/models/model.gguf", currentNctx]);
        if (n < 0) throw new Error(`llmjs_load_model failed (${n})`);
        // load_model builds a context with the defaults; this applies the KV
        // options the page asked for (a no-op when they already match)
        if (createContext(currentNctx) < 0) {
            throw new Error("could not create a context with the requested KV cache settings");
        }

        // optional multimodal projector: a second GGUF from the library, loaded
        // by llama.cpp's mtmd stack (llama-server --mmproj)
        let mmproj = null;
        if (mmproj_id) {
            mmproj = await getModel(mmproj_id);
            if (!mmproj) throw new Error(`unknown mmproj ${mmproj_id}`);
            if (mmproj.source === "file" && !mmproj.copied && !files.has(mmproj.id)) {
                throw new Error("the mmproj file is not attached - choose it again");
            }
            await loadIntoFS(mmproj, "/models/mmproj.gguf");
            const rc = module.ccall("llmjs_mmproj_load", "number", ["string"],
                                    ["/models/mmproj.gguf"]);
            if (rc <= 0) throw new Error("could not load the multimodal projector");
        } else {
            module.ccall("llmjs_mmproj_free", null, [], []);
        }

        loaded = true;
        loadedId = id;
        busy = false;
        post({
            event: action.INITIALIZED,
            model,
            bytes,
            n_ctx: module.ccall("llmjs_n_ctx", "number", [], []),
            ctx_config: ctxConfig(),
            threads: currentThreads,
            threaded_build: THREADED,
            hardware_threads: HARDWARE_THREADS,
            mmproj: (() => {
                const text = module.UTF8ToString(module._llmjs_mmproj_info());
                try { return JSON.parse(text); } catch (_) { return null; }
            })(),
            mmproj_name: mmproj ? mmproj.name : null,
            chat_template_name: module.UTF8ToString(module._llmjs_chat_template_name()),
            // the GGUF's own Jinja template, rendered by the engine inside the wasm build
            chat_template: module.UTF8ToString(module._llmjs_chat_template()),
            has_chat_template: module._llmjs_has_chat_template() !== 0,
        });
        post({ event: action.MODEL_UPDATED, model });
    } catch (err) {
        busy = false;
        if (err && err.name === "AbortError") {
            post({ event: action.PAUSED, id });
        } else {
            fail(err);
        }
    }
}

// Set by a STOP message. The generation loop below checks it between tokens,
// which is only possible because generation is stepped from JS instead of one
// long blocking call: inside a single ccall this worker could not read messages.
let cancelRequested = false;
let generating = false;

// mode "chat": a message list rendered through the chat template.
// mode "completion": the prompt is the prompt, no template involved.
// Both honour the same sampling limits and stop sequences.
async function run(opts) {
    const {
        mode = "completion",
        prompt = "",
        messages = null,
        stop = [],
        seed = -1,
        max_token_len = 64,
        top_k = 40,
        top_p = 0.9,
        temp = 0.7,
        repeat_last_n = 32,
        repeat_penalty = 1.1,
        context_size = 0,
        chat_template = "",
        chat_template_source = "",
        tools = "[]",
        template_kwargs = "{}",
        grammar = "",
        json_schema = "",
        images = [],
    } = opts;

    if (!loaded || !module) {
        fail(new Error("no model loaded"));
        return;
    }
    if (generating) {
        fail(new Error("a run is already in progress"));
        return;
    }

    // structured output: only when the page chose a schema. The tool-call path
    // sends neither, so it runs exactly as before. `grammar` above is a const
    // destructured binding, so the schema's result needs its own local.
    let grammarText = grammar;
    if (!grammarText && json_schema && json_schema.trim()) {
        if (typeof module._llmjs_grammar_from_schema !== "function") {
            fail(new Error("this module build has no JSON schema support; the threaded build "
                           + "(needs COOP/COEP) has it, or paste a GBNF grammar instead"));
            return;
        }
        const asGrammar = module.ccall("llmjs_grammar_from_schema", "string",
                                       ["string"], [json_schema]);
        if (!asGrammar) {
            const why = module.UTF8ToString(module._llmjs_last_error());
            fail(new Error(`could not use the JSON schema: ${why}`));
            return;
        }
        grammarText = asGrammar;
    }

    if (createContext(context_size || currentNctx) < 0) {
        fail(new Error("could not create the inference context"));
        return;
    }

    // stop sequences are matched in C++ against the decoded text, so a stop
    // string is never streamed to the page
    module.ccall("llmjs_set_stop_sequences", null, ["string"],
                 [JSON.stringify(Array.isArray(stop) ? stop : [])]);
    // data URLs for the images of this run; the chat template has one media
    // marker per image in the messages it renders
    module.ccall("llmjs_set_images", null, ["string"],
                 [JSON.stringify(Array.isArray(images) ? images : [])]);

    // arm the engine: everything promised by the parameters goes in here
    let rc;
    if (mode === "chat") {
        const list = Array.isArray(messages) && messages.length
            ? messages
            : [{ role: "user", content: prompt }];
        rc = module.ccall(
            "llmjs_gen_begin_chat",
            "number",
            ["string", "string", "string", "string", "number", "number", "number",
             "number", "number", "number", "number", "string"],
            [JSON.stringify(list), tools || "[]", chat_template_source, template_kwargs || "{}",
             max_token_len, temp, top_k, top_p, repeat_penalty, repeat_last_n, seed, grammarText],
        );
    } else {
        rc = module.ccall(
            "llmjs_gen_begin_raw",
            "number",
            ["string", "number", "number", "number", "number", "number",
             "number", "number", "string"],
            [prompt, max_token_len, temp, top_k, top_p, repeat_penalty,
             repeat_last_n, seed, grammarText],
        );
    }
    if (rc < 0) {
        // the engine knows why: an image larger than the context, a marker/image
        // mismatch, no projector … - "(-1)" alone tells the user nothing
        const why = module.UTF8ToString(module._llmjs_last_error());
        fail(new Error(`could not start generation (${rc})${why ? `: ${why}` : ""}`));
        return;
    }

    generating = true;
    cancelRequested = false;
    const t0 = performance.now();
    let tokens = 0;
    try {
        // one token per turn of the event loop: a STOP message is handled
        // during the await, and the loop then exits with the text so far
        while (!cancelRequested) {
            if (module.ccall("llmjs_gen_step", "number", [], []) <= 0) break;
            await tick();
        }
        tokens = module.ccall("llmjs_gen_end", "number", [], []);
    } finally {
        generating = false;
    }

    const ms = performance.now() - t0;
    const interrupted = cancelRequested;
    cancelRequested = false;

    post({
        event: action.RUN_COMPLETED,
        mode,
        tokens,
        ms,
        interrupted,
        tokens_per_second: tokens > 0 ? (tokens * 1000) / ms : 0,
        // the final text, with any stop sequence removed: whatever the page
        // streamed may still contain the beginning of a stop string
        output: module.UTF8ToString(module._llmjs_last_output()),
        // prompt_tokens / reused_tokens / prompt_ms of this run
        run_stats: (() => {
            const text = module.UTF8ToString(module._llmjs_last_run_stats());
            if (!text) return null;
            try { return JSON.parse(text); } catch (_) { return null; }
        })(),
        // jinja / jinja:override / builtin:<name> / raw
        template_strategy: module.UTF8ToString(module._llmjs_chat_template_strategy()),
        // { content, reasoning_content, tool_calls } when the template's parser
        // recognised the reply, null when it did not
        parsed: (() => {
            const text = module.UTF8ToString(module._llmjs_parsed_output());
            if (!text) return null;
            try { return JSON.parse(text); } catch (_) { return null; }
        })(),
    });
}

self.addEventListener("message", (e) => {
    const msg = e.data;
    switch (msg.event) {
        case action.ADD_URL:
            (async () => {
                const id = await addUrlModel(msg);
                // a projector is not a model to run: download its bytes into the
                // library, then stop - it is picked in the engine parameters and
                // loaded next to whatever model is selected
                await start({ id, context_size: msg.context_size, kv: msg.kv, threads: msg.threads,
                              ingest_only: !!msg.mmproj });
            })().catch(fail);
            break;

        case action.ADD_FILE:
            (async () => {
                const id = await addFileModel(msg);
                if (msg.mmproj) return;
                await start({ id, context_size: msg.context_size, kv: msg.kv, threads: msg.threads });
            })().catch(fail);
            break;

        case action.LOAD:
            start(msg);
            break;

        case action.PAUSE:
            if (abort) abort.abort();
            break;

        case action.RESUME:
            if (lastRequest) start(lastRequest);
            else fail(new Error("nothing to resume"));
            break;

        case action.PREVIEW_TEMPLATE:
            (async () => {
                await ensureModule();
                const messages = JSON.stringify(msg.messages && msg.messages.length
                    ? msg.messages
                    : [{ role: "user", content: msg.prompt || "" }]);
                const text = renderTemplateEx(messages, msg.tools || "[]",
                    msg.template_source || "", msg.template_kwargs || "{}");
                post({ event: action.TEMPLATE_PREVIEW, text });
            })().catch(fail);
            break;

        case action.RELOAD_MODEL:
            (async () => {
                // a real reload: drop the model and the backend, then load again
                if (module) {
                    module._llmjs_free();
                    module._llmjs_init();
                }
                loaded = false;
                loadedId = null;
                await start({
                    id: msg.id,
                    file: files.get(msg.id) || null,
                    context_size: msg.context_size,
                    kv: msg.kv,
                    threads: msg.threads,
                    mmproj_id: msg.mmproj_id,
                    mmproj_file: msg.mmproj_id ? files.get(msg.mmproj_id) || null : null,
                });
            })().catch(fail);
            break;

        case action.HASH_MODEL:
            (async () => {
                const model = await getModel(msg.id);
                if (!model) throw new Error(`unknown model ${msg.id}`);
                const hashed = await hashModel(model);
                post({ event: action.MODEL_UPDATED, model: hashed });
            })().catch(fail);
            break;

        case action.DELETE_MODEL:
            (async () => {
                if (loaded && loadedId === msg.id) {
                    module._llmjs_free();
                    loaded = false;
                    loadedId = null;
                }
                files.delete(msg.id);
                const chunks = await deleteModel(msg.id);
                post({ event: action.DELETED, id: msg.id, chunks });
            })().catch(fail);
            break;

        case action.LIST_MODELS:
            listModels()
                .then((models) => post({ event: action.MODELS, models }))
                .catch(fail);
            break;

        case action.STOP:
            if (generating) cancelRequested = true;
            break;

        case action.RUN_MAIN:
            run(msg).catch(fail);
            break;
    }
}, false);