vxcontrol/langchaingo
1
0
Fork 0
mirror of https://github.com/vxcontrol/langchaingo.git synced 2026-07-19 21:54:17 -04:00
Code Issues Packages Projects Releases 10 Wiki Activity

embeddings: refactor embedding interfaces to enable creation from Chat and LLM (#374)

Browse Source 
* Refactor embedding interfaces to enable creation from Chat and LLM

* embedding: add test that verifies assignability
This commit is contained in:
Eli Bendersky
2023-11-30 07:24:31 -08:00
committed by Travis Cline
parent a1d183f61b
commit 3d030fcbb3
8 changed files with 127 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
embeddings/embedding.go
+78
View File
@@ -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
}
embeddings/ernie/ernie.go
+1 -2
View File
@@ -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.
embeddings/internal/embedderclient/embedderclient.go
-45
View File
@@ -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
}
embeddings/ollama/ollama.go
+1 -2
View File
@@ -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.
embeddings/openai/openai.go
+1 -2
View File
@@ -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.
embeddings/openai/openai_test.go
+23
View File
@@ -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
}
embeddings/options.go
+22
View File
@@ -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
}
}
embeddings/vertexai/vertexai_palm.go
+1 -2
View File
@@ -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.