Added save directory for embedder
This commit is contained in:
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "@ztimson/ai-utils",
|
"name": "@ztimson/ai-utils",
|
||||||
"version": "0.5.3",
|
"version": "0.5.4",
|
||||||
"description": "AI Utility library",
|
"description": "AI Utility library",
|
||||||
"author": "Zak Timson",
|
"author": "Zak Timson",
|
||||||
"license": "MIT",
|
"license": "MIT",
|
||||||
|
|||||||
@@ -3,8 +3,11 @@ import { parentPort } from 'worker_threads';
|
|||||||
|
|
||||||
let embedder: any;
|
let embedder: any;
|
||||||
|
|
||||||
parentPort?.on('message', async ({ id, text, model }) => {
|
parentPort?.on('message', async ({ id, text, model, path }) => {
|
||||||
if(!embedder) embedder = await pipeline('feature-extraction', 'Xenova/' + model, {quantized: true});
|
if(!embedder) embedder = await pipeline('feature-extraction', 'Xenova/' + model, {
|
||||||
|
quantized: true,
|
||||||
|
cache_dir: path,
|
||||||
|
});
|
||||||
const output = await embedder(text, { pooling: 'mean', normalize: true });
|
const output = await embedder(text, { pooling: 'mean', normalize: true });
|
||||||
const embedding = Array.from(output.data);
|
const embedding = Array.from(output.data);
|
||||||
parentPort?.postMessage({ id, embedding });
|
parentPort?.postMessage({ id, embedding });
|
||||||
|
|||||||
@@ -271,7 +271,12 @@ class LLM {
|
|||||||
return new Promise((resolve, reject) => {
|
return new Promise((resolve, reject) => {
|
||||||
const id = this.embedId++;
|
const id = this.embedId++;
|
||||||
this.embedQueue.set(id, { resolve, reject });
|
this.embedQueue.set(id, { resolve, reject });
|
||||||
this.embedWorker?.postMessage({ id, text, model: this.ai.options?.embedder || 'bge-small-en-v1.5' });
|
this.embedWorker?.postMessage({
|
||||||
|
id,
|
||||||
|
text,
|
||||||
|
model: this.ai.options?.embedder || 'bge-small-en-v1.5',
|
||||||
|
path: this.ai.options.path
|
||||||
|
});
|
||||||
});
|
});
|
||||||
};
|
};
|
||||||
const chunks = this.chunk(target, maxTokens, overlapTokens);
|
const chunks = this.chunk(target, maxTokens, overlapTokens);
|
||||||
|
|||||||
Reference in New Issue
Block a user