Normalize vectors on embedding creation instead of querying

- Normalizes only once instead of each time
- Embedding creation takes time anyway, while query should be as fast as possible

Author: philippgille

collection.go

@@ -213,13 +213,17 @@ func (c *Collection) AddDocument(ctx context.Context, doc Document) error {
 		m[k] = v
 	}
 
-	// Create embedding if they don't exist
+	// Create embedding if they don't exist, otherwise normalize if necessary
 	if len(doc.Embedding) == 0 {
 		embedding, err := c.embed(ctx, doc.Content)
 		if err != nil {
 			return fmt.Errorf("couldn't create embedding of document: %w", err)
 		}
 		doc.Embedding = embedding
+	} else {
+		if !isNormalized(doc.Embedding) {
+			doc.Embedding = normalizeVector(doc.Embedding)
+		}
 	}
 
 	c.documentsLock.Lock()

collection_test.go

@@ -13,7 +13,7 @@ func TestCollection_Add(t *testing.T) {
 	ctx := context.Background()
 	name := "test"
 	metadata := map[string]string{"foo": "bar"}
-	vectors := []float32{-0.1, 0.1, 0.2}
+	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
 		return vectors, nil
 	}
@@ -109,7 +109,7 @@ func TestCollection_Add_Error(t *testing.T) {
 	ctx := context.Background()
 	name := "test"
 	metadata := map[string]string{"foo": "bar"}
-	vectors := []float32{-0.1, 0.1, 0.2}
+	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
 		return vectors, nil
 	}
@@ -160,7 +160,7 @@ func TestCollection_AddConcurrently(t *testing.T) {
 	ctx := context.Background()
 	name := "test"
 	metadata := map[string]string{"foo": "bar"}
-	vectors := []float32{-0.1, 0.1, 0.2}
+	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
 		return vectors, nil
 	}
@@ -256,7 +256,7 @@ func TestCollection_AddConcurrently_Error(t *testing.T) {
 	ctx := context.Background()
 	name := "test"
 	metadata := map[string]string{"foo": "bar"}
-	vectors := []float32{-0.1, 0.1, 0.2}
+	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
 		return vectors, nil
 	}
