Skip to content

Commit

Permalink
fix: use json as default output and parse embeddings
Browse files Browse the repository at this point in the history
  • Loading branch information
henomis committed Sep 10, 2024
1 parent 3930e2b commit 6fc5365
Showing 1 changed file with 23 additions and 12 deletions.
35 changes: 23 additions & 12 deletions embedder/llamacpp/llamacpp.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,10 @@ package llamacppembedder

import (
"context"
"encoding/json"
"errors"
"os"
"os/exec"
"strconv"
"strings"

"github.com/henomis/lingoose/embedder"
)
Expand All @@ -16,6 +16,16 @@ type LlamaCppEmbedder struct {
modelPath string
}

type output struct {
Object string `json:"object"`
Data []data `json:"data"`
}
type data struct {
Object string `json:"object"`
Index int `json:"index"`
Embedding []float64 `json:"embedding"`
}

func New() *LlamaCppEmbedder {
return &LlamaCppEmbedder{
llamacppPath: "./llama.cpp/embedding",
Expand Down Expand Up @@ -61,7 +71,7 @@ func (l *LlamaCppEmbedder) embed(ctx context.Context, text string) (embedder.Emb
return nil, err
}

llamacppArgs := []string{"-m", l.modelPath, "-p", text}
llamacppArgs := []string{"-m", l.modelPath, "--embd-output-format", "json", "-p", text}
llamacppArgs = append(llamacppArgs, l.llamacppArgs...)

//nolint:gosec
Expand All @@ -74,14 +84,15 @@ func (l *LlamaCppEmbedder) embed(ctx context.Context, text string) (embedder.Emb
}

func parseEmbeddings(str string) (embedder.Embedding, error) {
strSlice := strings.Split(strings.TrimSpace(str), " ")
floatSlice := make([]float64, len(strSlice))
for i, s := range strSlice {
f, err := strconv.ParseFloat(s, 64)
if err != nil {
return nil, err
}
floatSlice[i] = f
var output output

Check failure on line 87 in embedder/llamacpp/llamacpp.go

View workflow job for this annotation

GitHub Actions / lint

shadow: declaration of "output" shadows declaration at line 19 (govet)
err := json.Unmarshal([]byte(str), &output)
if err != nil {
return nil, err
}
return floatSlice, nil

if len(output.Data) != 1 {
return nil, errors.New("no embeddings found")
}

return output.Data[0].Embedding, nil
}

0 comments on commit 6fc5365

Please sign in to comment.