From ac00be1bb47a9874b9d7ed21094544f917ac6e5d Mon Sep 17 00:00:00 2001 From: emme1t <149944796+emme1t@users.noreply.github.com> Date: Sat, 31 Oct 2026 18:33:42 -0700 Subject: [PATCH] fix: resolve auto models from paths with trailing separators --- FlagEmbedding/inference/auto_embedder.py | 5 +- FlagEmbedding/inference/auto_reranker.py | 5 +- tests/test_auto_model_local_paths.py | 84 ++++++++++++++++++++++++ 3 files changed, 90 insertions(+), 4 deletions(-) create mode 100644 tests/test_auto_model_local_paths.py diff --git a/FlagEmbedding/inference/auto_embedder.py b/FlagEmbedding/inference/auto_embedder.py index 6fad0330f..b8e1f09da 100644 --- a/FlagEmbedding/inference/auto_embedder.py +++ b/FlagEmbedding/inference/auto_embedder.py @@ -59,9 +59,10 @@ def from_finetuned( Returns: AbsEmbedder: The model class to load model, which is child class of :class:`AbsEmbedder`. """ - model_name = os.path.basename(model_name_or_path) + model_path = os.path.normpath(model_name_or_path) + model_name = os.path.basename(model_path) if model_name.startswith("checkpoint-"): - model_name = os.path.basename(os.path.dirname(model_name_or_path)) + model_name = os.path.basename(os.path.dirname(model_path)) if model_class is not None: _model_class = EMBEDDER_CLASS_MAPPING[EmbedderModelClass(model_class)] diff --git a/FlagEmbedding/inference/auto_reranker.py b/FlagEmbedding/inference/auto_reranker.py index a8b2d699e..5157ba0c9 100644 --- a/FlagEmbedding/inference/auto_reranker.py +++ b/FlagEmbedding/inference/auto_reranker.py @@ -46,9 +46,10 @@ def from_finetuned( Returns: AbsReranker: The reranker class to load model, which is child class of :class:`AbsReranker`. """ - model_name = os.path.basename(model_name_or_path) + model_path = os.path.normpath(model_name_or_path) + model_name = os.path.basename(model_path) if model_name.startswith("checkpoint-"): - model_name = os.path.basename(os.path.dirname(model_name_or_path)) + model_name = os.path.basename(os.path.dirname(model_path)) if model_class is not None: _model_class = RERANKER_CLASS_MAPPING[RerankerModelClass(model_class)] diff --git a/tests/test_auto_model_local_paths.py b/tests/test_auto_model_local_paths.py new file mode 100644 index 000000000..d1b6ff850 --- /dev/null +++ b/tests/test_auto_model_local_paths.py @@ -0,0 +1,84 @@ +import os + +import numpy as np +import pytest +from tokenizers import Tokenizer, models, pre_tokenizers, processors +from transformers import ( + BertConfig, + BertForSequenceClassification, + BertModel, + PreTrainedTokenizerFast, +) + +from FlagEmbedding import FlagAutoModel, FlagAutoReranker + + +def run_inference(model, kind): + if kind == "embedder": + return model.encode(["hello world", "hello"]) + return np.asarray(model.compute_score([("hello", "world"), ("hello", "hello")])) + + +@pytest.fixture(scope="module", params=["embedder", "reranker"]) +def local_model(request, tmp_path_factory): + """Save small, randomly initialized models without downloading checkpoints.""" + kind = request.param + name = "bge-base-en-v1.5" if kind == "embedder" else "bge-reranker-base" + model_dir = tmp_path_factory.mktemp(kind) / name + checkpoint_dir = model_dir / "checkpoint-10" + vocabulary = { + token: i + for i, token in enumerate( + ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "hello", "world"] + ) + } + backend = Tokenizer(models.WordPiece(vocabulary, unk_token="[UNK]")) + backend.pre_tokenizer = pre_tokenizers.Whitespace() + backend.post_processor = processors.TemplateProcessing( + single="[CLS] $A [SEP]", + pair="[CLS] $A [SEP] $B:1 [SEP]:1", + special_tokens=[("[CLS]", 2), ("[SEP]", 3)], + ) + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=backend, + unk_token="[UNK]", + pad_token="[PAD]", + cls_token="[CLS]", + sep_token="[SEP]", + model_max_length=32, + ) + config = BertConfig( + vocab_size=len(vocabulary), + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + max_position_embeddings=32, + num_labels=1, + ) + model_type = BertModel if kind == "embedder" else BertForSequenceClassification + model = model_type(config) + for path in (model_dir, checkpoint_dir): + model.save_pretrained(path) + tokenizer.save_pretrained(path) + factory = FlagAutoModel if kind == "embedder" else FlagAutoReranker + reference = factory.from_finetuned(str(model_dir), devices="cpu", use_fp16=False) + expected = run_inference(reference, kind) + assert np.isfinite(expected).all() + return kind, factory, model_dir, checkpoint_dir, expected + + +@pytest.mark.parametrize("checkpoint", [False, True]) +@pytest.mark.parametrize("suffix", sorted({"", "/", "//", os.sep, os.sep * 2})) +def test_auto_model_loads_local_directory_with_trailing_separator( + local_model, checkpoint, suffix +): + kind, factory, model_dir, checkpoint_dir, expected = local_model + path = str(checkpoint_dir if checkpoint else model_dir) + suffix + + model = factory.from_finetuned(path, devices="cpu", use_fp16=False) + + assert model.model_name_or_path == path + np.testing.assert_allclose( + run_inference(model, kind), expected, rtol=1e-6, atol=1e-6 + )