mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 14:14:50 +00:00
fix: minor embed compat fixes
This commit is contained in:
1 parent
dd3876fdcb
commit
869d87452b
6 files changed
+121
-2
No files matched your search
@@ -202,7 +202,7 @@ def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int,
|
||||
``(chunks_seen, chunks_written)``.
|
||||
"""
|
||||
store = VectorCreator.create_vectorstore(
|
||||
"faiss", source_id=source_id, embeddings_key="embeddings"
|
||||
"faiss", source_id=source_id, embeddings_key=settings.EMBEDDINGS_KEY
|
||||
)
|
||||
chunks: List[Dict[str, Any]] = store.get_chunks() or []
|
||||
if dry_run or not chunks:
|
||||
@@ -217,7 +217,10 @@ def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int,
|
||||
# index and touches nothing on disk. The existing index is only replaced by
|
||||
# ``save_local`` below, so a failure while embedding leaves it intact.
|
||||
rebuilt = VectorCreator.create_vectorstore(
|
||||
"faiss", source_id=source_id, embeddings_key="embeddings", docs_init=docs
|
||||
"faiss",
|
||||
source_id=source_id,
|
||||
embeddings_key=settings.EMBEDDINGS_KEY,
|
||||
docs_init=docs,
|
||||
)
|
||||
rebuilt.save_local()
|
||||
return len(chunks), len(docs)
|
||||
|
||||
@@ -13,6 +13,11 @@ from application.vectorstore.model_registry import (
|
||||
)
|
||||
|
||||
|
||||
def _embeddings_name_is_explicit() -> bool:
|
||||
"""True when ``EMBEDDINGS_NAME`` was configured rather than defaulted."""
|
||||
return "EMBEDDINGS_NAME" in getattr(settings, "model_fields_set", set())
|
||||
|
||||
|
||||
class RemoteEmbeddings:
|
||||
"""
|
||||
Wrapper for remote embeddings API (OpenAI-compatible).
|
||||
@@ -47,10 +52,18 @@ class RemoteEmbeddings:
|
||||
model contributes its own context window, so a request that the server
|
||||
would reject -- or silently truncate -- is clipped here instead of
|
||||
being sent and paid for.
|
||||
|
||||
That fallback needs the name to mean something. For a remote server it
|
||||
is only a label forwarded as the ``model`` field, so the settings
|
||||
default must not lend the server mpnet's 384-token window: a name
|
||||
nobody chose describes nothing, and clipping on it would silently
|
||||
discard most of every chunk.
|
||||
"""
|
||||
configured = settings.EMBEDDINGS_MAX_INPUT_TOKENS
|
||||
if configured and configured > 0:
|
||||
return configured
|
||||
if not _embeddings_name_is_explicit():
|
||||
return None
|
||||
model_limit = max_input_tokens_for(self.model_name)
|
||||
return model_limit if model_limit and model_limit > 0 else None
|
||||
|
||||
|
||||
@@ -94,6 +94,27 @@ def _spec_for(model_name: str) -> EmbeddingModel:
|
||||
)
|
||||
|
||||
|
||||
def _pad_to_longest_in_batch(model: Any) -> None:
|
||||
"""Undo a fixed padding width baked into a model's ``tokenizer.json``."""
|
||||
# FastEmbed enables padding only when the tokenizer declares none, so a
|
||||
# fixed ``length`` survives loading. Shorter inputs are then padded to that
|
||||
# width while longer ones keep their own, the batch is ragged, and the ONNX
|
||||
# tensor build fails. mpnet ships ``length: 128``; granite does not.
|
||||
# Mean pooling masks pad tokens, so the vectors are unaffected.
|
||||
tokenizer = getattr(getattr(model, "model", None), "tokenizer", None)
|
||||
padding = getattr(tokenizer, "padding", None)
|
||||
if not isinstance(padding, dict) or padding.get("length") is None:
|
||||
return
|
||||
tokenizer.enable_padding(
|
||||
direction=padding.get("direction", "right"),
|
||||
pad_id=padding.get("pad_id", 0),
|
||||
pad_type_id=padding.get("pad_type_id", 0),
|
||||
pad_token=padding.get("pad_token", "<pad>"),
|
||||
length=None,
|
||||
pad_to_multiple_of=padding.get("pad_to_multiple_of"),
|
||||
)
|
||||
|
||||
|
||||
class EmbeddingsWrapper:
|
||||
"""Runs an embedding model locally through FastEmbed.
|
||||
|
||||
@@ -131,6 +152,7 @@ class EmbeddingsWrapper:
|
||||
f"{exc}. Known models: {', '.join(known_names())}."
|
||||
) from exc
|
||||
|
||||
_pad_to_longest_in_batch(self.model)
|
||||
self.dimension = self.spec.dimension or self._probe_dimension()
|
||||
logger.info("Embeddings model ready (dimension=%d)", self.dimension)
|
||||
|
||||
|
||||
@@ -99,6 +99,14 @@ class TestFaissRebuild:
|
||||
assert [d.page_content for d in docs] == ["alpha", "beta"]
|
||||
assert [d.metadata for d in docs] == [{"i": 0}, {"i": 1}]
|
||||
|
||||
def test_embeddings_key_comes_from_settings(self, stores):
|
||||
"""A placeholder here is sent as the server's bearer token."""
|
||||
_, _, factory = stores
|
||||
with patch.object(reembed.settings, "EMBEDDINGS_KEY", "sk-real", create=True):
|
||||
reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
|
||||
keys = [call.kwargs.get("embeddings_key") for call in factory.call_args_list]
|
||||
assert keys == ["sk-real", "sk-real"]
|
||||
|
||||
def test_existing_index_is_not_deleted(self, stores):
|
||||
"""The rebuild must not destroy the old index before the new one exists."""
|
||||
existing, rebuilt, _ = stores
|
||||
|
||||
@@ -165,3 +165,54 @@ class TestLengthSortedBatching:
|
||||
wrapper, _ = self._wrapper(fake_fastembed, 2)
|
||||
out = wrapper.embed_documents(["aa", "b", "aa", "ccc"])
|
||||
assert out == [[2.0], [1.0], [2.0], [3.0]]
|
||||
|
||||
|
||||
class TestTokenizerPadding:
|
||||
"""A fixed padding width in ``tokenizer.json`` makes mixed batches ragged.
|
||||
|
||||
FastEmbed calls ``enable_padding`` only when the tokenizer declares none,
|
||||
so mpnet's fixed ``length: 128`` survives loading. Any batch mixing an
|
||||
input longer than 128 tokens with a shorter one then produces rows of
|
||||
different widths and ONNX rejects the tensor.
|
||||
"""
|
||||
|
||||
def _tokenizer(self, padding):
|
||||
tokenizer = MagicMock()
|
||||
tokenizer.padding = padding
|
||||
return tokenizer
|
||||
|
||||
def test_fixed_width_padding_is_reset_to_batch_longest(self, fake_fastembed):
|
||||
_, instance = fake_fastembed
|
||||
tokenizer = self._tokenizer(
|
||||
{
|
||||
"length": 128,
|
||||
"pad_id": 1,
|
||||
"pad_token": "<pad>",
|
||||
"pad_type_id": 0,
|
||||
"direction": "right",
|
||||
"pad_to_multiple_of": None,
|
||||
}
|
||||
)
|
||||
instance.model.tokenizer = tokenizer
|
||||
|
||||
EmbeddingsWrapper(MPNET.name)
|
||||
|
||||
kwargs = tokenizer.enable_padding.call_args.kwargs
|
||||
assert kwargs["length"] is None, "padding must follow the longest input"
|
||||
# The model's own pad token must survive the reset.
|
||||
assert kwargs["pad_id"] == 1
|
||||
assert kwargs["pad_token"] == "<pad>"
|
||||
|
||||
def test_dynamic_padding_is_left_alone(self, fake_fastembed):
|
||||
_, instance = fake_fastembed
|
||||
tokenizer = self._tokenizer({"length": None, "pad_id": 0, "pad_token": "<pad>"})
|
||||
instance.model.tokenizer = tokenizer
|
||||
|
||||
EmbeddingsWrapper(GRANITE_97M.name)
|
||||
|
||||
tokenizer.enable_padding.assert_not_called()
|
||||
|
||||
def test_tokenizer_that_cannot_be_reached_is_not_fatal(self, fake_fastembed):
|
||||
_, instance = fake_fastembed
|
||||
instance.model = None
|
||||
EmbeddingsWrapper(GRANITE_97M.name)
|
||||
@@ -105,10 +105,31 @@ class TestInputLimitResolution:
|
||||
from application.vectorstore import base
|
||||
|
||||
monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", None)
|
||||
monkeypatch.setattr(base, "_embeddings_name_is_explicit", lambda: True)
|
||||
assert self._remote("granite-311m")._resolve_input_limit() == 32768
|
||||
# mpnet genuinely stops at 384; sending more is paid for and discarded.
|
||||
assert self._remote("all-mpnet-base-v2")._resolve_input_limit() == 384
|
||||
|
||||
def test_default_model_name_lends_the_server_no_ceiling(self, monkeypatch):
|
||||
"""An unset EMBEDDINGS_NAME describes nothing about the remote server.
|
||||
|
||||
The name is only forwarded as the ``model`` field. Letting the
|
||||
settings default contribute mpnet's 384-token window would clip every
|
||||
chunk on a server that may well serve a 32k-context model.
|
||||
"""
|
||||
from application.vectorstore import base
|
||||
|
||||
monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", None)
|
||||
monkeypatch.setattr(base, "_embeddings_name_is_explicit", lambda: False)
|
||||
assert self._remote("all-mpnet-base-v2")._resolve_input_limit() is None
|
||||
|
||||
def test_explicit_setting_still_wins_over_an_unset_name(self, monkeypatch):
|
||||
from application.vectorstore import base
|
||||
|
||||
monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", 512)
|
||||
monkeypatch.setattr(base, "_embeddings_name_is_explicit", lambda: False)
|
||||
assert self._remote("all-mpnet-base-v2")._resolve_input_limit() == 512
|
||||
|
||||
def test_unknown_model_stays_unlimited(self, monkeypatch):
|
||||
from application.vectorstore import base
|
||||
|
||||
@@ -119,6 +140,7 @@ class TestInputLimitResolution:
|
||||
from application.vectorstore import base
|
||||
|
||||
monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", 0)
|
||||
monkeypatch.setattr(base, "_embeddings_name_is_explicit", lambda: True)
|
||||
assert self._remote("granite-97m")._resolve_input_limit() == 32768
|
||||
|
||||
def test_dimension_is_taken_from_the_registry(self):
|
||||
|
||||
Reference in new issue
Block a user