Merge branch 'minor-embeddings-changes' into 'develop'

Published embeddings image

See merge request baserow/baserow!3802

Author: bram2w

embeddings/Dockerfile

@@ -1,31 +1,30 @@
-FROM python:3.13-slim AS builder
+FROM python:3.11-slim-bookworm AS builder
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    git \
-    && rm -rf /var/lib/apt/lists/*
+    git && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /build
+COPY _download_model.py .
 
 RUN pip install --no-cache-dir \
-    optimum[onnxruntime]==1.27.0 \
-    transformers==4.53.0
+    huggingface_hub==0.31.0 \
+    transformers==4.53.0 && \
+    python _download_model.py && \
+    rm _download_model.py
 
-COPY _download_model.py /tmp/download_model.py
-RUN python /tmp/download_model.py && rm /tmp/download_model.py
+FROM python:3.11-slim-bookworm
 
-FROM python:3.13-slim
+COPY --from=builder /model /model
 
 RUN pip install --no-cache-dir \
-    optimum[onnxruntime]==1.27.0 \
+    onnxruntime==1.20.1 \
     transformers==4.53.0 \
     starlette==0.48.0 \
-    uvicorn==0.37.0 \
-    && rm -rf /root/.cache
-
-COPY --from=builder /model /model
-
-COPY app.py /app/app.py
+    uvicorn==0.37.0 && \
+    rm -rf /root/.cache
 
 WORKDIR /app
+COPY app.py .
 
 EXPOSE 80
-
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "80"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "80"]

embeddings/README.md

@@ -36,11 +36,29 @@ The service uses a multi-stage Docker build:
 -   **Batch support**: Process multiple texts in a single request
 -   **Health checks**: Built-in health endpoint for monitoring
 
+## Docker
+
+### Run locally
+
+```
+docker run -p 8080:80 baserow/embeddings:1.0.0
+```
+
+### Build for publish
+
+```
+docker buildx build \
+  --platform linux/amd64,linux/arm64 \
+  -t baserow/embeddings:1.0.0 \
+  -t baserow/embeddings:latest \
+  --push .
+```
+
 ## API
 
 ### Endpoints
 
-#### `POST /embed`
+#### `POST http://localhost:8080/embed`
 
 Generate embeddings for one or more texts.
 

embeddings/_download_model.py

@@ -1,18 +1,27 @@
 #!/usr/bin/env python3
 """Download and convert the embedding model to ONNX format."""
 
-from optimum.onnxruntime import ORTModelForFeatureExtraction
+from huggingface_hub import snapshot_download
 from transformers import AutoTokenizer
+import os
+import shutil
 
-MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+MODEL_ID = "optimum/all-MiniLM-L6-v2"
 OUTPUT_DIR = "/model"
 
-print(f"Downloading {MODEL_NAME} (ONNX format)...")
-model = ORTModelForFeatureExtraction.from_pretrained(MODEL_NAME)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+print(f"Downloading {MODEL_ID} from Hugging Face...")
+local_dir = snapshot_download(repo_id=MODEL_ID, allow_patterns=["*.onnx", "*.json", "*.txt"])
 
-print(f"Saving to {OUTPUT_DIR}...")
-model.save_pretrained(OUTPUT_DIR)
+print(f"Saving tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 tokenizer.save_pretrained(OUTPUT_DIR)
 
-print("Done!")
+print("Copying ONNX model files...")
+for item in os.listdir(local_dir):
+    src = os.path.join(local_dir, item)
+    dst = os.path.join(OUTPUT_DIR, item)
+    if os.path.isfile(src):
+        shutil.copy(src, dst)
+
+print("Done! ONNX model is ready at", OUTPUT_DIR)

embeddings/app.py

@@ -1,56 +1,54 @@
 from starlette.applications import Starlette
 from starlette.responses import JSONResponse
 from starlette.routing import Route
-from optimum.onnxruntime import ORTModelForFeatureExtraction
+import onnxruntime as ort
 from transformers import AutoTokenizer
 import numpy as np
+import os
 
-# Load ONNX model and tokenizer once at startup
-model = ORTModelForFeatureExtraction.from_pretrained("/model")
-tokenizer = AutoTokenizer.from_pretrained("/model")
+# Load ONNX model directly
+MODEL_DIR = "/model"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
 
+# Create inference session
+model_path = os.path.join(MODEL_DIR, "model.onnx")
+session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
 
 def mean_pooling(token_embeddings, attention_mask):
-    """Mean pooling to get sentence embeddings"""
     input_mask_expanded = np.expand_dims(attention_mask, -1)
-    input_mask_expanded = np.broadcast_to(
-        input_mask_expanded, token_embeddings.shape
-    ).astype(float)
-
+    input_mask_expanded = np.broadcast_to(input_mask_expanded, token_embeddings.shape).astype(float)
     sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
     sum_mask = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
-
     return sum_embeddings / sum_mask
 
-
 async def embed(request):
     try:
         data = await request.json()
-    except:
+    except Exception:
         return JSONResponse({"error": "Invalid JSON"}, status_code=400)
 
     if not data or "texts" not in data:
         return JSONResponse({"error": "Missing texts field"}, status_code=400)
 
     texts = data["texts"]
-
     if isinstance(texts, str):
         texts = [texts]
 
     encoded = tokenizer(texts, padding=True, truncation=True, return_tensors="np")
-    outputs = model(**encoded)
-    embeddings = mean_pooling(outputs.last_hidden_state, encoded["attention_mask"])
+    ort_inputs = {k: v for k, v in encoded.items()}
 
-    # Normalize embeddings
+    # Run model
+    ort_outputs = session.run(None, ort_inputs)
+    token_embeddings = ort_outputs[0]
+
+    embeddings = mean_pooling(token_embeddings, encoded["attention_mask"])
     embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
 
     return JSONResponse({"embeddings": embeddings.tolist()})
 
-
 async def health(request):
     return JSONResponse({"status": "healthy"})
 
-
 app = Starlette(
     routes=[
         Route("/embed", embed, methods=["POST"]),