chore: revert 824abe6

Anush008 · Sep 19, 2023 · 472b119 · 472b119
1 parent 544afcb
commit 472b119
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/src/fastembed.ts b/src/fastembed.ts
@@ -289,7 +289,21 @@ export class FlagEmbedding extends Embedding {
         output.last_hidden_state.dims as number[]
       );
 
-      const embeddings = lastHiddenState.map((sentence) => sentence[0]);
+      const embeddings = lastHiddenState.map((layer, layerIdx) => {
+        const weightedSum = layer.reduce((acc, tokenEmbedding, idx) => {
+          const attentionWeight = maskArray[layerIdx][idx];
+          return acc.map(
+            (val, i) => val + tokenEmbedding[i] * Number(attentionWeight)
+          );
+        }, new Array(layer[0].length).fill(0));
+
+        const inputMaskSum = maskArray[layerIdx].reduce(
+          (acc, attentionWeight) => acc + Number(attentionWeight),
+          0
+        );
+
+        return weightedSum.map((val) => val / (inputMaskSum + 1e-9));
+      });
 
       yield embeddings.map(normalize);
     }