From 111af2d66e5430d71bfd096342e9bff16e039988 Mon Sep 17 00:00:00 2001 From: Christoph Haas Date: Sun, 13 Sep 2026 22:10:01 +0200 Subject: [PATCH] Switch to bge-small-en-v1.5 for asymmetric query/passage embeddings Replace all-MiniLM-L6-v2 with bge-small-en-v1.5, which separates short search queries from package descriptions via the retrieval instruction. Requests now default to query mode and opt out with type=passage for documents. Uses the export's pooled, L2-normalized sentence_embedding output instead of manual mean pooling. --- Dockerfile | 6 +-- README.md | 122 +++++++++++++++++++---------------------------------- main.go | 50 ++++++++++------------ 3 files changed, 68 insertions(+), 110 deletions(-) diff --git a/Dockerfile b/Dockerfile index dadbfbe..54914a4 100644 --- a/Dockerfile +++ b/Dockerfile @@ -41,9 +41,9 @@ RUN apt-get update && \ rm -rf /var/lib/apt/lists/* # Download model and tokenizer files -RUN curl -sL -o /app/model.onnx https://huggingface.co/onnx-community/all-MiniLM-L6-v2-ONNX/resolve/main/onnx/model.onnx && \ - curl -sL -o /app/model.onnx_data https://huggingface.co/onnx-community/all-MiniLM-L6-v2-ONNX/resolve/main/onnx/model.onnx_data && \ - curl -sL -o /app/tokenizer.json https://huggingface.co/onnx-community/all-MiniLM-L6-v2-ONNX/resolve/main/tokenizer.json && \ +RUN curl -sL -o /app/model.onnx https://huggingface.co/onnx-community/bge-small-en-v1.5-ONNX/resolve/main/onnx/model.onnx && \ + curl -sL -o /app/model.onnx_data https://huggingface.co/onnx-community/bge-small-en-v1.5-ONNX/resolve/main/onnx/model.onnx_data && \ + curl -sL -o /app/tokenizer.json https://huggingface.co/onnx-community/bge-small-en-v1.5-ONNX/resolve/main/tokenizer.json && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* diff --git a/README.md b/README.md index 48e03a2..066afe0 100644 --- a/README.md +++ b/README.md @@ -1,100 +1,64 @@ # Vector Service -A minimal Docker container that provides text embedding vectors using the all-MiniLM-L6-v2 model. The service accepts POST requests with text and returns a 384-dimensional embedding vector. +A minimal Go container that provides text embedding vectors using the +[bge-small-en-v1.5](https://huggingface.co/BAAI/bge-small-en-v1.5) model +(384 dimensions). The service accepts POST requests and returns an embedding +vector, ready to be consumed by pgvector. -## Features +## Model -- **Small image size**: ~373 MB (much smaller than typical Python-based solutions) -- **Fast inference**: Uses ONNX Runtime for efficient model execution -- **API authentication**: Optional API secret protection -- **Pre-converted ONNX**: Uses ready-to-use ONNX model from HuggingFace +bge-small-en-v1.5 uses **asymmetric** embeddings: queries and passages are +embedded differently, which is what lets a short search query align with a +package description. To support this, the request carries a `type`: -## Usage +- **default (omitted)** — treated as a **query**: the text is prefixed with + the bge retrieval instruction + (`Represent this sentence for searching relevant passages: `). This keeps + existing callers that send a bare `{"text": ...}` working unchanged. +- `"passage"` — embedded as-is (used for documents/descriptions). -### Build the container +## API + +`POST /vector` ```bash -podman build -t vector-service . +curl -X POST http://localhost:8080/vector \ + -H "Content-Type: application/json" \ + -d '{"text": "sqlite editor", "type": "query"}' ``` -Or with Docker: +```json +{ "vector": [0.023, -0.118, ...] } +``` + +Endpoints: + +- `POST /vector` — embed text. +- `GET /ok` — health check. +- `GET /version` — build revision. + +Optional Bearer auth via the `API_SECRET` environment variable. + +## Build ```bash docker build -t vector-service . ``` -The build process will: -1. Download the pre-converted ONNX model from `onnx-community/all-MiniLM-L6-v2-ONNX` on HuggingFace -2. Install only the runtime dependencies (ONNX Runtime + NumPy) -3. Create a minimal image (~373 MB) +The build downloads the pre-converted ONNX model +(`onnx-community/bge-small-en-v1.5-ONNX`) and the ONNX Runtime shared library, +then compiles the Go binary. -### Run the service - -Without authentication: -```bash -podman run --rm -p 8080:8080 -d vector-service -``` - -With API secret authentication: -```bash -podman run --rm -p 8080:8080 -e API_SECRET=your-secret-key -d vector-service -``` - -### Query the service - -Send a POST request with JSON body: +## Run ```bash -curl -X POST http://localhost:8080/vector \ - -H "Content-Type: application/json" \ - -H "Authorization: Bearer your-secret-key" \ - -d '{"text": "Hello World"}' +docker run --rm -p 8080:8080 -e API_SECRET=your-secret-key vector-service ``` -Response: -```json -{ - "vector": [0.1868536774709355, 0.8120285351760685, ...] -} -``` +## Technical details -The vector has 384 dimensions. - -### Without authentication - -If you didn't set `API_SECRET`, you can query without the Authorization header: - -```bash -curl -X POST http://localhost:8080/vector \ - -H "Content-Type: application/json" \ - -d '{"text": "Hello World"}' -``` - -## Files - -- `Dockerfile`: Single-stage build with pre-downloaded ONNX model -- `server.py`: HTTP server with ONNX inference - -## Technical Details - -### Model -- **Model**: all-MiniLM-L6-v2 (80 MB on disk) -- **Source**: Pre-converted ONNX from [onnx-community/all-MiniLM-L6-v2-ONNX](https://huggingface.co/onnx-community/all-MiniLM-L6-v2-ONNX) -- **Dimensions**: 384 -- **Format**: ONNX (pre-converted) - -### Dependencies -- Runtime: Python 3.11, ONNX Runtime, NumPy -- No build-time dependencies needed (uses pre-converted model) - -### Image Size Breakdown -- Model files (ONNX + ONNX data + tokenizer + vocab): ~95 MB -- Python runtime and dependencies: ~278 MB -- Total: ~373 MB - -## Notes - -- The build will download the pre-converted ONNX model from HuggingFace (~95 MB total) -- Much faster builds since no PyTorch or model conversion is needed -- The ONNX model includes an external data file (`model.onnx_data`) which is normal for larger models -- For production use, consider adding rate limiting and HTTPS +- **Model**: bge-small-en-v1.5 (ONNX), 384 dimensions +- **Pooling/normalization**: the export's `sentence_embedding` output is + mean-pooled and L2-normalized +- **Runtime**: Go + ONNX Runtime (no Python) +- **Files**: `main.go` (server), `Dockerfile` diff --git a/main.go b/main.go index e6e8a3f..10f0af4 100644 --- a/main.go +++ b/main.go @@ -23,6 +23,10 @@ var ( version = "dev" ) +// bge-small-en-v1.5 uses asymmetric query/passage embeddings: queries are +// prefixed with this instruction, passages (package descriptions) are not. +const queryInstruction = "Represent this sentence for searching relevant passages: " + func init() { // Initialize ONNX Runtime libPath := os.Getenv("ONNXRUNTIME_SHARED_LIBRARY_PATH") @@ -42,11 +46,12 @@ func init() { log.Fatalf("Failed to load tokenizer: %v", err) } - // Load ONNX model + // Load ONNX model. bge-small's export already provides a pooled, + // L2-normalized "sentence_embedding" output, so no manual pooling. session, err = onnxruntime_go.NewDynamicAdvancedSession( modelPath, []string{"input_ids", "attention_mask", "token_type_ids"}, - []string{"last_hidden_state"}, + []string{"sentence_embedding"}, nil) if err != nil { log.Fatalf("Failed to load ONNX model: %v", err) @@ -104,6 +109,7 @@ func vectorHandler(w http.ResponseWriter, r *http.Request) { var req struct { Text string `json:"text"` + Type string `json:"type"` } if err := json.NewDecoder(r.Body).Decode(&req); err != nil { http.Error(w, "Bad request", http.StatusBadRequest) @@ -115,8 +121,17 @@ func vectorHandler(w http.ResponseWriter, r *http.Request) { return } + // bge expects queries to be prefixed with the retrieval instruction; + // passages (package descriptions) are embedded as-is. The default is + // "query" so existing callers sending a bare {"text": ...} keep + // working unchanged; documents must opt out with type=passage. + text := req.Text + if req.Type != "passage" { + text = queryInstruction + text + } + // Encode text - inputIds, attentionMask, tokenTypeIds := encode(req.Text) + inputIds, attentionMask, tokenTypeIds := encode(text) // Create input tensors inputShape := onnxruntime_go.NewShape(1, maxLen) @@ -142,8 +157,8 @@ func vectorHandler(w http.ResponseWriter, r *http.Request) { } defer tokenTypeIdsTensor.Destroy() - // Create output tensor - outputShape := onnxruntime_go.NewShape(1, maxLen, embeddingSize) + // Create output tensor: "sentence_embedding" has shape [1, 384]. + outputShape := onnxruntime_go.NewShape(1, embeddingSize) outputTensor, err := onnxruntime_go.NewEmptyTensor[float32](outputShape) if err != nil { http.Error(w, fmt.Sprintf("Failed to create output tensor: %v", err), http.StatusInternalServerError) @@ -161,32 +176,11 @@ func vectorHandler(w http.ResponseWriter, r *http.Request) { return } - // Get embeddings + // The sentence embedding is already mean-pooled and L2-normalized. embeddings := outputTensor.GetData() - - // Mean pooling over sequence length (exclude padding) - var sum [384]float32 - count := 0 - for i := 0; i < int(maxLen); i++ { - if attentionMask[i] == 1 { - for j := 0; j < int(embeddingSize); j++ { - sum[j] += embeddings[i*int(embeddingSize)+j] - } - count++ - } - } - - var sentenceEmbedding [384]float32 - if count > 0 { - for j := 0; j < int(embeddingSize); j++ { - sentenceEmbedding[j] = sum[j] / float32(count) - } - } - - // Convert to slice for JSON vector := make([]float32, embeddingSize) for i := 0; i < int(embeddingSize); i++ { - vector[i] = sentenceEmbedding[i] + vector[i] = embeddings[i] } w.Header().Set("Content-Type", "application/json")