Add "normalized" parameter to skip check if normalization is known

Author: philippgille

README.md

@@ -98,7 +98,7 @@ func main() {
     db := chromem.NewDB()
 
     // Create collection. GetCollection, GetOrCreateCollection, DeleteCollection also available!
-    collection, _ := db.CreateCollection("all-my-documents", nil, nil)
+    collection, _ := db.CreateCollection("all-my-documents", nil, nil, nil)
 
     // Add docs to the collection. Update and delete will be added in the future.
     // Can be multi-threaded with AddConcurrently()!

collection.go

@@ -22,11 +22,12 @@ type Collection struct {
 	documents        map[string]*Document
 	documentsLock    sync.RWMutex
 	embed            EmbeddingFunc
+	normalized       *bool
 }
 
 // We don't export this yet to keep the API surface to the bare minimum.
 // Users create collections via [Client.CreateCollection].
-func newCollection(name string, metadata map[string]string, embed EmbeddingFunc, dir string) (*Collection, error) {
+func newCollection(name string, metadata map[string]string, embed EmbeddingFunc, normalized *bool, dir string) (*Collection, error) {
 	// We copy the metadata to avoid data races in case the caller modifies the
 	// map after creating the collection while we range over it.
 	m := make(map[string]string, len(metadata))
@@ -37,9 +38,10 @@ func newCollection(name string, metadata map[string]string, embed EmbeddingFunc,
 	c := &Collection{
 		Name: name,
 
-		metadata:  m,
-		documents: make(map[string]*Document),
-		embed:     embed,
+		metadata:   m,
+		documents:  make(map[string]*Document),
+		embed:      embed,
+		normalized: normalized,
 	}
 
 	// Persistence
@@ -301,7 +303,7 @@ func (c *Collection) Query(ctx context.Context, queryText string, nResults int,
 	}
 
 	// For the remaining documents, calculate cosine similarity.
-	docSim, err := calcDocSimilarity(ctx, queryVectors, filteredDocs)
+	docSim, err := calcDocSimilarity(ctx, queryVectors, filteredDocs, c.normalized)
 	if err != nil {
 		return nil, fmt.Errorf("couldn't calculate cosine similarity: %w", err)
 	}

collection_test.go

@@ -20,7 +20,7 @@ func TestCollection_Add(t *testing.T) {
 
 	// Create collection
 	db := NewDB()
-	c, err := db.CreateCollection(name, metadata, embeddingFunc)
+	c, err := db.CreateCollection(name, metadata, embeddingFunc, nil)
 	if err != nil {
 		t.Fatal("expected no error, got", err)
 	}
@@ -116,7 +116,7 @@ func TestCollection_Add_Error(t *testing.T) {
 
 	// Create collection
 	db := NewDB()
-	c, err := db.CreateCollection(name, metadata, embeddingFunc)
+	c, err := db.CreateCollection(name, metadata, embeddingFunc, nil)
 	if err != nil {
 		t.Fatal("expected no error, got", err)
 	}
@@ -167,7 +167,7 @@ func TestCollection_AddConcurrently(t *testing.T) {
 
 	// Create collection
 	db := NewDB()
-	c, err := db.CreateCollection(name, metadata, embeddingFunc)
+	c, err := db.CreateCollection(name, metadata, embeddingFunc, nil)
 	if err != nil {
 		t.Fatal("expected no error, got", err)
 	}
@@ -263,7 +263,7 @@ func TestCollection_AddConcurrently_Error(t *testing.T) {
 
 	// Create collection
 	db := NewDB()
-	c, err := db.CreateCollection(name, metadata, embeddingFunc)
+	c, err := db.CreateCollection(name, metadata, embeddingFunc, nil)
 	if err != nil {
 		t.Fatal("expected no error, got", err)
 	}
@@ -316,7 +316,7 @@ func TestCollection_Count(t *testing.T) {
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
 		return []float32{-0.1, 0.1, 0.2}, nil
 	}
-	c, err := db.CreateCollection(name, metadata, embeddingFunc)
+	c, err := db.CreateCollection(name, metadata, embeddingFunc, nil)
 	if err != nil {
 		t.Fatal("expected no error, got", err)
 	}
@@ -407,7 +407,7 @@ func benchmarkCollection_Query(b *testing.B, n int, withContent bool) {
 	// Create collection
 	db := NewDB()
 	name := "test"
-	c, err := db.CreateCollection(name, nil, embeddingFunc)
+	c, err := db.CreateCollection(name, nil, embeddingFunc, &trueVal)
 	if err != nil {
 		b.Fatal("expected no error, got", err)
 	}

db.go

@@ -142,14 +142,17 @@ func NewPersistentDB(path string) (*DB, error) {
 //   - metadata: Optional metadata to associate with the collection.
 //   - embeddingFunc: Optional function to use to embed documents.
 //     Uses the default embedding function if not provided.
-func (db *DB) CreateCollection(name string, metadata map[string]string, embeddingFunc EmbeddingFunc) (*Collection, error) {
+//   - normalized: Optional flag to indicate if the embeddings of the collection
+//     are normalized (when you add embeddings yourself, or the embeddings created
+//     by the embeddingFunc). If nil it will be autodetected.
+func (db *DB) CreateCollection(name string, metadata map[string]string, embeddingFunc EmbeddingFunc, normalized *bool) (*Collection, error) {
 	if name == "" {
 		return nil, errors.New("collection name is empty")
 	}
 	if embeddingFunc == nil {
 		embeddingFunc = NewEmbeddingFuncDefault()
 	}
-	collection, err := newCollection(name, metadata, embeddingFunc, db.persistDirectory)
+	collection, err := newCollection(name, metadata, embeddingFunc, normalized, db.persistDirectory)
 	if err != nil {
 		return nil, fmt.Errorf("couldn't create collection: %w", err)
 	}
@@ -213,12 +216,15 @@ func (db *DB) GetCollection(name string, embeddingFunc EmbeddingFunc) *Collectio
 //   - metadata: Optional metadata to associate with the collection.
 //   - embeddingFunc: Optional function to use to embed documents.
 //     Uses the default embedding function if not provided.
-func (db *DB) GetOrCreateCollection(name string, metadata map[string]string, embeddingFunc EmbeddingFunc) (*Collection, error) {
+//   - normalized: Optional flag to indicate if the embeddings of the collection
+//     are normalized (when you add embeddings yourself, or the embeddings created
+//     by the embeddingFunc). If nil it will be autodetected.
+func (db *DB) GetOrCreateCollection(name string, metadata map[string]string, embeddingFunc EmbeddingFunc, normalized *bool) (*Collection, error) {
 	// No need to lock here, because the methods we call do that.
 	collection := db.GetCollection(name, embeddingFunc)
 	if collection == nil {
 		var err error
-		collection, err = db.CreateCollection(name, metadata, embeddingFunc)
+		collection, err = db.CreateCollection(name, metadata, embeddingFunc, normalized)
 		if err != nil {
 			return nil, fmt.Errorf("couldn't create collection: %w", err)
 		}

db_test.go

@@ -18,7 +18,7 @@ func TestDB_CreateCollection(t *testing.T) {
 	db := NewDB()
 
 	t.Run("OK", func(t *testing.T) {
-		c, err := db.CreateCollection(name, metadata, embeddingFunc)
+		c, err := db.CreateCollection(name, metadata, embeddingFunc, nil)
 		if err != nil {
 			t.Fatal("expected no error, got", err)
 		}
@@ -70,7 +70,7 @@ func TestDB_CreateCollection(t *testing.T) {
 	})
 
 	t.Run("NOK - Empty name", func(t *testing.T) {
-		_, err := db.CreateCollection("", metadata, embeddingFunc)
+		_, err := db.CreateCollection("", metadata, embeddingFunc, nil)
 		if err == nil {
 			t.Fatal("expected error, got nil")
 		}
@@ -89,7 +89,7 @@ func TestDB_ListCollections(t *testing.T) {
 	// Create initial collection
 	db := NewDB()
 	// We ignore the return value. CreateCollection is tested elsewhere.
-	_, err := db.CreateCollection(name, metadata, embeddingFunc)
+	_, err := db.CreateCollection(name, metadata, embeddingFunc, nil)
 	if err != nil {
 		t.Fatal("expected no error, got", err)
 	}
@@ -155,7 +155,7 @@ func TestDB_GetCollection(t *testing.T) {
 	// Create initial collection
 	db := NewDB()
 	// We ignore the return value. CreateCollection is tested elsewhere.
-	_, err := db.CreateCollection(name, metadata, embeddingFunc)
+	_, err := db.CreateCollection(name, metadata, embeddingFunc, nil)
 	if err != nil {
 		t.Fatal("expected no error, got", err)
 	}
@@ -207,15 +207,15 @@ func TestDB_GetOrCreateCollection(t *testing.T) {
 		// Create collection so that the GetOrCreateCollection() call below only
 		// gets it.
 		// We ignore the return value. CreateCollection is tested elsewhere.
-		_, err := db.CreateCollection(name, metadata, embeddingFunc)
+		_, err := db.CreateCollection(name, metadata, embeddingFunc, nil)
 		if err != nil {
 			t.Fatal("expected no error, got", err)
 		}
 
 		// Call GetOrCreateCollection() with the same name to only get it. We pass
 		// nil for the metadata and embeddingFunc so we can check that the returned
 		// collection is the original one, and not a new one.
-		c, err := db.GetOrCreateCollection(name, nil, nil)
+		c, err := db.GetOrCreateCollection(name, nil, embeddingFunc, nil)
 		if err != nil {
 			t.Fatal("expected no error, got", err)
 		}
@@ -257,7 +257,7 @@ func TestDB_GetOrCreateCollection(t *testing.T) {
 		db := NewDB()
 
 		// Call GetOrCreateCollection()
-		c, err := db.GetOrCreateCollection(name, metadata, embeddingFunc)
+		c, err := db.GetOrCreateCollection(name, metadata, embeddingFunc, nil)
 		if err != nil {
 			t.Fatal("expected no error, got", err)
 		}
@@ -307,7 +307,7 @@ func TestDB_DeleteCollection(t *testing.T) {
 	// Create initial collection
 	db := NewDB()
 	// We ignore the return value. CreateCollection is tested elsewhere.
-	_, err := db.CreateCollection(name, metadata, embeddingFunc)
+	_, err := db.CreateCollection(name, metadata, embeddingFunc, nil)
 	if err != nil {
 		t.Fatal("expected no error, got", err)
 	}
@@ -339,7 +339,7 @@ func TestDB_Reset(t *testing.T) {
 	// Create initial collection
 	db := NewDB()
 	// We ignore the return value. CreateCollection is tested elsewhere.
-	_, err := db.CreateCollection(name, metadata, embeddingFunc)
+	_, err := db.CreateCollection(name, metadata, embeddingFunc, nil)
 	if err != nil {
 		t.Fatal("expected no error, got", err)
 	}

examples/rag-wikipedia-ollama/main.go

@@ -19,6 +19,9 @@ const (
 	embeddingModel = "nomic-embed-text"
 )
 
+// The nomic-embed-text-v1.5 model doesn't return normalized embeddings
+var normalized = false
+
 func main() {
 	ctx := context.Background()
 
@@ -49,7 +52,7 @@ func main() {
 	// variable to be set.
 	// For this example we choose to use a locally running embedding model though.
 	// It requires Ollama to serve its API at "http://localhost:11434/api".
-	collection, err := db.GetOrCreateCollection("Wikipedia", nil, chromem.NewEmbeddingFuncOllama(embeddingModel))
+	collection, err := db.GetOrCreateCollection("Wikipedia", nil, chromem.NewEmbeddingFuncOllama(embeddingModel), &normalized)
 	if err != nil {
 		panic(err)
 	}

examples/semantic-search-arxiv-openai/main.go

@@ -16,6 +16,9 @@ import (
 
 const searchTerm = "semantic search with vector databases"
 
+// OpenAI embeddings are already normalized.
+var normalized = true
+
 func main() {
 	ctx := context.Background()
 
@@ -30,7 +33,7 @@ func main() {
 	// We pass nil as embedding function to use the default (OpenAI text-embedding-3-small),
 	// which is very good and cheap. It requires the OPENAI_API_KEY environment
 	// variable to be set.
-	collection, err := db.GetOrCreateCollection("arXiv cs.CL 2023", nil, nil)
+	collection, err := db.GetOrCreateCollection("arXiv cs.CL 2023", nil, nil, &normalized)
 	if err != nil {
 		panic(err)
 	}

query.go

@@ -95,7 +95,7 @@ func documentMatchesFilters(document *Document, where, whereDocument map[string]
 	return true
 }
 
-func calcDocSimilarity(ctx context.Context, queryVectors []float32, docs []*Document) ([]docSim, error) {
+func calcDocSimilarity(ctx context.Context, queryVectors []float32, docs []*Document, isNormalized *bool) ([]docSim, error) {
 	similarities := make([]docSim, 0, len(docs))
 	similaritiesLock := sync.Mutex{}
 
@@ -145,7 +145,7 @@ func calcDocSimilarity(ctx context.Context, queryVectors []float32, docs []*Docu
 					return
 				}
 
-				sim, err := cosineSimilarity(queryVectors, doc.Embedding)
+				sim, err := cosineSimilarity(queryVectors, doc.Embedding, isNormalized)
 				if err != nil {
 					setSharedErr(fmt.Errorf("couldn't calculate similarity for document '%s': %w", doc.ID, err))
 					return

vector.go

@@ -7,19 +7,33 @@ import (
 
 const isNormalizedPrecisionTolerance = 1e-6
 
+var (
+	falseVal = false
+	trueVal  = true
+)
+
 // cosineSimilarity calculates the cosine similarity between two vectors.
-// Vectors are normalized first.
+// Pass isNormalized=true if the vectors are already normalized, false
+// to normalize them, and nil to autodetect.
 // The resulting value represents the similarity, so a higher value means the
 // vectors are more similar.
-func cosineSimilarity(a, b []float32) (float32, error) {
+func cosineSimilarity(a, b []float32, isNormalized *bool) (float32, error) {
 	// The vectors must have the same length
 	if len(a) != len(b) {
 		return 0, errors.New("vectors must have the same length")
 	}
 
-	if !isNormalized(a) || !isNormalized(b) {
+	if isNormalized == nil {
+		if !checkNormalized(a) || !checkNormalized(b) {
+			isNormalized = &falseVal
+		} else {
+			isNormalized = &trueVal
+		}
+	}
+	if !*isNormalized {
 		a, b = normalizeVector(a), normalizeVector(b)
 	}
+
 	var dotProduct float32
 	for i := range a {
 		dotProduct += a[i] * b[i]
@@ -44,8 +58,8 @@ func normalizeVector(v []float32) []float32 {
 	return res
 }
 
-// isNormalized checks if the vector is normalized.
-func isNormalized(v []float32) bool {
+// checkNormalized checks if the vector is normalized.
+func checkNormalized(v []float32) bool {
 	var sqSum float64
 	for _, val := range v {
 		sqSum += float64(val) * float64(val)