Spaces:

Mike0021
/

pulpie-web-demo

Running

File size: 18,594 Bytes

import { AutoTokenizer, env } from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.8.1";
import * as ort from "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.22.0/dist/ort.webgpu.min.mjs";

const MODEL_ID = "Mike0021/pulpie-orange-small-onnx";
const MODEL_URL = `https://huggingface.co/${MODEL_ID}/resolve/main/model.onnx`;
const MODEL_BYTES = 847167095;
const MAX_TOKENS = 8192;
const BOS_ID = 128000;
const EOS_ID = 128001;
const SEP_ID = 128256;

env.allowRemoteModels = true;
env.allowLocalModels = false;
env.useBrowserCache = true;

ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.22.0/dist/";
ort.env.wasm.numThreads = 1;
ort.env.wasm.simd = true;

const els = {
  modelStatus: document.querySelector("#modelStatus"),
  loadLabel: document.querySelector("#loadLabel"),
  loadPercent: document.querySelector("#loadPercent"),
  loadProgress: document.querySelector("#loadProgress"),
  urlInput: document.querySelector("#urlInput"),
  fetchButton: document.querySelector("#fetchButton"),
  extractButton: document.querySelector("#extractButton"),
  examples: document.querySelector("#examples"),
  htmlInput: document.querySelector("#htmlInput"),
  keptStat: document.querySelector("#keptStat"),
  droppedStat: document.querySelector("#droppedStat"),
  processingStat: document.querySelector("#processingStat"),
  modelLoadStat: document.querySelector("#modelLoadStat"),
  originalMeta: document.querySelector("#originalMeta"),
  cleanMeta: document.querySelector("#cleanMeta"),
  blockMeta: document.querySelector("#blockMeta"),
  originalOutput: document.querySelector("#originalOutput"),
  cleanOutput: document.querySelector("#cleanOutput"),
  blocksOutput: document.querySelector("#blocksOutput"),
};

const state = {
  tokenizer: null,
  session: null,
  provider: "wasm",
  loadingPromise: null,
  modelLoadMs: 0,
  busy: false,
};

const examples = [
  {
    name: "News article",
    html: `<html>
  <body>
    <header><a>Subscribe</a> <a>Markets</a> <a>Sports</a></header>
    <main>
      <article>
        <h1>Researchers release a compact content extraction model</h1>
        <p>The encoder reads a long page in one pass and marks boilerplate for removal.</p>
        <p>Benchmarks compare throughput, model size, and extraction quality.</p>
        <ul>
          <li>Small checkpoint</li>
          <li>Long context</li>
          <li>Block-level labels</li>
        </ul>
      </article>
    </main>
    <footer>Copyright 2026. Privacy policy. Cookie settings.</footer>
  </body>
</html>`,
  },
  {
    name: "Docs page",
    html: `<div class="layout">
  <nav>Home Install API Changelog</nav>
  <article>
    <h1>Install the extractor</h1>
    <p>Install the package, tokenize simplified HTML blocks, then classify separator tokens.</p>
    <pre>pip install pulpie onnxruntime</pre>
    <p>The output can be reconstructed as clean Markdown.</p>
  </article>
  <aside>Related links: pricing, contact, status</aside>
</div>`,
  },
  {
    name: "Blog post",
    html: `<body>
  <div class="cookie">We use cookies to improve this site. Accept all.</div>
  <article>
    <h1>Building fast browser ML demos</h1>
    <p>Static Spaces can host complete machine learning demos without a Python server.</p>
    <p>The browser downloads model weights once and keeps them in its cache for later visits.</p>
    <blockquote>Good progress indicators make large downloads feel predictable.</blockquote>
  </article>
  <section class="newsletter">Get weekly updates in your inbox.</section>
</body>`,
  },
  {
    name: "Product page",
    html: `<main>
  <header>Logo Docs Blog Sign in</header>
  <section>
    <h1>Pulpie Orange</h1>
    <p>A compact EuroBERT model for separating article text from page chrome.</p>
    <p>It scores simplified HTML blocks and keeps the pieces that read like main content.</p>
  </section>
  <footer>Terms Security Contact</footer>
</main>`,
  },
];

