From f2c66b0cb8056264b65881c195101bf647735035 Mon Sep 17 00:00:00 2001
From: ztimson <zaktimson@gmail.com>
Date: Wed, 11 Feb 2026 20:23:50 -0500
Subject: [PATCH] Updated default embedder

---
 package.json    | 2 +-
 src/ai.ts       | 2 ++
 src/embedder.ts | 8 ++++----
 src/llm.ts      | 2 +-
 4 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/package.json b/package.json
index 984986e..04a2e2e 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
 	"name": "@ztimson/ai-utils",
-	"version": "0.5.1",
+	"version": "0.5.2",
 	"description": "AI Utility library",
 	"author": "Zak Timson",
 	"license": "MIT",
diff --git a/src/ai.ts b/src/ai.ts
index 8aeba6c..c681dad 100644
--- a/src/ai.ts
+++ b/src/ai.ts
@@ -10,6 +10,8 @@ export type AbortablePromise<T> = Promise<T> & {
 export type AiOptions = {
 	/** Path to models */
 	path?: string;
+	/** Embedding model */
+	embedder?: string; // all-MiniLM-L6-v2, bge-small-en-v1.5, bge-large-en-v1.5
 	/** Large language models, first is default */
 	llm?: Omit<LLMRequest, 'model'> & {
 		models: {[model: string]: AnthropicConfig | OllamaConfig | OpenAiConfig};
diff --git a/src/embedder.ts b/src/embedder.ts
index 9b2c852..16dae3d 100644
--- a/src/embedder.ts
+++ b/src/embedder.ts
@@ -1,11 +1,11 @@
 import { pipeline } from '@xenova/transformers';
 import { parentPort } from 'worker_threads';
 
-let model: any;
+let embedder: any;
 
-parentPort?.on('message', async ({ id, text }) => {
-	if(!model) model = await pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2');
-	const output = await model(text, { pooling: 'mean', normalize: true });
+parentPort?.on('message', async ({ id, text, model }) => {
+	if(!embedder) embedder = await pipeline('feature-extraction', 'Xenova/' + model);
+	const output = await embedder(text, { pooling: 'mean', normalize: true });
 	const embedding = Array.from(output.data);
 	parentPort?.postMessage({ id, embedding });
 });
diff --git a/src/llm.ts b/src/llm.ts
index 763d721..5db4d7e 100644
--- a/src/llm.ts
+++ b/src/llm.ts
@@ -271,7 +271,7 @@ class LLM {
 			return new Promise((resolve, reject) => {
 				const id = this.embedId++;
 				this.embedQueue.set(id, { resolve, reject });
-				this.embedWorker?.postMessage({ id, text });
+				this.embedWorker?.postMessage({ id, text, model: this.ai.options?.embedder || 'bge-small-en-v1.5' });
 			});
 		};
 		const chunks = this.chunk(target, maxTokens, overlapTokens);