From fae2c1515f76c12d77b69c163db5e0fe11654075 Mon Sep 17 00:00:00 2001 From: jraedisch Date: Fri, 3 Apr 2026 02:49:01 +0200 Subject: [PATCH] match mlx_embeddings float32 pooling precision in verify.py mlx_embeddings.mean_pooling casts the attention mask to float32 before accumulating the pooling sum. verify.py was staying in bfloat16, losing precision during summation over many tokens. Cast h and mask to float32 for pooling and dense projection to match the server-side computation. Co-Authored-By: Claude Opus 4.6 (1M context) --- verify.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/verify.py b/verify.py index c4c7a86..9aee3a7 100644 --- a/verify.py +++ b/verify.py @@ -166,18 +166,17 @@ class EmbeddingModel: with torch.no_grad(): out = self.model(**inputs) - # Mean pooling, then dense projection (matches SentenceTransformers / mlx-embeddings ≥0.1.0) + # Mean pooling in float32 (matches mlx_embeddings mean_pooling which casts mask to float32) h = out.last_hidden_state - mask = inputs["attention_mask"].unsqueeze(-1).to(h.dtype) - emb = (h * mask).sum(dim=1) / mask.sum(dim=1) + mask = inputs["attention_mask"].unsqueeze(-1).float() + emb = (h.float() * mask).sum(dim=1) / mask.sum(dim=1) emb = emb.squeeze() # Dense projection (identity activation, no bias) - emb = emb @ self.dense0.T - emb = emb @ self.dense1.T + emb = emb @ self.dense0.float().T + emb = emb @ self.dense1.float().T # L2 normalize - emb = emb.float() emb = emb / emb.norm() return emb.numpy().astype(np.float32) -- 2.51.2