upgrade to nomic-v2-moe + dimensions support

app.py

@@ -1,13 +1,15 @@
 """Embedding Server (sentence-transformers) for HuggingFace Spaces."""
 import os
+import numpy as np
 from fastapi import FastAPI
 from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer
 
-MODEL_NAME = os.environ.get("MODEL_NAME", "intfloat/multilingual-e5-small")
+MODEL_NAME = os.environ.get("MODEL_NAME", "nomic-ai/nomic-embed-text-v2-moe")
 print(f"[Embedding] Loading model: {MODEL_NAME}...", flush=True)
-model = SentenceTransformer(MODEL_NAME)
-print("[Embedding] Model loaded.", flush=True)
+model = SentenceTransformer(MODEL_NAME, trust_remote_code=True)
+NATIVE_DIMS = model.get_sentence_embedding_dimension()
+print(f"[Embedding] Model loaded. Native dimensions: {NATIVE_DIMS}", flush=True)
 
 app = FastAPI()
 
@@ -17,33 +19,52 @@ class EmbedRequest(BaseModel):
     texts: list[str] | None = None
     model: str | None = None
     normalize: bool = True
-    prefix: str | None = None  # e5 models require "query: " or "passage: " prefix
+    prefix: str | None = None
+    dimensions: int | None = None
+
+
+def _process_embeddings(embeddings: np.ndarray, dimensions: int | None) -> np.ndarray:
+    """Truncate to target dimensions and re-normalize."""
+    if dimensions and dimensions < embeddings.shape[1]:
+        embeddings = embeddings[:, :dimensions]
+        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
+        embeddings = embeddings / norms
+    return embeddings
+
+
+def _encode(input_texts: list[str], req: EmbedRequest) -> dict:
+    if req.prefix:
+        input_texts = [req.prefix + t for t in input_texts]
+    embeddings = model.encode(input_texts, convert_to_numpy=True,
+                              normalize_embeddings=req.normalize)
+    embeddings = _process_embeddings(embeddings, req.dimensions)
+    return {
+        "embeddings": embeddings.tolist(),
+        "model": MODEL_NAME,
+        "dimensions": embeddings.shape[1],
+        "tokens": len(input_texts) * 32,
+    }
 
 
 @app.get("/health")
 def health():
-    return {"status": "ok", "model": MODEL_NAME}
+    return {
+        "status": "ok",
+        "model": MODEL_NAME,
+        "model_name": MODEL_NAME,
+        "native_dimensions": NATIVE_DIMS,
+    }
 
 
 @app.post("/embed")
 def embed(req: EmbedRequest):
-    # Accept both "text" (single/list) and "texts" (list) fields
     if req.texts:
         input_texts = req.texts
     elif req.text:
         input_texts = [req.text] if isinstance(req.text, str) else req.text
     else:
         return {"error": "Provide 'text' or 'texts' field"}, 400
-    # Apply prefix if provided (e5 models require "query: " or "passage: ")
-    if req.prefix:
-        input_texts = [req.prefix + t for t in input_texts]
-    embeddings = model.encode(input_texts, normalize_embeddings=req.normalize)
-    return {
-        "embeddings": embeddings.tolist(),
-        "model": MODEL_NAME,
-        "dimensions": embeddings.shape[1],
-        "tokens": len(input_texts) * 32,
-    }
+    return _encode(input_texts, req)
 
 
 @app.post("/embed_batch")
@@ -54,12 +75,4 @@ def embed_batch(req: EmbedRequest):
         input_texts = [req.text] if isinstance(req.text, str) else req.text
     else:
         return {"error": "Provide 'text' or 'texts' field"}, 400
-    if req.prefix:
-        input_texts = [req.prefix + t for t in input_texts]
-    embeddings = model.encode(input_texts, normalize_embeddings=req.normalize)
-    return {
-        "embeddings": embeddings.tolist(),
-        "model": MODEL_NAME,
-        "dimensions": embeddings.shape[1],
-        "tokens": len(input_texts) * 32,
-    }
+    return _encode(input_texts, req)