@@ -313,8 +313,9 @@ func TestCollection_Count(t *testing.T) {
 	db := NewDB()
 	name := "test"
 	metadata := map[string]string{"foo": "bar"}
+	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
-		return []float32{-0.1, 0.1, 0.2}, nil
+		return vectors, nil
 	}
 	c, err := db.CreateCollection(name, metadata, embeddingFunc)
 	if err != nil {

db.go

@@ -13,6 +13,9 @@ import (
 // EmbeddingFunc is a function that creates embeddings for a given text.
 // chromem-go will use OpenAI`s "text-embedding-3-small" model by default,
 // but you can provide your own function, using any model you like.
+// The function must return a *normalized* vector, i.e. the length of the vector
+// must be 1. OpenAI's and Mistral's embedding models do this by default. Some
+// others like Nomic's "nomic-embed-text-v1.5" don't.
 type EmbeddingFunc func(ctx context.Context, text string) ([]float32, error)
 
 // DB is the chromem-go database. It holds collections, which hold documents.

db_test.go

@@ -10,7 +10,7 @@ func TestDB_CreateCollection(t *testing.T) {
 	// Values in the collection
 	name := "test"
 	metadata := map[string]string{"foo": "bar"}
-	vectors := []float32{-0.1, 0.1, 0.2}
+	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
 		return vectors, nil
 	}
@@ -81,7 +81,7 @@ func TestDB_ListCollections(t *testing.T) {
 	// Values in the collection
 	name := "test"
 	metadata := map[string]string{"foo": "bar"}
-	vectors := []float32{-0.1, 0.1, 0.2}
+	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
 		return vectors, nil
 	}
@@ -147,7 +147,7 @@ func TestDB_GetCollection(t *testing.T) {
 	// Values in the collection
 	name := "test"
 	metadata := map[string]string{"foo": "bar"}
-	vectors := []float32{-0.1, 0.1, 0.2}
+	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
 		return vectors, nil
 	}
@@ -196,7 +196,7 @@ func TestDB_GetOrCreateCollection(t *testing.T) {
 	// Values in the collection
 	name := "test"
 	metadata := map[string]string{"foo": "bar"}
-	vectors := []float32{-0.1, 0.1, 0.2}
+	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
 		return vectors, nil
 	}
@@ -299,7 +299,7 @@ func TestDB_DeleteCollection(t *testing.T) {
 	// Values in the collection
 	name := "test"
 	metadata := map[string]string{"foo": "bar"}
-	vectors := []float32{-0.1, 0.1, 0.2}
+	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
 		return vectors, nil
 	}
@@ -331,7 +331,7 @@ func TestDB_Reset(t *testing.T) {
 	// Values in the collection
 	name := "test"
 	metadata := map[string]string{"foo": "bar"}
-	vectors := []float32{-0.1, 0.1, 0.2}
+	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
 		return vectors, nil
 	}

document_test.go

@@ -10,7 +10,7 @@ func TestDocument_New(t *testing.T) {
 	ctx := context.Background()
 	id := "test"
 	metadata := map[string]string{"foo": "bar"}
-	vectors := []float32{-0.1, 0.1, 0.2}
+	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 	content := "hello world"
 	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
 		return vectors, nil

embed_compat.go

@@ -9,9 +9,13 @@ const (
 // NewEmbeddingFuncMistral returns a function that creates embeddings for a text
 // using the Mistral API.
 func NewEmbeddingFuncMistral(apiKey string) EmbeddingFunc {
+	// Mistral embeddings are normalized, see section "Distance Measures" on
+	// https://docs.mistral.ai/guides/embeddings/.
+	normalized := true
+
 	// The Mistral API docs don't mention the `encoding_format` as optional,
 	// but it seems to be, just like OpenAI. So we reuse the OpenAI function.
-	return NewEmbeddingFuncOpenAICompat(baseURLMistral, apiKey, embeddingModelMistral)
+	return NewEmbeddingFuncOpenAICompat(baseURLMistral, apiKey, embeddingModelMistral, &normalized)
 }
 
 const baseURLJina = "https://api.jina.ai/v1"
@@ -28,7 +32,7 @@ const (
 // NewEmbeddingFuncJina returns a function that creates embeddings for a text
 // using the Jina API.
 func NewEmbeddingFuncJina(apiKey string, model EmbeddingModelJina) EmbeddingFunc {
-	return NewEmbeddingFuncOpenAICompat(baseURLJina, apiKey, string(model))
+	return NewEmbeddingFuncOpenAICompat(baseURLJina, apiKey, string(model), nil)
 }
 
 const baseURLMixedbread = "https://api.mixedbread.ai"
@@ -49,7 +53,7 @@ const (
 // NewEmbeddingFuncMixedbread returns a function that creates embeddings for a text
 // using the mixedbread.ai API.
 func NewEmbeddingFuncMixedbread(apiKey string, model EmbeddingModelMixedbread) EmbeddingFunc {
-	return NewEmbeddingFuncOpenAICompat(baseURLMixedbread, apiKey, string(model))
+	return NewEmbeddingFuncOpenAICompat(baseURLMixedbread, apiKey, string(model), nil)
 }
 
 const baseURLLocalAI = "http://localhost:8080/v1"
@@ -64,5 +68,5 @@ const baseURLLocalAI = "http://localhost:8080/v1"
 // But other embedding models are supported as well. See the LocalAI documentation
 // for details.
 func NewEmbeddingFuncLocalAI(model string) EmbeddingFunc {
-	return NewEmbeddingFuncOpenAICompat(baseURLLocalAI, "", model)
+	return NewEmbeddingFuncOpenAICompat(baseURLLocalAI, "", model, nil)
 }

embed_ollama.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"sync"
 )
 
 // TODO: Turn into const and use as default, but allow user to pass custom URL
@@ -28,6 +29,9 @@ func NewEmbeddingFuncOllama(model string) EmbeddingFunc {
 	// and it might have to be a long timeout, depending on the text length.
 	client := &http.Client{}
 
+	var checkedNormalized bool
+	checkNormalized := sync.Once{}
+
 	return func(ctx context.Context, text string) ([]float32, error) {
 		// Prepare the request body.
 		reqBody, err := json.Marshal(map[string]string{
@@ -74,6 +78,18 @@ func NewEmbeddingFuncOllama(model string) EmbeddingFunc {
 			return nil, errors.New("no embeddings found in the response")
 		}
 
-		return embeddingResponse.Embedding, nil
+		v := embeddingResponse.Embedding
+		checkNormalized.Do(func() {
+			if isNormalized(v) {
+				checkedNormalized = true
+			} else {
+				checkedNormalized = false
+			}
+		})
+		if !checkedNormalized {
+			v = normalizeVector(v)
+		}
+
+		return v, nil
 	}
 }

embed_ollama_test.go

@@ -25,7 +25,7 @@ func TestNewEmbeddingFuncOllama(t *testing.T) {
 	if err != nil {
 		t.Fatal("unexpected error:", err)
 	}
-	wantRes := []float32{-0.1, 0.1, 0.2}
+	wantRes := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 
 	// Mock server
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

embed_openai.go

@@ -9,6 +9,7 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"sync"
 )
 
 const BaseURLOpenAI = "https://api.openai.com/v1"
@@ -39,7 +40,9 @@ func NewEmbeddingFuncDefault() EmbeddingFunc {
 // NewEmbeddingFuncOpenAI returns a function that creates embeddings for a text
 // using the OpenAI API.
 func NewEmbeddingFuncOpenAI(apiKey string, model EmbeddingModelOpenAI) EmbeddingFunc {
-	return NewEmbeddingFuncOpenAICompat(BaseURLOpenAI, apiKey, string(model))
+	// OpenAI embeddings are normalized
+	normalized := true
+	return NewEmbeddingFuncOpenAICompat(BaseURLOpenAI, apiKey, string(model), &normalized)
 }
 
 // NewEmbeddingFuncOpenAICompat returns a function that creates embeddings for a text
@@ -48,12 +51,20 @@ func NewEmbeddingFuncOpenAI(apiKey string, model EmbeddingModelOpenAI) Embedding
 //   - LitLLM: https://github.com/BerriAI/litellm
 //   - Ollama: https://github.com/ollama/ollama/blob/main/docs/openai.md
 //   - etc.
-func NewEmbeddingFuncOpenAICompat(baseURL, apiKey, model string) EmbeddingFunc {
+//
+// The `normalized` parameter indicates whether the vectors returned by the embedding
+// model are already normalized, as is the case for OpenAI's and Mistral's models.
+// The flag is optional. If it's nil, it will be autodetected on the first request
+// (which bears a small risk that the vector just happens to have a length of 1).
+func NewEmbeddingFuncOpenAICompat(baseURL, apiKey, model string, normalized *bool) EmbeddingFunc {
 	// We don't set a default timeout here, although it's usually a good idea.
 	// In our case though, the library user can set the timeout on the context,
 	// and it might have to be a long timeout, depending on the text length.
 	client := &http.Client{}
 
+	var checkedNormalized bool
+	checkNormalized := sync.Once{}
+
 	return func(ctx context.Context, text string) ([]float32, error) {
 		// Prepare the request body.
 		reqBody, err := json.Marshal(map[string]string{
@@ -101,6 +112,24 @@ func NewEmbeddingFuncOpenAICompat(baseURL, apiKey, model string) EmbeddingFunc {
 			return nil, errors.New("no embeddings found in the response")
 		}
 
-		return embeddingResponse.Data[0].Embedding, nil
+		v := embeddingResponse.Data[0].Embedding
+		if normalized != nil {
+			if *normalized {
+				return v, nil
+			}
+			return normalizeVector(v), nil
+		}
+		checkNormalized.Do(func() {
+			if isNormalized(v) {
+				checkedNormalized = true
+			} else {
+				checkedNormalized = false
+			}
+		})
+		if !checkedNormalized {
+			v = normalizeVector(v)
+		}
+
+		return v, nil
 	}
 }

embed_openai_test.go

@@ -33,7 +33,7 @@ func TestNewEmbeddingFuncOpenAICompat(t *testing.T) {
 	if err != nil {
 		t.Fatal("unexpected error:", err)
 	}
-	wantRes := []float32{-0.1, 0.1, 0.2}
+	wantRes := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
 
 	// Mock server
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -75,7 +75,7 @@ func TestNewEmbeddingFuncOpenAICompat(t *testing.T) {
 	defer ts.Close()
 	baseURL := ts.URL + baseURLSuffix
 
-	f := chromem.NewEmbeddingFuncOpenAICompat(baseURL, apiKey, model)
+	f := chromem.NewEmbeddingFuncOpenAICompat(baseURL, apiKey, model, nil)
 	res, err := f(context.Background(), input)
 	if err != nil {
 		t.Fatal("expected nil, got", err)