function setStatus(text, mode = "") {
  els.modelStatus.textContent = text;
  els.modelStatus.className = `status-pill ${mode}`.trim();
}

function setBusy(isBusy) {
  state.busy = isBusy;
  els.extractButton.disabled = isBusy;
  els.fetchButton.disabled = isBusy;
}

function formatBytes(bytes) {
  const units = ["B", "KB", "MB", "GB"];
  let value = bytes;
  let unit = 0;
  while (value >= 1024 && unit < units.length - 1) {
    value /= 1024;
    unit += 1;
  }
  return `${value.toFixed(unit === 0 ? 0 : 1)} ${units[unit]}`;
}

function formatTime(ms) {
  if (!Number.isFinite(ms) || ms <= 0) return "0 ms";
  return ms < 1000 ? `${Math.round(ms)} ms` : `${(ms / 1000).toFixed(1)} s`;
}

function updateProgress(label, loaded, total = MODEL_BYTES) {
  const pct = total ? Math.min(100, (loaded / total) * 100) : 0;
  els.loadLabel.textContent = label;
  els.loadPercent.textContent = `${Math.round(pct)}%`;
  els.loadProgress.style.width = `${pct}%`;
}

function tokenizerProgress(info) {
  if (info.status === "progress") {
    const total = info.total || 1;
    updateProgress(`Tokenizer: ${info.file}`, info.loaded || 0, total);
  } else if (info.status === "download") {
    updateProgress(`Tokenizer: ${info.file}`, 0, 1);
  }
}

async function downloadModel() {
  updateProgress("Downloading model.onnx", 0, MODEL_BYTES);
  const response = await fetch(MODEL_URL);
  if (!response.ok) {
    throw new Error(`Model download failed with HTTP ${response.status}`);
  }

  const total = Number(response.headers.get("content-length")) || MODEL_BYTES;
  if (!response.body) {
    const buffer = await response.arrayBuffer();
    updateProgress("Downloaded model.onnx", buffer.byteLength, buffer.byteLength);
    return buffer;
  }

  const reader = response.body.getReader();
  const bytes = new Uint8Array(total);
  let loaded = 0;
  const chunks = [];

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    if (loaded + value.length <= bytes.length) {
      bytes.set(value, loaded);
    } else {
      chunks.push(value);
    }
    loaded += value.length;
    updateProgress(
      `Downloading model.onnx (${formatBytes(loaded)} / ${formatBytes(total)})`,
      loaded,
      total,
    );
  }

  updateProgress("Creating ONNX Runtime session", total, total);
  if (chunks.length === 0 && loaded === bytes.length) return bytes.buffer;

  const merged = new Uint8Array(loaded);
  merged.set(bytes.subarray(0, Math.min(loaded, bytes.length)));
  let offset = bytes.length;
  for (const chunk of chunks) {
    merged.set(chunk, offset);
    offset += chunk.length;
  }
  return merged.buffer;
}

async function createSession(modelBuffer) {
  const baseOptions = {
    graphOptimizationLevel: "all",
    executionMode: "sequential",
    logSeverityLevel: 2,
  };

  let hasWebGpuAdapter = false;
  if ("gpu" in navigator && typeof navigator.gpu.requestAdapter === "function") {
    try {
      hasWebGpuAdapter = Boolean(await navigator.gpu.requestAdapter());
    } catch {
      hasWebGpuAdapter = false;
    }
  }

  if (hasWebGpuAdapter) {
    try {
      const session = await ort.InferenceSession.create(modelBuffer, {
        ...baseOptions,
        executionProviders: ["webgpu"],
      });
      state.provider = "webgpu";
      return session;
    } catch (error) {
      console.warn("WebGPU session failed; falling back to WASM.", error);
    }
  }

  state.provider = "wasm";
  return ort.InferenceSession.create(modelBuffer, {
    ...baseOptions,
    executionProviders: ["wasm"],
  });
}

