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

@@ -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 && \
+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/all-MiniLM-L6-v2-ONNX/resolve/main/onnx/model.onnx_data && \
+    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/all-MiniLM-L6-v2-ONNX/resolve/main/tokenizer.json && \
+    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/*
 

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)
+bge-small-en-v1.5 uses **asymmetric** embeddings: queries and passages are
-- **Fast inference**: Uses ONNX Runtime for efficient model execution
+embedded differently, which is what lets a short search query align with a
-- **API authentication**: Optional API secret protection
+package description. To support this, the request carries a `type`:
-- **Pre-converted ONNX**: Uses ready-to-use ONNX model from HuggingFace
 
-## 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:
+The build downloads the pre-converted ONNX model
-1. Download the pre-converted ONNX model from `onnx-community/all-MiniLM-L6-v2-ONNX` on HuggingFace
+(`onnx-community/bge-small-en-v1.5-ONNX`) and the ONNX Runtime shared library,
-2. Install only the runtime dependencies (ONNX Runtime + NumPy)
+then compiles the Go binary.
-3. Create a minimal image (~373 MB)
 
-### Run the service
+## Run
-
-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:
 
 ```bash
-curl -X POST http://localhost:8080/vector \
+docker run --rm -p 8080:8080 -e API_SECRET=your-secret-key vector-service
-  -H "Content-Type: application/json" \
-  -H "Authorization: Bearer your-secret-key" \
-  -d '{"text": "Hello World"}'
 ```
 
-Response:
+## Technical details
-```json
-{
-  "vector": [0.1868536774709355, 0.8120285351760685, ...]
-}
-```
 
-The vector has 384 dimensions.
+- **Model**: bge-small-en-v1.5 (ONNX), 384 dimensions
-
+- **Pooling/normalization**: the export's `sentence_embedding` output is
-### Without authentication
+  mean-pooled and L2-normalized
-
+- **Runtime**: Go + ONNX Runtime (no Python)
-If you didn't set `API_SECRET`, you can query without the Authorization header:
+- **Files**: `main.go` (server), `Dockerfile`
-
-```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

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
+	// Create output tensor: "sentence_embedding" has shape [1, 384].
-	outputShape := onnxruntime_go.NewShape(1, maxLen, embeddingSize)
+	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")