Fit embeddings to Vectorize dimension

app.py

@@ -16,6 +16,7 @@ AGENT_TOKEN = os.environ.get("JIMS_RENDER_AGENT_TOKEN") or os.environ.get("JIMS_
 MODEL_NAME = os.environ.get("JIMS_EMBEDDING_MODEL", "intfloat/multilingual-e5-small")
 ARTIFACT_ID = os.environ.get("JIMS_ACTIVE_ARTIFACT_ID", "hf_space_encoder")
 HASH_FALLBACK = os.environ.get("JIMS_EMBEDDING_HASH_FALLBACK_ENABLED", "true").lower() == "true"
+TARGET_DIMENSIONS = max(int(os.environ.get("JIMS_EMBEDDING_DIMENSIONS", "768") or "768"), 1)
 
 model = None
 model_error = ""
@@ -45,14 +46,23 @@ def normalize(values: list[float]) -> list[float]:
     return [round(value / norm, 8) for value in values]
 
 
-def hash_embed(text: str, dimensions: int = 768) -> list[float]:
+def fit_dimensions(values: list[float]) -> list[float]:
+    if len(values) == TARGET_DIMENSIONS:
+        return normalize(values)
+    if len(values) > TARGET_DIMENSIONS:
+        return normalize(values[:TARGET_DIMENSIONS])
+    return normalize([*values, *([0.0] * (TARGET_DIMENSIONS - len(values)))])
+
+
+def hash_embed(text: str) -> list[float]:
+    dimensions = TARGET_DIMENSIONS
     vector = [0.0] * dimensions
     for token in text.lower().split():
         digest = hashlib.sha256(token.encode("utf-8")).digest()
         idx = int.from_bytes(digest[:2], "big") % dimensions
         sign = 1.0 if digest[2] % 2 == 0 else -1.0
         vector[idx] += sign
-    return normalize(vector)
+    return fit_dimensions(vector)
 
 
 def prefixed(text: str, purpose: str) -> str:
@@ -84,7 +94,7 @@ def embed_texts(texts: list[str], purpose: str) -> tuple[list[list[float]], bool
         return [hash_embed(text) for text in texts], True, "hash_fallback"
     try:
         vectors = loaded.encode([prefixed(text[:16000], purpose) for text in texts], normalize_embeddings=True).tolist()
-        return [[float(value) for value in vector] for vector in vectors], False, MODEL_NAME
+        return [fit_dimensions([float(value) for value in vector]) for vector in vectors], False, MODEL_NAME
     except Exception as exc:
         if not HASH_FALLBACK:
             raise HTTPException(status_code=500, detail=str(exc)) from exc
@@ -98,7 +108,13 @@ def root() -> dict[str, str]:
 
 @app.get("/health")
 def health() -> dict[str, Any]:
-    return {"status": "ok", "service": "embedding-service", "model": MODEL_NAME, "model_loaded": model is not None}
+    return {
+        "status": "ok",
+        "service": "embedding-service",
+        "model": MODEL_NAME,
+        "dimension": TARGET_DIMENSIONS,
+        "model_loaded": model is not None,
+    }
 
 
 @app.get("/ready")
@@ -143,7 +159,13 @@ def encode(payload: dict[str, Any]) -> dict[str, Any]:
 
 @app.get("/v1/artifact/current")
 def current_artifact() -> dict[str, Any]:
-    return {"artifact_id": ARTIFACT_ID, "model": MODEL_NAME, "loaded": model is not None, "error": model_error}
+    return {
+        "artifact_id": ARTIFACT_ID,
+        "model": MODEL_NAME,
+        "dimension": TARGET_DIMENSIONS,
+        "loaded": model is not None,
+        "error": model_error,
+    }
 
 
 @app.post("/v1/reload-artifact", dependencies=[Depends(verify_token)])