Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve query sorting #53

Merged
merged 2 commits into from
Mar 17, 2024
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Improve query sorting
slices.SortFunc is recommended over sort.Slice and more performant
philippgille committed Mar 17, 2024
commit 31ee6bf325c184b44aa0383d7271c1c0870a6779
24 changes: 12 additions & 12 deletions collection.go
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
package chromem

import (
"cmp"
"context"
"errors"
"fmt"
"os"
"path/filepath"
"slices"
"sort"
"sync"
)

@@ -324,29 +324,29 @@ func (c *Collection) QueryEmbedding(ctx context.Context, queryEmbedding []float3
}

// For the remaining documents, calculate cosine similarity.
docSim, err := calcDocSimilarity(ctx, queryEmbedding, filteredDocs)
docSims, err := calcDocSimilarity(ctx, queryEmbedding, filteredDocs)
if err != nil {
return nil, fmt.Errorf("couldn't calculate cosine similarity: %w", err)
}

// Sort by similarity
sort.Slice(docSim, func(i, j int) bool {
// The `less` function would usually use `<`, but we want to sort descending.
return docSim[i].similarity > docSim[j].similarity
slices.SortFunc(docSims, func(i, j docSim) int {
// i, j; for descending order
return cmp.Compare(j.similarity, i.similarity)
})

// Return the top nResults or len(docSim), whichever is smaller
if len(docSim) < nResults {
nResults = len(docSim)
if len(docSims) < nResults {
nResults = len(docSims)
}
res := make([]Result, 0, nResults)
for i := 0; i < nResults; i++ {
res = append(res, Result{
ID: docSim[i].docID,
Metadata: c.documents[docSim[i].docID].Metadata,
Embedding: c.documents[docSim[i].docID].Embedding,
Content: c.documents[docSim[i].docID].Content,
Similarity: docSim[i].similarity,
ID: docSims[i].docID,
Metadata: c.documents[docSims[i].docID].Metadata,
Embedding: c.documents[docSims[i].docID].Embedding,
Content: c.documents[docSims[i].docID].Content,
Similarity: docSims[i].similarity,
})
}