Spaces:
Sleeping
Sleeping
| <html> | |
| <head> | |
| <meta charset="UTF-8"> | |
| <style> | |
| body { font-family: sans-serif; margin: 0; padding: 6px; font-size: 14px; color: #333; } | |
| #status { color: #666; margin-bottom: 4px; } | |
| #progress { width: 100%; height: 6px; display: none; margin-bottom: 4px; } | |
| #error { color: #c00; line-height: 1.5; } | |
| </style> | |
| </head> | |
| <body> | |
| <div id="status"></div> | |
| <progress id="progress" max="100" value="0"></progress> | |
| <div id="error"></div> | |
| <script type="module"> | |
| const MODEL_REPO = "juanbascur/specter2-proximity-onnx"; | |
| const BATCH_SIZE = 64; | |
| // ── Streamlit component API ────────────────────────────────────────────── | |
| function sendValue(value) { | |
| window.parent.postMessage({ type: "streamlit:setComponentValue", value }, "*"); | |
| } | |
| function setHeight(h) { | |
| window.parent.postMessage({ type: "streamlit:setFrameHeight", height: h }, "*"); | |
| } | |
| window.parent.postMessage({ type: "streamlit:componentReady", apiVersion: 1 }, "*"); | |
| setHeight(20); | |
| window.addEventListener("message", async (event) => { | |
| if (event.data.type !== "streamlit:render") return; | |
| const { papers, run } = event.data.args; | |
| if (run) await runEmbeddings(papers); | |
| }); | |
| // ── Main ───────────────────────────────────────────────────────────────── | |
| async function runEmbeddings(papers) { | |
| const statusEl = document.getElementById("status"); | |
| const progressEl = document.getElementById("progress"); | |
| const errorEl = document.getElementById("error"); | |
| errorEl.textContent = ""; | |
| // WebGPU support check | |
| if (!navigator.gpu) { | |
| errorEl.innerHTML = | |
| "It seems your browser does not support WebGPU. " + | |
| "Chrome, Edge and Opera support it by default. " + | |
| "Firefox and Safari need additional configuration. " + | |
| "Alternatively, you can run the embeddings on the web server " + | |
| "by checking <b>Use fallback resources</b>, but this takes much longer."; | |
| setHeight(100); | |
| sendValue({ error: "webgpu_not_supported" }); | |
| return; | |
| } | |
| try { | |
| progressEl.style.display = "block"; | |
| setHeight(50); | |
| statusEl.textContent = "Loading model…"; | |
| const { AutoTokenizer, AutoModel } = await import( | |
| "https://cdn.jsdelivr.net/npm/@huggingface/transformers@3" | |
| ); | |
| const tokenizer = await AutoTokenizer.from_pretrained(MODEL_REPO); | |
| const model = await AutoModel.from_pretrained(MODEL_REPO, { | |
| device: "webgpu", | |
| dtype: "fp16", | |
| }); | |
| const total = papers.length; | |
| const allCLS = new Float32Array(total * 768); | |
| for (let start = 0; start < total; start += BATCH_SIZE) { | |
| const batch = papers.slice(start, start + BATCH_SIZE); | |
| const texts = batch.map(p => p.title + " [SEP] " + (p.abstract || "")); | |
| const inputs = tokenizer(texts, { padding: true, truncation: true, max_length: 512 }); | |
| const output = await model(inputs); | |
| // Extract CLS token (position 0) from last_hidden_state [B, L, H] | |
| const hs = output.last_hidden_state; | |
| const [B, L, H] = hs.dims; | |
| const data = hs.data; | |
| for (let b = 0; b < B; b++) { | |
| const src = b * L * H; | |
| const dst = (start + b) * 768; | |
| allCLS.set(data.slice(src, src + H), dst); | |
| } | |
| const done = Math.min(start + BATCH_SIZE, total); | |
| progressEl.value = (done / total) * 100; | |
| statusEl.textContent = `Encoding ${done}/${total}…`; | |
| } | |
| progressEl.value = 100; | |
| statusEl.textContent = "Done."; | |
| // Encode as base64 for efficient transfer to Python | |
| const bytes = new Uint8Array(allCLS.buffer); | |
| let binary = ""; | |
| for (let i = 0; i < bytes.length; i++) binary += String.fromCharCode(bytes[i]); | |
| const b64 = btoa(binary); | |
| sendValue({ embeddings_b64: b64, n_papers: total, n_dims: 768 }); | |
| } catch (err) { | |
| errorEl.textContent = "Error: " + err.message; | |
| setHeight(60); | |
| sendValue({ error: err.message }); | |
| } | |
| } | |
| </script> | |
| </body> | |
| </html> | |