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,
+        )
     )