From 3d030fcbb3c316b9d74969fa56b40e9d0a37c562 Mon Sep 17 00:00:00 2001 From: Eli Bendersky Date: Thu, 30 Nov 2023 07:24:31 -0800 Subject: [PATCH] embeddings: refactor embedding interfaces to enable creation from Chat and LLM (#374) * Refactor embedding interfaces to enable creation from Chat and LLM * embedding: add test that verifies assignability --- embeddings/embedding.go | 78 +++++++++++++++++++ embeddings/ernie/ernie.go | 3 +- .../internal/embedderclient/embedderclient.go | 45 ----------- embeddings/ollama/ollama.go | 3 +- embeddings/openai/openai.go | 3 +- embeddings/openai/openai_test.go | 23 ++++++ embeddings/options.go | 22 ++++++ embeddings/vertexai/vertexai_palm.go | 3 +- 8 files changed, 127 insertions(+), 53 deletions(-) delete mode 100644 embeddings/internal/embedderclient/embedderclient.go create mode 100644 embeddings/options.go diff --git a/embeddings/embedding.go b/embeddings/embedding.go index c796524a..7a7ab20d 100644 --- a/embeddings/embedding.go +++ b/embeddings/embedding.go @@ -13,6 +13,51 @@ type Embedder interface { EmbedQuery(ctx context.Context, text string) ([]float32, error) } +// EmbedderClient is the interface LLM clients implement for embeddings. +type EmbedderClient interface { + CreateEmbedding(ctx context.Context, texts []string) ([][]float32, error) +} + +type EmbedderImpl struct { + client EmbedderClient + + StripNewLines bool + BatchSize int +} + +func NewEmbedder(client EmbedderClient, opts ...Option) (*EmbedderImpl, error) { + e := &EmbedderImpl{ + client: client, + StripNewLines: defaultStripNewLines, + BatchSize: defaultBatchSize, + } + + for _, opt := range opts { + opt(e) + } + return e, nil +} + +// EmbedQuery embeds a single text. +func (ei *EmbedderImpl) EmbedQuery(ctx context.Context, text string) ([]float32, error) { + if ei.StripNewLines { + text = strings.ReplaceAll(text, "\n", " ") + } + + emb, err := ei.client.CreateEmbedding(ctx, []string{text}) + if err != nil { + return nil, err + } + + return emb[0], nil +} + +// EmbedDocuments creates one vector embedding for each of the texts. +func (ei *EmbedderImpl) EmbedDocuments(ctx context.Context, texts []string) ([][]float32, error) { + texts = MaybeRemoveNewLines(texts, ei.StripNewLines) + return BatchedEmbed(ctx, ei.client, texts, ei.BatchSize) +} + func MaybeRemoveNewLines(texts []string, removeNewLines bool) []string { if !removeNewLines { return texts @@ -43,3 +88,36 @@ func BatchTexts(texts []string, batchSize int) [][]string { return batchedTexts } + +// BatchedEmbed creates embeddings for the given input texts, batching them +// into batches of batchSize if needed. +func BatchedEmbed(ctx context.Context, embedder EmbedderClient, texts []string, batchSize int) ([][]float32, error) { + batchedTexts := BatchTexts(texts, batchSize) + + emb := make([][]float32, 0, len(texts)) + for _, texts := range batchedTexts { + curTextEmbeddings, err := embedder.CreateEmbedding(ctx, texts) + if err != nil { + return nil, err + } + // If the size of this batch is 1, don't average/combine the vectors. + if len(texts) == 1 { + emb = append(emb, curTextEmbeddings[0]) + continue + } + + textLengths := make([]int, 0, len(texts)) + for _, text := range texts { + textLengths = append(textLengths, len(text)) + } + + combined, err := CombineVectors(curTextEmbeddings, textLengths) + if err != nil { + return nil, err + } + + emb = append(emb, combined) + } + + return emb, nil +} diff --git a/embeddings/ernie/ernie.go b/embeddings/ernie/ernie.go index eaaf6226..9e2a3e27 100644 --- a/embeddings/ernie/ernie.go +++ b/embeddings/ernie/ernie.go @@ -4,7 +4,6 @@ import ( "context" "github.com/tmc/langchaingo/embeddings" - "github.com/tmc/langchaingo/embeddings/internal/embedderclient" "github.com/tmc/langchaingo/llms/ernie" ) @@ -44,7 +43,7 @@ func NewErnie(opts ...Option) (*Ernie, error) { // EmbedDocuments use ernie Embedding-V1. func (e *Ernie) EmbedDocuments(ctx context.Context, texts []string) ([][]float32, error) { texts = embeddings.MaybeRemoveNewLines(texts, e.stripNewLines) - return embedderclient.BatchedEmbed(ctx, e.client, texts, e.batchSize) + return embeddings.BatchedEmbed(ctx, e.client, texts, e.batchSize) } // EmbedQuery use ernie Embedding-V1. diff --git a/embeddings/internal/embedderclient/embedderclient.go b/embeddings/internal/embedderclient/embedderclient.go deleted file mode 100644 index 4db5c34f..00000000 --- a/embeddings/internal/embedderclient/embedderclient.go +++ /dev/null @@ -1,45 +0,0 @@ -package embedderclient - -import ( - "context" - - "github.com/tmc/langchaingo/embeddings" -) - -// EmbedderClient is the interface LLM clients implement for embeddings. -type EmbedderClient interface { - CreateEmbedding(ctx context.Context, texts []string) ([][]float32, error) -} - -// BatchedEmbed creates embeddings for the given input texts, batching them -// into batches of batchSize if needed. -func BatchedEmbed(ctx context.Context, embedder EmbedderClient, texts []string, batchSize int) ([][]float32, error) { - batchedTexts := embeddings.BatchTexts(texts, batchSize) - - emb := make([][]float32, 0, len(texts)) - for _, texts := range batchedTexts { - curTextEmbeddings, err := embedder.CreateEmbedding(ctx, texts) - if err != nil { - return nil, err - } - // If the size of this batch is 1, don't average/combine the vectors. - if len(texts) == 1 { - emb = append(emb, curTextEmbeddings[0]) - continue - } - - textLengths := make([]int, 0, len(texts)) - for _, text := range texts { - textLengths = append(textLengths, len(text)) - } - - combined, err := embeddings.CombineVectors(curTextEmbeddings, textLengths) - if err != nil { - return nil, err - } - - emb = append(emb, combined) - } - - return emb, nil -} diff --git a/embeddings/ollama/ollama.go b/embeddings/ollama/ollama.go index 98eabe47..ab522e71 100644 --- a/embeddings/ollama/ollama.go +++ b/embeddings/ollama/ollama.go @@ -5,7 +5,6 @@ import ( "strings" "github.com/tmc/langchaingo/embeddings" - "github.com/tmc/langchaingo/embeddings/internal/embedderclient" "github.com/tmc/langchaingo/llms/ollama" ) @@ -32,7 +31,7 @@ func NewOllama(opts ...Option) (Ollama, error) { // EmbedDocuments creates one vector embedding for each of the texts. func (e Ollama) EmbedDocuments(ctx context.Context, texts []string) ([][]float32, error) { texts = embeddings.MaybeRemoveNewLines(texts, e.StripNewLines) - return embedderclient.BatchedEmbed(ctx, e.client, texts, e.BatchSize) + return embeddings.BatchedEmbed(ctx, e.client, texts, e.BatchSize) } // EmbedQuery embeds a single text. diff --git a/embeddings/openai/openai.go b/embeddings/openai/openai.go index da9b0b90..309d4033 100644 --- a/embeddings/openai/openai.go +++ b/embeddings/openai/openai.go @@ -5,7 +5,6 @@ import ( "strings" "github.com/tmc/langchaingo/embeddings" - "github.com/tmc/langchaingo/embeddings/internal/embedderclient" "github.com/tmc/langchaingo/llms/openai" ) @@ -32,7 +31,7 @@ func NewOpenAI(opts ...Option) (OpenAI, error) { // EmbedDocuments creates one vector embedding for each of the texts. func (e OpenAI) EmbedDocuments(ctx context.Context, texts []string) ([][]float32, error) { texts = embeddings.MaybeRemoveNewLines(texts, e.StripNewLines) - return embedderclient.BatchedEmbed(ctx, e.client, texts, e.BatchSize) + return embeddings.BatchedEmbed(ctx, e.client, texts, e.BatchSize) } // EmbedQuery embeds a single text. diff --git a/embeddings/openai/openai_test.go b/embeddings/openai/openai_test.go index 0331c8c0..db1e2e56 100644 --- a/embeddings/openai/openai_test.go +++ b/embeddings/openai/openai_test.go @@ -7,6 +7,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/tmc/langchaingo/embeddings" "github.com/tmc/langchaingo/llms/openai" ) @@ -103,3 +104,25 @@ func TestOpenaiEmbeddingsWithAzureAPI(t *testing.T) { require.NoError(t, err) assert.Len(t, embeddings, 1) } + +func TestUseLLMAndChatAsEmbedderClient(t *testing.T) { + t.Parallel() + + if openaiKey := os.Getenv("OPENAI_API_KEY"); openaiKey == "" { + t.Skip("OPENAI_API_KEY not set") + } + + llm, err := openai.New() + require.NoError(t, err) + + embedderFromLLM, err := embeddings.NewEmbedder(llm) + require.NoError(t, err) + var _ embeddings.Embedder = embedderFromLLM + + chat, err := openai.NewChat() + require.NoError(t, err) + + embedderFromChat, err := embeddings.NewEmbedder(chat) + require.NoError(t, err) + var _ embeddings.Embedder = embedderFromChat +} diff --git a/embeddings/options.go b/embeddings/options.go new file mode 100644 index 00000000..0c73a1ce --- /dev/null +++ b/embeddings/options.go @@ -0,0 +1,22 @@ +package embeddings + +const ( + defaultBatchSize = 512 + defaultStripNewLines = true +) + +type Option func(p *EmbedderImpl) + +// WithStripNewLines is an option for specifying the should it strip new lines. +func WithStripNewLines(stripNewLines bool) Option { + return func(p *EmbedderImpl) { + p.StripNewLines = stripNewLines + } +} + +// WithBatchSize is an option for specifying the batch size. +func WithBatchSize(batchSize int) Option { + return func(p *EmbedderImpl) { + p.BatchSize = batchSize + } +} diff --git a/embeddings/vertexai/vertexai_palm.go b/embeddings/vertexai/vertexai_palm.go index ff3d1b2d..79a9c077 100644 --- a/embeddings/vertexai/vertexai_palm.go +++ b/embeddings/vertexai/vertexai_palm.go @@ -5,7 +5,6 @@ import ( "strings" "github.com/tmc/langchaingo/embeddings" - "github.com/tmc/langchaingo/embeddings/internal/embedderclient" "github.com/tmc/langchaingo/llms/vertexai" ) @@ -32,7 +31,7 @@ func NewVertexAIPaLM(opts ...Option) (*VertexAIPaLM, error) { // EmbedDocuments creates one vector embedding for each of the texts. func (e VertexAIPaLM) EmbedDocuments(ctx context.Context, texts []string) ([][]float32, error) { texts = embeddings.MaybeRemoveNewLines(texts, e.StripNewLines) - return embedderclient.BatchedEmbed(ctx, e.client, texts, e.BatchSize) + return embeddings.BatchedEmbed(ctx, e.client, texts, e.BatchSize) } // EmbedQuery embeds a single text.