chaas/vector-service
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages 1 Wiki Activity Actions

Switch to bge-small-en-v1.5 for asymmetric query/passage embeddings
All checks were successful
Build and push container image / build-and-push (push) Successful in 4m10s
Details

Browse source 
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.
This commit is contained in:
Christoph Haas 2026-09-13 22:10:01 +02:00
parent 1b4d24cada
commit 111af2d66e
3 changed files with 68 additions and 110 deletions
Download patch file Download diff file Expand all files Collapse all files

50
main.go
View file

@ -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")