feat: embed pooling

application/core/settings.py

@@ -62,6 +62,12 @@ class Settings(BaseSettings):
     # processes beat one many-threaded process on the same cores.
     EMBEDDINGS_THREADS: Optional[int] = None
     EMBEDDINGS_CACHE_DIR: Optional[str] = None  # where FastEmbed caches model artifacts
+    # How a local model turns token vectors into one vector, and whether the
+    # result is L2-normalised. Both are normally read from the model's own
+    # repository; set these only for a repository that declares neither, or to
+    # override what it declares. "cls" or "mean".
+    EMBEDDINGS_POOLING: Optional[str] = None
+    EMBEDDINGS_NORMALIZE: Optional[bool] = None
     GITHUB_INGEST_MAX_FILE_BYTES: int = 1048576  # skip repo blobs larger than this (0 = no cap)
     GITHUB_INGEST_MAX_WORKERS: int = 8  # parallel file fetches per GitHub repo ingest
     # Optional directory of operator-supplied model YAMLs, loaded after the

application/vectorstore/embeddings_local.py

@@ -14,11 +14,14 @@ fp32 index of the same model.
 
 Models are described in :mod:`application.vectorstore.model_registry`. A name
 the registry does not know is treated as a Hugging Face repository, which is
-what someone configuring an arbitrary model expects.
+what someone configuring an arbitrary model expects; how to run it is read
+from the repository itself rather than assumed.
 """
 
+import json
 import logging
 import threading
+from dataclasses import replace
 from typing import Any, List, Optional
 
 from application.core.settings import settings
@@ -32,10 +35,19 @@ logger = logging.getLogger(__name__)
 _registered: set = set()
 _register_lock = threading.Lock()
 
-# Assumed layout for a model the registry does not describe.
+# Last-resort layout for a repository that declares nothing about itself.
 _FALLBACK_ONNX_FILE = "onnx/model.onnx"
 _FALLBACK_POOLING = "mean"
 
+# Sentence-transformers records how a model turns token vectors into one
+# vector, and whether it normalises the result, as files in the repository.
+# Reading them is the difference between running a model and running something
+# that merely shares its weights: mean-pooling a CLS model returns vectors at
+# cosine ~0.95 to the correct ones -- close enough to look like it works, far
+# enough to degrade retrieval, and silent either way.
+_POOLING_CONFIG = "1_Pooling/config.json"
+_MODULES_CONFIG = "modules.json"
+
 
 def _pooling_type(pooling: str):
     """Map our ``"cls"``/``"mean"`` spelling onto FastEmbed's enum."""
@@ -82,6 +94,108 @@ def _register(model: EmbeddingModel) -> None:
         _registered.add(model.repo)
 
 
