Update index.py

2025-05-18 19:27:14 +02:00 · 2025-05-18 19:27:14 +02:00 · 50042f6655
commit 50042f6655
parent 89d167a2f8
1 changed files with 3 additions and 1 deletions
--- a/index.py
+++ b/index.py
@ -98,7 +98,9 @@ def main():
    )

    model = BGEM3FlagModel(MODEL_NAME, device="cpu")
-    emb = model.encode(new_docs, batch_size=BATCH, return_dict=False)
+    emb_out = model.encode(new_docs, batch_size=BATCH)
+    # FlagEmbedding renvoie soit un ndarray, soit un dict {"embedding": …}
+    emb = emb_out["embedding"] if isinstance(emb_out, dict) else emb_out
    emb = emb.astype("float32")
    emb /= np.linalg.norm(emb, axis=1, keepdims=True) + 1e-12