From d85046ca38d8358485a947aae8c7584b2c80f3d4 Mon Sep 17 00:00:00 2001 From: Bogdan Abaev Date: Mon, 7 Sep 2026 16:49:41 -0700 Subject: [PATCH] more effective chunking of a large section Instead of breaking up section into pieces of at least X size (high number of smallersections), try to split section into chunks of similar sizes, meaning there is less chunks but they are larger. Overall, it is better for retrival quality and, importantly, for the amount of storage used --- chrome/content/zotero/xpcom/embeddings.js | 40 ++-- .../zotero/xpcom/utilities_internal.js | 172 +++++++++++++----- test/tests/embeddingsTest.js | 102 +++++++++++ 3 files changed, 249 insertions(+), 65 deletions(-) diff --git a/chrome/content/zotero/xpcom/embeddings.js b/chrome/content/zotero/xpcom/embeddings.js index 6d0c1c0991..7a02d33b97 100644 --- a/chrome/content/zotero/xpcom/embeddings.js +++ b/chrome/content/zotero/xpcom/embeddings.js @@ -1687,7 +1687,7 @@ Zotero.Embeddings.Indexing = new function () { const QUEUE_SLICE_SIZE = 32; // Bump when chunking changes, so stored attachment rows are rebuilt (see // _getAttachmentSourceHash()) - const CHUNKER_VERSION = 1; + const CHUNKER_VERSION = 2; // The inference process's memory arena only grows: fragmentation from // varying batch shapes accumulates and is never returned to the OS @@ -2342,19 +2342,22 @@ Zotero.Embeddings.Indexing = new function () { return indexable; } - // The embeddable chunks of an attachment's full text: its outline - // sections (extracted and cached by Zotero.SDT), split to fit the model - // window. When structured extraction yields nothing, the flat text falls - // back to note-style paragraph chunking -- searchable, just without - // section locations. Null when there's no embeddable text at all. - // - // Only the chunks' source references are stored -- block ranges and - // offsets for section chunks, flat-text offsets for fallback ones -- and - // the preview text is re-derived from them later (see - // getMatchingChunks()). That's also why a stale-processor pack won't do - // here: a background regeneration would shift the blocks out from under - // the stored references as soon as they were written. - async function _getAttachmentChunks(item) { + /** + * The embeddable chunks of an attachment's full text: its outline + * sections (extracted by Zotero.SDT), split to fit the model window. A + * document with no structured text falls back to paragraph chunking of + * its plain text, without section locations. + * + * Each chunk locates its text in the document by block and offset rather + * than copying it, so this waits for a current extraction instead of + * accepting one from an older processor, whose blocks it would point + * past. + * + * @param {Zotero.Item} item - A PDF, EPUB or snapshot attachment + * @return {Promise} - Null when there's no embeddable + * text at all + */ + this.getAttachmentChunks = async function (item) { let result = await Zotero.SDT.getSections(item.id, { allowStale: false }); let sections = result.ok ? _toIndexableSections(result.sections) : []; // The word minimum applies to the document as a whole, not each @@ -2391,7 +2394,7 @@ Zotero.Embeddings.Indexing = new function () { sectionPart: index + 1, sectionParts: chunks.length })); - } + }; // The stored source hash of each of the given items that has rows, read // in one query per chunk rather than one per item (every start @@ -2597,7 +2600,7 @@ Zotero.Embeddings.Indexing = new function () { await Zotero.SDT.ensure(item.id); // Counted while the pack is fresh, through the same derivation // the embedder uses so the two can't disagree - let chunks = await _getAttachmentChunks(item); + let chunks = await Zotero.Embeddings.Indexing.getAttachmentChunks(item); await _storeChunkCount(item.id, hash, chunks ? chunks.length : 0); progress.done++; await _tick(); @@ -2676,7 +2679,7 @@ Zotero.Embeddings.Indexing = new function () { // Derive each entry's chunks. Notes are split to fit the model's // context window; attachments are extracted (see - // _getAttachmentChunks()) and split section by section. A title and + // getAttachmentChunks()) and split section by section. A title and // abstract, or an annotation's passage and comment, fit the window in // almost all cases, so they're embedded as a single chunk and the // pipeline truncates the rare outlier. @@ -2691,7 +2694,8 @@ Zotero.Embeddings.Indexing = new function () { return 0; } if (entry.item.isAttachment()) { - entry.chunks = await _getAttachmentChunks(entry.item); + entry.chunks = await Zotero.Embeddings.Indexing + .getAttachmentChunks(entry.item); // Nothing embeddable anywhere in the attachment (missing // file, password-protected, no text layer). Recorded below // anyway, so the item counts as processed and isn't looked diff --git a/chrome/content/zotero/xpcom/utilities_internal.js b/chrome/content/zotero/xpcom/utilities_internal.js index c62f966864..fb81d25558 100644 --- a/chrome/content/zotero/xpcom/utilities_internal.js +++ b/chrome/content/zotero/xpcom/utilities_internal.js @@ -3641,9 +3641,9 @@ Zotero.Utilities.Internal.Chunking = new function () { * Split a text into passages that each fit the budget. Text that already * fits comes back as a single passage. * - * Paragraphs are the topic units, so two never share a passage unless one - * was too small to stand alone; a block over the budget is split into even - * pieces at sentence boundaries. + * Text over the budget is divided into as few and as even passages as + * possible, cut at paragraph boundaries; a paragraph over the budget on + * its own is split at sentence boundaries. * * @param {String} text * @param {Object} [metrics] - The character measure when omitted @@ -3671,31 +3671,8 @@ Zotero.Utilities.Internal.Chunking = new function () { return [_sliceUnits(text, paragraphs, totalSize)]; } - // Group paragraphs into blocks, combining any too small to stand alone - // with those that follow. A paragraph reaching the minimum on its own - // becomes its own block, keeping distinct subjects apart. - let groups = []; - let pending = []; - let pendingSize = 0; - for (let paragraph of paragraphs) { - pending.push(paragraph); - pendingSize += paragraph.size; - if (pendingSize >= metrics.minSize) { - groups.push(pending); - pending = []; - pendingSize = 0; - } - } - // A trailing group under the minimum joins the previous block rather - // than standing alone; the block is split evenly below anyway - if (pending.length) { - if (groups.length) { - groups[groups.length - 1].push(...pending); - } - else { - groups.push(pending); - } - } + let groups = _partitionEvenly(paragraphs, totalSize, budget, joinSize); + _absorbUndersized(groups, metrics.minSize, joinSize); let chunks = []; for (let group of groups) { @@ -3710,6 +3687,72 @@ Zotero.Utilities.Internal.Chunking = new function () { return chunks; } + // Paragraphs divided into as few and as even groups as possible: filling + // each to the budget instead would leave a short remainder at the end. + // Recomputing the target from what's left spreads the slack rather than + // accumulating it, the same way _splitBlockEvenly() spreads sentences. + function _partitionEvenly(paragraphs, totalSize, budget, joinSize) { + let groups = []; + let current = []; + let currentSize = 0; + let remainingSize = totalSize; + let remainingGroups = Math.ceil(totalSize / budget); + for (let paragraph of paragraphs) { + let withNext = currentSize + joinSize + paragraph.size; + // Close at whichever boundary lands nearer the target. On the last + // group only the budget closes it. + let closeHere = false; + if (current.length) { + if (withNext > budget) { + closeHere = true; + } + else if (remainingGroups > 1) { + let target = remainingSize / remainingGroups; + closeHere = Math.abs(withNext - target) > Math.abs(currentSize - target); + } + } + if (closeHere) { + groups.push(current); + remainingSize -= currentSize + joinSize; + remainingGroups = Math.max(1, remainingGroups - 1); + current = []; + currentSize = 0; + } + currentSize += (current.length ? joinSize : 0) + paragraph.size; + current.push(paragraph); + } + if (current.length) { + groups.push(current); + } + return groups; + } + + // A group under the minimum merges into its smaller neighbor: an item + // scores as its best chunk, so a fragment standing alone would inflate + // that score. The merged group can exceed the budget, which leaves it to + // the sentence splitter -- the only way to divide a paragraph that fills + // the budget by itself. + function _absorbUndersized(groups, minSize, joinSize) { + let i = 0; + while (groups.length > 1 && i < groups.length) { + if (_sumSizes(groups[i], joinSize) >= minSize) { + i++; + continue; + } + let before = i > 0 ? _sumSizes(groups[i - 1], joinSize) : Infinity; + let after = i < groups.length - 1 + ? _sumSizes(groups[i + 1], joinSize) + : Infinity; + if (before <= after) { + groups[i - 1].push(...groups[i]); + } + else { + groups[i + 1].unshift(...groups[i]); + } + groups.splice(i, 1); + } + } + /** * Split a document's outline sections (see Zotero.SDT.getSections()) into * chunks that each fit the budget. Sections play the role paragraphs play @@ -3720,9 +3763,10 @@ Zotero.Utilities.Internal.Chunking = new function () { * too-small merging in both directions, so each stands as its own chunk * rather than mixing into the running text. * - * `embedText` is the chunk's text prefixed with its section's outline path - * ("Methods > Participants"), giving a fragment the context of its - * headings at the cost of part of the budget; `text` stays the plain piece. + * `embedText` weaves each section's outline path ("Methods > Participants") + * into the chunk at the point that section's text begins, so a chunk + * spanning merged sections carries the heading of each rather than one + * heading for all of them; `text` stays the plain piece. * * Each chunk records where it lives: the block range it covers, with * offsets into the first and last block, so the text can be re-derived @@ -3804,29 +3848,21 @@ Zotero.Utilities.Internal.Chunking = new function () { let chunks = []; for (let group of groups) { - // A merged group takes its first section's outline path -- the - // heading its text starts under - let outlinePath = group.entries[0].section.outlinePath || ''; - let prefix = outlinePath ? outlinePath + '\n\n' : ''; - let prefixSize = prefix ? count(prefix) : 0; - // A pathological outline path that would eat a real share of the - // window hurts more than it helps - if (prefixSize > budget / 4) { - prefix = ''; - prefixSize = 0; - } // The group's source string: its sections' texts joined, with the // sections' paragraphs shifted to their place in it and every // block's extent recorded, so each chunk's slice can be mapped - // back to the blocks it covers + // back to the blocks it covers. Where each section starts is kept + // too, for the headings woven in below. let text = ''; let paragraphs = []; let blocks = []; + let sections = []; for (let entry of group.entries) { if (text) { text += '\n\n'; } let base = text.length; + sections.push({ start: base, path: entry.section.outlinePath || '' }); for (let paragraph of entry.paragraphs) { paragraphs.push({ ...paragraph, @@ -3850,7 +3886,29 @@ Zotero.Utilities.Internal.Chunking = new function () { } text += entry.section.text; } - let pieces = _chunkParagraphs(text, paragraphs, budget - prefixSize, metrics); + // A heading labels its own section only, so it runs to where the + // next one starts. One repeating the heading before it says + // nothing, and one long enough to eat a real share of the window + // hurts more than it helps. + let headings = sections + .map((section, j) => ({ + ...section, + end: sections[j + 1] ? sections[j + 1].start : text.length, + size: count(section.path + '\n\n') + })) + .filter((heading, j) => heading.path + && heading.path !== sections[j - 1]?.path + && count(heading.path) <= budget / 4); + // Every heading the group could contribute is held back from the + // budget, since which of them a piece carries isn't known until + // the pieces exist. Past a quarter of the window only the first + // is kept, so the reservation can't crowd out the text. + let headingSize = headings.reduce((sum, heading) => sum + heading.size, 0); + if (headingSize > budget / 4) { + headings = headings.slice(0, 1); + headingSize = headings.length ? headings[0].size : 0; + } + let pieces = _chunkParagraphs(text, paragraphs, budget - headingSize, metrics); for (let i = 0; i < pieces.length; i++) { let piece = pieces[i]; // The blocks the piece's extent overlaps. A piece boundary @@ -3861,11 +3919,31 @@ Zotero.Utilities.Internal.Chunking = new function () { block => block.end > piece.start && block.start < piece.end); let first = covered[0]; let last = covered[covered.length - 1]; + // Each heading the piece reaches, at the point its section + // starts -- at the top for the section the piece opens in + let embedText = ''; + let embedSize = piece.size; + let emitted = ''; + let cursor = piece.start; + for (let heading of headings) { + if (heading.end <= piece.start || heading.start >= piece.end + || heading.path === emitted) { + continue; + } + let at = Math.max(heading.start, piece.start); + embedText += text.slice(cursor, at) + heading.path + '\n\n'; + embedSize += heading.size; + emitted = heading.path; + cursor = at; + } + embedText += text.slice(cursor, piece.end); chunks.push({ text: piece.text, - embedText: prefix + piece.text, - size: piece.size + prefixSize, - outlinePath, + embedText, + size: embedSize, + outlinePath: headings.find( + heading => heading.end > piece.start && heading.start < piece.end + )?.path || '', startBlock: first ? first.index : null, endBlock: last ? last.index : null, startOffset: first ? piece.start - first.start : null, diff --git a/test/tests/embeddingsTest.js b/test/tests/embeddingsTest.js index bdd38e13c6..fefde60f79 100644 --- a/test/tests/embeddingsTest.js +++ b/test/tests/embeddingsTest.js @@ -478,6 +478,11 @@ describe("Zotero.Embeddings", function () { // chunkText() returns { text, tokens, start, end }; most assertions // here are about the text var texts = chunks => chunks.map(chunk => chunk.text); + // A paragraph of `count` sentences, each about 12 estimated tokens. + // Sentences start with a capital: segmentation doesn't break on a + // period followed by lowercase. + var sentences = (tag, count) => Array.from({ length: count }, + (x, i) => `${tag} sentence number ${i} with several words in it.`).join(' '); var stubs = []; beforeEach(function () { @@ -535,6 +540,37 @@ describe("Zotero.Embeddings", function () { assert.notInclude(chunks[1], 'alpha'); }); + it("should divide oversized text into even pieces at paragraph boundaries", async function () { + // Four paragraphs of about 150 tokens, 600 in all: over the budget, + // but only just, so the even division is two pieces of two + // paragraphs. Filling each piece to the budget instead would put + // three in the first and leave one on its own. + let chunks = Zotero.Embeddings.Chunking.chunkText( + ['Alpha', 'Bravo', 'Charlie', 'Delta'].map(tag => sentences(tag, 12)).join('\n\n')); + assert.lengthOf(chunks, 2); + // Paragraphs stay whole, and the two pieces come out even + assert.include(chunks[0].text, 'Alpha sentence'); + assert.include(chunks[0].text, 'Bravo sentence'); + assert.notInclude(chunks[0].text, 'Charlie'); + assert.include(chunks[1].text, 'Charlie sentence'); + assert.include(chunks[1].text, 'Delta sentence'); + assert.closeTo(chunks[0].tokens, chunks[1].tokens, chunks[0].tokens * 0.2); + }); + + it("should absorb a tail too small to stand alone rather than leaving it a chunk", async function () { + // A paragraph filling most of the budget and a short one after it: + // keeping the paragraph boundary would leave the tail as a runt, so + // the two are divided at sentence boundaries instead + let chunks = Zotero.Embeddings.Chunking.chunkText( + sentences('Alpha', 32) + '\n\n' + sentences('Bravo', 8)); + assert.lengthOf(chunks, 2); + for (let chunk of chunks) { + assert.isAtLeast(chunk.tokens, Zotero.Utilities.Internal.Chunking.MIN_TOKENS); + assert.isAtMost(chunk.tokens, BUDGET); + } + assert.closeTo(chunks[0].tokens, chunks[1].tokens, chunks[0].tokens * 0.35); + }); + it("should combine paragraphs too small to embed on their own", async function () { // A heading and a date, then a substantial paragraph, then a second // substantial paragraph -- the shape of an annotations note @@ -653,6 +689,43 @@ describe("Zotero.Embeddings", function () { assert.equal(chunks[0].outlinePath, 'Results > Field studies'); }); + it("should carry each merged section's own heading into the embedded text", async function () { + // A stub too small to stand alone merges into the section after + // it, so one chunk holds text from both. Labelling the whole + // chunk with the stub's heading would describe almost none of it. + let chunks = Zotero.Embeddings.Chunking.chunkSections([ + sdtSection('Funding', 0, [words('alpha', 12)]), + sdtSection('Methods', 1, wordBlocks('bravo', 4, 30)) + ]); + assert.lengthOf(chunks, 1); + let embedded = chunks[0].embedText; + // Each heading sits with the text it belongs to + assert.isBelow(embedded.indexOf('Funding'), embedded.indexOf('alpha0')); + assert.isBelow(embedded.indexOf('alpha0'), embedded.indexOf('Methods')); + assert.isBelow(embedded.indexOf('Methods'), embedded.indexOf('bravo0')); + // The plain text stays free of them, since block offsets index it + assert.notInclude(chunks[0].text, 'Funding'); + assert.notInclude(chunks[0].text, 'Methods'); + // Both headings are counted against the window + assert.isAbove(chunks[0].tokens, + Zotero.Embeddings.Chunking.estimateTokens(chunks[0].text)); + }); + + it("should label a chunk with the section it starts in", async function () { + // Two small sections merge, then a third large one splits: the + // pieces after the first belong to the section they sit in + let chunks = Zotero.Embeddings.Chunking.chunkSections([ + sdtSection('Preface', 0, [words('alpha', 12)]), + sdtSection('Discussion', 1, wordBlocks('bravo', 8, 60)) + ]); + assert.isAbove(chunks.length, 1); + assert.equal(chunks[0].outlinePath, 'Preface'); + for (let chunk of chunks.slice(1)) { + assert.equal(chunk.outlinePath, 'Discussion'); + assert.notInclude(chunk.embedText, 'Preface'); + } + }); + 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 @@ -1987,6 +2060,35 @@ describe("Zotero.Embeddings", function () { } }); + it("should chunk an attachment that has nothing stored yet", async function () { + this.timeout(60000); + let item = await createDataObject('item', { title: 'Parent of derived attachment' }); + let attachment = await importPDFAttachment(item); + let stubs = [ + sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), + sinon.stub(Zotero.SDT, 'getSections').resolves({ + ok: true, + sections: [ + sdtSection('Results', 0, ['Owls hunt at night. '.repeat(200)]), + sdtSection('Discussion', 1, ['Hawks hunt by day. '.repeat(200)]) + ] + }) + ]; + try { + assert.isEmpty(await Zotero.Embeddings.getChunks(attachment.id)); + let chunks = await Zotero.Embeddings.Indexing.getAttachmentChunks(attachment); + assert.isAbove(chunks.length, 1); + for (let chunk of chunks) { + assert.isAbove(chunk.tokens, 0); + assert.isNotEmpty(chunk.text); + } + assert.include(chunks.map(chunk => chunk.outlinePath), 'Results'); + } + finally { + stubs.forEach(stub => stub.restore()); + } + }); + it("should index recently touched attachments first", async function () { this.timeout(60000); // One attachment per signal, each under its own parent. All are