+def _read_repo_json(repo: str, filename: str) -> Optional[dict]:
+    """Fetch one small JSON from ``repo``, or ``None`` when it is not there.
+
+    Reads through the Hugging Face hub cache, so a warmed image finds it
+    offline. Every failure -- absent file, no network, malformed JSON -- is the
+    same answer to the caller: this repository does not tell us.
+    """
+    try:
+        from huggingface_hub import hf_hub_download
+
+        with open(hf_hub_download(repo_id=repo, filename=filename), encoding="utf-8") as handle:
+            return json.load(handle)
+    except Exception as exc:
+        logger.debug("No %s for %s (%s)", filename, repo, exc)
+        return None
+
+
+def _describe_from_repo(repo: str) -> Optional[EmbeddingModel]:
+    """Build a spec from a repository's sentence-transformers metadata.
+
+    Args:
+        repo: Hugging Face repository id.
+
+    Returns:
+        The described model, or ``None`` when the repository carries no
+        metadata to read -- leaving the caller to fall back to assumptions.
+
+    Raises:
+        RuntimeError: If the model has a Dense projection head. FastEmbed runs
+            the transformer and pools it, and nothing else, so the projection
+            would be skipped and the vectors come out both the wrong width and
+            in a different space. There is no correct way to run it here.
+    """
+    pooling_config = _read_repo_json(repo, _POOLING_CONFIG)
+    if pooling_config is None:
+        return None
+
+    kinds = {
+        str(module.get("type", "")).rsplit(".", 1)[-1]
+        for module in (_read_repo_json(repo, _MODULES_CONFIG) or [])
+        if isinstance(module, dict)
+    }
+    if "Dense" in kinds:
+        raise RuntimeError(
+            f"Embedding model {repo!r} has a Dense projection layer, which FastEmbed "
+            "cannot run: its vectors would be the wrong width and in a different "
+            "space than the model was trained to produce. Choose a model without "
+            "one, or serve this one over EMBEDDINGS_BASE_URL."
+        )
+
+    if pooling_config.get("pooling_mode_cls_token"):
+        pooling = "cls"
+    elif pooling_config.get("pooling_mode_mean_tokens"):
+        pooling = "mean"
+    else:
+        # max, mean_sqrt_len, weighted-mean: FastEmbed offers none of them, so
+        # there is nothing to describe and guessing is what we are avoiding.
+        logger.warning(
+            "Embedding model %s uses a pooling mode FastEmbed cannot reproduce (%s).",
+            repo,
+            ", ".join(sorted(k for k, v in pooling_config.items() if v is True)) or "unknown",
+        )
+        return None
+
+    # ``Normalize`` appears in modules.json only when the model L2-normalises.
+    # Its absence is a fact, not missing data: a dot-product model is trained
+    # on unnormalised vectors and normalising re-ranks its results.
+    normalize = "Normalize" in kinds
+    dimension = int(pooling_config.get("word_embedding_dimension") or 0)
+    logger.info(
+        "Embedding model %s declares %s pooling, normalize=%s, dimension=%s.",
+        repo,
+        pooling,
+        normalize,
+        dimension or "unknown",
+    )
+    return EmbeddingModel(
+        name=repo,
+        dimension=dimension,
+        max_input_tokens=512,
+        pooling=pooling,
+        normalize=normalize,
+        repo=repo,
+        onnx_file=_FALLBACK_ONNX_FILE,
+    )
+
+
+def _apply_overrides(spec: EmbeddingModel) -> EmbeddingModel:
+    """Let ``EMBEDDINGS_POOLING``/``EMBEDDINGS_NORMALIZE`` win over any source."""
+    pooling = getattr(settings, "EMBEDDINGS_POOLING", None)
+    normalize = getattr(settings, "EMBEDDINGS_NORMALIZE", None)
+    changes = {}
+    if isinstance(pooling, str) and pooling.strip().lower() in ("cls", "mean"):
+        changes["pooling"] = pooling.strip().lower()
+    if isinstance(normalize, bool):
+        changes["normalize"] = normalize
+    if not changes:
+        return spec
+    logger.info("Overriding %s from settings: %s", spec.repo, changes)
+    return replace(spec, **changes)
+
+
 def _spec_for(model_name: str) -> EmbeddingModel:
     """Registry entry for ``model_name``, or a best-effort one for a raw repo.
 
@@ -89,29 +203,39 @@ def _spec_for(model_name: str) -> EmbeddingModel:
         model_name: Configured ``EMBEDDINGS_NAME``.
 
     Returns:
-        A registry entry. For an unregistered name the returned entry assumes
-        the standard Hugging Face ONNX layout and mean pooling, and carries
-        ``dimension = 0`` so the caller knows to probe for the real width.
+        The registry entry, else one read from the repository's own
+        sentence-transformers metadata, else a last-resort entry assuming the
+        standard ONNX layout with mean pooling and ``dimension = 0`` so the
+        caller knows to probe for the real width. Settings overrides win over
+        all three.
     """
     spec = resolve(model_name)
     if spec is not None:
