Add support for the larger Chinese BGE model

fastembed/common/preprocessor_utils.py

@@ -47,7 +47,9 @@ def _valid_context(value: Any) -> int | None:
     return value
 
 
-def _resolve_max_context(tokenizer_config: dict[str, Any], model_dir: Path) -> int:
+def _resolve_max_context(
+    tokenizer_config: dict[str, Any], model_dir: Path, max_context: int | None = None
+) -> int:
     """Pick the truncation limit, preferring the stricter of the two tokenizer config keys.
 
     `config.json:max_position_embeddings` deliberately is not used as a fallback: it is the size
@@ -62,6 +64,11 @@ def _resolve_max_context(tokenizer_config: dict[str, Any], model_dir: Path) -> i
         )
         if context is not None
     ]
+    if max_context is not None:
+        context = _valid_context(max_context)
+        if context is None:
+            raise ValueError("max_context must be a positive integer truncation limit")
+        candidates.append(context)
     if not candidates:
         raise ValueError(
             f"Could not determine the maximum context length for {model_dir}. Set a positive "
@@ -71,7 +78,10 @@ def _resolve_max_context(tokenizer_config: dict[str, Any], model_dir: Path) -> i
     return min(candidates)
 
 
-def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
+def load_tokenizer(
+    model_dir: Path, max_context: int | None = None
+) -> tuple[Tokenizer, dict[str, int]]:
+    """Load a tokenizer, optionally capping it at a model's verified context limit."""
     tokenizer_path = model_dir / "tokenizer.json"
     if not tokenizer_path.exists():
         raise ValueError(f"Could not find tokenizer.json in {model_dir}")
@@ -90,7 +100,7 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
     with open(str(tokenizer_config_path)) as tokenizer_config_file:
         tokenizer_config = json.load(tokenizer_config_file)
 
-    max_context = _resolve_max_context(tokenizer_config, model_dir)
+    max_context = _resolve_max_context(tokenizer_config, model_dir, max_context=max_context)
 
     tokens_map = load_special_tokens(model_dir)
 

fastembed/text/onnx_embedding.py

@@ -1,9 +1,11 @@
 from typing import Any, Iterable, Sequence, Type
+from pathlib import Path
 
 from fastembed.common.model_description import DenseModelDescription, ModelSource
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.types import Device, NumpyArray, OnnxProvider
 from fastembed.common.utils import define_cache_dir, normalize
+from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
 from fastembed.text.text_embedding_base import TextEmbeddingBase
 
@@ -92,6 +94,18 @@
         ),
         model_file="model_optimized.onnx",
     ),
+    DenseModelDescription(
+        model="BAAI/bge-large-zh-v1.5",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), Chinese, 512 input tokens truncation, "
+            "Prefixes for queries/documents: not so necessary, 2023 year."
+        ),
+        license="mit",
+        size_in_GB=1.30,
+        sources=ModelSource(hf="Xenova/bge-large-zh-v1.5"),
+        model_file="onnx/model.onnx",
+    ),
     DenseModelDescription(
         model="mixedbread-ai/mxbai-embed-large-v1",
         dim=1024,
@@ -208,6 +222,12 @@
 class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
     """Implementation of the Flag Embedding model."""
 
+    def _load_tokenizer(self, model_dir: Path) -> None:
+        if self.model_description.model != "BAAI/bge-large-zh-v1.5":
+            return super()._load_tokenizer(model_dir)
+        # This export has an unknown-length sentinel, but BGE's usable context is 512 tokens.
+        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir, max_context=512)
+
     @classmethod
     def _list_supported_models(cls) -> list[DenseModelDescription]:
         """

tests/test_bge_large_zh.py

@@ -0,0 +1,116 @@
+"""Reference checks for the larger Chinese BGE model.
+
+Canonical values come from the original PyTorch model using the following snippet.
+
+    from transformers import AutoModel, AutoTokenizer
+    import torch
+
+    name = "BAAI/bge-large-zh-v1.5"
+    tokenizer = AutoTokenizer.from_pretrained(name)
+    model = AutoModel.from_pretrained(name).eval()
+    inputs = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt")
+    with torch.no_grad():
+        vectors = torch.nn.functional.normalize(
+            model(**inputs).last_hidden_state[:, 0], p=2, dim=1
+        )
+    print(vectors[:, :8].tolist())
+"""
+
+import os
+
+import numpy as np
+import pytest
+
+from fastembed import TextEmbedding
+from tests.utils import delete_model_cache, is_manual_run, should_test_model
+
+MODEL_NAME = "BAAI/bge-large-zh-v1.5"
+REFERENCE_CASES = [
+    (
+        "hello world",
+        [
+            -0.025394167751073837,
+            -0.0016842425102367997,
+            -0.094914011657238,
+            -0.007070891559123993,
+            0.0027618478052318096,
+            0.050447918474674225,
+            -0.0512826070189476,
+            0.021190207451581955,
+        ],
+    ),
+    (
+        "你好世界",
+        [
+            0.04448169842362404,
+            -0.023043543100357056,
+            -0.06824296712875366,
+            -0.005314783658832312,
+            -0.025446804240345955,
+            0.016195841133594513,
+            -0.06172054260969162,
+            0.028417708352208138,
+        ],
+    ),
+    (
+        "北京是中国的首都。",
+        [
+            -0.07171106338500977,
+            0.028986480087041855,
+            0.0031650469172745943,
+            0.02360720932483673,
+            0.00862447265535593,
+            0.05719650164246559,
+            -0.0070158555172383785,
+            -0.037890322506427765,
+        ],
+    ),
+    (
+        "为这个句子生成表示以用于检索相关文章： 中国的首都是哪里？",
+        [
+            -0.06612120568752289,
+            0.03332896530628204,
+            0.028995418921113014,
+            0.03479088470339775,
+            0.005401355680078268,
+            0.07183656841516495,
+            0.001583364442922175,
+            -0.029327189549803734,
+        ],
+    ),
+]
+
+
+@pytest.fixture(scope="module")
+def model():
+    description = next(
+        item for item in TextEmbedding._list_supported_models() if item.model == MODEL_NAME
+    )
+    if not should_test_model(description, "", os.getenv("CI"), is_manual_run()):
+        pytest.skip("The large-model check runs in the scheduled model matrix")
+    embedding = TextEmbedding(MODEL_NAME, threads=2, cuda=False)
+    yield embedding
+    if os.getenv("CI"):
+        delete_model_cache(embedding.model._model_dir)
+
+
+@pytest.mark.parametrize("text,expected", REFERENCE_CASES)
+def test_reference_embeddings(model, text, expected):
+    passage = next(iter(model.passage_embed([text])))
+    query = next(iter(model.query_embed(text)))
+    assert passage.shape == query.shape == (1024,)
+    assert np.allclose(passage[:8], expected, atol=1e-3)
+    assert np.allclose(query[:8], expected, atol=1e-3)
+    assert np.linalg.norm(passage) == pytest.approx(1.0, abs=1e-6)
+
+
+def test_mixed_length_parallel_embeddings(model):
+    texts = [text for text, _ in REFERENCE_CASES]
+    sequential = np.array(list(model.embed(texts, batch_size=4)))
+    parallel = np.array(list(model.embed(texts, batch_size=1, parallel=2)))
+    assert sequential.shape == parallel.shape == (len(texts), 1024)
+    assert np.allclose(sequential, parallel, atol=1e-5)
+
+
+def test_long_chinese_input_is_truncated(model):
+    assert model.token_count("中国" * 600) == 512

tests/test_preprocessor_utils.py

@@ -288,6 +288,51 @@ def test_absent_max_context_keys_raise(make_model_dir) -> None:
         load_tokenizer(model_dir)
 
 
+@pytest.mark.parametrize(
+    "configured,limit,expected",
+    [(HF_SENTINEL, 512, 512), (128, 512, 128), (512, 256, 256)],
+)
+def test_explicit_context_limit_never_extends_a_configured_limit(
+    make_model_dir, configured, limit, expected
+) -> None:
+    model_dir = make_model_dir(
+        tokenizer_config={"model_max_length": configured, "max_length": None}
+    )
+    tokenizer, _ = load_tokenizer(model_dir, max_context=limit)
+    assert tokenizer.truncation["max_length"] == expected
+
+
+@pytest.mark.parametrize("limit", [0, -1, True, HF_SENTINEL])
+def test_unusable_explicit_context_limit_raises(make_model_dir, limit) -> None:
+    with pytest.raises(ValueError, match="max_context must be a positive integer"):
+        load_tokenizer(make_model_dir(), max_context=limit)
+
+
+@pytest.mark.parametrize(
+    "model_name,configured,expected",
+    [
+        ("BAAI/bge-large-zh-v1.5", HF_SENTINEL, 512),
+        ("baai/BGE-LARGE-ZH-v1.5", HF_SENTINEL, 512),
+        ("BAAI/bge-large-zh-v1.5", 128, 128),
+        (BASE_MODEL, 256, 256),
+    ],
+)
+def test_model_tokenizer_context_without_onnx_weights(
+    make_model_dir, model_name, configured, expected
+) -> None:
+    model_dir = make_model_dir(
+        tokenizer_config={"model_max_length": configured, "max_length": None}
+    )
+    embedding = TextEmbedding(
+        model_name, lazy_load=True, specific_model_path=str(model_dir), cuda=False
+    )
+
+    embedding.token_count("hello world")
+
+    assert embedding.model.tokenizer.truncation["max_length"] == expected
+    assert getattr(embedding.model, "model", None) is None
+
+
 @pytest.fixture(scope="module")
 def token_id(make_model_dir):
     """Resolve vocabulary ids by name, so the cases below carry no magic numbers."""

tests/test_text_onnx_embeddings.py

@@ -28,6 +28,15 @@
     "BAAI/bge-base-en-v1.5": np.array(
         [0.01129394, 0.05493144, 0.02615099, 0.00328772, 0.02996045]
     ),
+    "BAAI/bge-large-zh-v1.5": np.array(
+        [
+            -0.025394167751073837,
+            -0.0016842425102367997,
+            -0.094914011657238,
+            -0.007070891559123993,
+            0.0027618478052318096,
+        ]
+    ),
     "BAAI/bge-large-en-v1.5": np.array(
         [0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]
     ),