@mastra/fastembed 1.3.4 → 1.3.5-alpha.1

package/dist/index.js

@@ -1,12 +1,12 @@
 import fsp from "fs/promises";
 import os from "os";
 import path from "path";
-import { AddedToken, Tokenizer } from "@anush008/tokenizers";
 import fs from "fs";
 import { Readable } from "stream";
 import { pipeline } from "stream/promises";
-import * as ort from "onnxruntime-node";
+import { AddedToken, Tokenizer } from "@anush008/tokenizers";
 import { downloadFile } from "@huggingface/hub";
+import * as ort from "onnxruntime-node";
 //#region ../../../../../setup-pnpm/node_modules/.bin/store/v11/links/@/zod/4.6.7/4c9f15c99caaca8b83ba8ab68f247a75121c86fb7e498876ed36daccd0f3d3ce/node_modules/zod/v4/core/util.js
 function getEnumValues(entries) {
 	const numericValues = Object.values(entries).filter((v) => typeof v === "number");
@@ -18387,12 +18387,19 @@ function normalize(v) {
 	const epsilon = 1e-12;
 	return v.map((val) => val / Math.max(norm, epsilon));
 }
-function getEmbeddings(data, dimensions) {
-	const [x, _y, z] = dimensions;
-	return new Array(x).fill(void 0).map((_, index) => {
-		const startIndex = index * _y * z;
-		const endIndex = startIndex + z;
-		return data.slice(startIndex, endIndex);
+function getEmbeddings(data, dimensions, masks) {
+	const [batchSize, sequenceLength, hiddenSize] = dimensions;
+	return Array.from({ length: batchSize }, (_, batchIndex) => {
+		const start = batchIndex * sequenceLength * hiddenSize;
+		if (!masks) return data.slice(start, start + hiddenSize);
+		const embedding = new Array(hiddenSize).fill(0);
+		let tokenCount = 0;
+		for (let tokenIndex = 0; tokenIndex < sequenceLength; tokenIndex++) {
+			if (masks[batchIndex][tokenIndex] === 0n) continue;
+			tokenCount++;
+			for (let dimension = 0; dimension < hiddenSize; dimension++) embedding[dimension] = embedding[dimension] + data[start + tokenIndex * hiddenSize + dimension];
+		}
+		return embedding.map((value) => value / Math.max(tokenCount, 1));
 	});
 }
 var Embedding = class {};
@@ -18409,7 +18416,7 @@ function loadTokenizerFromDir(modelDir, maxLength) {
 	const tokenizerFilePath = path.join(modelDir.toString(), "tokenizer_config.json");
 	if (!fs.existsSync(tokenizerFilePath)) throw new Error(`Tokenizer config file not found at ${tokenizerFilePath}`);
 	const tokenizerConfig = JSON.parse(fs.readFileSync(tokenizerFilePath, "utf-8"));
-	maxLength = Math.min(maxLength, tokenizerConfig["model_max_length"]);
+	maxLength = Math.min(maxLength, tokenizerConfig["model_max_length"] ?? maxLength);
 	const tokensMapPath = path.join(modelDir.toString(), "special_tokens_map.json");
 	if (!fs.existsSync(tokensMapPath)) throw new Error(`Tokens map file not found at ${tokensMapPath}`);
 	const tokensMap = JSON.parse(fs.readFileSync(tokensMapPath, "utf-8"));
@@ -18544,7 +18551,7 @@ var FlagEmbedding = class FlagEmbedding extends Embedding {
 			};
 			if (this.model === "fast-multilingual-e5-large") delete inputs.token_type_ids;
 			const lastHiddenState = (await this.session.run(inputs)).last_hidden_state;
-			yield getEmbeddings(lastHiddenState.data, lastHiddenState.dims).map(normalize);
+			yield getEmbeddings(lastHiddenState.data, lastHiddenState.dims, this.model === "fast-all-MiniLM-L6-v2" || this.model === "fast-multilingual-e5-large" ? maskArray : void 0).map(normalize);
 		}
 	}
 	passageEmbed(texts, batchSize = 256) {
@@ -18619,29 +18626,46 @@ var SparseTextEmbedding = class SparseTextEmbedding extends SparseEmbedding {
 		return { tokenizer: loadTokenizerFromDir(modelDir, maxLength) };
 	}
 	static async retrieveModel(model, cacheDir, _showDownloadProgress = true) {
-		if (!fs.existsSync(cacheDir)) fs.mkdirSync(cacheDir, { mode: 493 });
+		if (!fs.existsSync(cacheDir)) fs.mkdirSync(cacheDir, {
+			recursive: true,
+			mode: 493
+		});
 		const modelDir = path.join(cacheDir.toString(), model.replace("/", "_"));
 		if (fs.existsSync(modelDir)) return modelDir;
-		fs.mkdirSync(modelDir, { mode: 493 });
-		for (const fileName of [
+		const filesToDownload = [
 			"onnx/model.onnx",
 			"tokenizer.json",
 			"tokenizer_config.json",
 			"config.json",
 			"special_tokens_map.json"
-		]) {
-			const outputPath = path.join(modelDir, fileName);
-			const outputDir = path.dirname(outputPath);
-			if (!fs.existsSync(outputDir)) fs.mkdirSync(outputDir, {
+		];
+		const tmpDir = fs.mkdtempSync(`${modelDir}.tmp-`);
+		try {
+			for (const fileName of filesToDownload) {
+				const outputPath = path.join(tmpDir, fileName);
+				fs.mkdirSync(path.dirname(outputPath), {
+					recursive: true,
+					mode: 493
+				});
+				const downloaded = await downloadFile({
+					repo: model,
+					path: fileName
+				});
+				if (!downloaded) throw new Error(`Model file not found in ${model}: ${fileName}`);
+				await pipeline(Readable.fromWeb(downloaded.stream()), fs.createWriteStream(outputPath));
+			}
+			try {
+				fs.renameSync(tmpDir, modelDir);
+			} catch (error) {
+				const code = error.code;
+				if ((code === "EEXIST" || code === "ENOTEMPTY") && fs.existsSync(modelDir)) return modelDir;
+				throw error;
+			}
+		} finally {
+			fs.rmSync(tmpDir, {
 				recursive: true,
-				mode: 493
-			});
-			const downloaded = await downloadFile({
-				repo: model,
-				path: fileName
+				force: true
 			});
-			if (!downloaded) throw new Error(`Model file not found in ${model}: ${fileName}`);
-			await pipeline(Readable.fromWeb(downloaded.stream()), fs.createWriteStream(outputPath));
 		}
 		return modelDir;
 	}