From fbb0495f054efced655a3174a6d55cf5b7e7e15d Mon Sep 17 00:00:00 2001 From: Bogdan Abaev Date: Tue, 1 Sep 2026 09:36:06 -0700 Subject: [PATCH] approximate chunk token sizes instead of tokenizer Using tokenizer to determine exact token size of a chunk puts a lot of work on the main thread during indexing only to find if a chunk is over the min limit and is under the max. There is not much value in precision so we estimate the token value per chunk with a language table (e.g. english ~4 chars/token, CJK ~1 char/token, etc.) If needed precision ends up being necessary, we can run a tokenizer on a small subset of the text in document and derive chars/token ration from that instead of actually tokenizing the entire document. --- chrome/content/zotero/xpcom/embeddings.js | 233 ++++++------------ .../zotero/xpcom/utilities_internal.js | 81 ++++-- test/tests/embeddingsTest.js | 190 +++++++------- test/tests/utilities_internalTest.js | 12 + 4 files changed, 224 insertions(+), 292 deletions(-) diff --git a/chrome/content/zotero/xpcom/embeddings.js b/chrome/content/zotero/xpcom/embeddings.js index 1e082096e4..71b6608f86 100644 --- a/chrome/content/zotero/xpcom/embeddings.js +++ b/chrome/content/zotero/xpcom/embeddings.js @@ -1498,205 +1498,108 @@ Zotero.Embeddings = new function () { /** - * Zotero.Embeddings.Chunking -- splitting a text into passages small enough to - * each be embedded on their own, so a long text's later paragraphs are - * searchable instead of being averaged into one vector or truncated away by the - * pipeline. The indexer applies this to notes (chunkText()) and to attachment - * full text (chunkSections()). Chunk size is bounded by the shared passage - * geometry's BUDGET_TOKENS (see Zotero.Utilities.Internal.Chunking), capped - * by the model's window, and counted in the model's own tokens. + * Zotero.Embeddings.Chunking -- the shared chunker fitted to the active model. + * Splitting a text into passages small enough to each be embedded on their own + * keeps a long text's later paragraphs searchable instead of averaged into one + * vector or truncated away. The geometry and the splitting live in + * Zotero.Utilities.Internal.Chunking; this module caps the budget to the + * model's window and reports passage sizes as estimated tokens. */ Zotero.Embeddings.Chunking = new function () { - // The shared passage geometry (see Zotero.Utilities.Internal.Chunking), - // counted here in the model's own tokens: BUDGET_TOKENS caps a chunk, - // within what the model's window allows (see _getMetrics()); MIN_TOKENS - // is the least worth embedding alone -- an item scores as its best chunk, - // so fragments would inflate the score while losing their context; and - // OVERLAP_TOKENS is carried across a split, so a thought spanning the - // boundary is searchable in both halves. + // The shared geometry: BUDGET_TOKENS caps a chunk within the model's + // window (see _getBudget()); MIN_TOKENS is the least worth embedding + // alone -- an item scores as its best chunk, so fragments would inflate + // the score; OVERLAP_TOKENS is carried across a split, so a thought + // spanning the boundary is searchable in both halves. const { BUDGET_TOKENS, MIN_TOKENS, OVERLAP_TOKENS } = Zotero.Utilities.Internal.Chunking; - // Most characters per encode() call. The multilingual models' - // SentencePiece Unigram tokenizer is quadratic in input length (its - // Metaspace pre-tokenizer doesn't split at whitespace), so a single long - // encode can take seconds. Segments this size stay in the linear regime, - // and token counts are additive across whitespace boundaries, so the - // segmented measurement is exact. - const TOKENIZER_SEGMENT_CHARS = 1000; + // Window taken by the special tokens wrapped around every input -- two + // for the BERT-family models here; the budget's headroom absorbs more + const SPECIAL_TOKENS = 2; - // Tokenizer instances by model name. Failures aren't cached, so a later - // indexing run retries the load. - let _tokenizers = new Map(); + // Headroom held back from the window: the character estimate can + // undercount, and a chunk that overshoots loses its tail to the + // pipeline's truncation + const BUDGET_SAFETY = 0.9; - // Counters and metrics cached per tokenizer: each costs a few encode() - // calls to derive, which adds up when many short texts are measured one - // at a time. A tokenizer belongs to one model (see getTokenizer()), so - // the model facts baked into its metrics can't go stale. - let _metrics = new WeakMap(); - let _counters = new WeakMap(); - - /** - * The active model's tokenizer, built from the files in the runtime's - * model cache with the transformers.js implementation Firefox ships -- - * the same code, over the same files, that the inference process - * tokenizes with. - * - * A failure throws rather than falling back to an approximate count: - * approximated chunk boundaries would persist uncorrected (an item's - * source hash covers its text, not the chunker), and with the tokenizer - * files broken nothing would embed anyway. - * - * @return {Promise} - */ - this.getTokenizer = function () { - let name = Zotero.Embeddings.getModelName(); - if (!_tokenizers.has(name)) { - let promise = (async () => { - let { PreTrainedTokenizer } = ChromeUtils.importESModule( - 'chrome://global/content/ml/transformers.js' - ); - let decoder = new TextDecoder(); - let [tokenizerJSON, tokenizerConfig] = await Promise.all( - ['tokenizer.json', 'tokenizer_config.json'].map( - async file => JSON.parse(decoder.decode( - await Zotero.Embeddings.getModelFile(file) - )) - ) - ); - return new PreTrainedTokenizer(tokenizerJSON, tokenizerConfig); - })(); - // Allow a later run to retry after a failed load - promise.catch(() => { - if (_tokenizers.get(name) === promise) { - _tokenizers.delete(name); - } - }); - _tokenizers.set(name, promise); - } - return _tokenizers.get(name); - }; - - // A counter of a text's tokens that leaves out the special tokens the - // tokenizer wraps every input in, so counts of several pieces add up to - // the count of the pieces joined -- the additivity that lets the shared - // chunker measure each paragraph once and size every grouping by summing. - function _getCounter(tokenizer) { - let counter = _counters.get(tokenizer); - if (!counter) { - let specialTokens = tokenizer.encode('').length; - // Measured in segments (see _splitForTokenizer()), so a long - // paragraph never hits the tokenizer's quadratic regime in a - // single call - counter = text => _splitForTokenizer(text).reduce( - (sum, segment) => sum + tokenizer.encode(segment).length - specialTokens, - 0 - ); - _counters.set(tokenizer, counter); - } - return counter; + // The shared ceiling capped to the model's window, less the special + // tokens and the passage prefix embedPassages() prepends -- not part of + // the text the chunker sees, but in the window once a chunk is embedded + function _getBudget() { + let prefix = Zotero.Embeddings.getPassagePrefix(); + let prefixTokens = prefix + ? Math.ceil(prefix.length / Zotero.Utilities.Internal.Chunking.getCharsPerToken(prefix)) + : 0; + return Math.floor( + (Math.min(BUDGET_TOKENS, Zotero.Embeddings.getModelMaxTokens()) + - SPECIAL_TOKENS - prefixTokens) * BUDGET_SAFETY + ); } - // The shared chunker's metrics for the active model. joinSize is what one - // '\n\n' join adds, measured rather than assumed (a SentencePiece model - // can spend a token on collapsed whitespace). The budget is the shared - // ceiling, capped by the model's window, less the special tokens and the - // passage prefix embedPassages() prepends -- neither is part of the text - // the chunker sees, but both take up window once the chunk is embedded. - function _getMetrics(tokenizer) { - let cached = _metrics.get(tokenizer); - if (cached) { - return cached; - } - let count = _getCounter(tokenizer); - let specialTokens = tokenizer.encode('').length; - let prefix = Zotero.Embeddings.getPassagePrefix(); - let metrics = { - count, - joinSize: Math.max(0, count('a\n\na') - 2 * count('a')), - budget: Math.min(BUDGET_TOKENS, Zotero.Embeddings.getModelMaxTokens()) - - specialTokens - (prefix ? count(prefix) : 0), + // The shared chunker's metrics for the active model: counts estimated + // from characters at the sample's chars-per-token scale. The estimate + // errs both ways -- a misjudged small chunk merges or stands alone, an + // overshooting one is truncated by the pipeline. + function _getMetrics(sample) { + let charsPerToken = Zotero.Utilities.Internal.Chunking.getCharsPerToken(sample); + return { + count: text => text.length / charsPerToken, + joinSize: 2 / charsPerToken, + budget: _getBudget(), minSize: MIN_TOKENS, overlap: OVERLAP_TOKENS }; - _metrics.set(tokenizer, metrics); - return metrics; - } - - // A text in segments of at most TOKENIZER_SEGMENT_CHARS, each ending - // right before a whitespace character, so the next segment carries it and - // the tokenizer sees every word with its leading space. Text with no - // whitespace in a whole window (e.g., unsegmented CJK) is cut mid-run, - // which can miscount by a token per boundary -- noise against the budget. - function _splitForTokenizer(text) { - if (text.length <= TOKENIZER_SEGMENT_CHARS) { - return [text]; - } - let segments = []; - let start = 0; - while (start < text.length) { - let end = Math.min(start + TOKENIZER_SEGMENT_CHARS, text.length); - if (end < text.length) { - for (let i = end; i > start; i--) { - if (/\s/.test(text[i])) { - end = i; - break; - } - } - } - segments.push(text.slice(start, end)); - start = end; - } - return segments; } /** - * Split a text into chunks that each fit the active model's context - * window, measured in the model's own tokens (see + * Split a text into chunks that each fit the active model's window, + * sized in estimated tokens (see * Zotero.Utilities.Internal.Chunking.chunkText()). * * @param {String} text - * @return {Promise} - [{ text, tokens, start, end }] + * @return {Object[]} - [{ text, tokens, start, end }] */ - this.chunkText = async function (text) { - let metrics = _getMetrics(await this.getTokenizer()); + this.chunkText = function (text) { + let metrics = _getMetrics(text); return _asTokens(Zotero.Utilities.Internal.Chunking.chunkText(text, metrics)); }; /** - * Token count of a text against the active model, without the special - * tokens the tokenizer wraps every input in -- the same additive measure - * the chunkers report on their chunks. + * Estimated token count of a text -- the same character measure the + * chunkers size their chunks with, no model involved. * * @param {String} text - * @return {Promise} + * @return {Number} */ - this.countTokens = async function (text) { - let tokenizer = await this.getTokenizer(); - return _getCounter(tokenizer)(text); + this.estimateTokens = function (text) { + return Math.round( + text.length / Zotero.Utilities.Internal.Chunking.getCharsPerToken(text) + ); }; /** * Split a document's outline sections (see Zotero.SDT.getSections()) into - * chunks that each fit the active model's context window, measured in the - * model's own tokens, and prefixed for embedding with their section's - * outline path (see Zotero.Utilities.Internal.Chunking.chunkSections()). + * chunks that each fit the active model's window, sized in estimated + * tokens and prefixed for embedding with their section's outline path + * (see Zotero.Utilities.Internal.Chunking.chunkSections()). * * @param {Object[]} sections - * @return {Promise} - [{ text, embedText, tokens, outlinePath, + * @return {Object[]} - [{ text, embedText, tokens, outlinePath, * startBlock, endBlock, startOffset, endOffset, pageIndex, pageLabel, * position, sectionPart, sectionParts, auxiliary }], where tokens * counts embedText */ - this.chunkSections = async function (sections) { - let metrics = _getMetrics(await this.getTokenizer()); + this.chunkSections = function (sections) { + let sample = sections.map(section => section.text).join('\n\n'); + let metrics = _getMetrics(sample); return _asTokens(Zotero.Utilities.Internal.Chunking.chunkSections(sections, metrics)); }; // Chunks as this module reports them: the shared chunker sizes passages in - // whatever it was handed, and what it was handed here is tokens + // whatever it was handed, and what it was handed here is estimated tokens function _asTokens(chunks) { - return chunks.map(({ size, ...chunk }) => ({ ...chunk, tokens: size })); + return chunks.map(({ size, ...chunk }) => ({ ...chunk, tokens: Math.round(size) })); } }; @@ -2065,8 +1968,10 @@ Zotero.Embeddings.Indexing = new function () { // The title is the note's first line, so its words are among the // note's own -- a title with enough of them settles it without // stripping the body's HTML - if (_hasEmbeddableText(row.title) - || _hasEmbeddableText(_htmlToText(row.note, true))) { + if (_hasEmbeddableText(row.title)) { + add(row); + } + else if (_hasEmbeddableText(_htmlToText(row.note, true))) { add(row); } } @@ -2408,7 +2313,7 @@ Zotero.Embeddings.Indexing = new function () { // to rank if (sections.length && _hasEmbeddableText(sections.map(section => section.text).join(' '))) { - let chunks = await Zotero.Embeddings.Chunking.chunkSections(sections); + let chunks = Zotero.Embeddings.Chunking.chunkSections(sections); // An auxiliary chunk stands alone, so it's held to the same word // minimum as any standalone text -- a bare "Figure 1" gives the // model nothing to rank @@ -2431,7 +2336,7 @@ Zotero.Embeddings.Indexing = new function () { // Flat text has no sections, so the whole document plays that role: // the part numbering says where in it a chunk falls, and the chunk's // extent in the flat text is its source reference - let chunks = await Zotero.Embeddings.Chunking.chunkText(text); + let chunks = Zotero.Embeddings.Chunking.chunkText(text); return chunks.map((chunk, index) => ({ text: chunk.text, tokens: chunk.tokens, @@ -2675,14 +2580,14 @@ Zotero.Embeddings.Indexing = new function () { } else if (entry.item.isNote()) { // chunkText() already carries each chunk's token count - entry.chunks = await Zotero.Embeddings.Chunking.chunkText(entry.text); + entry.chunks = Zotero.Embeddings.Chunking.chunkText(entry.text); } else { // Item text is embedded whole, so it never passes through a // chunker and has to be counted here entry.chunks = [{ text: entry.text, - tokens: await Zotero.Embeddings.Chunking.countTokens(entry.text) + tokens: Zotero.Embeddings.Chunking.estimateTokens(entry.text) }]; } entry.vectors = new Array(entry.chunks.length); diff --git a/chrome/content/zotero/xpcom/utilities_internal.js b/chrome/content/zotero/xpcom/utilities_internal.js index 375730a961..c62f966864 100644 --- a/chrome/content/zotero/xpcom/utilities_internal.js +++ b/chrome/content/zotero/xpcom/utilities_internal.js @@ -3357,14 +3357,12 @@ Zotero.Utilities.Internal.onDragItems = function (event, itemIDs, dragImage = ev * Zotero.Utilities.Internal.Chunking -- splitting a text into passages. * * Passages are measured in characters unless the caller supplies its own - * metrics (see Zotero.Embeddings.Chunking, which measures in a model's - * tokens): { count(text), joinSize (charged per join, so pieces' sizes add up + * metrics (see Zotero.Embeddings.Chunking, which fits the budget to a model's + * window): { count(text), joinSize (charged per join, so pieces' sizes add up * to the joined size), budget (most a passage may reach), minSize (least * worth standing alone), overlap (carried into the next passage where one is * split mid-thought) }. The character default (see getCharacterMetrics()) - * approximates the same geometry, so it suits cuts made and discarded in one - * sitting; boundaries that get stored should be measured exactly, with the - * caller's own metrics. + * scales the same geometry by what a token is worth in the text's scripts. * * Every passage is a slice of its source, located by its start/end extent. */ @@ -3381,21 +3379,59 @@ Zotero.Utilities.Internal.Chunking = new function () { this.BUDGET_TOKENS = BUDGET_TOKENS; this.MIN_TOKENS = MIN_TOKENS; this.OVERLAP_TOKENS = OVERLAP_TOKENS; - // What a token is worth in characters in either kind of script - const CHARS_PER_TOKEN = 4; - const CJK_CHARS_PER_TOKEN = 1; + // What a token is worth in characters, by script class. Alphabetic + // scripts pack several characters into a token; CJK is written without + // spaces, so a character carries about what a word of an alphabetic + // script does; abjads and Indic scripts fragment in between. The + // unclassified rest (digits, punctuation, symbols) fragments more than + // alphabetic prose. + const SCRIPT_CHARS_PER_TOKEN = [ + [/\p{Script=Latin}/gu, 4], + [/[\p{Script=Cyrillic}\p{Script=Greek}]/gu, 3], + [/[\p{Script=Arabic}\p{Script=Hebrew}]/gu, 2.5], + [/[\p{Script=Devanagari}\p{Script=Bengali}\p{Script=Gurmukhi}\p{Script=Gujarati}\p{Script=Tamil}\p{Script=Telugu}\p{Script=Kannada}\p{Script=Malayalam}\p{Script=Sinhala}\p{Script=Thai}\p{Script=Lao}\p{Script=Khmer}\p{Script=Myanmar}]/gu, 1.5], + [/[\p{Script=Han}\p{Script=Hiragana}\p{Script=Katakana}\p{Script=Hangul}]/gu, 1] + ]; + const DEFAULT_CHARS_PER_TOKEN = 3; + // How much of a text its script mix is read from + const SCRIPT_SAMPLE_CHARS = 4000; + + /** + * What a token of the text is worth in characters: the share-weighted + * average of the script classes in a leading sample. An estimate -- real + * tokenizers vary by model and by text -- so a consumer sizing against a + * hard window should hold back some headroom. + * + * @param {String} text + * @return {Number} + */ + this.getCharsPerToken = function (text) { + let sample = text.slice(0, SCRIPT_SAMPLE_CHARS).replace(/\s/g, ''); + if (!sample) { + return DEFAULT_CHARS_PER_TOKEN; + } + let weighted = 0; + let classified = 0; + for (let [pattern, charsPerToken] of SCRIPT_CHARS_PER_TOKEN) { + let matched = sample.match(pattern)?.length || 0; + weighted += matched * charsPerToken; + classified += matched; + } + weighted += (sample.length - classified) * DEFAULT_CHARS_PER_TOKEN; + return weighted / sample.length; + }; /** * A measure counting characters, for a consumer with no tokenizer -- what * the chunkers here default to when no metrics are given. The budgets - * mirror the token geometry above, so passages come out about the size a - * model-driven chunker would make them. + * mirror the token geometry above at the text's chars-per-token scale, so + * passages come out about the size a model-driven chunker would make them. * - * @param {String} text - Read for the script it's written in + * @param {String} text - Read for the scripts it's written in * @return {Object} - Metrics */ this.getCharacterMetrics = function (text) { - let scale = _isMostlyCJK(text) ? CJK_CHARS_PER_TOKEN : CHARS_PER_TOKEN; + let scale = this.getCharsPerToken(text); return { count: value => value.length, // The two newlines paragraphs are joined with @@ -3412,17 +3448,6 @@ Zotero.Utilities.Internal.Chunking = new function () { return Zotero.Utilities.Internal.Chunking.getCharacterMetrics(text); } - // CJK is written without spaces, so a character of it carries about what - // a word of an alphabetic script does - function _isMostlyCJK(text) { - let sample = text.slice(0, 4000).replace(/\s/g, ''); - if (!sample) { - return false; - } - let cjk = sample.match(/[\p{Script=Han}\p{Script=Hiragana}\p{Script=Katakana}\p{Script=Hangul}]/gu); - return !!cjk && cjk.length * 2 > sample.length; - } - // The paragraphs of a text, each counted exactly once -- the only place // chunking pays to measure it -- and each carrying its trimmed extent function _measureParagraphs(text, count) { @@ -3500,11 +3525,17 @@ Zotero.Utilities.Internal.Chunking = new function () { return units; } + // Built once: constructing a segmenter is expensive next to running one, + // and a text with no paragraph breaks segments a block at a time + let _sentenceSegmenter = null; + // The sentences of a range as they stand, whatever their size function _segmentSentences(source, start, end, count) { let units = []; - let segmenter = new Intl.Segmenter(undefined, { granularity: 'sentence' }); - for (let { segment, index } of segmenter.segment(source.slice(start, end))) { + if (!_sentenceSegmenter) { + _sentenceSegmenter = new Intl.Segmenter(undefined, { granularity: 'sentence' }); + } + for (let { segment, index } of _sentenceSegmenter.segment(source.slice(start, end))) { let unit = _measureRange(source, start + index, start + index + segment.length, count); if (unit) { units.push(unit); diff --git a/test/tests/embeddingsTest.js b/test/tests/embeddingsTest.js index 15118b859d..0be8f9bf02 100644 --- a/test/tests/embeddingsTest.js +++ b/test/tests/embeddingsTest.js @@ -10,17 +10,6 @@ describe("Zotero.Embeddings", function () { testMean = await calibrateTestModel(); }); - // Stands in for a model's tokenizer, which the test environment has no - // downloaded model to provide: one token per whitespace-separated word, plus - // the two special tokens a real tokenizer wraps every input in, so counts - // here mean what they mean in production - function wordTokenizer() { - return { - encode: text => ['', ...text.split(/\s+/).filter(Boolean), ''], - decode: ids => ids.filter(id => id !== '' && id !== '').join(' ') - }; - } - // A Zotero.SDT.getSections() section, built from its blocks -- which are // what indexing reads, the section's own text and span being derived. // A block is its text, or an object adding flowClass/reference/location. @@ -475,14 +464,17 @@ describe("Zotero.Embeddings", function () { describe("#chunkText()", function () { // bge has no passage prefix, so the window less the two special tokens - // that wrap every input is what a chunk's own text gets (see MODELS). - // bge's window is under the chunking ceiling, so the window governs here. - const BUDGET = 512 - 2; - // The shared BUDGET_TOKENS geometry less those same special tokens - const CEILING = 768 - 2; - var fakeTokenizer = wordTokenizer(); - // A chunk's own tokens, the way chunking counts them - var contentTokens = text => fakeTokenizer.encode(text).length - 2; + // that wrap every input, less the headroom held back for the character + // estimate, is what a chunk's own text gets (see MODELS and + // _getBudget()). bge's window is under the chunking ceiling, so the + // window governs here. + const BUDGET = Math.floor((512 - 2) * 0.9); + // The shared BUDGET_TOKENS geometry under the same derivation + const CEILING = Math.floor((768 - 2) * 0.9); + // A chunk's own tokens, the way chunking estimates them: characters + // at the chars-per-token scale of the whole text being chunked + var estTokens = (text, whole) => text.length + / Zotero.Utilities.Internal.Chunking.getCharsPerToken(whole || text); // chunkText() returns { text, tokens, start, end }; most assertions // here are about the text var texts = chunks => chunks.map(chunk => chunk.text); @@ -498,60 +490,57 @@ describe("Zotero.Embeddings", function () { }); it("should return text that fits the window as a single chunk", async function () { - stubs.push(sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(fakeTokenizer)); - let single = await Zotero.Embeddings.Chunking.chunkText('A short title'); + let single = Zotero.Embeddings.Chunking.chunkText('A short title'); assert.deepEqual(texts(single), ['A short title']); - // Each chunk carries the count measured on the way - assert.equal(single[0].tokens, contentTokens('A short title')); - // Right up to the budget it's still one chunk, and one token past it - // splits -- the budget being the window less the special tokens that - // wrap every input and the model's passage prefix - let words = n => Array.from({ length: n }, (x, i) => `word${i}`).join(' '); - assert.lengthOf(await Zotero.Embeddings.Chunking.chunkText(words(BUDGET)), 1); - assert.isAbove((await Zotero.Embeddings.Chunking.chunkText(words(BUDGET + 1))).length, 1); + // Each chunk carries the count estimated on the way + assert.equal(single[0].tokens, Math.round(estTokens('A short title'))); + // Right up to the budget it's still one chunk, and a token past it + // splits -- pure letters, so the estimate is exactly four + // characters per token + let letters = tokens => 'a'.repeat(tokens * 4); + assert.lengthOf(Zotero.Embeddings.Chunking.chunkText(letters(BUDGET)), 1); + assert.isAbove((Zotero.Embeddings.Chunking.chunkText(letters(BUDGET + 1))).length, 1); }); it("should bound a chunk by the chunking ceiling, not the model's window", async function () { - stubs.push(sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(fakeTokenizer)); // A model that accepts far more at once doesn't get one vector per // note: that would average a note's subjects together, which is // what scoring an item by its best chunk exists to avoid stubs.push(sinon.stub(Zotero.Embeddings, 'getModelMaxTokens').returns(8192)); let words = (tag, n) => Array.from({ length: n }, (x, i) => `${tag}${i}`).join(' '); - // 2100 tokens: comfortably inside the window, past the ceiling - let chunks = await Zotero.Embeddings.Chunking.chunkText( + // ~1400 estimated tokens a paragraph: comfortably inside the + // window, past the ceiling + let chunks = Zotero.Embeddings.Chunking.chunkText( [words('alpha', 700), words('bravo', 700), words('charlie', 700)].join('\n\n') ); assert.isAbove(chunks.length, 1); for (let chunk of chunks) { - assert.isAtMost(contentTokens(chunk.text), CEILING); + assert.isAtMost(chunk.tokens, CEILING); } }); it("shouldn't put two substantial paragraphs in one chunk", async function () { - stubs.push(sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(fakeTokenizer)); // Two paragraphs on different subjects, each well under the // window but together over it. Packing them by size alone would // leave a chunk straddling both. - let a = Array.from({ length: 300 }, (x, i) => `alpha${i}`).join(' '); - let b = Array.from({ length: 300 }, (x, i) => `bravo${i}`).join(' '); - let chunks = texts(await Zotero.Embeddings.Chunking.chunkText(`${a}\n\n${b}`)); + let a = Array.from({ length: 150 }, (x, i) => `alpha${i}`).join(' '); + let b = Array.from({ length: 150 }, (x, i) => `bravo${i}`).join(' '); + let chunks = texts(Zotero.Embeddings.Chunking.chunkText(`${a}\n\n${b}`)); assert.lengthOf(chunks, 2); // Neither chunk mixes the two subjects assert.include(chunks[0], 'alpha0'); - assert.include(chunks[0], 'alpha299'); + assert.include(chunks[0], 'alpha149'); assert.notInclude(chunks[0], 'bravo'); assert.include(chunks[1], 'bravo0'); assert.notInclude(chunks[1], 'alpha'); }); it("should combine paragraphs too small to embed on their own", async function () { - stubs.push(sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(fakeTokenizer)); // A heading and a date, then a substantial paragraph, then a second // substantial paragraph -- the shape of an annotations note - let big1 = Array.from({ length: 300 }, (x, i) => `alpha${i}`).join(' '); - let big2 = Array.from({ length: 300 }, (x, i) => `bravo${i}`).join(' '); - let chunks = texts(await Zotero.Embeddings.Chunking.chunkText( + let big1 = Array.from({ length: 150 }, (x, i) => `alpha${i}`).join(' '); + let big2 = Array.from({ length: 150 }, (x, i) => `bravo${i}`).join(' '); + let chunks = texts(Zotero.Embeddings.Chunking.chunkText( `Annotations\n(11/12/2024)\n${big1}\n\n${big2}` )); assert.lengthOf(chunks, 2); @@ -567,31 +556,29 @@ describe("Zotero.Embeddings", function () { }); it("should split an oversized paragraph into even pieces at sentence boundaries", async function () { - stubs.push(sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(fakeTokenizer)); - // One paragraph of 60 ten-token sentences -- 600 tokens, over the - // window, with no paragraph breaks to split at + // One paragraph of 60 sentences -- ~780 estimated tokens, over + // the window, with no paragraph breaks to split at let sentences = Array.from({ length: 60 }, (x, i) => `Sentence ${i} has some words about subject number ${i}.`); - let chunks = texts(await Zotero.Embeddings.Chunking.chunkText(sentences.join(' '))); + let chunks = Zotero.Embeddings.Chunking.chunkText(sentences.join(' ')); assert.lengthOf(chunks, 2); - let sizes = chunks.map(contentTokens); - for (let size of sizes) { - assert.isAtMost(size, BUDGET); + for (let chunk of chunks) { + assert.isAtMost(chunk.tokens, BUDGET); // Filling the first piece to the budget would leave a short - // remainder; even pieces are ~300 plus the overlap - assert.isAbove(size, 250); + // remainder; even pieces are ~390 plus the overlap + assert.isAbove(chunk.tokens, 250); } // No sentence was dropped - let joined = chunks.join('\n'); + let joined = texts(chunks).join('\n'); for (let sentence of sentences) { assert.include(joined, sentence); } // Adjacent pieces of one paragraph still overlap - assert.include(chunks[1], sentences[29]); + assert.isTrue(sentences.some(sentence => chunks[0].text.includes(sentence) + && chunks[1].text.includes(sentence))); }); it("shouldn't leave an undersized piece at the end of a split", async function () { - stubs.push(sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(fakeTokenizer)); // A block only ever closes on a sentence boundary, so each piece // lands a little under its target. Without spreading that slack over // the pieces still to come, it accumulates into an extra runt piece @@ -599,9 +586,10 @@ describe("Zotero.Embeddings", function () { for (let count of [45, 64, 83, 97, 140]) { let sentences = Array.from({ length: count }, (x, i) => `Sentence ${i} has a few more words in it about subject ${i}.`); - let chunks = texts(await Zotero.Embeddings.Chunking.chunkText(sentences.join(' '))); - let sizes = chunks.map(contentTokens); - let total = contentTokens(sentences.join(' ')); + let text = sentences.join(' '); + let chunks = Zotero.Embeddings.Chunking.chunkText(text); + let sizes = chunks.map(chunk => chunk.tokens); + let total = estTokens(text); // No more pieces than the window requires assert.equal(chunks.length, Math.ceil(total / (BUDGET - 48)), `piece count for ${count} sentences (sizes ${sizes.join(', ')})`); @@ -615,12 +603,10 @@ describe("Zotero.Embeddings", function () { }); describe("#chunkSections()", function () { - var fakeTokenizer = wordTokenizer(); var stubs = []; beforeEach(function () { stubs.push(sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5')); - stubs.push(sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(fakeTokenizer)); }); afterEach(function () { @@ -636,9 +622,9 @@ describe("Zotero.Embeddings", function () { (y, j) => `${tag}${i * per + j}`).join(' ')); it("shouldn't put two substantial sections in one chunk", async function () { - let chunks = await Zotero.Embeddings.Chunking.chunkSections([ - sdtSection('Introduction', 0, wordBlocks('alpha', 5, 40)), - sdtSection('Methods', 5, wordBlocks('bravo', 5, 40)) + let chunks = Zotero.Embeddings.Chunking.chunkSections([ + sdtSection('Introduction', 0, wordBlocks('alpha', 5, 30)), + sdtSection('Methods', 5, wordBlocks('bravo', 5, 30)) ]); assert.lengthOf(chunks, 2); // Neither chunk mixes the two sections, and each points back at @@ -648,7 +634,7 @@ describe("Zotero.Embeddings", function () { assert.equal(chunks[0].startBlock, 0); assert.equal(chunks[0].endBlock, 4); assert.equal(chunks[0].startOffset, 0); - assert.equal(chunks[0].endOffset, wordBlocks('alpha', 5, 40)[4].length); + assert.equal(chunks[0].endOffset, wordBlocks('alpha', 5, 30)[4].length); assert.include(chunks[1].text, 'bravo0'); assert.notInclude(chunks[1].text, 'alpha0'); assert.equal(chunks[1].startBlock, 5); @@ -656,8 +642,8 @@ describe("Zotero.Embeddings", function () { }); it("should prefix the embedded text with the section's outline path", async function () { - let chunks = await Zotero.Embeddings.Chunking.chunkSections([ - sdtSection('Results > Field studies', 2, [words('alpha', 200)]) + let chunks = Zotero.Embeddings.Chunking.chunkSections([ + sdtSection('Results > Field studies', 2, [words('alpha', 150)]) ]); assert.lengthOf(chunks, 1); // What gets embedded carries the heading context; the display @@ -670,11 +656,11 @@ describe("Zotero.Embeddings", function () { it("should combine sections too small to embed on their own", async function () { // Front matter before the first heading rides along with the // section that follows it, the way small paragraphs do in a note - let chunks = await Zotero.Embeddings.Chunking.chunkSections([ + let chunks = Zotero.Embeddings.Chunking.chunkSections([ sdtSection('', 0, ['Title page']), sdtSection('', 1, ['Copyright notice']), - sdtSection('Introduction', 2, wordBlocks('alpha', 8, 25)), - sdtSection('Methods', 10, wordBlocks('bravo', 10, 20)) + sdtSection('Introduction', 2, wordBlocks('alpha', 8, 18)), + sdtSection('Methods', 10, wordBlocks('bravo', 10, 15)) ]); assert.lengthOf(chunks, 2); assert.include(chunks[0].text, 'Title page'); @@ -688,8 +674,8 @@ describe("Zotero.Embeddings", function () { }); it("should join a trailing small section to the previous chunk", async function () { - let chunks = await Zotero.Embeddings.Chunking.chunkSections([ - sdtSection('Body', 0, wordBlocks('alpha', 10, 20)), + let chunks = Zotero.Embeddings.Chunking.chunkSections([ + sdtSection('Body', 0, wordBlocks('alpha', 10, 15)), sdtSection('Appendix', 10, ['Short appendix note.', 'A closing line.']) ]); assert.lengthOf(chunks, 1); @@ -706,7 +692,7 @@ describe("Zotero.Embeddings", function () { (x, i) => `Sentence ${i} has some words about subject number ${i}.`); let block = sentences.join(' '); let position = { pageIndex: 4, rects: [[10, 20, 300, 40]] }; - let chunks = await Zotero.Embeddings.Chunking.chunkSections([ + let chunks = Zotero.Embeddings.Chunking.chunkSections([ sdtSection('Discussion', 3, [{ text: block, pageIndex: 4, @@ -751,7 +737,7 @@ describe("Zotero.Embeddings", function () { pageLabel: String(i + 1), position: { pageIndex: i, rects: [[10, 20, 300, 40]] } })); - let chunks = await Zotero.Embeddings.Chunking.chunkSections([ + let chunks = Zotero.Embeddings.Chunking.chunkSections([ sdtSection('Discussion', 0, blocks) ]); assert.isAbove(chunks.length, 1); @@ -770,7 +756,7 @@ describe("Zotero.Embeddings", function () { }); it("should keep auxiliary sections as standalone chunks", async function () { - let chunks = await Zotero.Embeddings.Chunking.chunkSections([ + let chunks = Zotero.Embeddings.Chunking.chunkSections([ // A small body section, a tiny caption, then a substantial // body section sdtSection('Results', 0, ['A short opening paragraph.']), @@ -779,7 +765,7 @@ describe("Zotero.Embeddings", function () { ['Figure 3: Owl migration routes across the Baltic.']), auxiliary: true }, - sdtSection('Results', 2, wordBlocks('alpha', 8, 25)) + sdtSection('Results', 2, wordBlocks('alpha', 8, 18)) ]); assert.lengthOf(chunks, 2); // The caption is a chunk of its own, however small... @@ -800,8 +786,8 @@ describe("Zotero.Embeddings", function () { }); it("shouldn't fold a trailing small body section into an auxiliary chunk", async function () { - let chunks = await Zotero.Embeddings.Chunking.chunkSections([ - sdtSection('Body', 0, wordBlocks('alpha', 10, 20)), + let chunks = Zotero.Embeddings.Chunking.chunkSections([ + sdtSection('Body', 0, wordBlocks('alpha', 10, 15)), { ...sdtSection('Body', 10, ['Figure 1: A caption with enough words to keep.']), @@ -819,8 +805,8 @@ describe("Zotero.Embeddings", function () { }); it("should mark an unsplit section as its only piece", async function () { - let chunks = await Zotero.Embeddings.Chunking.chunkSections([ - sdtSection('Body', 0, [words('alpha', 200)]) + let chunks = Zotero.Embeddings.Chunking.chunkSections([ + sdtSection('Body', 0, [words('alpha', 150)]) ]); assert.lengthOf(chunks, 1); assert.equal(chunks[0].sectionPart, 1); @@ -1197,7 +1183,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()) ]; try { await Zotero.Embeddings.Indexing.startIndexing(); @@ -1242,7 +1227,6 @@ describe("Zotero.Embeddings", function () { // would kick off a model switch), so name one to keep the // window and passage prefix chunking reads consistent with it sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()) ]; try { await Zotero.Embeddings.Indexing.startIndexing(); @@ -1297,7 +1281,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()) ]; try { await Zotero.Embeddings.Indexing.startIndexing(); @@ -1353,7 +1336,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()) ]; try { await Zotero.Embeddings.Indexing.startIndexing(); @@ -1406,7 +1388,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()) ]; try { await Zotero.Embeddings.Indexing.startIndexing(); @@ -1449,7 +1430,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()), // The extraction itself is sdt.js's concern (see sdtTest.js); // what's under test is what indexing does with the sections sinon.stub(Zotero.SDT, 'ensure').resolves(true), @@ -1569,7 +1549,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()), sinon.stub(Zotero.SDT, 'ensure').resolves(true), sinon.stub(Zotero.SDT, 'getSections').resolves({ ok: true, @@ -1620,7 +1599,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()), sinon.stub(Zotero.SDT, 'ensure').resolves(true), sinon.stub(Zotero.SDT, 'getSections').resolves({ ok: true, @@ -1693,7 +1671,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()), sinon.stub(Zotero.SDT, 'ensure').resolves(true), sinon.stub(Zotero.SDT, 'getSections').resolves({ ok: true, @@ -1766,7 +1743,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()), sinon.stub(Zotero.SDT, 'ensure').resolves(true), sinon.stub(Zotero.SDT, 'getSections').callsFake(async (itemID) => { if (itemID === big.id || itemID === small.id) { @@ -1825,7 +1801,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()), ensureStub, getSectionsStub ]; @@ -1877,7 +1852,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()), ensureStub, sinon.stub(Zotero.SDT, 'getSections').resolves({ ok: true, @@ -1920,7 +1894,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()), sinon.stub(Zotero.SDT, 'ensure').resolves(true), sinon.stub(Zotero.SDT, 'getSections').resolves({ ok: false, reason: 'failed' }) ]; @@ -1988,7 +1961,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer').resolves(wordTokenizer()), sinon.stub(Zotero.Embeddings, 'embedQuery').resolves(Float32Array.from(testMean)), sinon.stub(Zotero.SDT, 'ensure').resolves(true), getSectionsStub @@ -2041,8 +2013,6 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'isDownloaded').resolves(true), sinon.stub(Zotero.Embeddings, 'download').resolves(), sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), - sinon.stub(Zotero.Embeddings.Chunking, 'getTokenizer') - .resolves(wordTokenizer()) ]; let queries = []; let queryStub = sinon.stub(Zotero.DB, 'queryAsync') @@ -2111,29 +2081,43 @@ describe("Zotero.Embeddings", function () { assert.isTrue(await Zotero.Embeddings.isDownloaded()); }); - it("should chunk with the model's own tokenizer", async function () { + it("should chunk within the model's real window", async function () { this.timeout(1800000); await loadZoteroPane(); Zotero.Prefs.set('embeddings.model', 'bge-small-en-v1.5'); await Zotero.Embeddings.download(); - let tokenizer = await Zotero.Embeddings.Chunking.getTokenizer(); - assert.ok(tokenizer); + // The model's real tokenizer, built from the same files the + // inference process reads, with the transformers.js + // implementation Firefox ships + let { PreTrainedTokenizer } = ChromeUtils.importESModule( + 'chrome://global/content/ml/transformers.js' + ); + let decoder = new TextDecoder(); + let [tokenizerJSON, tokenizerConfig] = await Promise.all( + ['tokenizer.json', 'tokenizer_config.json'].map( + async file => JSON.parse(decoder.decode( + await Zotero.Embeddings.getModelFile(file) + )) + ) + ); + let tokenizer = new PreTrainedTokenizer(tokenizerJSON, tokenizerConfig); assert.isAbove(tokenizer.encode('a passage about owls').length, 3); - // A long text splits into chunks that each fit the real window + // A long text splits into chunks that each fit the real window, + // even though their sizes are estimated from characters let sentences = []; for (let i = 0; i < 100; i++) { sentences.push(`Sentence number ${i} concerns the ecology of temperate wetlands.`); } - let chunks = await Zotero.Embeddings.Chunking.chunkText(sentences.join(' ')); + let chunks = Zotero.Embeddings.Chunking.chunkText(sentences.join(' ')); assert.isAbove(chunks.length, 1); for (let chunk of chunks) { assert.isAtMost(tokenizer.encode(chunk.text).length, 512); - // The carried count is the model's own, less the special - // tokens encode() wraps every input in - assert.equal(chunk.tokens, - tokenizer.encode(chunk.text).length - tokenizer.encode('').length); + // The carried count is an estimate of the model's own count + let real = tokenizer.encode(chunk.text).length - tokenizer.encode('').length; + assert.isAbove(chunk.tokens, real / 2); + assert.isBelow(chunk.tokens, real * 2); } }); }); diff --git a/test/tests/utilities_internalTest.js b/test/tests/utilities_internalTest.js index 488f10376c..a25639a507 100644 --- a/test/tests/utilities_internalTest.js +++ b/test/tests/utilities_internalTest.js @@ -780,6 +780,18 @@ describe("Zotero.Utilities.Internal", function () { } }); + it("should scale a token's worth by the script of the text", function () { + let chunking = Zotero.Utilities.Internal.Chunking; + // Pure scripts land on their table entries... + assert.equal(chunking.getCharsPerToken('climate change'), 4); + assert.equal(chunking.getCharsPerToken('\u6c17\u5019\u5909\u52d5'.repeat(10)), 1); + assert.equal(chunking.getCharsPerToken('\u0438\u0437\u043c\u0435\u043d\u0435\u043d\u0438\u0435 \u043a\u043b\u0438\u043c\u0430\u0442\u0430'), 3); + // ...and a mixture lands between its parts + let mixed = chunking.getCharsPerToken('climate change \u6c17\u5019\u5909\u52d5\u6c17\u5019\u5909\u52d5'); + assert.isAbove(mixed, 1); + assert.isBelow(mixed, 4); + }); + it("should leave a text within the budget whole", function () { let chunks = chunk('A short paragraph.'); assert.lengthOf(chunks, 1);