Spaces:
Running
Running
| import { AutoTokenizer, CLIPTextModelWithProjection, env } from '@huggingface/transformers'; | |
| import { cachedFile, DATASET, TOKENIZER, TEXT_MODEL, DIMENSIONS } from './assets.js'; | |
| import { normalizeDatabase, rankImages } from './search.js'; | |
| import { PRODUCT_EMBEDDINGS, PRODUCT_ROWS, loadGallery } from './products.js'; | |
| env.allowLocalModels = false; | |
| env.backends.onnx.wasm.numThreads = 1; | |
| let instance; | |
| let productsPromise; | |
| let pending; | |
| let running = false; | |
| const send = message => self.postMessage(message); | |
| /** | |
| * Interior-photo search. The text tower is bep40/clip-vit-b32-multilingual-vi-text: | |
| * a CLIP-B/32 text config with mean pooling and the 768->512 projection baked in, | |
| * so Vietnamese queries land in the same 512-d space as the image matrix | |
| * (projection_dim: 512 in config.json, confirmed on the Hub). | |
| * | |
| * The tokenizer comes from the upstream sentence-transformers repo, NOT the fork: | |
| * the fork's tokenizer_config.json is a BERT WordPiece config ([CLS]/[SEP], | |
| * model_max_length 512) that does not match its own CLIP config | |
| * (max_position_embeddings 77). The fork supplies only the ONNX graph. | |
| */ | |
| function initialize() { | |
| return instance ??= (async () => { | |
| send({ status: 'loading', label: 'Đang tải mô hình tìm kiếm' }); | |
| const progress_callback = event => { | |
| if (event.status === 'progress' && event.file?.endsWith('.onnx')) { | |
| send({ status: 'loading', label: 'Đang tải mô hình tìm kiếm', progress: event.progress }); | |
| } | |
| }; | |
| const [tokenizer, model, metadata, buffer] = await Promise.all([ | |
| AutoTokenizer.from_pretrained(TOKENIZER), | |
| // onnx/ contains exactly one graph, text_model.onnx (130.6 MB), and no | |
| // *_quantized variant -- verified on the Hub. Any dtype other than 'fp32' | |
| // would ask for a file that does not exist, so pin 'fp32' and device 'wasm'. | |
| CLIPTextModelWithProjection.from_pretrained(TEXT_MODEL, { dtype: 'fp32', device: 'wasm', progress_callback }), | |
| cachedFile(DATASET + 'image-embeddings.json').then(buffer => JSON.parse(new TextDecoder().decode(buffer))), | |
| cachedFile(DATASET + 'image-embeddings_25k-512-32bit.bin'), | |
| ]); | |
| const embeddings = new Float32Array(buffer); | |
| if (embeddings.length !== metadata.length * DIMENSIONS) { | |
| throw new Error('The image database and embeddings do not match.'); | |
| } | |
| normalizeDatabase(embeddings, DIMENSIONS); | |
| send({ status: 'loading', label: 'Đang khởi động tìm kiếm', progress: 100 }); | |
| await model(tokenizer('a photograph', { padding: true, truncation: true })); | |
| send({ status: 'ready', count: metadata.length }); | |
| return { tokenizer, model, metadata, embeddings }; | |
| })().catch(error => { instance = null; throw error; }); | |
| } | |
| /** Expand IEEE 754 binary16 into binary32; JS has no widely-supported Float16Array. */ | |
| function halfToFloat(view, count) { | |
| const out = new Float32Array(count); | |
| for (let i = 0; i < count; i++) { | |
| const bits = view.getUint16(i * 2, true); | |
| const sign = (bits & 0x8000) ? -1 : 1; | |
| const exponent = (bits >> 10) & 0x1f; | |
| const fraction = bits & 0x03ff; | |
| let value; | |
| if (exponent === 0) value = fraction * 2 ** -24; // subnormal | |
| else if (exponent === 0x1f) value = fraction ? NaN : Infinity; // inf / NaN | |
| else value = (1 + fraction / 1024) * 2 ** (exponent - 15); | |
| out[i] = sign * value; | |
| } | |
| return out; | |
| } | |
| /** | |
| * Malloca product photos, ranked in the same text tower -- so a Vietnamese query | |
| * works directly, with no translation model. The matrix is L2-normalized offline, | |
| * so ranking is a plain dot product. | |
| * | |
| * `meta.rows` is [{ sku, image, index }] written by tools/build_product_embeddings.py, | |
| * with `index` the row offset in the matrix. It is the single source of truth for | |
| * the image -> product mapping; nothing here guesses gallery structure. | |
| */ | |
| function initializeProducts() { | |
| return productsPromise ??= (async () => { | |
| const on_progress = (loaded, total) => { | |
| send({ status: 'products-loading', progress: total ? Math.round((loaded / total) * 100) : 0 }); | |
| }; | |
| const [meta, buffer] = await Promise.all([ | |
| fetch(PRODUCT_ROWS).then(response => { | |
| if (!response.ok) throw new Error(`Không tải được product-emb.json (${response.status}).`); | |
| return response.json(); | |
| }), | |
| cachedFile(PRODUCT_EMBEDDINGS, on_progress).catch(error => { | |
| throw new Error(`Chưa có ma trận embedding sản phẩm. Chạy tools/run_embedding_job.py trước. (${error.message})`); | |
| }), | |
| ]); | |
| const { count, dim, dtype } = meta; | |
| if (dim !== DIMENSIONS) throw new Error(`Sai kích thước vector: ${dim} != ${DIMENSIONS}`); | |
| if (!Array.isArray(meta.rows) || meta.rows.length !== count) { | |
| throw new Error(`product-emb.json: ${meta.rows?.length ?? 0} rows cho ${count} embeddings.`); | |
| } | |
| const bytes = dtype === 'float16' ? 2 : 4; | |
| const expect = count * dim * bytes; | |
| if (buffer.byteLength !== expect) { | |
| throw new Error(`Ma trận embedding không khớp: ${buffer.byteLength} B, cần ${expect} B.`); | |
| } | |
| const data = dtype === 'float16' | |
| ? halfToFloat(new DataView(buffer), count * dim) | |
| : new Float32Array(buffer.slice(0)); | |
| normalizeDatabase(data, dim); | |
| return { data, dim, count, rows: meta.rows }; | |
| })().catch(error => { productsPromise = null; throw error; }); | |
| } | |
| async function drain() { | |
| if (running) return; | |
| running = true; | |
| try { | |
| const { tokenizer, model, metadata, embeddings } = await initialize(); | |
| while (pending) { | |
| const request = pending; | |
| pending = null; | |
| const started = performance.now(); | |
| try { | |
| const { text_embeds } = await model(tokenizer(request.text, { padding: true, truncation: true })); | |
| // A newer query has already arrived: don't animate an obsolete result. | |
| if (pending) continue; | |
| const output = rankImages(text_embeds.data, embeddings, DIMENSIONS) | |
| .map(({ index, score }) => ({ ...metadata[index], score })); | |
| send({ status: 'complete', id: request.id, text: request.text, output, elapsed: performance.now() - started }); | |
| } catch (error) { | |
| send({ status: 'error', id: request.id, message: error.message }); | |
| } | |
| } | |
| } catch (error) { | |
| send({ status: 'error', id: pending?.id, message: error.message }); | |
| pending = null; | |
| } finally { running = false; } | |
| } | |
| async function embedText(text) { | |
| const { tokenizer, model } = await initialize(); | |
| const { text_embeds } = await model(tokenizer(text, { padding: true, truncation: true })); | |
| const query = text_embeds.data; | |
| let norm = 0; | |
| for (const value of query) norm += value * value; | |
| return { query, norm: Math.sqrt(norm) || 1 }; | |
| } | |
| /** Rank Malloca products. Independent of the interior index. */ | |
| async function searchProducts(text, id, limit) { | |
| const started = performance.now(); | |
| try { | |
| const { data, dim, count, rows } = await initializeProducts(); | |
| const { query, norm } = await embedText(text); | |
| // A product may have several photos; keep its best-scoring row. | |
| const best = new Map(); | |
| for (const { index, sku, image } of rows) { | |
| const offset = index * dim; | |
| let dot = 0; | |
| for (let j = 0; j < dim; j++) dot += query[j] * data[offset + j]; | |
| const score = dot / norm; | |
| const seen = best.get(sku); | |
| if (!seen || score > seen.score) best.set(sku, { score, image }); | |
| } | |
| const output = Array.from(best, ([sku, { score, image }]) => ({ sku, score, image })) | |
| .sort((a, b) => b.score - a.score) | |
| .slice(0, limit); | |
| send({ status: 'products-complete', id, text, output, count, elapsed: performance.now() - started }); | |
| } catch (error) { | |
| send({ status: 'products-error', id, message: error.message }); | |
| } | |
| } | |
| self.addEventListener('message', ({ data }) => { | |
| if (data.type === 'init') { | |
| initialize().catch(error => send({ status: 'error', message: error.message })); | |
| } else if (data.type === 'search') { | |
| pending = { id: data.id, text: data.text }; | |
| drain(); | |
| } else if (data.type === 'search-products') { | |
| searchProducts(data.text, data.id, data.limit ?? 24); | |
| } else if (data.type === 'clear') { | |
| pending = null; | |
| } | |
| }); | |