mirror of
https://github.com/zotero/zotero.git
synced 2026-10-08 03:08:19 +00:00
Embed passages through a verified llama endpoint
Passages can be embedded by a local llama.cpp server serving the same model, configured from Settings -> Advanced. The server is trusted only because its vectors match the local model's, never by name: - Zotero.Embeddings.Endpoint: verify() embeds calibration passages locally and remotely (plus one chunk-length text) and stores a typed verdict (ok, unreachable, unauthorized, not-embeddings, width-mismatch, low-agreement, context-too-small) keyed to the URL and model version. Passages route to the endpoint only while a matching verdict says ok; queries always embed locally. - Every batch carries a sentinel text whose local vector is cached, so a server switched to another model or pooling is caught on that batch and nothing wrong is stored. Three consecutive failures skip the endpoint for the rest of the run; a run start rechecks a verified endpoint. - MODELS gains `serving` (GGUF repo and quant) for both bekko models; getCommand() renders the llama.cpp command with a 4096 context. - Preferences: status line and Configure Endpoint dialog with the URL, the facts the server must match, a copyable command, typed errors, and a Stop Using Endpoint button. Also fixes two tokenizer defects in the runtime's transformers.js that the endpoint comparison exposed, both affecting every locally embedded bekko vector: - Metaspace's leading word marker was never prepended (the runtime needs the legacy add_prefix_space flag). `leadingSpace: true` in MODELS adds it in embedMany(). Short texts had drifted to cosine 0.83-0.95 against the reference. - Newlines and repeated spaces were tokenized differently from the reference. Whitespace runs now collapse to one space in both embedding paths.
This commit is contained in:
Bogdan Abaev
parent
9948d4b2c2
commit
0e2fdcec8d
8 changed files with 1020 additions and 16 deletions
179
chrome/content/zotero/preferences/embeddingsEndpoint.xhtml
Normal file
179
chrome/content/zotero/preferences/embeddingsEndpoint.xhtml
Normal file
|
|
@ -0,0 +1,179 @@
|
|||
<?xml version="1.0"?>
|
||||
<!--
|
||||
***** BEGIN LICENSE BLOCK *****
|
||||
|
||||
Copyright © 2026 Corporation for Digital Scholarship
|
||||
Vienna, Virginia, USA
|
||||
https://www.zotero.org
|
||||
|
||||
This file is part of Zotero.
|
||||
|
||||
Zotero is free software: you can redistribute it and/or modify
|
||||
it under the terms of the GNU Affero General Public License as published by
|
||||
the Free Software Foundation, either version 3 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
Zotero is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU Affero General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU Affero General Public License
|
||||
along with Zotero. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
***** END LICENSE BLOCK *****
|
||||
-->
|
||||
|
||||
<?xml-stylesheet href="chrome://global/skin/" type="text/css"?>
|
||||
<?xml-stylesheet href="chrome://zotero/skin/zotero.css" type="text/css"?>
|
||||
<?xml-stylesheet href="chrome://zotero/skin/preferences.css"?>
|
||||
<?xml-stylesheet href="chrome://zotero-platform/content/zotero.css"?>
|
||||
|
||||
<window
|
||||
xmlns="http://www.mozilla.org/keymaster/gatekeeper/there.is.only.xul"
|
||||
xmlns:html="http://www.w3.org/1999/xhtml"
|
||||
data-l10n-id="preferences-advanced-semantic-search-endpoint-dialog"
|
||||
onload="Zotero_EmbeddingsEndpointDialog.init()">
|
||||
<dialog
|
||||
id="zotero-embeddings-endpoint"
|
||||
orient="vertical"
|
||||
buttons="cancel,accept,extra2"
|
||||
style="padding: 20px 15px; min-width: 620px">
|
||||
|
||||
<script src="chrome://global/content/globalOverlay.js"/>
|
||||
<script src="chrome://zotero/content/include.js"/>
|
||||
|
||||
<html:link rel="localization" href="preferences.ftl"/>
|
||||
|
||||
<vbox id="zotero-embeddings-endpoint-container">
|
||||
<label data-l10n-id="preferences-advanced-semantic-search-endpoint-intro"/>
|
||||
<separator class="thin"/>
|
||||
|
||||
<label data-l10n-id="preferences-advanced-semantic-search-endpoint-url" control="endpoint-url"/>
|
||||
<html:input type="text" id="endpoint-url" placeholder="http://localhost:8080/v1/embeddings"/>
|
||||
<description class="endpoint-note"
|
||||
data-l10n-id="preferences-advanced-semantic-search-endpoint-privacy"/>
|
||||
<separator class="thin"/>
|
||||
|
||||
<label data-l10n-id="preferences-advanced-semantic-search-endpoint-requirements"/>
|
||||
<html:div class="endpoint-facts">
|
||||
<label data-l10n-id="preferences-advanced-semantic-search-endpoint-model"/>
|
||||
<label id="endpoint-model"/>
|
||||
<label data-l10n-id="preferences-advanced-semantic-search-endpoint-pooling"/>
|
||||
<label id="endpoint-pooling"/>
|
||||
<label data-l10n-id="preferences-advanced-semantic-search-endpoint-file"/>
|
||||
<label id="endpoint-file"/>
|
||||
</html:div>
|
||||
<separator class="thin"/>
|
||||
|
||||
<label data-l10n-id="preferences-advanced-semantic-search-endpoint-llama"/>
|
||||
<hbox align="center">
|
||||
<html:textarea id="endpoint-command" class="endpoint-command" readonly="readonly" rows="2"/>
|
||||
<button id="endpoint-copy" data-l10n-id="preferences-advanced-semantic-search-endpoint-copy"/>
|
||||
</hbox>
|
||||
<label id="endpoint-command-url" class="endpoint-note"/>
|
||||
<separator class="thin"/>
|
||||
|
||||
<hbox id="endpoint-progress" align="center" hidden="true">
|
||||
<image class="zotero-spinner-16"/>
|
||||
<label data-l10n-id="preferences-advanced-semantic-search-endpoint-verifying"/>
|
||||
</hbox>
|
||||
<description id="endpoint-error" class="endpoint-error" hidden="true"/>
|
||||
</vbox>
|
||||
|
||||
<script>
|
||||
<![CDATA[
|
||||
var Zotero_EmbeddingsEndpointDialog = {
|
||||
init() {
|
||||
this.io = window.arguments[0];
|
||||
this.urlInput = document.getElementById('endpoint-url');
|
||||
this.dialog = document.getElementById('zotero-embeddings-endpoint');
|
||||
this.urlInput.value = Zotero.Prefs.get('embeddings.endpoint') || '';
|
||||
document.l10n.setAttributes(this.dialog.getButton('accept'),
|
||||
'preferences-advanced-semantic-search-endpoint-accept');
|
||||
// Offered only while there is an endpoint to remove
|
||||
let remove = this.dialog.getButton('extra2');
|
||||
document.l10n.setAttributes(remove, 'preferences-advanced-semantic-search-endpoint-remove');
|
||||
remove.hidden = !this.urlInput.value;
|
||||
document.addEventListener('dialogextra2', () => this.remove());
|
||||
|
||||
let serving = Zotero.Embeddings.getServing();
|
||||
document.getElementById('endpoint-model').value = Zotero.Embeddings.getModelName();
|
||||
document.getElementById('endpoint-pooling').value = Zotero.Embeddings.getPooling();
|
||||
document.getElementById('endpoint-file').value = `${serving.gguf} (${serving.quant})`;
|
||||
let { command, url } = Zotero.Embeddings.Endpoint.getCommand();
|
||||
document.getElementById('endpoint-command').value = command;
|
||||
document.l10n.setAttributes(document.getElementById('endpoint-command-url'),
|
||||
'preferences-advanced-semantic-search-endpoint-then-use', { url });
|
||||
document.getElementById('endpoint-copy').addEventListener('command', () => {
|
||||
Zotero.Utilities.Internal.copyTextToClipboard(command);
|
||||
});
|
||||
|
||||
// Verification is asynchronous: hold the dialog open until it
|
||||
// has an answer
|
||||
document.addEventListener('dialogaccept', (event) => {
|
||||
event.preventDefault();
|
||||
this.accept();
|
||||
});
|
||||
this.urlInput.addEventListener('input', () => this.showError(null));
|
||||
this.urlInput.focus();
|
||||
},
|
||||
|
||||
// Stop using the endpoint. Nothing stored is affected: the vectors
|
||||
// it produced match the model's, and a later endpoint verifies anew.
|
||||
remove() {
|
||||
Zotero.Prefs.clear('embeddings.endpoint');
|
||||
this.io.ok = true;
|
||||
window.close();
|
||||
},
|
||||
|
||||
async accept() {
|
||||
let url = this.urlInput.value.trim();
|
||||
// An emptied URL turns the endpoint off; nothing to verify
|
||||
if (!url) {
|
||||
this.remove();
|
||||
return;
|
||||
}
|
||||
this.setBusy(true);
|
||||
this.showError(null);
|
||||
let verdict;
|
||||
try {
|
||||
verdict = await Zotero.Embeddings.Endpoint.verify(url);
|
||||
}
|
||||
catch (e) {
|
||||
Zotero.logError(e);
|
||||
verdict = { state: 'unreachable' };
|
||||
}
|
||||
this.setBusy(false);
|
||||
if (verdict.state != 'ok') {
|
||||
this.showError(verdict);
|
||||
return;
|
||||
}
|
||||
Zotero.Prefs.set('embeddings.endpoint', url);
|
||||
this.io.ok = true;
|
||||
window.close();
|
||||
},
|
||||
|
||||
setBusy(busy) {
|
||||
document.getElementById('endpoint-progress').hidden = !busy;
|
||||
this.dialog.getButton('accept').disabled = busy;
|
||||
this.dialog.getButton('cancel').disabled = busy;
|
||||
this.dialog.getButton('extra2').disabled = busy;
|
||||
this.urlInput.disabled = busy;
|
||||
},
|
||||
|
||||
showError(verdict) {
|
||||
let box = document.getElementById('endpoint-error');
|
||||
box.hidden = !verdict;
|
||||
if (!verdict) {
|
||||
return;
|
||||
}
|
||||
document.l10n.setAttributes(box,
|
||||
`preferences-advanced-semantic-search-endpoint-error-${verdict.state}`,
|
||||
{ pooling: Zotero.Embeddings.getPooling() });
|
||||
}
|
||||
};
|
||||
]]>
|
||||
</script>
|
||||
</dialog>
|
||||
</window>
|
||||
|
|
@ -110,6 +110,15 @@ Zotero_Preferences.Advanced = {
|
|||
Zotero.Embeddings.Indexing.stopIndexing();
|
||||
});
|
||||
|
||||
document.getElementById('semantic-search-endpoint-configure').addEventListener('command', () => {
|
||||
this.openSemanticSearchEndpointDialog();
|
||||
});
|
||||
// The endpoint's stored verdict is read lazily; have it in memory
|
||||
// before the status line first renders
|
||||
Zotero.Embeddings.Endpoint.load().then(() => {
|
||||
this.updateSemanticSearchUI(Zotero.Embeddings.Indexing.getStatus());
|
||||
});
|
||||
|
||||
// Render current state, then compute up-to-date per-library counts
|
||||
this.updateSemanticSearchUI(Zotero.Embeddings.Indexing.getStatus());
|
||||
Zotero.Embeddings.Indexing.refreshStatus();
|
||||
|
|
@ -161,6 +170,8 @@ Zotero_Preferences.Advanced = {
|
|||
return;
|
||||
}
|
||||
|
||||
this.updateSemanticSearchEndpointUI(status.endpoint);
|
||||
|
||||
// Phase / status message
|
||||
let phaseLabel = document.getElementById('semantic-search-phase');
|
||||
let hasRemaining = status.items.done < status.items.total
|
||||
|
|
@ -235,6 +246,39 @@ Zotero_Preferences.Advanced = {
|
|||
},
|
||||
|
||||
|
||||
// Whether the active model can be served at all, and how the configured
|
||||
// server stands (see Zotero.Embeddings.Endpoint.getStatus())
|
||||
updateSemanticSearchEndpointUI: function (endpoint) {
|
||||
let row = document.getElementById('semantic-search-endpoint-row');
|
||||
row.hidden = !Zotero.Embeddings.Endpoint.isSupported();
|
||||
if (row.hidden) {
|
||||
return;
|
||||
}
|
||||
let label = document.getElementById('semantic-search-endpoint-status');
|
||||
let id = {
|
||||
off: 'off',
|
||||
unknown: null,
|
||||
unverified: 'unverified',
|
||||
ok: 'valid',
|
||||
unreachable: 'unreachable'
|
||||
}[endpoint.state];
|
||||
if (id === null) {
|
||||
label.removeAttribute('data-l10n-id');
|
||||
label.value = '';
|
||||
return;
|
||||
}
|
||||
document.l10n.setAttributes(label, `preferences-advanced-semantic-search-endpoint-${id || 'invalid'}`);
|
||||
},
|
||||
|
||||
|
||||
openSemanticSearchEndpointDialog: function () {
|
||||
let io = { ok: false };
|
||||
window.openDialog('chrome://zotero/content/preferences/embeddingsEndpoint.xhtml',
|
||||
'zotero-preferences-embeddingsEndpoint', 'chrome,modal,centerscreen', io);
|
||||
this.updateSemanticSearchUI(Zotero.Embeddings.Indexing.getStatus());
|
||||
},
|
||||
|
||||
|
||||
// Key/value rows of pipeline diagnostics for developers, so the labels
|
||||
// are plain English rather than localized
|
||||
_updateSemanticSearchDiagnostics: function (diagnostics, eta) {
|
||||
|
|
|
|||
|
|
@ -343,6 +343,11 @@
|
|||
data-l10n-id="preferences-advanced-semantic-search-stop"
|
||||
hidden="true"/>
|
||||
</hbox>
|
||||
<hbox id="semantic-search-endpoint-row" align="center" hidden="true">
|
||||
<label id="semantic-search-endpoint-status"/>
|
||||
<button id="semantic-search-endpoint-configure"
|
||||
data-l10n-id="preferences-advanced-semantic-search-endpoint-configure"/>
|
||||
</hbox>
|
||||
<html:div id="semantic-search-progress">
|
||||
<html:div id="semantic-search-items-row" class="semantic-search-progress-row">
|
||||
<label data-l10n-id="preferences-advanced-semantic-search-items"/>
|
||||
|
|
|
|||
|
|
@ -64,13 +64,19 @@ Zotero.Embeddings = new function () {
|
|||
// maxTokens: 512, // context window; longer text is chunked to fit
|
||||
// dims: 256, // optional: keep only the first N dimensions of every vector
|
||||
// // (Matryoshka) -- only for a model trained to truncate
|
||||
// leadingSpace: true, // optional: the tokenizer marks the first word with a space
|
||||
// // the runtime fails to add (see embedMany())
|
||||
// serving: { // optional: how to serve the same model from llama.cpp
|
||||
// gguf: 'user/repo-GGUF', // (see Zotero.Embeddings.Endpoint); a model without it
|
||||
// quant: 'F16' // can't be embedded through an endpoint
|
||||
// },
|
||||
// l10nID: '...', // optional Fluent id for the menu
|
||||
// label: '...' // optional plain-English menu label, for a model that isn't
|
||||
// // shipped. The menu prefers l10nID, then label, then modelId.
|
||||
// }
|
||||
const MODELS = {
|
||||
'bekko-embedding-v1-a8m': {
|
||||
revision: 1,
|
||||
revision: 2,
|
||||
modelId: 'hotchpotch/bekko-embedding-v1-a8m',
|
||||
// The repo's default artifact is onnx/model.onnx (fp32 layers,
|
||||
// int8 embedding table); it has no model_quantized.onnx
|
||||
|
|
@ -80,10 +86,12 @@ Zotero.Embeddings = new function () {
|
|||
passagePrefix: '',
|
||||
maxTokens: 8192,
|
||||
dims: 256,
|
||||
leadingSpace: true,
|
||||
serving: { gguf: 'hotchpotch/bekko-embedding-v1-a8m-GGUF', quant: 'F16' },
|
||||
l10nID: 'preferences-advanced-semantic-search-multilingual'
|
||||
},
|
||||
'bekko-embedding-v1-a25m': {
|
||||
revision: 1,
|
||||
revision: 2,
|
||||
modelId: 'hotchpotch/bekko-embedding-v1-a25m',
|
||||
// The repo's default artifact is onnx/model.onnx (fp32 layers,
|
||||
// int8 embedding table); it has no model_quantized.onnx
|
||||
|
|
@ -93,6 +101,8 @@ Zotero.Embeddings = new function () {
|
|||
passagePrefix: '',
|
||||
maxTokens: 8192,
|
||||
dims: 256,
|
||||
leadingSpace: true,
|
||||
serving: { gguf: 'hotchpotch/bekko-embedding-v1-a25m-GGUF', quant: 'F16' },
|
||||
label: "better but slower multilingual"
|
||||
},
|
||||
'bge-small-zh-v1.5': {
|
||||
|
|
@ -217,6 +227,23 @@ Zotero.Embeddings = new function () {
|
|||
return _getModel().passagePrefix;
|
||||
};
|
||||
|
||||
/**
|
||||
* How the active model combines token vectors: 'cls' or 'mean'.
|
||||
* @return {String}
|
||||
*/
|
||||
this.getPooling = function () {
|
||||
return _getModel().pooling;
|
||||
};
|
||||
|
||||
/**
|
||||
* How the active model can be served from llama.cpp (see MODELS), or
|
||||
* null for a model that can't be.
|
||||
* @return {Object|null} - { gguf, quant }
|
||||
*/
|
||||
this.getServing = function () {
|
||||
return _getModel().serving || null;
|
||||
};
|
||||
|
||||
/**
|
||||
* Read one of the active model's files (e.g. its tokenizer) from the
|
||||
* runtime's model cache, fetching it from the model hub if the runtime
|
||||
|
|
@ -331,6 +358,7 @@ Zotero.Embeddings = new function () {
|
|||
}
|
||||
await Zotero.DB.queryAsync("DROP TABLE IF EXISTS embeddings.itemEmbeddings");
|
||||
await Zotero.DB.queryAsync("DROP TABLE IF EXISTS embeddings.itemEmbeddingsMeta");
|
||||
Zotero.Embeddings.Endpoint.reset();
|
||||
await Zotero.DB.queryAsync("DROP TABLE IF EXISTS embeddings.modelCalibration");
|
||||
await Zotero.DB.queryAsync("DROP TABLE IF EXISTS embeddings.itemChunkCounts");
|
||||
// One row per chunk of an item's text, its vector centered and
|
||||
|
|
@ -914,6 +942,17 @@ Zotero.Embeddings = new function () {
|
|||
if (!texts.length) {
|
||||
return [];
|
||||
}
|
||||
texts = texts.map(_normalizeInput);
|
||||
// A Metaspace tokenizer with prepend_scheme "always" marks the first
|
||||
// word with a space like any other, and the model was trained on that.
|
||||
// The runtime's tokenizer (transformers.js 3.5.1) only prepends it
|
||||
// when the legacy add_prefix_space flag is also set, which a current
|
||||
// tokenizer.json omits -- so the text gets its space here. Without it
|
||||
// the vectors drift from the model's, badly for short texts (cosine
|
||||
// 0.83-0.95 for titles and queries, 0.99+ for chunks).
|
||||
if (_getModel().leadingSpace) {
|
||||
texts = texts.map(text => ' ' + text);
|
||||
}
|
||||
let engine;
|
||||
let run = async () => {
|
||||
engine = await _getEngine();
|
||||
|
|
@ -945,6 +984,15 @@ Zotero.Embeddings = new function () {
|
|||
return vectors.map(vector => _finish(new Float32Array(vector)));
|
||||
};
|
||||
|
||||
// Text as every embedding path sees it, local or served: runs of
|
||||
// whitespace, newlines included, become one space. Line breaks and
|
||||
// double spaces in extracted text are layout, not language -- and the
|
||||
// runtime's tokenizer reads them differently from the reference
|
||||
// implementation, which moved vectors by up to 0.1 in cosine.
|
||||
function _normalizeInput(text) {
|
||||
return text.replace(/\s+/g, ' ').trim();
|
||||
}
|
||||
|
||||
// A model's raw output into its embedding: cut to its `dims` if it
|
||||
// truncates, and unit length -- centering subtracts a mean measured over
|
||||
// unit vectors (see center())
|
||||
|
|
@ -956,11 +1004,25 @@ Zotero.Embeddings = new function () {
|
|||
return _normalize(vector);
|
||||
}
|
||||
|
||||
// POST { input: [...] } to an OpenAI-style /v1/embeddings endpoint (as
|
||||
// served by llama.cpp, Ollama, TEI and the hosted APIs), which answers
|
||||
// { data: [{ index, embedding }] }. Ordered by index, since the spec
|
||||
// doesn't promise the array comes back in input order.
|
||||
async function _embedViaEndpoint(endpoint, texts) {
|
||||
/**
|
||||
* Embed texts through an OpenAI-style /v1/embeddings endpoint, as served
|
||||
* by llama.cpp: POST { input } and read { model, data: [{ index,
|
||||
* embedding }] }. Rows are ordered by
|
||||
* index, since the spec doesn't promise input order. The vectors come
|
||||
* back finished the way the local engine's do (see _finish()), so the
|
||||
* two are comparable.
|
||||
*
|
||||
* @param {String} endpoint - URL
|
||||
* @param {String[]} texts - With any prefix already applied
|
||||
* @return {Promise<Object>} - { vectors, model, width }: the finished
|
||||
* vectors, the model string the server reported, and the width of
|
||||
* the vectors as served
|
||||
* @throws {Zotero.HTTP.UnexpectedStatusException} on a non-2xx response
|
||||
* @throws {Zotero.Embeddings.EndpointResponseError} on a 2xx response that
|
||||
* isn't an embeddings response for these inputs
|
||||
*/
|
||||
this.embedViaEndpoint = async function (endpoint, texts) {
|
||||
texts = texts.map(_normalizeInput);
|
||||
Zotero.debug(`Embeddings: embedding batch of ${texts.length} via endpoint`);
|
||||
let xmlhttp = await Zotero.HTTP.request('POST', endpoint, {
|
||||
body: JSON.stringify({ input: texts }),
|
||||
|
|
@ -980,12 +1042,18 @@ Zotero.Embeddings = new function () {
|
|||
let received = vectors
|
||||
? `${vectors.length} vectors`
|
||||
: JSON.stringify(xmlhttp.response).substring(0, 200);
|
||||
throw new Error(`Embeddings: endpoint returned ${received} `
|
||||
throw new this.EndpointResponseError(`Endpoint returned ${received} `
|
||||
+ `for ${texts.length} inputs`);
|
||||
}
|
||||
Zotero.debug(`Embeddings: batch of ${texts.length} done`);
|
||||
return vectors.map(vector => _finish(new Float32Array(vector)));
|
||||
}
|
||||
return {
|
||||
vectors: vectors.map(vector => _finish(new Float32Array(vector))),
|
||||
model: typeof xmlhttp.response.model == 'string' ? xmlhttp.response.model : null,
|
||||
width: vectors[0].length
|
||||
};
|
||||
};
|
||||
|
||||
this.EndpointResponseError = class extends Error {};
|
||||
|
||||
// The last embedded query, reused across the scoring passes a single
|
||||
// search triggers (per-row membership cutoffs plus the merged ranking)
|
||||
|
|
@ -1045,21 +1113,31 @@ Zotero.Embeddings = new function () {
|
|||
this.embedPassages = async function (texts) {
|
||||
let passagePrefix = _getModel().passagePrefix;
|
||||
texts = texts.map(text => passagePrefix + text);
|
||||
// Passages can route to an external endpoint serving the same model;
|
||||
// queries always embed locally. A failed batch embeds locally instead
|
||||
// of waiting on a retry -- the next batch tries the endpoint again.
|
||||
let endpoint = Zotero.Prefs.get('embeddings.endpoint');
|
||||
// Passages can route to an endpoint verified to serve the same model
|
||||
// (see Zotero.Embeddings.Endpoint); queries always embed locally. A
|
||||
// failed batch embeds locally instead of waiting on a retry -- the
|
||||
// next batch tries the endpoint again.
|
||||
let endpoint = await this.Endpoint.getActive();
|
||||
if (!endpoint) {
|
||||
return this.embedMany(texts);
|
||||
}
|
||||
// Every batch carries a sentinel text whose local vector is known, so
|
||||
// a server that stops matching the model is caught on that batch and
|
||||
// never gets a vector stored (see Endpoint.checkBatch())
|
||||
let sentinel = await this.Endpoint.getSentinel();
|
||||
let served;
|
||||
try {
|
||||
return await _embedViaEndpoint(endpoint, texts);
|
||||
served = await this.embedViaEndpoint(endpoint.url, [...texts, sentinel.text]);
|
||||
}
|
||||
catch (e) {
|
||||
Zotero.logError(e);
|
||||
this.Endpoint.recordFailure(e);
|
||||
Zotero.debug('Embeddings: endpoint failed -- embedding this batch locally');
|
||||
return this.embedMany(texts);
|
||||
}
|
||||
this.Endpoint.recordSuccess();
|
||||
let checked = await this.Endpoint.checkBatch(endpoint, served, sentinel.vector, texts.length);
|
||||
return checked ? served.vectors.slice(0, texts.length) : this.embedMany(texts);
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
@ -2900,6 +2978,7 @@ Zotero.Embeddings.Indexing = new function () {
|
|||
boosts: [..._threadBoosts]
|
||||
}
|
||||
},
|
||||
endpoint: Zotero.Embeddings.Endpoint.getStatus(),
|
||||
error: _lastError ? (_lastError.message || String(_lastError)) : null
|
||||
};
|
||||
};
|
||||
|
|
@ -3114,6 +3193,8 @@ Zotero.Embeddings.Indexing = new function () {
|
|||
// the first run for a given model version pays for this; every
|
||||
// later one finds the numbers already in the database.
|
||||
await Zotero.Embeddings.ensureCalibration();
|
||||
// A verified endpoint may have changed or gone away since
|
||||
await Zotero.Embeddings.Endpoint.recheck();
|
||||
|
||||
_downloadProgress = null;
|
||||
let shouldStop = () => _stopping;
|
||||
|
|
@ -3672,3 +3753,366 @@ Zotero.Embeddings.Calibration = new function () {
|
|||
}
|
||||
};
|
||||
|
||||
|
||||
/**
|
||||
* Zotero.Embeddings.Endpoint -- embedding passages through a server that
|
||||
* serves the active model, on this machine or another.
|
||||
*
|
||||
* A server is trusted because its vectors match the local model's, never
|
||||
* because of its name: verify() embeds a fixed set of texts both ways and
|
||||
* compares them, and the verdict is stored keyed to the URL and the model
|
||||
* version it was measured for. Passages route to the endpoint only while a
|
||||
* matching verdict says so (see Zotero.Embeddings.embedPassages()); queries
|
||||
* always embed locally, and so does any batch the endpoint fails.
|
||||
*/
|
||||
Zotero.Embeddings.Endpoint = new function () {
|
||||
// Least per-text cosine between served and local vectors for the two to
|
||||
// count as the same model. Measured for bekko-a8m against llama.cpp at
|
||||
// BF16 on 162 real library chunks: none below 0.9998. A Q8 GGUF of the
|
||||
// same weights sits around 0.997; the same model with the wrong pooling
|
||||
// lands near 0.7, a different model below 0.5.
|
||||
const AGREEMENT_MIN = 0.98;
|
||||
// Texts the probe embeds both ways, drawn from the calibration corpus
|
||||
const PROBE_SHORT = 8;
|
||||
const PROBE_LONG = 12;
|
||||
// The probe's longest text, in characters -- a chunk near the indexer's
|
||||
// worst case, so a server whose window is too small for real chunks is
|
||||
// found out here rather than mid-run
|
||||
const PROBE_PADDED_CHARS = 8000;
|
||||
// Context and batch size baked into the command shown to the user, for
|
||||
// the same reason
|
||||
const SERVER_CONTEXT = 4096;
|
||||
// Consecutive failed requests before the endpoint is left alone for the
|
||||
// rest of the run, rather than waiting out a timeout per batch
|
||||
const MAX_FAILURES = 3;
|
||||
const META_KEY = 'endpoint';
|
||||
|
||||
// The stored verdict: undefined until read, null when there is none
|
||||
let _verdict;
|
||||
// The sentinel text and its local vector, for the active model version
|
||||
let _sentinel = null;
|
||||
// Consecutive failed requests
|
||||
let _failures = 0;
|
||||
// Whether the endpoint is being skipped for the rest of this run
|
||||
let _suspended = false;
|
||||
|
||||
/**
|
||||
* Whether the active model can be served at all (see `serving` in MODELS).
|
||||
* @return {Boolean}
|
||||
*/
|
||||
this.isSupported = function () {
|
||||
return !!Zotero.Embeddings.getServing();
|
||||
};
|
||||
|
||||
/**
|
||||
* The llama.cpp command that serves the active model with the settings
|
||||
* its vectors depend on, and the URL it serves at.
|
||||
*
|
||||
* @return {Object|null} - { command, url }, or null for a model that
|
||||
* can't be served
|
||||
*/
|
||||
this.getCommand = function () {
|
||||
let serving = Zotero.Embeddings.getServing();
|
||||
if (!serving) {
|
||||
return null;
|
||||
}
|
||||
return {
|
||||
command: `llama serve -hf ${serving.gguf}:${serving.quant} --embeddings `
|
||||
+ `--pooling ${Zotero.Embeddings.getPooling()} `
|
||||
+ `-c ${SERVER_CONTEXT} -ub ${SERVER_CONTEXT} -b ${SERVER_CONTEXT} --port 8080`,
|
||||
url: 'http://localhost:8080/v1/embeddings'
|
||||
};
|
||||
};
|
||||
|
||||
/**
|
||||
* Embed a fixed set of texts locally and through the endpoint, compare
|
||||
* them, and store the verdict for the current model version.
|
||||
*
|
||||
* @param {String} url
|
||||
* @return {Promise<Object>} - { state, url, modelVersion, serverModel,
|
||||
* agreement, time }. state is 'ok', or why not: 'unreachable' (no
|
||||
* usable response), 'unauthorized' (the server wants credentials,
|
||||
* which aren't supported), 'not-embeddings' (not an OpenAI-style
|
||||
* embeddings endpoint), 'width-mismatch' (a different model),
|
||||
* 'low-agreement' (a different model, or the wrong pooling), or
|
||||
* 'context-too-small' (the server can't take a chunk-length text).
|
||||
* agreement is the least per-text cosine, once there are vectors to
|
||||
* compare.
|
||||
*/
|
||||
this.verify = async function (url) {
|
||||
let E = Zotero.Embeddings;
|
||||
_suspended = false;
|
||||
_failures = 0;
|
||||
let verdict = {
|
||||
state: null, url, modelVersion: E.getModelVersion(),
|
||||
serverModel: null, agreement: null, time: Date.now()
|
||||
};
|
||||
let fail = async (state, detail) => {
|
||||
Zotero.debug(`Embeddings: endpoint ${url} failed verification -- ${state}: ${detail}`);
|
||||
verdict.state = state;
|
||||
await _save(verdict);
|
||||
return verdict;
|
||||
};
|
||||
let prefix = E.getPassagePrefix();
|
||||
let { regular, padded } = _probeTexts();
|
||||
let local = await E.embedMany([...regular, padded].map(text => prefix + text));
|
||||
let localPadded = local.pop();
|
||||
|
||||
let served;
|
||||
try {
|
||||
served = await E.embedViaEndpoint(url, regular.map(text => prefix + text));
|
||||
}
|
||||
catch (e) {
|
||||
return fail(_classifyFailure(e), e.message);
|
||||
}
|
||||
verdict.serverModel = served.model;
|
||||
let dims = local[0].length;
|
||||
if (served.width < dims) {
|
||||
return fail('width-mismatch', `served ${served.width} dimensions, model has ${dims}`);
|
||||
}
|
||||
verdict.agreement = Math.min(...served.vectors.map((vector, i) => E.cosine(vector, local[i])));
|
||||
if (verdict.agreement < AGREEMENT_MIN) {
|
||||
return fail('low-agreement', `agreement ${verdict.agreement.toFixed(3)}`);
|
||||
}
|
||||
// A server whose window is too small for a chunk rejects it, or embeds
|
||||
// what fits and says nothing
|
||||
try {
|
||||
let { vectors } = await E.embedViaEndpoint(url, [prefix + padded]);
|
||||
let agreement = E.cosine(vectors[0], localPadded);
|
||||
if (agreement < AGREEMENT_MIN) {
|
||||
return fail('context-too-small', `long text agreement ${agreement.toFixed(3)}`);
|
||||
}
|
||||
}
|
||||
catch (e) {
|
||||
// A server that took the other texts and rejects this one is
|
||||
// short of window, not of network
|
||||
if (e instanceof Zotero.HTTP.UnexpectedStatusException) {
|
||||
return fail('context-too-small', e.message);
|
||||
}
|
||||
return fail(_classifyFailure(e), e.message);
|
||||
}
|
||||
verdict.state = 'ok';
|
||||
await _save(verdict);
|
||||
return verdict;
|
||||
};
|
||||
|
||||
/**
|
||||
* The endpoint passages route to: the configured URL, when a stored
|
||||
* verdict says it serves the active model.
|
||||
*
|
||||
* @return {Promise<Object|null>} - { url, serverModel }
|
||||
*/
|
||||
this.getActive = async function () {
|
||||
let url = Zotero.Prefs.get('embeddings.endpoint');
|
||||
if (!url || _suspended) {
|
||||
return null;
|
||||
}
|
||||
let verdict = await _load();
|
||||
if (!_applies(verdict, url) || verdict.state != 'ok') {
|
||||
return null;
|
||||
}
|
||||
return { url, serverModel: verdict.serverModel };
|
||||
};
|
||||
|
||||
/**
|
||||
* A short text sent along with every batch, and the active model's own
|
||||
* vector for it, embedded once per model version. Checking the server's
|
||||
* vector for it costs no local inference per batch.
|
||||
*
|
||||
* @return {Promise<Object>} - { text, vector }
|
||||
*/
|
||||
this.getSentinel = async function () {
|
||||
let modelVersion = Zotero.Embeddings.getModelVersion();
|
||||
if (_sentinel?.modelVersion !== modelVersion) {
|
||||
let text = Zotero.Embeddings.getPassagePrefix() + _probeTexts().regular[0];
|
||||
let [vector] = await Zotero.Embeddings.embedMany([text]);
|
||||
_sentinel = { modelVersion, text, vector };
|
||||
}
|
||||
return _sentinel;
|
||||
};
|
||||
|
||||
/**
|
||||
* Whether a served batch can be stored: the server still reports the
|
||||
* model it was verified with, and its vector for the text at `index`
|
||||
* matches the local model's. Anything else invalidates the endpoint.
|
||||
*
|
||||
* @param {Object} endpoint - As getActive() returned it
|
||||
* @param {Object} served - As Zotero.Embeddings.embedViaEndpoint() returned it
|
||||
* @param {Float32Array} local - The local vector for that text
|
||||
* @param {Number} index
|
||||
* @return {Promise<Boolean>}
|
||||
*/
|
||||
this.checkBatch = async function (endpoint, served, local, index) {
|
||||
if (served.model !== endpoint.serverModel) {
|
||||
await this.invalidate(`served model changed from '${endpoint.serverModel}' to '${served.model}'`);
|
||||
return false;
|
||||
}
|
||||
let agreement = Zotero.Embeddings.cosine(served.vectors[index], local);
|
||||
if (agreement < AGREEMENT_MIN) {
|
||||
await this.invalidate(`served vectors stopped matching (agreement ${agreement.toFixed(3)})`);
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
};
|
||||
|
||||
/**
|
||||
* Note a failed request. After MAX_FAILURES in a row the endpoint is
|
||||
* skipped until the next run (see recheck()) or verification.
|
||||
* @param {Error} e
|
||||
*/
|
||||
this.recordFailure = function (e) {
|
||||
if (++_failures >= MAX_FAILURES && !_suspended) {
|
||||
_suspended = true;
|
||||
Zotero.debug(`Embeddings: endpoint skipped for the rest of this run -- `
|
||||
+ `${MAX_FAILURES} requests failed in a row, the last with: ${e.message}`);
|
||||
}
|
||||
};
|
||||
|
||||
this.recordSuccess = function () {
|
||||
_failures = 0;
|
||||
};
|
||||
|
||||
/**
|
||||
* At the start of a run: try the endpoint again if it was skipped, and
|
||||
* confirm a verified one still serves the model, on one text.
|
||||
*/
|
||||
this.recheck = async function () {
|
||||
_suspended = false;
|
||||
_failures = 0;
|
||||
let active = await this.getActive();
|
||||
if (!active) {
|
||||
return;
|
||||
}
|
||||
let sentinel = await this.getSentinel();
|
||||
let served;
|
||||
try {
|
||||
served = await Zotero.Embeddings.embedViaEndpoint(active.url, [sentinel.text]);
|
||||
}
|
||||
catch (e) {
|
||||
_suspended = true;
|
||||
Zotero.debug(`Embeddings: endpoint skipped for this run -- unreachable: ${e.message}`);
|
||||
return;
|
||||
}
|
||||
await this.checkBatch(active, served, sentinel.vector, 0);
|
||||
};
|
||||
|
||||
/**
|
||||
* Stop routing to the endpoint until it's verified again.
|
||||
* @param {String} detail - For the debug log
|
||||
*/
|
||||
this.invalidate = async function (detail) {
|
||||
let verdict = await _load();
|
||||
if (!verdict) {
|
||||
return;
|
||||
}
|
||||
Zotero.debug(`Embeddings: endpoint no longer trusted -- ${detail}`);
|
||||
await _save(Object.assign(verdict, { state: 'invalid' }));
|
||||
};
|
||||
|
||||
/**
|
||||
* The endpoint as the preferences show it. Synchronous, from the verdict
|
||||
* last read; 'unknown' until one has been.
|
||||
*
|
||||
* @return {Object} - { url, state, serverModel }: state is 'off' with no
|
||||
* URL configured, 'unverified' with a URL no stored verdict covers,
|
||||
* 'unreachable' while a verified endpoint is being skipped this run,
|
||||
* or the verdict's own state
|
||||
*/
|
||||
this.getStatus = function () {
|
||||
let url = Zotero.Prefs.get('embeddings.endpoint') || '';
|
||||
let status = { url, state: 'off', serverModel: null };
|
||||
if (!url) {
|
||||
return status;
|
||||
}
|
||||
if (_verdict === undefined) {
|
||||
status.state = 'unknown';
|
||||
}
|
||||
else if (!_applies(_verdict, url)) {
|
||||
status.state = 'unverified';
|
||||
}
|
||||
else {
|
||||
Object.assign(status, {
|
||||
state: _verdict.state == 'ok' && _suspended ? 'unreachable' : _verdict.state,
|
||||
serverModel: _verdict.serverModel
|
||||
});
|
||||
}
|
||||
return status;
|
||||
};
|
||||
|
||||
/**
|
||||
* Read the stored verdict into memory, so getStatus() can answer.
|
||||
* @return {Promise<Object|null>}
|
||||
*/
|
||||
this.load = function () {
|
||||
return _load();
|
||||
};
|
||||
|
||||
/**
|
||||
* Forget the verdict read into memory, after the table holding it was
|
||||
* rebuilt (see Zotero.Embeddings.initDB())
|
||||
*/
|
||||
this.reset = function () {
|
||||
_verdict = undefined;
|
||||
_sentinel = null;
|
||||
_suspended = false;
|
||||
_failures = 0;
|
||||
};
|
||||
|
||||
// Whether a stored verdict is about this URL and the active model
|
||||
function _applies(verdict, url) {
|
||||
return !!verdict && verdict.url === url
|
||||
&& verdict.modelVersion === Zotero.Embeddings.getModelVersion();
|
||||
}
|
||||
|
||||
async function _load() {
|
||||
if (_verdict === undefined) {
|
||||
await Zotero.Embeddings.initDB();
|
||||
let json = await Zotero.DB.valueQueryAsync(
|
||||
"SELECT value FROM embeddings.itemEmbeddingsMeta WHERE key=?", [META_KEY]);
|
||||
_verdict = json ? JSON.parse(json) : null;
|
||||
}
|
||||
return _verdict;
|
||||
}
|
||||
|
||||
async function _save(verdict) {
|
||||
await Zotero.Embeddings.initDB();
|
||||
await Zotero.DB.queryAsync(
|
||||
"REPLACE INTO embeddings.itemEmbeddingsMeta (key, value) VALUES (?, ?)",
|
||||
[META_KEY, JSON.stringify(verdict)]);
|
||||
_verdict = verdict;
|
||||
}
|
||||
|
||||
// What went wrong with a request, as a verdict state
|
||||
function _classifyFailure(e) {
|
||||
if (e instanceof Zotero.HTTP.UnexpectedStatusException) {
|
||||
if (e.status == 401 || e.status == 403) {
|
||||
return 'unauthorized';
|
||||
}
|
||||
if (e.status == 404 || e.status == 405) {
|
||||
return 'not-embeddings';
|
||||
}
|
||||
return 'unreachable';
|
||||
}
|
||||
if (e instanceof Zotero.Embeddings.EndpointResponseError) {
|
||||
return 'not-embeddings';
|
||||
}
|
||||
return 'unreachable';
|
||||
}
|
||||
|
||||
// Short and long passages from the calibration corpus, and one long
|
||||
// enough to overrun a small server window, its tail distinct so that a
|
||||
// silently truncated embedding no longer matches the local one
|
||||
function _probeTexts() {
|
||||
let { short, long } = Zotero.Embeddings.Calibration.getCorpus();
|
||||
let regular = [
|
||||
...short.slice(0, PROBE_SHORT).map(pair => pair.passage),
|
||||
...long.slice(0, PROBE_LONG).map(pair => pair.passage)
|
||||
];
|
||||
let padded = '';
|
||||
for (let i = 0; padded.length < PROBE_PADDED_CHARS; i++) {
|
||||
padded += long[i % long.length].passage + ' ';
|
||||
}
|
||||
padded += 'The closing sentence names the Antikythera mechanism and the lighthouse at Alexandria.';
|
||||
return { regular, padded };
|
||||
}
|
||||
};
|
||||
|
|
|
|||
|
|
@ -135,6 +135,51 @@ preferences-advanced-semantic-search-switch-button = Change Mode
|
|||
preferences-advanced-semantic-search-disable-title = Disable Best-Match Search?
|
||||
preferences-advanced-semantic-search-disable-text = Disabling will delete the search index and downloaded data.
|
||||
preferences-advanced-semantic-search-disable-button = Disable
|
||||
preferences-advanced-semantic-search-endpoint-configure =
|
||||
.label = Configure Endpoint…
|
||||
preferences-advanced-semantic-search-endpoint-off =
|
||||
.value = Embedding endpoint: not configured
|
||||
preferences-advanced-semantic-search-endpoint-unverified =
|
||||
.value = Embedding endpoint: not verified
|
||||
preferences-advanced-semantic-search-endpoint-valid =
|
||||
.value = Embedding endpoint: valid
|
||||
preferences-advanced-semantic-search-endpoint-unreachable =
|
||||
.value = Embedding endpoint: unreachable
|
||||
preferences-advanced-semantic-search-endpoint-invalid =
|
||||
.value = Embedding endpoint: not valid
|
||||
preferences-advanced-semantic-search-endpoint-dialog =
|
||||
.title = Configure Embedding Endpoint
|
||||
preferences-advanced-semantic-search-endpoint-intro =
|
||||
.value = Send the embedding workload to a local or remote server:
|
||||
preferences-advanced-semantic-search-endpoint-url =
|
||||
.value = Server URL:
|
||||
preferences-advanced-semantic-search-endpoint-privacy = A remote server receives the text of your library.
|
||||
preferences-advanced-semantic-search-endpoint-requirements =
|
||||
.value = The server must serve the model Zotero is using:
|
||||
preferences-advanced-semantic-search-endpoint-model =
|
||||
.value = Model:
|
||||
preferences-advanced-semantic-search-endpoint-pooling =
|
||||
.value = Pooling:
|
||||
preferences-advanced-semantic-search-endpoint-file =
|
||||
.value = File:
|
||||
preferences-advanced-semantic-search-endpoint-llama =
|
||||
.value = Run with llama.cpp:
|
||||
preferences-advanced-semantic-search-endpoint-then-use =
|
||||
.value = Then use { $url }
|
||||
preferences-advanced-semantic-search-endpoint-copy =
|
||||
.label = Copy
|
||||
preferences-advanced-semantic-search-endpoint-verifying =
|
||||
.value = Checking that the server’s vectors match the model’s…
|
||||
preferences-advanced-semantic-search-endpoint-accept =
|
||||
.label = Verify and Use
|
||||
preferences-advanced-semantic-search-endpoint-remove =
|
||||
.label = Stop Using Endpoint
|
||||
preferences-advanced-semantic-search-endpoint-error-unreachable = The server didn’t respond.
|
||||
preferences-advanced-semantic-search-endpoint-error-unauthorized = The server requires authentication, which isn’t supported.
|
||||
preferences-advanced-semantic-search-endpoint-error-not-embeddings = That URL isn’t an OpenAI-style embeddings endpoint.
|
||||
preferences-advanced-semantic-search-endpoint-error-width-mismatch = The server is serving a different model.
|
||||
preferences-advanced-semantic-search-endpoint-error-low-agreement = The server’s vectors don’t match this model’s. Check that it serves the model above with { $pooling } pooling.
|
||||
preferences-advanced-semantic-search-endpoint-error-context-too-small = The server can’t embed long passages. Start it with a context of at least 4096 tokens.
|
||||
preferences-advanced-language-and-region-title = Language and Region
|
||||
preferences-advanced-enable-bidi-ui =
|
||||
.label = Enable bidirectional text editing utilities
|
||||
|
|
|
|||
|
|
@ -123,6 +123,9 @@ pref("extensions.zotero.embeddings.model", "");
|
|||
pref("extensions.zotero.embeddings.indexingPaused", false);
|
||||
// Also index the full text of PDF/EPUB/snapshot attachments
|
||||
pref("extensions.zotero.embeddings.indexFulltext", false);
|
||||
// An OpenAI-style /v1/embeddings URL to send indexing to, used only once
|
||||
// verified to serve the active model (see Zotero.Embeddings.Endpoint)
|
||||
pref("extensions.zotero.embeddings.endpoint", "");
|
||||
|
||||
// Notes
|
||||
pref("extensions.zotero.note.fontFamily", "-apple-system, BlinkMacSystemFont, \"Segoe UI\", \"Helvetica Neue\", Helvetica, Arial, sans-serif");
|
||||
|
|
|
|||
|
|
@ -154,4 +154,33 @@
|
|||
|
||||
#semantic-search-libraries {
|
||||
padding-top: 10px;
|
||||
}
|
||||
}
|
||||
|
||||
#semantic-search-endpoint-row {
|
||||
margin: 4px 0;
|
||||
}
|
||||
|
||||
#zotero-embeddings-endpoint-container {
|
||||
.endpoint-note {
|
||||
opacity: 0.7;
|
||||
font-size: 0.9em;
|
||||
}
|
||||
|
||||
.endpoint-facts {
|
||||
display: grid;
|
||||
grid-template-columns: max-content 1fr;
|
||||
column-gap: 8px;
|
||||
margin: 4px 0 4px 12px;
|
||||
}
|
||||
|
||||
.endpoint-command {
|
||||
flex: 1;
|
||||
font-family: monospace;
|
||||
font-size: 0.9em;
|
||||
resize: none;
|
||||
}
|
||||
|
||||
.endpoint-error {
|
||||
color: var(--accent-red);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1167,6 +1167,35 @@ describe("Zotero.Embeddings", function () {
|
|||
stubs.forEach(stub => stub.restore());
|
||||
}
|
||||
});
|
||||
it("should collapse whitespace and give a model's tokenizer the leading space the runtime omits", async function () {
|
||||
let seen = [];
|
||||
let engine = fakeEngine(async (engine, { args: [texts] }) => {
|
||||
seen.push(...texts);
|
||||
return texts.map(() => Array.from({ length: 384 }, (_, i) => Math.sin(i + 1)));
|
||||
});
|
||||
let modelName = sinon.stub(Zotero.Embeddings, 'getModelName').returns('bekko-embedding-v1-a8m');
|
||||
let stubs = [
|
||||
sinon.stub(Zotero.ML, 'createEngine').resolves(engine),
|
||||
sinon.stub(Zotero.ML, 'shutdown').resolves(),
|
||||
sinon.stub(Zotero.ML, 'getOptimalConcurrency').returns(2),
|
||||
modelName,
|
||||
sinon.stub(Zotero.Embeddings, 'getModelVersion').returns('test-spacing/1')
|
||||
];
|
||||
try {
|
||||
await Zotero.Embeddings.embedMany(['some text', 'more', 'line one\nline two\n']);
|
||||
assert.deepEqual(seen, [' some text', ' more', ' line one line two']);
|
||||
// A model whose tokenizer doesn't mark the first word is left alone
|
||||
await Zotero.Embeddings.shutdownEngine({ modelChanged: false });
|
||||
seen = [];
|
||||
modelName.returns('bge-small-zh-v1.5');
|
||||
await Zotero.Embeddings.embedMany(['some text']);
|
||||
assert.deepEqual(seen, ['some text']);
|
||||
}
|
||||
finally {
|
||||
await Zotero.Embeddings.shutdownEngine({ modelChanged: false });
|
||||
stubs.forEach(stub => stub.restore());
|
||||
}
|
||||
});
|
||||
it("should keep only the first dims of a model that truncates", async function () {
|
||||
// 384 raw dimensions from the engine; bekko-a25m stores 256 of them
|
||||
let raw = Array.from({ length: 384 }, (_, i) => Math.sin(i + 1));
|
||||
|
|
@ -1314,6 +1343,232 @@ describe("Zotero.Embeddings", function () {
|
|||
});
|
||||
});
|
||||
|
||||
describe("Endpoint", function () {
|
||||
const URL = 'http://localhost:8080/v1/embeddings';
|
||||
// Deterministic unit vectors, one per text, in the width the served
|
||||
// model stores (bekko keeps 256 of its 384)
|
||||
let vectorFor = (text, dims = 256) => {
|
||||
let state = 0;
|
||||
for (let i = 0; i < text.length; i++) {
|
||||
state = (state * 31 + text.charCodeAt(i)) % 2147483647;
|
||||
}
|
||||
let vector = new Float32Array(dims);
|
||||
for (let d = 0; d < dims; d++) {
|
||||
state = (state * 1103515245 + 12345) % 2147483648;
|
||||
vector[d] = state / 2147483648 - 0.5;
|
||||
}
|
||||
let norm = Math.sqrt(vector.reduce((sum, val) => sum + val * val, 0));
|
||||
return vector.map(val => val / norm);
|
||||
};
|
||||
// A server: answers each request from `remoteFor(text)`, or throws
|
||||
// what `failWith(texts)` returns
|
||||
let serve = ({ remoteFor = vectorFor, model = 'served', failWith = () => null, shape = null } = {}) => {
|
||||
let calls = [];
|
||||
return {
|
||||
calls,
|
||||
stub: sinon.stub(Zotero.HTTP, 'request').callsFake(async (method, url, options) => {
|
||||
let { input } = JSON.parse(options.body);
|
||||
calls.push({ url, input, headers: options.headers });
|
||||
let error = failWith(input);
|
||||
if (error) {
|
||||
throw error;
|
||||
}
|
||||
if (shape) {
|
||||
return { status: 200, response: shape };
|
||||
}
|
||||
return {
|
||||
status: 200,
|
||||
response: {
|
||||
model,
|
||||
data: input.map((text, index) => ({ index, embedding: Array.from(remoteFor(text)) }))
|
||||
}
|
||||
};
|
||||
})
|
||||
};
|
||||
};
|
||||
let status = code => new Zotero.HTTP.UnexpectedStatusException({ status: code }, URL, `HTTP ${code}`);
|
||||
let stubs;
|
||||
beforeEach(async function () {
|
||||
stubs = [
|
||||
sinon.stub(Zotero.Embeddings, 'getModelName').returns('bekko-embedding-v1-a8m'),
|
||||
sinon.stub(Zotero.Embeddings, 'getModelVersion').returns('test-endpoint/1'),
|
||||
sinon.stub(Zotero.Embeddings, 'embedMany').callsFake(async texts => texts.map(text => vectorFor(text)))
|
||||
];
|
||||
await Zotero.DB.queryAsync("DELETE FROM embeddings.itemEmbeddingsMeta WHERE key='endpoint'");
|
||||
Zotero.Embeddings.Endpoint.reset();
|
||||
});
|
||||
afterEach(async function () {
|
||||
stubs.forEach(stub => stub.restore());
|
||||
Zotero.HTTP.request.restore?.();
|
||||
Zotero.Prefs.clear('embeddings.endpoint');
|
||||
await Zotero.DB.queryAsync("DELETE FROM embeddings.itemEmbeddingsMeta WHERE key='endpoint'");
|
||||
});
|
||||
|
||||
it("should describe how to serve the model", function () {
|
||||
let { command, url } = Zotero.Embeddings.Endpoint.getCommand();
|
||||
assert.include(command, 'hotchpotch/bekko-embedding-v1-a8m-GGUF:F16');
|
||||
assert.include(command, '--pooling mean');
|
||||
assert.include(url, 'localhost');
|
||||
});
|
||||
|
||||
it("should accept a server whose vectors match the model's", async function () {
|
||||
let server = serve();
|
||||
let verdict = await Zotero.Embeddings.Endpoint.verify(URL);
|
||||
assert.equal(verdict.state, 'ok');
|
||||
assert.equal(verdict.serverModel, 'served');
|
||||
assert.closeTo(verdict.agreement, 1, 1e-5);
|
||||
// The regular texts in one request, the chunk-length one alone
|
||||
assert.lengthOf(server.calls, 2);
|
||||
assert.isAbove(server.calls[0].input.length, 10);
|
||||
assert.isAbove(server.calls[1].input[0].length, 8000);
|
||||
});
|
||||
|
||||
it("should reject a server serving something else", async function () {
|
||||
serve({ remoteFor: text => vectorFor(text + ' but different') });
|
||||
let verdict = await Zotero.Embeddings.Endpoint.verify(URL);
|
||||
assert.equal(verdict.state, 'low-agreement');
|
||||
assert.isBelow(verdict.agreement, 0.5);
|
||||
});
|
||||
|
||||
it("should reject vectors narrower than the model's", async function () {
|
||||
serve({ remoteFor: text => vectorFor(text, 100) });
|
||||
let verdict = await Zotero.Embeddings.Endpoint.verify(URL);
|
||||
assert.equal(verdict.state, 'width-mismatch');
|
||||
});
|
||||
|
||||
it("should tell an unreachable server, a wrong endpoint, and one wanting credentials apart", async function () {
|
||||
serve({ failWith: () => new Error('connection refused') });
|
||||
assert.equal((await Zotero.Embeddings.Endpoint.verify(URL)).state, 'unreachable');
|
||||
Zotero.HTTP.request.restore();
|
||||
serve({ failWith: () => status(404) });
|
||||
assert.equal((await Zotero.Embeddings.Endpoint.verify(URL)).state, 'not-embeddings');
|
||||
Zotero.HTTP.request.restore();
|
||||
serve({ shape: { choices: [] } });
|
||||
assert.equal((await Zotero.Embeddings.Endpoint.verify(URL)).state, 'not-embeddings');
|
||||
Zotero.HTTP.request.restore();
|
||||
serve({ failWith: () => status(401) });
|
||||
assert.equal((await Zotero.Embeddings.Endpoint.verify(URL)).state, 'unauthorized');
|
||||
});
|
||||
|
||||
it("should reject a server that can't take a chunk-length text", async function () {
|
||||
// llama.cpp rejects it
|
||||
serve({ failWith: texts => (texts[0].length > 8000 ? status(500) : null) });
|
||||
assert.equal((await Zotero.Embeddings.Endpoint.verify(URL)).state, 'context-too-small');
|
||||
Zotero.HTTP.request.restore();
|
||||
// Or embeds what fits and says nothing
|
||||
serve({ remoteFor: text => vectorFor(text.length > 8000 ? text.slice(0, 2000) : text) });
|
||||
assert.equal((await Zotero.Embeddings.Endpoint.verify(URL)).state, 'context-too-small');
|
||||
});
|
||||
|
||||
it("should route passages only to a verified endpoint for this model", async function () {
|
||||
let server = serve();
|
||||
// Configured but unverified: local
|
||||
Zotero.Prefs.set('embeddings.endpoint', URL);
|
||||
await Zotero.Embeddings.embedPassages(['a passage']);
|
||||
assert.lengthOf(server.calls, 0);
|
||||
assert.equal(Zotero.Embeddings.Endpoint.getStatus().state, 'unverified');
|
||||
|
||||
await Zotero.Embeddings.Endpoint.verify(URL);
|
||||
let vectors = await Zotero.Embeddings.embedPassages(['a passage']);
|
||||
assert.lengthOf(server.calls, 3);
|
||||
// The batch plus its sentinel went out; only the batch came back
|
||||
assert.lengthOf(server.calls[2].input, 2);
|
||||
assert.lengthOf(vectors, 1);
|
||||
assert.closeTo(Zotero.Embeddings.cosine(vectors[0], vectorFor('a passage')), 1, 1e-5);
|
||||
assert.equal(Zotero.Embeddings.Endpoint.getStatus().state, 'ok');
|
||||
|
||||
// Verified for another URL: local
|
||||
Zotero.Prefs.set('embeddings.endpoint', 'http://elsewhere/v1/embeddings');
|
||||
await Zotero.Embeddings.embedPassages(['a passage']);
|
||||
assert.lengthOf(server.calls, 3);
|
||||
assert.equal(Zotero.Embeddings.Endpoint.getStatus().state, 'unverified');
|
||||
|
||||
// Verified for another model version: local
|
||||
Zotero.Prefs.set('embeddings.endpoint', URL);
|
||||
stubs[1].returns('test-endpoint/2');
|
||||
await Zotero.Embeddings.embedPassages(['a passage']);
|
||||
assert.lengthOf(server.calls, 3);
|
||||
});
|
||||
|
||||
it("should stop trusting a server whose model changes mid-run", async function () {
|
||||
let server = serve();
|
||||
Zotero.Prefs.set('embeddings.endpoint', URL);
|
||||
await Zotero.Embeddings.Endpoint.verify(URL);
|
||||
Zotero.HTTP.request.restore();
|
||||
server = serve({ model: 'something-else' });
|
||||
let [vector] = await Zotero.Embeddings.embedPassages(['a passage']);
|
||||
// The batch embedded locally, and the endpoint is out
|
||||
assert.lengthOf(server.calls, 1);
|
||||
assert.closeTo(Zotero.Embeddings.cosine(vector, vectorFor('a passage')), 1, 1e-5);
|
||||
let status = Zotero.Embeddings.Endpoint.getStatus();
|
||||
assert.equal(status.state, 'invalid');
|
||||
await Zotero.Embeddings.embedPassages(['another']);
|
||||
assert.lengthOf(server.calls, 1);
|
||||
});
|
||||
|
||||
it("should stop trusting a server whose vectors drift under the same name", async function () {
|
||||
serve();
|
||||
Zotero.Prefs.set('embeddings.endpoint', URL);
|
||||
await Zotero.Embeddings.Endpoint.verify(URL);
|
||||
Zotero.HTTP.request.restore();
|
||||
let server = serve({ remoteFor: text => vectorFor(text + ' drifted') });
|
||||
let [vector] = await Zotero.Embeddings.embedPassages(['first', 'a longer second passage']);
|
||||
// Caught on the very batch through its sentinel: nothing served was returned
|
||||
assert.lengthOf(server.calls, 1);
|
||||
assert.closeTo(Zotero.Embeddings.cosine(vector, vectorFor('first')), 1, 1e-5);
|
||||
let status = Zotero.Embeddings.Endpoint.getStatus();
|
||||
assert.equal(status.state, 'invalid');
|
||||
});
|
||||
|
||||
it("should leave a server alone after repeated failures until the next run", async function () {
|
||||
serve();
|
||||
Zotero.Prefs.set('embeddings.endpoint', URL);
|
||||
await Zotero.Embeddings.Endpoint.verify(URL);
|
||||
Zotero.HTTP.request.restore();
|
||||
let server = serve({ failWith: () => new Error('timeout') });
|
||||
for (let i = 0; i < 5; i++) {
|
||||
await Zotero.Embeddings.embedPassages([`passage ${i}`]);
|
||||
}
|
||||
assert.lengthOf(server.calls, 3);
|
||||
let status = Zotero.Embeddings.Endpoint.getStatus();
|
||||
assert.equal(status.state, 'unreachable');
|
||||
|
||||
// The next run tries again, and a server that answers is back
|
||||
Zotero.HTTP.request.restore();
|
||||
server = serve();
|
||||
await Zotero.Embeddings.Endpoint.recheck();
|
||||
assert.equal(Zotero.Embeddings.Endpoint.getStatus().state, 'ok');
|
||||
await Zotero.Embeddings.embedPassages(['again']);
|
||||
assert.lengthOf(server.calls, 2);
|
||||
});
|
||||
|
||||
it("should recheck a verified server when a run starts", async function () {
|
||||
serve();
|
||||
Zotero.Prefs.set('embeddings.endpoint', URL);
|
||||
await Zotero.Embeddings.Endpoint.verify(URL);
|
||||
Zotero.HTTP.request.restore();
|
||||
serve({ failWith: () => new Error('connection refused') });
|
||||
await Zotero.Embeddings.Endpoint.recheck();
|
||||
assert.equal(Zotero.Embeddings.Endpoint.getStatus().state, 'unreachable');
|
||||
Zotero.HTTP.request.restore();
|
||||
serve({ model: 'swapped' });
|
||||
await Zotero.Embeddings.Endpoint.recheck();
|
||||
assert.equal(Zotero.Embeddings.Endpoint.getStatus().state, 'invalid');
|
||||
});
|
||||
|
||||
it("should embed a failed batch locally and keep the endpoint", async function () {
|
||||
let server = serve();
|
||||
Zotero.Prefs.set('embeddings.endpoint', URL);
|
||||
await Zotero.Embeddings.Endpoint.verify(URL);
|
||||
Zotero.HTTP.request.restore();
|
||||
server = serve({ failWith: () => status(500) });
|
||||
let [vector] = await Zotero.Embeddings.embedPassages(['a passage']);
|
||||
assert.closeTo(Zotero.Embeddings.cosine(vector, vectorFor('a passage')), 1, 1e-5);
|
||||
assert.lengthOf(server.calls, 1);
|
||||
assert.equal(Zotero.Embeddings.Endpoint.getStatus().state, 'ok');
|
||||
});
|
||||
});
|
||||
|
||||
describe("Indexing", function () {
|
||||
it("should clear embeddings when the model changes", async function () {
|
||||
let stubs = [
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue