From 56b343b0e653e3d3996d02aaaaaaed93ed5be65d Mon Sep 17 00:00:00 2001 From: IlyaKhalafi Date: Wed, 14 Oct 2026 08:48:08 +0000 Subject: [PATCH] Add larger Chinese BGE embeddings Closes #289 --- fastembed/common/preprocessor_utils.py | 16 +++- fastembed/text/onnx_embedding.py | 20 +++++ tests/test_bge_large_zh.py | 116 +++++++++++++++++++++++++ tests/test_preprocessor_utils.py | 45 ++++++++++ tests/test_text_onnx_embeddings.py | 9 ++ 5 files changed, 203 insertions(+), 3 deletions(-) create mode 100644 tests/test_bge_large_zh.py diff --git a/fastembed/common/preprocessor_utils.py b/fastembed/common/preprocessor_utils.py index fb853924..a70877b0 100644 --- a/fastembed/common/preprocessor_utils.py +++ b/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) diff --git a/fastembed/text/onnx_embedding.py b/fastembed/text/onnx_embedding.py index cae94ac5..96983360 100644 --- a/fastembed/text/onnx_embedding.py +++ b/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]: """ diff --git a/tests/test_bge_large_zh.py b/tests/test_bge_large_zh.py new file mode 100644 index 00000000..75c5b9e0 --- /dev/null +++ b/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 diff --git a/tests/test_preprocessor_utils.py b/tests/test_preprocessor_utils.py index d04f0312..e7f5e4ca 100644 --- a/tests/test_preprocessor_utils.py +++ b/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.""" diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 179e8c25..013e137d 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/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] ),