Ask your own PDF, with nothing uploaded

Sandbox

// Ask questions about your own PDF, answered by a model running in this tab.
// Nothing downloads until you click Load. The PDF is read in the page and never uploaded.
// mount(el) builds the whole thing inside el and returns a cleanup function.

const WEBLLM = 'https://esm.run/@mlc-ai/[email protected]';
const PDFJS = 'https://cdn.jsdelivr.net/npm/[email protected]/build/pdf.min.mjs';
const PDFJS_WORKER = 'https://cdn.jsdelivr.net/npm/[email protected]/build/pdf.worker.min.mjs';
const M16 = 'Qwen2.5-1.5B-Instruct-q4f16_1-MLC', M32 = 'Qwen2.5-1.5B-Instruct-q4f32_1-MLC';

// The model holds about 3,000 words, so the PDF is cut into ~150-word passages. A second, small
// model turns each passage into numbers that capture its meaning. The main model searches
// (keywords + meaning), reads a page, then answers. The code adds the page number.
const EMBED = 'snowflake-arctic-embed-s-q0f32-MLC-b4';
const QUERY_PREFIX = 'Represent this sentence for searching relevant passages: ';
const W_BM25 = 0.5, W_COS = 0.5;
const PDF_OPTS = { disableFontFace: true, useSystemFonts: false, isEvalSupported: false, useWorkerFetch: false };
const SYSTEM = `You answer questions about a document by using tools. You cannot see the document until you use a tool. Reply with exactly one JSON object each turn, one of:
{"action":"search","query":"..."}  searches the document and returns the top 3 hits, each with its page number.
{"action":"read_page","page":N}  returns the full text of page N.
{"action":"answer","answer":"...","evidence_page":N}  gives your final answer.
Rules: the question is written in everyday words but the document may use more formal or technical wording, so reword the question into words the document would use before you search. Read the page that looks most relevant before you answer. Answer briefly, in one sentence, using only text from the document. If the document does not say, answer "Not in the document."`;
const CHUNK = 150, OVERLAP = 20, SNIPPET = 400, PAGE_CHARS = 2500, MAX_STEPS = 5, CTX_CHARS = 12000, EMBED_CHARS = 1000, EMBED_BATCH = 16;
const STOP = new Set('the a an of to in and or is are was be by for on at as it this that with from what which who how does do must shall can'.split(' '));
const tok = (s) => (s.toLowerCase().match(/[a-z0-9]+/g) || []).filter((w) => !STOP.has(w)).map((w) => (w.length > 3 ? w.replace(/(ing|es|s)$/, '') : w));

const S_SEARCH = { type: 'object', properties: { action: { const: 'search' }, query: { type: 'string' } }, required: ['action', 'query'] };
const S_READ = { type: 'object', properties: { action: { const: 'read_page' }, page: { type: 'integer' } }, required: ['action', 'page'] };
const S_ANS = { type: 'object', properties: { action: { const: 'answer' }, answer: { type: 'string' }, evidence_page: { type: 'integer' } }, required: ['action', 'answer', 'evidence_page'] };
const SCHEMA_ANY = JSON.stringify({ anyOf: [S_SEARCH, S_READ, S_ANS] }), SCHEMA_ANSWER = JSON.stringify(S_ANS);
const SCHEMA_FLAT = JSON.stringify({ type: 'object', properties: { action: { enum: ['search', 'read_page', 'answer'] }, query: { type: 'string' }, page: { type: 'integer' }, answer: { type: 'string' }, evidence_page: { type: 'integer' } }, required: ['action', 'query', 'page', 'answer', 'evidence_page'] });

const BOX = 'display:block;box-sizing:border-box;width:100%;padding:8px;background:transparent;color:var(--ink);border:1px solid var(--ink-faint);border-radius:var(--radius-sm);font:13px/1.5 var(--font-mono)';
const NOTE = 'margin:8px 0 0;font:13px/1.5 var(--font-mono);color:var(--ink-muted)';

