From 836e4b7cbc34488c64c1fdd5711906e96f5aca98 Mon Sep 17 00:00:00 2001 From: Bogdan Abaev Date: Wed, 1 Jul 2026 16:39:07 -0700 Subject: [PATCH] local semantic search on abstracts - added environment to run embedding models locally (transformers.js, ONNX Runtime WASM binary, etc.). The actual inference execution happens in a separate worker environment (worker.js) - added local itemEmbeddings table to store embeddings locally - in advanced preferences, one can select two options for semantic search model: english and multilingual. English model (bge-small-en-v1.5) is better for english-only corpus but multilingual (multilingual-e5-small) is necessary to handle abstracts with any other language than english. We can add more language-specific models as needed. - when the model is selected, Zotero.Embeddings.download will download the model (quantized ~100mb) and store it locally. - Zotero.Embeddings.Indexing will start a process to index all regular items with title+abstract. It happens in batches and takes some time. The progress will appear in the advanced preferences pane. Embeddings are inserted into itemEmbeddings SQL table. For now, the table is local only, no syncing is involved. - when embedding model pref is set to "Disabled", the model is deleted and embeddings table is cleared. - when an embedding model is selected, quick search dropdown has a new "Similarity" mode, which will run semantic search on the current scope of items. - semantic search does not clearly define what counts as "relevant" and what is "not relevant". In addition, it will change depending on the library and query. So we cannot semantically filter out items the way it is done via SQL. Semantic search returns the ranking but items cannot be sorted because it is done by the itemTree based on column selection. So in "similarity" quicksearch mode, there is also a dropdown to select how many top relevant items to keep (top 5 - top 100). It allows the user to keep the most relevant items depending on the context, without conflicting with itemTree sorting. - semantic search happens in-memory. On a large 5K library it's fast, but we could consider sqlite-vec extension if needed. --- .gitattributes | 1 + .../content/zotero/collectionViewItemTree.jsx | 7 + .../zotero/elements/quickSearchTextbox.js | 82 +- .../preferences/preferences_advanced.js | 86 + .../preferences/preferences_advanced.xhtml | 29 +- .../content/zotero/xpcom/collectionTreeRow.js | 49 +- chrome/content/zotero/xpcom/embeddings.js | 1031 + chrome/content/zotero/xpcom/schema.js | 4 + chrome/content/zotero/zotero.mjs | 1 + chrome/content/zotero/zoteroPane.js | 1 + chrome/locale/en-US/zotero/preferences.ftl | 18 + chrome/locale/en-US/zotero/zotero.ftl | 3 + defaults/preferences/zotero.js | 9 + .../ort-wasm-simd-threaded.jsep.mjs | 125 + .../ort-wasm-simd-threaded.jsep.wasm | 3 + resource/embeddings/transformers.js | 49046 ++++++++++++++++ resource/embeddings/worker.js | 125 + resource/schema/userdata.sql | 8 + scss/elements/_quickSearchTextbox.scss | 30 + scss/preferences/_advanced.scss | 4 + 20 files changed, 50647 insertions(+), 15 deletions(-) create mode 100644 chrome/content/zotero/xpcom/embeddings.js create mode 100644 resource/embeddings/ort-wasm-simd-threaded.jsep.mjs create mode 100644 resource/embeddings/ort-wasm-simd-threaded.jsep.wasm create mode 100644 resource/embeddings/transformers.js create mode 100644 resource/embeddings/worker.js diff --git a/.gitattributes b/.gitattributes index 4fb3d77151..c0373931d8 100644 --- a/.gitattributes +++ b/.gitattributes @@ -3,3 +3,4 @@ app/win/zotero.exe.tar.xz filter=lfs diff=lfs merge=lfs -text app/mac/updater.tar.xz filter=lfs diff=lfs merge=lfs -text app/win/updater.exe.tar.xz filter=lfs diff=lfs merge=lfs -text app/linux/updater.tar.xz filter=lfs diff=lfs merge=lfs -text +resource/embeddings/*.wasm filter=lfs diff=lfs merge=lfs -text diff --git a/chrome/content/zotero/collectionViewItemTree.jsx b/chrome/content/zotero/collectionViewItemTree.jsx index 30ad42f7c8..7ae2f48629 100644 --- a/chrome/content/zotero/collectionViewItemTree.jsx +++ b/chrome/content/zotero/collectionViewItemTree.jsx @@ -770,6 +770,13 @@ class CollectionViewItemTreeRowProvider extends ItemTreeRowProvider { refresh = true; madeChanges = true; } + // Under an active similarity quick search, handle removals with a full refresh too + else if (['remove', 'delete', 'trash'].includes(action) + && collectionTreeRows.some(row => row.isSimilaritySearch())) { + this.itemTree.invalidateRowCache(ids); + refresh = true; + madeChanges = true; + } else if ((action == 'remove' && !collectionTreeRows.some(row => row.isLibrary(true))) || action == 'delete' || action == 'trash' || (action == 'removeDuplicatesMaster' && this.viewMode == 'duplicates')) { diff --git a/chrome/content/zotero/elements/quickSearchTextbox.js b/chrome/content/zotero/elements/quickSearchTextbox.js index 79f51869e1..0095132a0a 100644 --- a/chrome/content/zotero/elements/quickSearchTextbox.js +++ b/chrome/content/zotero/elements/quickSearchTextbox.js @@ -45,11 +45,17 @@ `, ['chrome://zotero/locale/zotero.dtd']); } - _searchModes = { - titleCreatorYear: Zotero.getString('quickSearch.mode.titleCreatorYear'), - fields: Zotero.getString('quickSearch.mode.fieldsAndTags'), - everything: Zotero.getString('quickSearch.mode.everything') - }; + get _searchModes() { + let modes = { + titleCreatorYear: Zotero.getString('quickSearch.mode.titleCreatorYear'), + fields: Zotero.getString('quickSearch.mode.fieldsAndTags'), + everything: Zotero.getString('quickSearch.mode.everything') + }; + if (Zotero.Embeddings.isEnabled()) { + modes.similarity = Zotero.getString('quickSearch-mode-similarity'); + } + return modes; + } _searchModePopup = null; @@ -122,8 +128,36 @@ } }); wrapper.appendChild(advancedButton); + this._advancedButton = advancedButton; } - + + // Dropdown selecting how many results the similarity mode keeps; + // shown in place of the Advanced Search button + let topKList = document.createXULElement('menulist'); + topKList.id = 'zotero-tb-search-topk'; + topKList.hidden = true; + document.l10n.setAttributes(topKList, 'quicksearch-semantic-topk'); + let topKPopup = document.createXULElement('menupopup'); + for (let n of [5, 10, 25, 50, 100]) { + let item = document.createXULElement('menuitem'); + item.label = String(n); + item.value = String(n); + topKPopup.append(item); + } + topKList.append(topKPopup); + topKList.value = String(Zotero.Prefs.get('search.quicksearch-semantic-topK')); + topKList.addEventListener('command', (event) => { + // Don't trigger a quick search via the oncommand handler + event.stopPropagation(); + Zotero.Prefs.set('search.quicksearch-semantic-topK', parseInt(topKList.value)); + // Re-run the current search with the new top-K + if (this.value) { + this.dispatchEvent(new Event('command')); + } + }); + wrapper.appendChild(topKList); + this._topKList = topKList; + this.deck = this.firstElementChild; this.querySelector('.advanced-collapse-button').addEventListener('command', (event) => { @@ -158,6 +192,15 @@ popup.id = "search-mode-popup"; popup.toggleAttribute("needsgutter", true); + this._populateSearchModePopup(popup); + // Rebuild the menu if the available modes changed since it was built + // (e.g. semantic search was enabled or disabled in the preferences) + popup.addEventListener('popupshowing', () => this._syncSearchModePopup()); + + return this._searchModePopup = popup; + } + + _populateSearchModePopup(popup) { for (let [mode, label] of Object.entries(this._searchModes)) { let item = document.createXULElement('menuitem'); item.setAttribute('type', 'radio'); @@ -175,8 +218,22 @@ popup.append(item); } - - return this._searchModePopup = popup; + } + + _syncSearchModePopup() { + let popup = this._searchModePopup; + if (!popup) { + return; + } + let modes = Object.keys(this._searchModes); + let current = [...popup.children].map(item => item.value); + if (current.length === modes.length && current.every((mode, i) => mode === modes[i])) { + return; + } + popup.replaceChildren(); + this._populateSearchModePopup(popup); + let active = Zotero.Prefs.get('search.quicksearch-mode'); + popup.querySelector(`menuitem[value="${active}"]`)?.setAttribute('checked', 'true'); } onCollectionSelected() { @@ -192,10 +249,19 @@ mode = 'fields'; } + this._syncSearchModePopup(); this.searchModePopup.querySelector(`menuitem[value="${mode}"]`) .setAttribute('checked', 'true'); document.l10n.setAttributes(this.searchTextbox.inputField, "quicksearch-input", { placeholder: this._searchModes[mode] }); + // Advanced Search doesn't apply to semantic search, so swap its + // button for the similarity result-count dropdown + let isSimilarity = mode === 'similarity'; + if (this._advancedButton) { + this._advancedButton.hidden = isSimilarity; + } + this._topKList.hidden = !isSimilarity; + let advancedSearchDeck = document.getElementById('zotero-advanced-search-pane-deck'); if (advancedSearchDeck) { let state = advancedSearchDeck.state; diff --git a/chrome/content/zotero/preferences/preferences_advanced.js b/chrome/content/zotero/preferences/preferences_advanced.js index 1021c26d27..05ae78d5aa 100644 --- a/chrome/content/zotero/preferences/preferences_advanced.js +++ b/chrome/content/zotero/preferences/preferences_advanced.js @@ -57,6 +57,92 @@ Zotero_Preferences.Advanced = { document.getElementById('zotero-prefpane-advanced-enable-local-api').addEventListener('synctopreference', () => { this.updateLocalAPIUI(); }); + + this.initSemanticSearch(); + }, + + + initSemanticSearch: function () { + // Populate the model menu from the model registry. The preferences + // code watches bound menulists for added items and selects the one + // matching the current pref value. + let modelPopup = document.querySelector('#semantic-search-model > menupopup'); + for (let { name, l10nID } of Zotero.Embeddings.getAvailableModels()) { + let menuitem = document.createXULElement('menuitem'); + menuitem.setAttribute('value', name); + document.l10n.setAttributes(menuitem, l10nID); + modelPopup.append(menuitem); + } + + // Live progress updates from the background indexer + this._semanticSearchListener = status => this.updateSemanticSearchUI(status); + Zotero.Embeddings.Indexing.addProgressListener(this._semanticSearchListener); + + document.getElementById('zotero-prefpane-advanced').addEventListener('unload', () => { + Zotero.Embeddings.Indexing.removeProgressListener(this._semanticSearchListener); + }); + + document.getElementById('semantic-search-resume').addEventListener('command', () => { + Zotero.Embeddings.Indexing.startIndexing(); + }); + + document.getElementById('semantic-search-stop').addEventListener('command', () => { + Zotero.Embeddings.Indexing.stopIndexing(); + }); + + // Render current state, then compute up-to-date per-library counts + this.updateSemanticSearchUI(Zotero.Embeddings.Indexing.getStatus()); + Zotero.Embeddings.Indexing.refreshStatus(); + }, + + + updateSemanticSearchUI: function (status) { + let statusBox = document.getElementById('semantic-search-status'); + statusBox.hidden = !status.enabled; + if (!status.enabled) { + return; + } + + // Phase / status message + let phaseLabel = document.getElementById('semantic-search-phase'); + let hasRemaining = status.libraries.some(lib => lib.indexed < lib.eligible); + if (status.error) { + document.l10n.setAttributes(phaseLabel, 'preferences-advanced-semantic-search-error', { error: status.error }); + } + else if (status.phase === 'downloading') { + document.l10n.setAttributes(phaseLabel, 'preferences-advanced-semantic-search-downloading'); + } + else if (status.phase === 'indexing') { + document.l10n.setAttributes(phaseLabel, 'preferences-advanced-semantic-search-indexing'); + } + else { + document.l10n.setAttributes(phaseLabel, + (status.paused || hasRemaining) + ? 'preferences-advanced-semantic-search-paused' + : 'preferences-advanced-semantic-search-idle'); + } + + // Offer a manual restart when enabled but not currently indexing and + // indexing is stopped or there's outstanding work (or the last run + // errored out) + document.getElementById('semantic-search-resume').hidden + = status.indexing || !(status.error || status.paused || hasRemaining); + + // Offer to stop indexing while it's running + document.getElementById('semantic-search-stop').hidden = !status.indexing; + + // Per-library "indexed / total" counts + let grid = document.getElementById('semantic-search-libraries'); + if (grid.childElementCount !== status.libraries.length * 2) { + grid.textContent = ''; + for (let i = 0; i < status.libraries.length; i++) { + grid.append(document.createXULElement('label'), document.createXULElement('label')); + } + } + status.libraries.forEach((lib, i) => { + grid.children[i * 2].setAttribute('value', lib.name); + grid.children[i * 2 + 1].setAttribute('value', `${lib.indexed} / ${lib.eligible}`); + }); }, diff --git a/chrome/content/zotero/preferences/preferences_advanced.xhtml b/chrome/content/zotero/preferences/preferences_advanced.xhtml index 0140a75edb..03eb77b692 100644 --- a/chrome/content/zotero/preferences/preferences_advanced.xhtml +++ b/chrome/content/zotero/preferences/preferences_advanced.xhtml @@ -292,7 +292,34 @@