From 921f4b63b813c93cd65ec3c1c8f44cbc5aa37c96 Mon Sep 17 00:00:00 2001 From: Gellert Ilya Date: Fri, 10 Apr 2026 19:25:08 +0300 Subject: [PATCH] fix: change output names, change work with tensors and add close to main.go --- backend/cmd/server/main.go | 2 ++ backend/internal/tokenizer/embedder.go | 21 +++++++++++++++++---- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/backend/cmd/server/main.go b/backend/cmd/server/main.go index f1ada00..3ecf9e0 100644 --- a/backend/cmd/server/main.go +++ b/backend/cmd/server/main.go @@ -80,6 +80,8 @@ func main() { logger.Fatal().Err(err).Msg("Failed to create embedded tokenizer") } + defer embeder.Close() + cardService := service.NewCardService(cardRepo, algoEngine, embeder, &logger) cardHandler := http.NewHandler(cardService, &logger) diff --git a/backend/internal/tokenizer/embedder.go b/backend/internal/tokenizer/embedder.go index c187b23..9cae0df 100644 --- a/backend/internal/tokenizer/embedder.go +++ b/backend/internal/tokenizer/embedder.go @@ -34,7 +34,7 @@ func NewOnnxEmbedder(tokenizerPath, onnxModelPath, onnxLibraryPath string, log * } inputNames := []string{"input_ids", "attention_mask", "token_type_ids"} - outputNames := []string{"embeddings"} + outputNames := []string{"last_hidden_state"} session, err := ort.NewDynamicAdvancedSession( onnxModelPath, @@ -128,8 +128,9 @@ func (e *OnnxEmbedder) Encode(text string) ([]float32, error) { } }(typeIDsTensor) - outShape := ort.NewShape(1, 384) - outData := make([]float32, 384) + outShape := ort.NewShape(1, seqLen, 384) + outData := make([]float32, seqLen*384) + outTensor, err := ort.NewTensor(outShape, outData) if err != nil { return nil, err @@ -149,7 +150,19 @@ func (e *OnnxEmbedder) Encode(text string) ([]float32, error) { return nil, fmt.Errorf("ONNX inference failed: %w", err) } - return outData, nil + pooled := make([]float32, 384) + + for i := int64(0); i < seqLen; i++ { + for j := int64(0); j < 384; j++ { + pooled[j] += outData[i*384+j] + } + } + + for j := 0; j < 384; j++ { + pooled[j] /= float32(seqLen) + } + + return pooled, nil } func (e *OnnxEmbedder) CosineSimilarity(v1, v2 []float32) float32 {