fix(static): 修复 float16 静态向量加权池化精度

fastembed/text/static_embedding.py

@@ -235,7 +235,11 @@ def _sum_rows(self, ids: NDArray[np.int64], lengths: NDArray[np.int64]) -> Numpy
             # a count matrix, which dominates for short texts.
             vectors = self.embeddings[rows]
             if self.weights is not None:
-                vectors = vectors * self.weights[ids, None]
+                vectors = np.multiply(
+                    vectors,
+                    self.weights[ids, None],
+                    dtype=np.result_type(vectors.dtype, self.weights.dtype, np.float32),
+                )
             return vectors.sum(axis=0, keepdims=True, dtype=np.float32)
         distinct_rows, row_index = np.unique(rows, return_inverse=True)
         if len(lengths) > 1 and len(lengths) * len(distinct_rows) > self.POOLING_CELLS:

tests/test_static_embeddings.py

@@ -1,10 +1,12 @@
 import json
 import os
 import warnings
-from typing import Iterator
+from collections.abc import Iterator
 
 import numpy as np
 import pytest
+from tokenizers import Tokenizer
+from tokenizers.models import WordLevel
 
 from fastembed.common.safetensors_utils import load_safetensors
 from fastembed.text.static_embedding import StaticTextEmbedding
@@ -101,6 +103,7 @@ def test_embedding_size() -> None:
 
 
 def _write_safetensors(path, tensors: dict[str, tuple[str, list[int], bytes]]) -> None:
+    """Write a minimal safetensors fixture without downloading a model."""
     header, data, offset = {}, b"", 0
     for name, (dtype, shape, raw) in tensors.items():
         header[name] = {
@@ -118,8 +121,43 @@ def _write_safetensors(path, tensors: dict[str, tuple[str, list[int], bytes]]) -
         f.write(data)
 
 
+@pytest.mark.parametrize("mmap", [False, True])
+@pytest.mark.parametrize("dtype", ["F16", "F32"])
+@pytest.mark.parametrize("vector_scale,weight", [(1e-4, 1e-4), (1e4, 10.0)])
+def test_weighted_pooling_matches_across_batch_sizes(tmp_path, mmap, dtype, vector_scale, weight):
+    """Keep weighted pooling finite and batch-independent for float16 and float32."""
+    tokenizer = Tokenizer(WordLevel({"[UNK]": 0, "hello": 1}, unk_token="[UNK]"))
+    tokenizer.save(str(tmp_path / "tokenizer.json"))
+    numpy_dtype = "<f2" if dtype == "F16" else "<f4"
+    embeddings = np.zeros((2, 256), dtype=numpy_dtype)
+    embeddings[1, :2] = [vector_scale, 2 * vector_scale]
+    weights = np.array([1.0, weight], dtype=numpy_dtype)
+    _write_safetensors(
+        tmp_path / "model.safetensors",
+        {
+            "embeddings": (dtype, list(embeddings.shape), embeddings.tobytes()),
+            "weights": (dtype, list(weights.shape), weights.tobytes()),
+        },
+    )
+    model = TextEmbedding(
+        CI_MODEL,
+        specific_model_path=str(tmp_path),
+        cache_dir=str(tmp_path / "cache"),
+        mmap=mmap,
+    )
+
+    single = np.stack(list(model.embed(["hello", "hello"], batch_size=1)))
+    batched = np.stack(list(model.embed(["hello", "hello"], batch_size=2)))
+
+    assert single.dtype == batched.dtype == np.float32
+    assert np.isfinite(single).all()
+    np.testing.assert_allclose(single, batched, rtol=1e-5, atol=1e-7)
+    np.testing.assert_allclose(single[0, :2], np.array([1.0, 2.0]) / np.sqrt(5), rtol=1e-5)
+
+
 @pytest.mark.parametrize("mmap", [True, False])
 def test_load_safetensors(tmp_path, mmap: bool) -> None:
+    """Decode supported tensor types equally with and without memory mapping."""
     values = np.arange(6, dtype=np.float32).reshape(2, 3) / 4
     bf16 = (values.view(np.uint32) >> 16).astype("<u2")
     path = tmp_path / "model.safetensors"