@amemhq/core 1.0.0 → 1.0.1

package/dist/index.cjs

@@ -49,6 +49,7 @@ __export(index_exports, {
   generateReviewBatch: () => generateReviewBatch,
   getEmbeddingDim: () => getEmbeddingDim,
   getEmbeddingModel: () => getEmbeddingModel,
+  getEmbeddingPooling: () => getEmbeddingPooling,
   getNote: () => getNote,
   invalidateNote: () => invalidateNote,
   isModelLoaded: () => isModelLoaded,
@@ -88,6 +89,25 @@ var DEFAULT_EMBEDDING_MODEL = "Xenova/paraphrase-multilingual-MiniLM-L12-v2";
 function getEmbeddingModel() {
   return process.env.AMEM_EMBED_MODEL?.trim() || DEFAULT_EMBEDDING_MODEL;
 }
+var CLS_POOLED_MODELS = /* @__PURE__ */ new Set([
+  "bge-m3",
+  "bge-base-zh-v1.5",
+  "bge-small-zh-v1.5",
+  "bge-base-en-v1.5",
+  "bge-small-en-v1.5",
+  "bge-large-en-v1.5",
+  "gte-multilingual-base",
+  "gte-modernbert-base",
+  "gte-large-en-v1.5",
+  "snowflake-arctic-embed-m",
+  "snowflake-arctic-embed-l"
+]);
+function getEmbeddingPooling() {
+  const explicit = process.env.AMEM_EMBED_POOLING?.trim().toLowerCase();
+  if (explicit === "mean" || explicit === "cls") return explicit;
+  const basename2 = getEmbeddingModel().split("/").pop()?.toLowerCase() ?? "";
+  return CLS_POOLED_MODELS.has(basename2) ? "cls" : "mean";
+}
 async function getExtractor() {
   const wanted = getEmbeddingModel();
   if (extractor && loadedModelName === wanted) return extractor;
@@ -108,21 +128,25 @@ async function getEmbeddingDim() {
   cachedDim = probe.length;
   return cachedDim;
 }
-function meanPoolingNormalize(output, attentionMask) {
+function poolNormalize(output, attentionMask, mode) {
   const seqLen = output.length;
   const dim = output[0].length;
   const pooled = new Array(dim).fill(0);
-  let maskSum = 0;
-  for (let i = 0; i < seqLen; i++) {
-    const m = attentionMask[i];
-    maskSum += m;
+  if (mode === "cls") {
+    for (let j = 0; j < dim; j++) pooled[j] = output[0][j];
+  } else {
+    let maskSum = 0;
+    for (let i = 0; i < seqLen; i++) {
+      const m = attentionMask[i];
+      maskSum += m;
+      for (let j = 0; j < dim; j++) {
+        pooled[j] += output[i][j] * m;
+      }
+    }
     for (let j = 0; j < dim; j++) {
-      pooled[j] += output[i][j] * m;
+      pooled[j] /= Math.max(maskSum, 1e-9);
     }
   }
-  for (let j = 0; j < dim; j++) {
-    pooled[j] /= Math.max(maskSum, 1e-9);
-  }
   let norm = 0;
   for (const v of pooled) norm += v * v;
   norm = Math.sqrt(norm);
@@ -130,7 +154,8 @@ function meanPoolingNormalize(output, attentionMask) {
 }
 async function encode(text) {
   const ext = await getExtractor();
-  const result = await ext(text, { pooling: "mean", normalize: true });
+  const pooling = getEmbeddingPooling();
+  const result = await ext(text, { pooling, normalize: true });
   if (result && result.data) {
     return Array.from(result.data);
   }
@@ -146,7 +171,7 @@ async function encode(text) {
       }
       raw.push(row);
     }
-    return meanPoolingNormalize(raw, new Array(seqLen).fill(1));
+    return poolNormalize(raw, new Array(seqLen).fill(1), pooling);
   }
   throw new Error("Unexpected embedding output shape");
 }
@@ -2349,6 +2374,7 @@ function isPlausibleUpdateTarget(newEmbedding, targetEmbedding, minSimilarity) {
   generateReviewBatch,
   getEmbeddingDim,
   getEmbeddingModel,
+  getEmbeddingPooling,
   getNote,
   invalidateNote,
   isModelLoaded,