diff --git a/fastembed/text/static_embedding.py b/fastembed/text/static_embedding.py index 853d8214..60f44cde 100644 --- a/fastembed/text/static_embedding.py +++ b/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: diff --git a/tests/test_static_embeddings.py b/tests/test_static_embeddings.py index 60f695e7..22aa0df3 100644 --- a/tests/test_static_embeddings.py +++ b/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 = " 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("