async function ensureModel() {
  if (state.session && state.tokenizer) return;
  if (state.loadingPromise) return state.loadingPromise;

  state.loadingPromise = (async () => {
    const start = performance.now();
    setStatus("Loading", "busy");
    els.modelLoadStat.textContent = "loading";

    state.tokenizer = await AutoTokenizer.from_pretrained(MODEL_ID, {
      progress_callback: tokenizerProgress,
    });

    const modelBuffer = await downloadModel();
    state.session = await createSession(modelBuffer);

    state.modelLoadMs = performance.now() - start;
    els.modelLoadStat.textContent = formatTime(state.modelLoadMs);
    updateProgress(`Model ready on ${state.provider.toUpperCase()}`, MODEL_BYTES, MODEL_BYTES);
    setStatus(state.provider.toUpperCase(), "ready");
  })();

  try {
    await state.loadingPromise;
  } finally {
    state.loadingPromise = null;
  }
}

function normalizeText(text) {
  return text.replace(/\s+/g, " ").trim();
}

function escapeHtml(value) {
  return value
    .replaceAll("&", "&amp;")
    .replaceAll("<", "&lt;")
    .replaceAll(">", "&gt;")
    .replaceAll('"', "&quot;");
}

function simpleTag(tagName) {
  const tag = tagName.toLowerCase();
  if (/^h[1-6]$/.test(tag) || ["p", "li", "blockquote", "pre"].includes(tag)) return tag;
  return "p";
}

function extractBlocks(rawHtml) {
  const doc = new DOMParser().parseFromString(rawHtml, "text/html");
  doc
    .querySelectorAll("script, style, noscript, template, svg, canvas, iframe, object, embed")
    .forEach((node) => node.remove());

  const selector = [
    "h1",
    "h2",
    "h3",
    "h4",
    "h5",
    "h6",
    "p",
    "li",
    "blockquote",
    "pre",
    "figcaption",
    "td",
    "th",
    "nav",
    "header",
    "footer",
    "aside",
    "section",
    "article",
    "main",
    "div",
  ].join(",");
  const strongBlockSelector = "h1,h2,h3,h4,h5,h6,p,li,blockquote,pre,figcaption,td,th";
  const nodes = Array.from(doc.body.querySelectorAll(selector));
  const seen = new Set();
  const blocks = [];

  for (const node of nodes) {
    const tag = node.tagName.toLowerCase();
    const text = normalizeText(node.textContent || "");
    if (text.length < 2) continue;
    if (seen.has(text) && text.length > 24) continue;

    const isStrong = node.matches(strongBlockSelector);
    const hasStrongChildren = Boolean(node.querySelector(strongBlockSelector));
    const directText = Array.from(node.childNodes)
      .filter((child) => child.nodeType === Node.TEXT_NODE)
      .map((child) => child.textContent || "")
      .join(" ");
    const directRatio = normalizeText(directText).length / Math.max(text.length, 1);

    if (!isStrong && hasStrongChildren && directRatio < 0.32) continue;
    if (!isStrong && text.length < 18) continue;

    const outTag = simpleTag(tag);
    const id = blocks.length;
    const simplified = `<${outTag} _item_id="${id}">${escapeHtml(text)}</${outTag}>`;
    blocks.push({
      id,
      tag: outTag,
      text,
      simplified,
      original: node.outerHTML,
      tokenIds: [],
      prediction: 0,
      score: 0,
    });
    seen.add(text);
  }

  if (blocks.length === 0) {
    const text = normalizeText(doc.body.textContent || rawHtml);
    if (text) {
      blocks.push({
        id: 0,
        tag: "p",
        text,
        simplified: `<p _item_id="0">${escapeHtml(text)}</p>`,
        original: rawHtml,
        tokenIds: [],
        prediction: 0,
        score: 0,
      });
    }
  }

  return blocks;
}