-        return spec
-    logger.info(
-        "Embedding model %r is not in the registry (known: %s); treating it as a "
-        "Hugging Face repository with %s and %s pooling.",
+        return _apply_overrides(spec)
+
+    described = _describe_from_repo(model_name)
+    if described is not None:
+        return _apply_overrides(described)
+
+    logger.warning(
+        "Embedding model %r is not in the registry (known: %s) and its repository "
+        "declares no pooling, so %s with L2 normalisation is assumed. If that is "
+        "wrong the vectors will be quietly poor rather than fail; set "
+        "EMBEDDINGS_POOLING=cls|mean and EMBEDDINGS_NORMALIZE to pin it.",
         model_name,
         ", ".join(known_names()),
-        _FALLBACK_ONNX_FILE,
         _FALLBACK_POOLING,
     )
-    return EmbeddingModel(
-        name=model_name,
-        dimension=0,
-        max_input_tokens=512,
-        pooling=_FALLBACK_POOLING,
-        normalize=True,
-        repo=model_name,
-        onnx_file=_FALLBACK_ONNX_FILE,
+    return _apply_overrides(
+        EmbeddingModel(
+            name=model_name,
+            dimension=0,
+            max_input_tokens=512,
+            pooling=_FALLBACK_POOLING,
+            normalize=True,
+            repo=model_name,
+            onnx_file=_FALLBACK_ONNX_FILE,
+        )
     )
 
 

docs/content/upgrading.mdx

@@ -50,6 +50,27 @@ The script processes sources independently: one failing source is reported and s
   Running [GraphRAG](/Sources/GraphRAG)? The script also rewrites `graph_nodes.name_embedding`, which seeds every graph traversal. Those vectors are written once at extraction time and share the chunk vectors' width, so leaving them in the old model's space degrades graph retrieval just as silently as the chunk vectors would — and needs no LLM re-extraction to fix.
 </Callout>
 
+### Custom local models
+
+A local model now runs through ONNX Runtime, so its repository must ship an ONNX
+export (`onnx/model.onnx`) or be one of FastEmbed's built-in models. Repositories
+with PyTorch weights only no longer load; `hkunlp/instructor-large`, previously
+supported by name, is one of them. Serve such a model over `EMBEDDINGS_BASE_URL`
+instead, or switch to a model with an export.
+
+How to run the model — pooling, and whether outputs are L2-normalised — is read
+from the repository's own `1_Pooling/config.json` and `modules.json`. Two cases
+need attention:
+
+- **Models with a Dense projection layer** (`sentence-transformers/LaBSE`,
+  `distiluse-base-multilingual-cased-v1`) are now **refused at startup**.
+  FastEmbed cannot apply the projection, so it would have produced vectors of the
+  wrong width in a different space. If you were running one, its stored vectors
+  were already wrong; move it to `EMBEDDINGS_BASE_URL` or pick another model.
+- **Repositories that declare nothing** fall back to mean pooling with
+  normalisation and log a warning. Pin the real values with `EMBEDDINGS_POOLING`
+  (`cls` or `mean`) and `EMBEDDINGS_NORMALIZE`.
+
 <Callout type="info">
   Staying on `all-mpnet-base-v2` is a supported choice — it remains in the model registry and in `setup.sh`. You only need this section if you want to move to granite.
 </Callout>

tests/vectorstore/test_embeddings_local.py

@@ -18,6 +18,19 @@ def _clear_registration():
     embeddings_local._registered.clear()
 
 
+@pytest.fixture(autouse=True)
+def _no_hub_reads():
+    """Keep unit tests off the network.
+
+    ``_spec_for`` now asks a repository how it pools; without this every test
+    naming an unregistered model would reach the Hugging Face hub. ``None`` is
+    the "declares nothing" answer, which is the behaviour these tests were
+    written against. Tests that exercise the metadata patch it themselves.
+    """
+    with patch.object(embeddings_local, "_read_repo_json", return_value=None):
+        yield
+
+
 @pytest.fixture
 def fake_fastembed():
     """Patch FastEmbed so no model is downloaded or run."""
@@ -263,3 +276,118 @@ class TestTokenizerPadding:
         _, instance = fake_fastembed
         instance.model = None
         EmbeddingsWrapper(GRANITE_97M.name)
+
+
+def _repo_json(pooling_file, modules_file):
+    """Stub ``_read_repo_json`` returning canned repository metadata."""
+
+    def read(repo, filename):
+        return pooling_file if filename == embeddings_local._POOLING_CONFIG else modules_file
+
+    return read
+
+
+class TestPoolingReadFromTheRepository:
+    """A model's pooling is a fact its repository states, not a default.
+
+    Assuming mean pooling for a CLS model returns vectors at cosine ~0.95 to
+    the correct ones: no error, no dimension mismatch, just quietly worse
+    retrieval. These cover the shapes seen on the hub.
+    """
+
+    def test_cls_pooling_is_read_rather_than_assumed(self):
+        with patch.object(
+            embeddings_local,
+            "_read_repo_json",
+            _repo_json(
+                {"pooling_mode_cls_token": True, "word_embedding_dimension": 384},
+                [{"type": "sentence_transformers.models.Transformer"},
+                 {"type": "sentence_transformers.models.Pooling"},
+                 {"type": "sentence_transformers.models.Normalize"}],
+            ),
+        ):
+            spec = embeddings_local._spec_for("BAAI/bge-small-en-v1.5")
+        assert spec.pooling == "cls"
+        assert spec.normalize is True
+        # Declared width, so no probe run is needed to learn it.
+        assert spec.dimension == 384
+
+    def test_missing_normalize_module_means_unnormalised(self):
+        """multi-qa-mpnet-base-dot-v1 is trained on unnormalised vectors."""
+        with patch.object(
+            embeddings_local,
+            "_read_repo_json",
+            _repo_json(
+                {"pooling_mode_cls_token": True, "word_embedding_dimension": 768},
+                [{"type": "sentence_transformers.models.Transformer"},
+                 {"type": "sentence_transformers.models.Pooling"}],
+            ),
+        ):
+            spec = embeddings_local._spec_for("sentence-transformers/multi-qa-mpnet-base-dot-v1")
+        assert spec.pooling == "cls"
+        assert spec.normalize is False
+
+    def test_dense_projection_head_is_refused(self):
+        """FastEmbed would skip the projection and emit the wrong vectors."""
+        with patch.object(
+            embeddings_local,
+            "_read_repo_json",
+            _repo_json(
+                {"pooling_mode_cls_token": True, "word_embedding_dimension": 768},
+                [{"type": "sentence_transformers.models.Transformer"},
+                 {"type": "sentence_transformers.models.Pooling"},
+                 {"type": "sentence_transformers.models.Dense"},
+                 {"type": "sentence_transformers.models.Normalize"}],
+            ),
+        ):
+            with pytest.raises(RuntimeError) as excinfo:
+                embeddings_local._spec_for("sentence-transformers/LaBSE")
+        message = str(excinfo.value)
+        assert "LaBSE" in message
+        assert "Dense" in message
+
+    def test_unsupported_pooling_mode_falls_back_rather_than_lying(self):
+        with patch.object(
+            embeddings_local,
+            "_read_repo_json",
+            _repo_json({"pooling_mode_max_tokens": True}, []),
+        ):
+            spec = embeddings_local._spec_for("some-org/max-pooled")
+        assert spec.pooling == embeddings_local._FALLBACK_POOLING
+        assert spec.dimension == 0
+
+    def test_repository_without_metadata_keeps_the_assumption(self):
+        spec = embeddings_local._spec_for("some-org/plain-onnx-export")
+        assert spec.pooling == embeddings_local._FALLBACK_POOLING
+        assert spec.normalize is True
+        assert spec.dimension == 0
+
+    def test_registry_wins_over_the_repository(self):
+        """A described model is never re-read; the registry is the answer."""
+        read = MagicMock()
+        with patch.object(embeddings_local, "_read_repo_json", read):
+            spec = embeddings_local._spec_for(MPNET.name)
+        assert spec is MPNET
+        read.assert_not_called()
+
+
+class TestPoolingOverrides:
+    def test_settings_override_what_the_repository_declares(self):
+        with patch.object(
+            embeddings_local,
+            "_read_repo_json",
+            _repo_json(
+                {"pooling_mode_mean_tokens": True, "word_embedding_dimension": 768},
+                [{"type": "sentence_transformers.models.Normalize"}],
+            ),
+        ):
+            with patch.object(embeddings_local.settings, "EMBEDDINGS_POOLING", "cls"), \
+                 patch.object(embeddings_local.settings, "EMBEDDINGS_NORMALIZE", False):
+                spec = embeddings_local._spec_for("some-org/mislabelled")
+        assert spec.pooling == "cls"
+        assert spec.normalize is False
+
+    def test_a_meaningless_override_is_ignored(self):
+        with patch.object(embeddings_local.settings, "EMBEDDINGS_POOLING", "banana"):
+            spec = embeddings_local._spec_for("some-org/plain-onnx-export")
+        assert spec.pooling == embeddings_local._FALLBACK_POOLING

tests/vectorstore/test_model_registry.py

@@ -1,5 +1,7 @@
 """The registry is the single source of truth for embedding-model facts."""
 
+import os
+
 import pytest
 
 from application.vectorstore import model_registry as reg
@@ -86,3 +88,30 @@ class TestHelpers:
             for key in (model.name, *model.aliases):
                 assert key.lower() not in seen, f"{key} claimed twice"
                 seen[key.lower()] = model
+
+
+class TestRegistryMatchesTheHub:
+    """Every registry entry restates facts the model's repository already holds.
+
+    Restating them is what makes an offline install work, but a restatement can
+    drift from its source and nothing else would notice: wrong pooling is not a
+    crash, only worse retrieval. Opt in with ``DOCSGPT_HUB_TESTS=1``; needs
+    network.
+    """
+
+    @pytest.mark.skipif(
+        os.environ.get("DOCSGPT_HUB_TESTS") != "1",
+        reason="set DOCSGPT_HUB_TESTS=1 to check the registry against the hub",
+    )
+    @pytest.mark.parametrize(
+        "model", [m for m in reg.MODELS if m.provider == "fastembed"]
+    )
+    def test_entry_matches_repository_metadata(self, model):
+        from application.vectorstore.embeddings_local import _describe_from_repo
+
+        described = _describe_from_repo(model.repo)
+        assert described is not None, f"{model.repo} declares no pooling metadata"
+        assert described.pooling == model.pooling
+        assert described.normalize == model.normalize
+        if described.dimension:
+            assert described.dimension == model.dimension