export function mount(el) {
  el.innerHTML = `
    <p class="msg" hidden>Your browser can't run this. Try Chrome or Safari 26 on a laptop.</p>
    <div data-notice style="margin:0 0 16px;padding:12px;border:1px solid var(--line);border-radius:var(--radius-sm);font:13px/1.5 var(--font-mono)">
      <div class="t-label" style="color:var(--ink-muted);margin-bottom:8px">Online first, then offline</div>
      <p style="margin:0 0 4px;color:var(--ink)"><span style="color:var(--ink-muted)">1. ONLINE</span> Load the models. It needs internet once, to download it.</p>
      <p style="margin:0 0 8px;color:var(--ink)"><span style="color:var(--ink-muted)">2. OFFLINE</span> Then switch off Wi-Fi, pick a PDF and ask.</p>
      <p style="margin:0;color:var(--error)">Don't refresh the page while offline. A refresh needs internet, and you'll lose the models until you're back online.</p>
    </div>
    <div data-load-box>
      <button class="btn btn-primary" type="button" data-load>Load models: about 1.1 GB</button>
      <p style="${NOTE}">Qwen2.5-1.5B (chat) and snowflake-arctic-embed-s (search). Downloads once, then runs on your GPU.</p>
    </div>
    <p class="readout" data-prog style="margin:12px 0 0"></p>
    <div data-main hidden>
      <label style="display:block;margin-top:16px;font:13px var(--font-mono);color:var(--ink-muted)">Choose a PDF
        <input type="file" accept="application/pdf" data-file style="display:block;margin-top:8px;max-width:100%;color:var(--ink);font:13px var(--font-mono)">
      </label>
      <p data-doc style="${NOTE}"></p>
      <div data-ask hidden>
        <textarea rows="2" placeholder="Ask a question about the PDF" aria-label="Question" style="${BOX};margin-top:16px;resize:vertical"></textarea>
        <p style="margin:12px 0"><button class="btn" type="button" data-go>Ask</button></p>
        <pre data-trace style="margin:0 0 8px;white-space:pre-wrap;overflow-wrap:anywhere;font:12px/1.5 var(--font-mono);color:var(--ink-muted)"></pre>
        <pre data-out aria-live="polite" style="margin:0;min-height:3em;white-space:pre-wrap;overflow-wrap:anywhere;font:13px/1.5 var(--font-mono);color:var(--ink)"></pre>
      </div>
      <div class="controls">
        <span class="readout">network requests since model ready: <output data-net style="min-width:0">0</output></span>
        <span data-online></span>
        <span data-warn style="color:var(--error)" hidden>Don't refresh while offline.</span>
      </div>
      <p style="${NOTE}">Turn off Wi-Fi, pick a PDF and ask. Watch the counter.</p>
    </div>`;

  const $ = (s) => el.querySelector(s);
  const [msg, notice, loadBox, load, prog, main, file, doc, askBox, q, go, trace, out, net, online, warn] =
    ['.msg', '[data-notice]', '[data-load-box]', '[data-load]', '[data-prog]', '[data-main]', '[data-file]', '[data-doc]', '[data-ask]', 'textarea', '[data-go]', '[data-trace]', '[data-out]', '[data-net]', '[data-online]', '[data-warn]'].map($);

  let engine = null, pdfjs = null, observer = null, count = 0, cancelled = false, adapter = null, model = '';
  let chunks = [], pages = [], anySchema = SCHEMA_ANY;

  const showOnline = () => {
    online.textContent = navigator.onLine ? 'online' : 'offline';
    warn.hidden = navigator.onLine || !engine;
  };
  window.addEventListener('online', showOnline);
  window.addEventListener('offline', showOnline);
  showOnline();

  (async () => {
    try { adapter = navigator.gpu && (await navigator.gpu.requestAdapter()); } catch {}
    if (!adapter) { msg.hidden = false; loadBox.hidden = true; notice.hidden = true; }
  })();

  const gpuStopped = (err) => /mapAsync|Instance reference/.test(err?.message || '') || err?.name === 'DeviceLostError' || /DeviceLost/.test(String(err?.constructor?.name));
  const explain = (err) => (gpuStopped(err) ? "Your device's GPU stopped the model. This is a known issue on some Android phones. Try a laptop." : `Error: ${err?.message || err}`);

  load.addEventListener('click', async () => {
    load.disabled = true;
    try {
      // Imported only now, after the click, so nothing downloads on page load.
      const [webllm, pdf] = await Promise.all([import(/* @vite-ignore */ WEBLLM), import(/* @vite-ignore */ PDFJS)]);
      // Fetch the pdf.js worker now so reading a PDF later makes no network request.
      const w = await (await fetch(PDFJS_WORKER)).text();
      pdf.GlobalWorkerOptions.workerSrc = URL.createObjectURL(new Blob([w], { type: 'text/javascript' }));
      pdfjs = pdf;
      model = adapter.features.has('shader-f16') ? M16 : M32;
      // Two models in one engine: every request below names which one it wants.
      const e = await webllm.CreateMLCEngine([model, EMBED], {
        initProgressCallback: (p) => (prog.textContent = p.progress < 1 ? `Loading: ${Math.round(p.progress * 100)}%` : 'Almost ready'),
      });
      if (cancelled) return e.unload();
      engine = e;
      // From here on, every request this page makes shows up in the counter.
      observer = new PerformanceObserver((list) => (net.textContent = count += list.getEntries().length));
      observer.observe({ type: 'resource', buffered: false });
      prog.textContent = 'Running: Qwen2.5-1.5B + snowflake-arctic-embed-s. Ready. Choose a PDF.';
      loadBox.hidden = true;
      main.hidden = false;
      showOnline();
    } catch (err) {
      prog.textContent = explain(err);
      load.disabled = false;
    }
  });

  const embed = async (texts) => (await engine.embeddings.create({ input: texts, model: EMBED })).data.map((d) => d.embedding);

  file.addEventListener('change', async () => {
    const f = file.files[0];
    if (!f) return;
    askBox.hidden = true;
    doc.textContent = 'Reading...';
    try {
      const t0 = performance.now();
      const pdf = await pdfjs.getDocument({ data: new Uint8Array(await f.arrayBuffer()), ...PDF_OPTS }).promise;
      pages = [];
      for (let i = 1; i <= pdf.numPages; i++) {
        const t = await (await pdf.getPage(i)).getTextContent();
        pages.push(t.items.map((x) => x.str).join(' ').replace(/\s+/g, ' ').trim());
      }
      chunks = [];
      pages.forEach((txt, i) => {
        const w = txt.split(/\s+/).filter(Boolean);
        for (let s = 0; s < w.length; s += CHUNK - OVERLAP) {
          const text = w.slice(s, s + CHUNK).join(' ');
          if (text) chunks.push({ page: i + 1, text, toks: tok(text) });
          if (s + CHUNK >= w.length) break;
        }
      });
      // Turn every passage into numbers once, so each question can search by meaning.
      for (let i = 0; i < chunks.length; i += EMBED_BATCH) {
        const vs = await embed(chunks.slice(i, i + EMBED_BATCH).map((c) => c.text.slice(0, EMBED_CHARS)));
        vs.forEach((v, k) => (chunks[i + k].vec = v));
        doc.textContent = `Indexing ${Math.min(i + EMBED_BATCH, chunks.length)}/${chunks.length} passages...`;
      }
      doc.textContent = `Read ${pdf.numPages} pages, indexed ${chunks.length} passages in ${((performance.now() - t0) / 1000).toFixed(1)} s.` + (pages.some((p) => !p) ? " Some pages have no text (scanned?). Those can't be searched yet." : '');
      askBox.hidden = !chunks.length;
    } catch (err) {
      doc.textContent = `Error: ${err?.message || err}`;
    }
  });

  const cosine = (a, b) => { let d = 0, na = 0, nb = 0; for (let i = 0; i < a.length; i++) { d += a[i] * b[i]; na += a[i] * a[i]; nb += b[i] * b[i]; } return d / (Math.sqrt(na * nb) || 1); };
  const minmax = (xs) => { const lo = Math.min(...xs), hi = Math.max(...xs); return xs.map((x) => (hi > lo ? (x - lo) / (hi - lo) : 0)); };

  // Hybrid search: BM25 keywords (half) plus meaning via embeddings (half). Returns the top 3 as {page, text}.
  const search = async (query, k1 = 1.5, b = 0.75) => {
    const q = String(query || ''), qt = [...new Set(tok(q))], N = chunks.length, avg = chunks.reduce((a, c) => a + c.toks.length, 0) / N || 1;
    const df = {};
    qt.forEach((t) => (df[t] = chunks.filter((c) => c.toks.includes(t)).length));
    const bm = chunks.map((c) => {
      const tf = {};
      c.toks.forEach((t) => (tf[t] = (tf[t] || 0) + 1));
      let sc = 0;
      qt.forEach((t) => { if (tf[t]) sc += (Math.log(1 + (N - df[t] + 0.5) / (df[t] + 0.5)) * tf[t] * (k1 + 1)) / (tf[t] + k1 * (1 - b + (b * c.toks.length) / avg)); });
      return sc;
    });
    const bmMax = Math.max(...bm) || 1;
    const [qv] = await embed([QUERY_PREFIX + q.slice(0, EMBED_CHARS)]);
    const cos = minmax(chunks.map((c) => cosine(qv, c.vec)));
    return chunks.map((c, i) => ({ page: c.page, text: c.text.slice(0, SNIPPET), score: W_BM25 * (bm[i] / bmMax) + W_COS * cos[i] }))
      .sort((x, y) => y.score - x.score).slice(0, 3);
  };

  // Page number comes from the code: among pages the model was shown, the one whose text shares
  // the most words with the answer (ties go to the model's own evidence_page).
  const cite = (answer, modelPage, shown) => {
    const A = new Set(tok(answer));
    let best = null, bs = -1;
    for (const [page, txt] of shown) {
      const T = new Set(tok(txt));
      let n = 0;
      A.forEach((t) => T.has(t) && n++);
      if (n > bs || (n === bs && page === modelPage)) { bs = n; best = page; }
    }
    return best;
  };
  const see = (shown, page, text) => shown.set(page, (shown.get(page) || '') + ' ' + text);
  // Guards for a small model: strip trailing whitespace, and pull the answer out of broken JSON.
  const clean = (raw) => String(raw).replace(/[\s\t]+$/, '').replace(/([}"])[\s\t]+$/, '$1').trim();
  const salvage = (raw) => {
    const a = /"answer"\s*:\s*"([^"]*)"/.exec(raw);
    if (!a) return null;
    const p = /"evidence_page"\s*:\s*(\d+)/.exec(raw);
    return { action: 'answer', answer: a[1], ...(p ? { evidence_page: Number(p[1]) } : {}) };
  };
  // Keep the running context under CTX_CHARS by dropping the oldest tool results first.
  const trim = (msgs) => {
    while (msgs.reduce((a, m) => a + m.content.length, 0) > CTX_CHARS) {
      const i = msgs.findIndex((m, k) => k >= 2 && m.tool);
      if (i < 0) break;
      msgs.splice(i, 1);
    }
    return msgs.map(({ role, content }) => ({ role, content }));
  };
  const llm = async (messages, schema) => {
    const run = (s) => engine.chat.completions.create({ stream: false, temperature: 0, max_tokens: 200, messages, model, response_format: { type: 'json_object', schema: s } });
    try { return (await run(schema)).choices[0].message.content || ''; }
    catch (e) { if (schema === SCHEMA_ANY) { anySchema = SCHEMA_FLAT; return (await run(SCHEMA_FLAT)).choices[0].message.content || ''; } throw e; }
  };

  go.addEventListener('click', async () => {
    go.disabled = true;
    out.textContent = trace.textContent = '';
    const line = (t) => (trace.textContent += t + '\n');
    const msgs = [{ role: 'system', content: SYSTEM }, { role: 'user', content: `Question: ${q.value}` }];
    const shown = new Map(), used = new Set();
    let retried = false;
    try {
      for (let n = 1; n <= MAX_STEPS; n++) {
        const last = n === MAX_STEPS;
        if (last) msgs.push({ role: 'user', content: 'This is your last step. Answer now with the action answer.' });
        let raw = '', obj = null;
        try { raw = await llm(trim(msgs), last ? SCHEMA_ANSWER : anySchema); obj = JSON.parse(clean(raw)); } catch (e) { if (gpuStopped(e)) throw e; }
        if (!obj && raw) obj = salvage(raw);
        if (!obj && !retried) {
          retried = true;
          line('invalid JSON, retrying once');
          msgs.push({ role: 'assistant', content: raw || '{}' }, { tool: true, role: 'user', content: 'Reply with one valid JSON object.' });
        } else if (!obj) {
          line('invalid reply');
          msgs.push({ role: 'assistant', content: raw || '{}' }, { tool: true, role: 'user', content: 'That was not valid. Reply with one JSON object.' });
        } else if (obj.action === 'answer' || last) {
          const answer = String(obj.answer ?? raw), p = cite(answer, obj.evidence_page, shown);
          // No page for "not found": a page number there would look like evidence.
          out.textContent = /not in the document/i.test(answer) ? answer : `${answer} (p. ${p ?? '?'})`;
          return;
        } else if (obj.action === 'search') {
          const key = String(obj.query ?? '').trim().toLowerCase().replace(/\s+/g, ' ');
          if (used.has(key)) {
            line(`searched again, blocked: ${obj.query}`);
            msgs.push({ role: 'assistant', content: raw }, { tool: true, role: 'user', content: 'You already searched that. Search with different words or answer.' });
          } else {
            used.add(key);
            const hits = await search(obj.query);
            hits.forEach((h) => see(shown, h.page, h.text));
            line(`searched: ${obj.query} -> ${hits.map((h) => 'p' + h.page).join(', ') || 'no hits'}`);
            msgs.push({ role: 'assistant', content: raw }, { tool: true, role: 'user', content: 'Search results:\n' + (hits.length ? hits.map((h, i) => `Hit ${i + 1} (page ${h.page}): ${h.text}`).join('\n\n') : 'No hits. Try different words.') });
          }
        } else if (obj.action === 'read_page') {
          const p = Math.trunc(Number(obj.page)), ok = p >= 1 && p <= pages.length;
          if (ok) see(shown, p, pages[p - 1].slice(0, PAGE_CHARS));
          line(ok ? `read p${p}` : `read p${obj.page} (no such page)`);
          msgs.push({ role: 'assistant', content: raw }, { tool: true, role: 'user', content: ok ? `Page ${p}:\n${pages[p - 1].slice(0, PAGE_CHARS)}` : `There is no page ${obj.page}. The document has ${pages.length} pages.` });
        } else {
          line('unknown action');
          msgs.push({ role: 'assistant', content: raw }, { tool: true, role: 'user', content: 'Unknown action. Use search, read_page or answer.' });
        }
      }
      out.textContent = '(no answer)';
    } catch (err) {
      out.textContent = explain(err);
    } finally {
      go.disabled = false;
    }
  });

  return () => {
    cancelled = true;
    observer?.disconnect();
    window.removeEventListener('online', showOnline);
    window.removeEventListener('offline', showOnline);
    engine?.unload();
  };
}