function packChunks(blocks) {
  const chunks = [];
  let ids = [BOS_ID];
  let sepPositions = [];
  let blockIndices = [];

  function flush() {
    if (blockIndices.length === 0) return;
    ids.push(EOS_ID);
    chunks.push({ ids, sepPositions, blockIndices });
    ids = [BOS_ID];
    sepPositions = [];
    blockIndices = [];
  }

  for (const block of blocks) {
    let tokenIds = block.tokenIds;
    const maxBlockTokens = MAX_TOKENS - 3;
    if (tokenIds.length > maxBlockTokens) {
      tokenIds = tokenIds.slice(0, maxBlockTokens);
    }

    if (ids.length + tokenIds.length + 2 > MAX_TOKENS) flush();
    ids.push(...tokenIds);
    ids.push(SEP_ID);
    sepPositions.push(ids.length - 1);
    blockIndices.push(block.id);
  }

  flush();
  return chunks;
}

function toBigIntTensor(values, dims) {
  return new ort.Tensor("int64", BigInt64Array.from(values, (value) => BigInt(value)), dims);
}

async function classifyBlocks(blocks) {
  for (const block of blocks) {
    block.tokenIds = state.tokenizer.encode(block.simplified, { add_special_tokens: false });
  }

  const chunks = packChunks(blocks);
  const predictions = new Array(blocks.length).fill(0);
  const scores = new Array(blocks.length).fill(0);
  let totalTokens = 0;

  for (const chunk of chunks) {
    totalTokens += chunk.ids.length;
    const dims = [1, chunk.ids.length];
    const feeds = {
      input_ids: toBigIntTensor(chunk.ids, dims),
      attention_mask: toBigIntTensor(new Array(chunk.ids.length).fill(1), dims),
    };

    if (state.session.inputNames.includes("token_type_ids")) {
      feeds.token_type_ids = toBigIntTensor(new Array(chunk.ids.length).fill(0), dims);
    }

    const output = await state.session.run(feeds);
    const logits = output.logits || output[state.session.outputNames[0]];
    const data = logits.data;

    for (let i = 0; i < chunk.sepPositions.length; i += 1) {
      const blockIndex = chunk.blockIndices[i];
      const sep = chunk.sepPositions[i];
      const other = Number(data[sep * 2]);
      const main = Number(data[sep * 2 + 1]);
      predictions[blockIndex] = main >= other ? 1 : 0;
      scores[blockIndex] = 1 / (1 + Math.exp(other - main));
    }
  }

  blocks.forEach((block, index) => {
    block.prediction = predictions[index];
    block.score = scores[index];
  });

  return { chunks: chunks.length, totalTokens };
}

function renderClean(blocks) {
  const kept = blocks.filter((block) => block.prediction === 1);
  els.cleanOutput.replaceChildren();

  if (kept.length === 0) {
    const empty = document.createElement("p");
    empty.className = "empty-state";
    empty.textContent = "No main content blocks were selected.";
    els.cleanOutput.append(empty);
    return "";
  }

  let cleanText = "";
  let list = null;

  for (const block of kept) {
    if (block.tag === "li") {
      if (!list) {
        list = document.createElement("ul");
        els.cleanOutput.append(list);
      }
      const li = document.createElement("li");
      li.textContent = block.text;
      list.append(li);
      cleanText += `- ${block.text}\n`;
      continue;
    }

    list = null;
    const element = /^h[1-6]$/.test(block.tag)
      ? document.createElement("h3")
      : document.createElement("p");
    element.textContent = block.text;
    els.cleanOutput.append(element);
    cleanText += `${block.text}\n\n`;
  }

  return cleanText.trim();
}

