src/worker.js

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';
import { searchLikeAvatar } from './product-search.js';
import { loadRecentAvatarInteriors, searchAvatarInteriors, RECENT_LIMIT } from './avatar-interiors.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;
}

/* ── Interiors published on vai-avatar ───────────────────────────────────── */

/**
 * The newest interior photographs from vai-avatar, kept warm.
 *
 * Warm for the same reason the catalogue is: a poll tick must not make a visitor
 * wait on a Hub download. `avatarInteriorSignature` is the ids, so a tick that
 * finds nothing new costs one fetch and no UI work -- which matters because this
 * is the only thing that tells the main thread a new image exists at all.
 */
let avatarInteriors = null;
let avatarInteriorSignature = '';
let recentProductSignature = '';

async function avatarInteriorItems() {
  if (avatarInteriors) return avatarInteriors;
  avatarInteriors = await loadRecentAvatarInteriors(RECENT_LIMIT);
  avatarInteriorSignature = avatarInteriors.map(item => item.avatarImageId).join(',');
  return avatarInteriors;
}

/**
 * Re-read vai-avatar's interiors and report whether anything changed.
 *
 * `changed` is computed here rather than in App.jsx so the comparison is made
 * against the same value the worker already holds, instead of App keeping its
 * own copy that could drift from what was actually sent.
 */
async function refreshAvatarInteriors() {
  const items = await loadRecentAvatarInteriors(RECENT_LIMIT);
  const next = items.map(item => item.avatarImageId).join(',');
  // The newest ADDED PRODUCTS ride the same message as the newest interiors, so
  // "mới" means the same thing for both kinds of photograph and App.jsx has one
  // state update to make. Without this, a product pasted on vai-avatar (e.g.
  // "Thêm SP" -> products_url_added.json) would stay invisible until someone
  // typed its exact sku -- the mirror of the interiors feature it is meant to
  // sit beside.
  const { items: catItems } = await catalogueItems().catch(() => ({ items: null }));
  const recentProducts = catItems ? recentProductsOf(catItems) : [];
  const rpKey = recentProducts.map(p => p.sku + '@' + (p.addedAt || '')).join(',');
  // `changed` is composite: a poll that only spins up a new product (or only
  // a new interior) still counts as a change worth re-painting. A poll with
  // neither touches nothing, which is what stops the field rebuilding every 30s.
  const changed = next !== avatarInteriorSignature || rpKey !== recentProductSignature;
  avatarInteriors = items;
  avatarInteriorSignature = next;
  recentProductSignature = rpKey;
  send({ status: 'avatar-updated', items, signature: next, changed, count: items.length, recentProducts });
}

/**
 * The newest added products, ready to ring like the newest interiors.
 *
 * `addedAt` comes from `_added_at` preserved by normalizeProduct (catalogue.js).
 * Products with no timestamp are not "recent"; the newest RECENT_PRODUCT_LIMIT
 * of those that have one are returned, newest first.
 */
const RECENT_PRODUCT_LIMIT = 12;
function recentProductsOf(items) {
  if (!items?.size) return [];
  const withTime = [];
  for (const product of items.values()) {
    if (!product?.addedAt) continue;
    if (!Array.isArray(product.images) || !product.images.some(src => typeof src === 'string' && /^(https?:\/\/|data:image\/)/i.test(src))) continue;
    withTime.push({ ...product, score: 1 });
  }
  withTime.sort((a, b) => String(b.addedAt || '').localeCompare(String(a.addedAt || '')));
  return withTime.slice(0, RECENT_PRODUCT_LIMIT);
}

async function searchProducts(text, id, limit) {
  const started = performance.now();
  try {
    send({ status: 'products-loading' });
    const { items, signature } = await catalogueItems();
    // Ranked with vai-avatar's own scorer (see product-search.js) so a query
    // typed here finds the same product, in the same order, as the same query
    // typed in vai-avatar's panel.
    //
    // Both scorers run and both sets are returned. `output` drives the field;
    // `legacy` is kept so the ranking change can be measured rather than
    // assumed -- it costs one extra pass over a few hundred products, and
    // searchCatalogue was the only ranking this Space had until now.
    const output = searchLikeAvatar(text, items, limit);
    const legacy = searchCatalogue(text, items, limit);
    // vai-avatar's newest interiors, matched against the SAME query. Carried on
    // the same message as the products so App.jsx needs one state update and the
    // three kinds of result can never arrive out of step with each other -- the
    // field would otherwise briefly show a new photograph beside results for a
    // query that did not produce it.
    const interiors = searchAvatarInteriors(text, await avatarInteriorItems());
    send({
      status: 'products-complete', id, text, output, legacy, interiors,
      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(() => {});
    // Same treatment for vai-avatar's interiors, so the "mới" panel can paint on
    // first load instead of after a Hub round trip. Non-fatal for the same
    // reason: the field works perfectly well without it.
    refreshAvatarInteriors().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 ?? 48);
  } 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 }));
    // The interiors re-read rides the same tick as the catalogue, so one poll
    // covers both live sources. Two independent intervals would drift apart and
    // the "mới" panel would show a stale set for part of every cycle.
    refreshAvatarInteriors().catch(() => {});
  } else if (data.type === 'clear') {
    pending = null;
  }
});