mirror of
https://github.com/zotero/zotero.git
synced 2026-09-08 22:21:04 +00:00
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.
This commit is contained in:
Bogdan Abaev
parent
cb27b757d7
commit
fbb0495f05
4 changed files with 224 additions and 292 deletions
|
|
@ -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<Object>}
|
||||
*/
|
||||
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<Object[]>} - [{ 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<Number>}
|
||||
* @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<Object[]>} - [{ 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);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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 => ['<s>', ...text.split(/\s+/).filter(Boolean), '</s>'],
|
||||
decode: ids => ids.filter(id => id !== '<s>' && id !== '</s>').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);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue