From da15d299e633c3de0a36cae9a094fe4980345094 Mon Sep 17 00:00:00 2001 From: ztimson Date: Thu, 19 Feb 2026 21:37:58 -0500 Subject: [PATCH] parallel embedding cap --- package.json | 2 +- src/llm.ts | 27 +++++++++++++++++---------- 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/package.json b/package.json index 855fa25..da52d32 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@ztimson/ai-utils", - "version": "0.7.1", + "version": "0.7.2", "description": "AI Utility library", "author": "Zak Timson", "license": "MIT", diff --git a/src/llm.ts b/src/llm.ts index 139e565..a951b3d 100644 --- a/src/llm.ts +++ b/src/llm.ts @@ -255,11 +255,12 @@ class LLM { /** * Create a vector representation of a string * @param {object | string} target Item that will be embedded (objects get converted) - * @param {number} maxTokens Chunking size. More = better context, less = more specific (Search by paragraphs or lines) - * @param {number} overlapTokens Includes previous X tokens to provide continuity to AI (In addition to max tokens) + * @param {maxTokens?: number, overlapTokens?: number, parellel?: number} opts Options for embedding such as chunk sizes and parallel processing * @returns {Promise[]>} Chunked embeddings */ - embedding(target: object | string, maxTokens = 500, overlapTokens = 50) { + async embedding(target: object | string, opts: {maxTokens?: number, overlapTokens?: number, parallel?: number} = {}) { + let {maxTokens = 500, overlapTokens = 50, parallel = 1} = opts; + const embed = (text: string): Promise => { return new Promise((resolve, reject) => { const worker = new Worker(join(dirname(fileURLToPath(import.meta.url)), 'embedder.js')); @@ -279,13 +280,19 @@ class LLM { worker.postMessage({text, model: this.ai.options?.embedder || 'bge-small-en-v1.5', modelDir: this.ai.options.path}); }); }; - const chunks = this.chunk(target, maxTokens, overlapTokens); - return Promise.all(chunks.map(async (text, index) => ({ - index, - embedding: await embed(text), - text, - tokens: this.estimateTokens(text), - }))); + let i = 0, chunks = this.chunk(target, maxTokens, overlapTokens), results: any[] = []; + const next: Function = () => { + const index = i++; + if(index >= chunks.length) return; + const text = chunks[index]; + return embed(text).then(embedding => { + results.push({index, embedding, text, tokens: this.estimateTokens(text)}); + return next(); + }) + } + + await Promise.all(Array(parallel).fill(null).map(() => next())); + return results.toSorted((a, b) => a.index - b.index); } /**