Spaces:
Running
Running
| import { AutoTokenizer, CLIPTextModelWithProjection, env } from '@huggingface/transformers'; | |
| import { cachedFile, DATASET, TOKENIZER, TEXT_MODEL, TEXT_MODEL_REVISION, DIMENSIONS } from './assets.js'; | |
| import { normalizeDatabase, rankImages } from './search.js'; | |
| import { loadCatalogue, searchCatalogue, tryLoadCatalogue } from './catalogue.js'; | |
| env.allowLocalModels = false; | |
| env.backends.onnx.wasm.numThreads = 1; | |
| let instance; | |
| 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'. | |
| // | |
| // `revision` is load-bearing, not a reproducibility nicety: without it the | |
| // graph is fetched and cached under the unversioned /resolve/main/ URL, so | |
| // browsers that cached the broken 68dc949 export keep failing with | |
| // [ShapeInferenceError] Inferred=1 Declared=768 on every reload. See | |
| // TEXT_MODEL_REVISION in assets.js. | |
| CLIPTextModelWithProjection.from_pretrained( | |
| TEXT_MODEL, | |
| { dtype: 'fp32', device: 'wasm', revision: TEXT_MODEL_REVISION, 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; }); | |
| } | |
| 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; | |
| // Keep more interiors than the field shows: merge.js then interleaves | |
| // the product hits into this list, so a top-60 request that yields 12 | |
| // products still fills a 96-tile field instead of stopping at 60. | |
| const output = rankImages(text_embeds.data, embeddings, DIMENSIONS, 96) | |
| .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; } | |
| } | |
| /** | |
| * Rank Malloca products against the live catalogue that vai-avatar writes. | |
| * | |
| * This is deliberately lexical rather than a baked CLIP matrix. A matrix built | |
| * from a snapshot can only ever know the products that existed when it was | |
| * built, so "Thêm SP" on vai-avatar would stay invisible until someone re-ran a | |
| * 6,372-image embed. Reading the same JSON files vai-avatar writes means a new | |
| * product is searchable on the next query, and a deleted one disappears on its | |
| * own -- no rebuild, no staleness window. | |
| * | |
| * The catalogue is kept warm so a keystroke does not refetch; `refresh-catalogue` | |
| * forces a re-read, and the periodic poll re-reads it without a query. | |
| */ | |
| let catalogue = null; | |
| function catalogueItems() { | |
| return catalogue ?? loadCatalogue().then(next => { catalogue = next; return next; }); | |
| } | |
| /** Re-read both source files. Shared by the poll loop and the manual refresh. */ | |
| async function refreshCatalogue(force = false) { | |
| const next = force ? await loadCatalogue() : await tryLoadCatalogue(); | |
| if (next) catalogue = next; | |
| return next; | |
| } | |
| async function searchProducts(text, id, limit) { | |
| const started = performance.now(); | |
| try { | |
| send({ status: 'products-loading' }); | |
| const { items, signature } = await catalogueItems(); | |
| const output = searchCatalogue(text, items, limit); | |
| send({ | |
| status: 'products-complete', id, text, output, | |
| signature, count: items.size, 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 })); | |
| // Warm the catalogue in the background so the first product query does not | |
| // wait on two dataset downloads. A failure here is non-fatal: it surfaces | |
| // on the first search, where it can be reported properly. | |
| catalogueItems().catch(() => {}); | |
| } 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 === 'refresh-catalogue') { | |
| refreshCatalogue(data.force === true) | |
| .then(next => next && send({ status: 'catalogue-updated', signature: next.signature, count: next.items.size, syncedAt: next.syncedAt })) | |
| .catch(error => send({ status: 'catalogue-error', message: error.message })); | |
| } else if (data.type === 'clear') { | |
| pending = null; | |
| } | |
| }); | |