diff --git a/src/fastembed.ts b/src/fastembed.ts index 6301e58..a530511 100644 --- a/src/fastembed.ts +++ b/src/fastembed.ts @@ -1,10 +1,10 @@ +import { AddedToken, Tokenizer } from "@anush008/tokenizers"; import fs, { PathLike } from "fs"; import https from "https"; +import * as ort from "onnxruntime-node"; import path from "path"; import Progress from "progress"; import tar from "tar"; -import { AddedToken, Tokenizer } from "@anush008/tokenizers"; -import * as ort from "onnxruntime-node"; export enum ExecutionProvider { CPU = "cpu", @@ -22,14 +22,14 @@ export enum EmbeddingModel { BGESmallENV15 = "fast-bge-small-en-v1.5", BGESmallZH = "fast-bge-small-zh-v1.5", MLE5Large = "fast-multilingual-e5-large", + CUSTOM = "custom", } -interface InitOptions { - model: EmbeddingModel; - executionProviders: ExecutionProvider[]; - maxLength: number; - cacheDir: string; - showDownloadProgress: boolean; +export interface InitOptionsBase { + executionProviders?: ExecutionProvider[]; + maxLength?: number; + cacheDir?: string; + showDownloadProgress?: boolean; } interface ModelInfo { @@ -86,7 +86,20 @@ function getEmbeddings( // return resultArray; // } +// Cas standard +export interface InitStandardOptions extends InitOptionsBase { + model: Exclude; + modelAbsoluteDirPath?: undefined; + modelName?: string; +} +// Cas custom +export interface InitCustomOptions extends InitOptionsBase { + model: EmbeddingModel.CUSTOM; + modelAbsoluteDirPath: fs.PathLike; + modelName: string; +} +export type InitOptions = InitStandardOptions | InitCustomOptions; abstract class Embedding { abstract listSupportedModels(): ModelInfo[]; @@ -111,28 +124,47 @@ export class FlagEmbedding extends Embedding { ) { super(); } - + static async init(options: InitStandardOptions): Promise; + static async init(options: InitCustomOptions): Promise; static async init({ model = EmbeddingModel.BGESmallENV15, executionProviders = [ExecutionProvider.CPU], maxLength = 512, cacheDir = "local_cache", showDownloadProgress = true, + modelAbsoluteDirPath = "", + modelName = "", }: Partial = {}) { - const modelDir = await FlagEmbedding.retrieveModel( - model, - cacheDir, - showDownloadProgress - ); + if (model === EmbeddingModel.CUSTOM) { + if (!modelAbsoluteDirPath) { + throw new Error( + "For custom model, modelAbsoluteDirPath is required in FlagEmbedding.init" + ); + } + if (!modelName) { + throw new Error( + "For custom model, modelName is required in FlagEmbedding.init" + ); + } + } + const modelDir = + model === EmbeddingModel.CUSTOM + ? modelAbsoluteDirPath + : await FlagEmbedding.retrieveModel( + model, + cacheDir, + showDownloadProgress + ); const tokenizer = this.loadTokenizer(modelDir, maxLength); - - const modelPath = path.join( - modelDir.toString(), + const defaultModelName = model === EmbeddingModel.MLE5Large || - model === EmbeddingModel.AllMiniLML6V2 + model === EmbeddingModel.AllMiniLML6V2 ? "model.onnx" - : "model_optimized.onnx", + : "model_optimized.onnx"; + const modelPath = path.join( + modelDir.toString(), + modelName || defaultModelName ); if (!fs.existsSync(modelPath)) { throw new Error(`Model file not found at ${modelPath}`); diff --git a/tests/fastembed_custom.test.ts b/tests/fastembed_custom.test.ts new file mode 100644 index 0000000..4381e3e --- /dev/null +++ b/tests/fastembed_custom.test.ts @@ -0,0 +1,169 @@ +import fs from "fs"; +import path from "path"; +import { beforeAll, describe, expect, test } from "vitest"; +import { EmbeddingModel, FlagEmbedding } from "../src"; +async function downloadFile(url: string, dest: string) { + const res = await fetch(url); + if (!res.ok) throw new Error(`Failed to fetch ${url}`); + const buffer = await res.arrayBuffer(); + fs.mkdirSync(path.dirname(dest), { recursive: true }); + fs.writeFileSync(dest, Buffer.from(buffer)); +} +describe("FastEmbed Custom Model Tests with Download files", () => { + beforeAll(async () => { + const files = [ + { + repoUrl: + "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/onnx/model_quint8_avx2.onnx", + outputPath: "../models/customs/mymodel.onnx", + }, + { + repoUrl: + "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/config.json", + outputPath: "../models/customs/config.json", + }, + { + repoUrl: + "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/special_tokens_map.json", + outputPath: "../models/customs/special_tokens_map.json", + }, + { + repoUrl: + "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/tokenizer.json", + outputPath: "../models/customs/tokenizer.json", + }, + { + repoUrl: + "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/tokenizer_config.json", + outputPath: "../models/customs/tokenizer_config.json", + }, + ]; + for (const element of files) { + await downloadFile( + element.repoUrl, + path.resolve(__dirname, element.outputPath) + ); + } + }); + + test("Init EmbeddingModel", async () => { + const pathModel = path.resolve(__dirname, "../models/customs"); + console.log("pathModel", pathModel); + const model = await FlagEmbedding.init({ + model: EmbeddingModel.CUSTOM, + modelAbsoluteDirPath: pathModel, + modelName: "mymodel.onnx", + }); + expect(model).toBeDefined(); + }); + + test("FlagEmbedding embed", async () => { + const pathModel = path.resolve(__dirname, "../models/customs"); + const flagEmbedding = await FlagEmbedding.init({ + model: EmbeddingModel.CUSTOM, + modelAbsoluteDirPath: pathModel, + modelName: "mymodel.onnx", + maxLength: 512, + }); + const embeddings = (await flagEmbedding.embed(["This is a test"]).next()) + .value!; + expect(embeddings).toBeDefined(); + expect(embeddings.length).toBe(1); + }); + + test("FlagEmbedding embed batch", async () => { + const pathModel = path.resolve(__dirname, "../models/customs"); + const flagEmbedding = await FlagEmbedding.init({ + model: EmbeddingModel.CUSTOM, + modelAbsoluteDirPath: pathModel, + modelName: "mymodel.onnx", + + maxLength: 512, + }); + const embeddingsBatch = flagEmbedding.embed([ + "This is a test", + "Some text", + "Some more test", + "This is a test", + "Some text", + "Some more test", + ]); + for await (const embeddings of embeddingsBatch) { + expect(embeddings).toBeDefined(); + expect(embeddings.length).toBe(6); + expect(embeddings[0].length).toBe(384); + } + }); + + test("FlagEmbedding embed small batch", async () => { + const pathModel = path.resolve(__dirname, "../models/customs"); + const flagEmbedding = await FlagEmbedding.init({ + model: EmbeddingModel.CUSTOM, + modelAbsoluteDirPath: pathModel, + modelName: "mymodel.onnx", + maxLength: 512, + }); + const embeddingsBatch = flagEmbedding.embed( + ["This is a test", "Some text"], + 1 + ); + for await (const embeddings of embeddingsBatch) { + expect(embeddings).toBeDefined(); + expect(embeddings.length).toBe(1); + expect(embeddings[0].length).toBe(384); + } + }); + + test("FlagEmbedding queryEmbed", async () => { + const pathModel = path.resolve(__dirname, "../models/customs"); + const flagEmbedding = await FlagEmbedding.init({ + model: EmbeddingModel.CUSTOM, + modelAbsoluteDirPath: pathModel, + modelName: "mymodel.onnx", + maxLength: 512, + }); + const embeddings = await flagEmbedding.queryEmbed("This is a test"); + expect(embeddings).toBeDefined(); + expect(embeddings.length).toBe(384); + }); + + test("FlagEmbedding passageEmbed", async () => { + const pathModel = path.resolve(__dirname, "../models/customs"); + const flagEmbedding = await FlagEmbedding.init({ + model: EmbeddingModel.CUSTOM, + modelAbsoluteDirPath: pathModel, + modelName: "mymodel.onnx", + + maxLength: 512, + }); + const embeddings = ( + await flagEmbedding.passageEmbed(["This is a test"]).next() + ).value!; + expect(embeddings).toBeDefined(); + expect(embeddings.length).toBe(1); + }); + + test("FlagEmbedding canonical values", async () => { + const pathModel = path.resolve(__dirname, "../models/customs"); + const flagEmbedding = await FlagEmbedding.init({ + model: EmbeddingModel.CUSTOM, + modelAbsoluteDirPath: pathModel, + modelName: "mymodel.onnx", + maxLength: 512, + }); + const expected = [ + 0.025276897475123405, 0.013033483177423477, 0.005586996208876371, + 0.04152565822005272, -0.018848471343517303, -0.05523142218589783, + 0.018086062744259834, -0.000535094877704978, -0.013765564188361168, + -0.016923097893595695, + ]; + + const embeddings = (await flagEmbedding.embed(["hello world"]).next()) + .value!; + console.log("embeddings", embeddings[0].slice(0, 10)); + expect(embeddings).toBeDefined(); + for (let i = 0; i < expected.length; i++) { + expect(embeddings[0][i]).toBeCloseTo(expected[i], 3); + } + }); +}); diff --git a/tsconfig.json b/tsconfig.json index 6adcd02..3e25af3 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -13,5 +13,6 @@ "skipLibCheck": true, "types": ["vitest/globals"] }, - "include": ["./src"] + "include": ["./src"], + "exclude": ["./models"] }