function renderBlocks(blocks) {
  els.blocksOutput.replaceChildren();
  for (const block of blocks) {
    const item = document.createElement("article");
    item.className = `block ${block.prediction === 1 ? "main" : "other"}`;

    const top = document.createElement("div");
    top.className = "block-top";
    const label = document.createElement("span");
    label.textContent = block.prediction === 1 ? "Content" : "Boilerplate";
    const score = document.createElement("span");
    score.textContent = `${Math.round(block.score * 100)}%`;
    top.append(label, score);

    const text = document.createElement("p");
    text.textContent = block.text;
    item.append(top, text);
    els.blocksOutput.append(item);
  }
}

function updateStats(blocks, timingMs, totalTokens, chunks) {
  const kept = blocks.filter((block) => block.prediction === 1).length;
  const dropped = blocks.length - kept;
  els.keptStat.textContent = String(kept);
  els.droppedStat.textContent = String(dropped);
  els.processingStat.textContent = formatTime(timingMs);
  els.blockMeta.textContent = `${blocks.length} blocks, ${totalTokens} tokens, ${chunks} chunk${chunks === 1 ? "" : "s"}`;
}

async function runExtraction() {
  const rawHtml = els.htmlInput.value.trim();
  if (!rawHtml) {
    setStatus("Add HTML", "error");
    return;
  }

  try {
    setBusy(true);
    setStatus("Running", "busy");
    els.originalOutput.textContent = rawHtml;
    els.originalMeta.textContent = `${rawHtml.length.toLocaleString()} chars`;

    await ensureModel();

    const started = performance.now();
    const blocks = extractBlocks(rawHtml);
    const { chunks, totalTokens } = await classifyBlocks(blocks);
    const cleanText = renderClean(blocks);
    const timingMs = performance.now() - started;

    renderBlocks(blocks);
    updateStats(blocks, timingMs, totalTokens, chunks);
    els.cleanMeta.textContent = `${cleanText.length.toLocaleString()} chars`;
    setStatus(state.provider.toUpperCase(), "ready");
  } catch (error) {
    console.error(error);
    setStatus("Error", "error");
    els.cleanOutput.replaceChildren();
    const message = document.createElement("p");
    message.className = "empty-state";
    message.textContent = error instanceof Error ? error.message : String(error);
    els.cleanOutput.append(message);
  } finally {
    setBusy(false);
  }
}

async function fetchUrl() {
  const url = els.urlInput.value.trim();
  if (!url) return;

  setBusy(true);
  setStatus("Fetching", "busy");
  try {
    let html;
    try {
      const response = await fetch(url);
      if (!response.ok) throw new Error(`HTTP ${response.status}`);
      html = await response.text();
    } catch {
      const proxyUrl = `https://api.allorigins.win/raw?url=${encodeURIComponent(url)}`;
      const response = await fetch(proxyUrl);
      if (!response.ok) throw new Error(`Proxy HTTP ${response.status}`);
      html = await response.text();
    }
    els.htmlInput.value = html;
    els.originalOutput.textContent = html;
    els.originalMeta.textContent = `${html.length.toLocaleString()} chars`;
    setStatus(state.session ? state.provider.toUpperCase() : "Fetched", state.session ? "ready" : "");
  } catch (error) {
    console.error(error);
    setStatus("Fetch error", "error");
  } finally {
    setBusy(false);
  }
}

function loadExample(example) {
  els.htmlInput.value = example.html;
  els.originalOutput.textContent = example.html;
  els.originalMeta.textContent = `${example.html.length.toLocaleString()} chars`;
  if (state.session && !state.busy) {
    runExtraction();
  } else {
    setStatus("Example loaded");
  }
}

function initExamples() {
  for (const example of examples) {
    const button = document.createElement("button");
    button.className = "example-button";
    button.type = "button";
    button.textContent = example.name;
    button.addEventListener("click", () => loadExample(example));
    els.examples.append(button);
  }
  loadExample(examples[0]);
}

els.extractButton.addEventListener("click", runExtraction);
els.fetchButton.addEventListener("click", fetchUrl);
els.urlInput.addEventListener("keydown", (event) => {
  if (event.key === "Enter") fetchUrl();
});

initExamples();