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