embeddings/huggingface: Creating embedder for the Huggingface hub

embeddings/huggingface/huggingface.go

@@ -0,0 +1,72 @@
+package huggingface
+
+import (
+	"context"
+	"strings"
+
+	"github.com/tmc/langchaingo/embeddings"
+	"github.com/tmc/langchaingo/llms/huggingface"
+)
+
+// Huggingface is the embedder using the Huggingface hub api.
+type Huggingface struct {
+	client *huggingface.LLM
+	Model  string
+	Task   string
+
+	StripNewLines bool
+	BatchSize     int
+}
+
+var _ embeddings.Embedder = &Huggingface{}
+
+func NewHuggingface(opts ...Option) (*Huggingface, error) {
+	v, err := applyOptions(opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	return v, nil
+}
+
+func (e *Huggingface) EmbedDocuments(ctx context.Context, texts []string) ([][]float64, error) {
+	batchedTexts := embeddings.BatchTexts(
+		embeddings.MaybeRemoveNewLines(texts, e.StripNewLines),
+		e.BatchSize,
+	)
+
+	emb := make([][]float64, 0, len(texts))
+	for _, texts := range batchedTexts {
+		curTextEmbeddings, err := e.client.CreateEmbedding(ctx, texts, e.Model, e.Task)
+		if err != nil {
+			return nil, err
+		}
+
+		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
+}
+
+func (e *Huggingface) EmbedQuery(ctx context.Context, text string) ([]float64, error) {
+	if e.StripNewLines {
+		text = strings.ReplaceAll(text, "\n", " ")
+	}
+
+	emb, err := e.client.CreateEmbedding(ctx, []string{text}, e.Model, e.Task)
+	if err != nil {
+		return nil, err
+	}
+
+	return emb[0], nil
+}

embeddings/huggingface/huggingface_test.go

@@ -0,0 +1,27 @@
+package huggingface
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestHuggingfaceEmbeddings(t *testing.T) {
+	t.Parallel()
+
+	if openaiKey := os.Getenv("HUGGINGFACEHUB_API_TOKEN"); openaiKey == "" {
+		t.Skip("HUGGINGFACEHUB_API_TOKEN not set")
+	}
+	e, err := NewHuggingface()
+	require.NoError(t, err)
+
+	_, err = e.EmbedQuery(context.Background(), "Hello world!")
+	require.NoError(t, err)
+
+	embeddings, err := e.EmbedDocuments(context.Background(), []string{"Hello world", "The world is ending", "good bye"})
+	require.NoError(t, err)
+	assert.Len(t, embeddings, 3)
+}

embeddings/huggingface/options.go

@@ -0,0 +1,73 @@
+package huggingface
+
+import (
+	"github.com/tmc/langchaingo/llms/huggingface"
+)
+
+const (
+	_defaultBatchSize     = 512
+	_defaultStripNewLines = true
+	_defaultModel         = "sentence-transformers/all-mpnet-base-v2"
+	_defaultTask          = "feature-extraction"
+)
+
+// Option is a function type that can be used to modify the client.
+type Option func(p *Huggingface)
+
+// WithModel is an option for providing the model name to use.
+func WithModel(model string) Option {
+	return func(p *Huggingface) {
+		p.Model = model
+	}
+}
+
+// WithTask is an option for providing the task to call the model with.
+func WithTask(task string) Option {
+	return func(p *Huggingface) {
+		p.Task = task
+	}
+}
+
+// WithClient is an option for providing the LLM client.
+func WithClient(client huggingface.LLM) Option {
+	return func(p *Huggingface) {
+		p.client = &client
+	}
+}
+
+// WithStripNewLines is an option for specifying the should it strip new lines.
+func WithStripNewLines(stripNewLines bool) Option {
+	return func(p *Huggingface) {
+		p.StripNewLines = stripNewLines
+	}
+}
+
+// WithBatchSize is an option for specifying the batch size.
+func WithBatchSize(batchSize int) Option {
+	return func(p *Huggingface) {
+		p.BatchSize = batchSize
+	}
+}
+
+func applyOptions(opts ...Option) (*Huggingface, error) {
+	o := &Huggingface{
+		StripNewLines: _defaultStripNewLines,
+		BatchSize:     _defaultBatchSize,
+		Model:         _defaultModel,
+		Task:          _defaultTask,
+	}
+
+	for _, opt := range opts {
+		opt(o)
+	}
+
+	if o.client == nil {
+		client, err := huggingface.New()
+		if err != nil {
+			return nil, err
+		}
+		o.client = client
+	}
+
+	return o, nil
+}

llms/huggingface/huggingfacellm.go

@@ -11,8 +11,9 @@ import (
 )
 
 var (
-	ErrEmptyResponse = errors.New("empty response")
-	ErrMissingToken  = errors.New("missing the Hugging Face API token. Set it in the HUGGINGFACEHUB_API_TOKEN environment variable") //nolint:lll
+	ErrEmptyResponse            = errors.New("empty response")
+	ErrMissingToken             = errors.New("missing the Hugging Face API token. Set it in the HUGGINGFACEHUB_API_TOKEN environment variable") //nolint:lll
+	ErrUnexpectedResponseLength = errors.New("unexpected length of response")
 )
 
 type LLM struct {
@@ -91,3 +92,29 @@ func New(opts ...Option) (*LLM, error) {
 		client: c,
 	}, nil
 }
+
+// CreateEmbedding creates embeddings for the given input texts.
+func (o *LLM) CreateEmbedding(
+	ctx context.Context,
+	inputTexts []string,
+	model string,
+	task string,
+) ([][]float64, error) {
+	embeddings, err := o.client.CreateEmbedding(ctx, model, task, &huggingfaceclient.EmbeddingRequest{
+		Inputs: inputTexts,
+		Options: map[string]any{
+			"use_gpu":        false,
+			"wait_for_model": true,
+		},
+	})
+	if err != nil {
+		return nil, err
+	}
+	if len(embeddings) == 0 {
+		return nil, ErrEmptyResponse
+	}
+	if len(inputTexts) != len(embeddings) {
+		return embeddings, ErrUnexpectedResponseLength
+	}
+	return embeddings, nil
+}

llms/huggingface/internal/huggingfaceclient/embeddings.go

@@ -0,0 +1,54 @@
+package huggingfaceclient
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+)
+
+type embeddingPayload struct {
+	Options map[string]any
+	Inputs  []string `json:"inputs"`
+}
+
+// nolint:lll
+func (c *Client) createEmbedding(ctx context.Context, model string, task string, payload *embeddingPayload) ([][]float32, error) {
+	body := map[string]any{
+		"inputs": payload.Inputs,
+	}
+	for key, value := range payload.Options {
+		body[key] = value
+	}
+
+	payloadBytes, err := json.Marshal(body)
+	if err != nil {
+		return nil, fmt.Errorf("marshal payload: %w", err)
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("%s/pipeline/%s/%s", c.url, task, model), bytes.NewReader(payloadBytes))
+	if err != nil {
+		return nil, fmt.Errorf("create request: %w", err)
+	}
+	req.Header.Set("Authorization", "Bearer "+c.Token)
+	req.Header.Set("Content-Type", "application/json")
+
+	r, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer r.Body.Close()
+
+	if r.StatusCode != http.StatusOK {
+		msg := fmt.Sprintf("API returned unexpected status code: %d", r.StatusCode)
+
+		return nil, fmt.Errorf("%s: %s", msg, "unable to create embeddings") // nolint:goerr113
+	}
+
+	var response [][]float32
+	if err := json.NewDecoder(r.Body).Decode(&response); err != nil {
+		return nil, fmt.Errorf("decode response: %w", err)
+	}
+
+	return response, nil
+}

llms/huggingface/internal/huggingfaceclient/huggingfaceclient.go

@@ -11,6 +11,8 @@ var (
 	ErrEmptyResponse = errors.New("empty response")
 )
 
+const huggingfaceAPIBaseURL = "https://api-inference.huggingface.co"
+
 type Client struct {
 	Token string
 	Model string
@@ -24,7 +26,7 @@ func New(token string, model string) (*Client, error) {
 	return &Client{
 		Token: token,
 		Model: model,
-		url:   hfInferenceAPI,
+		url:   huggingfaceAPIBaseURL,
 	}, nil
 }
 
@@ -73,3 +75,43 @@ func (c *Client) RunInference(ctx context.Context, request *InferenceRequest) (*
 		Text: text,
 	}, nil
 }
+
+// EmbeddingRequest is a request to create an embedding.
+type EmbeddingRequest struct {
+	Options map[string]any `json:"options"`
+	Inputs  []string       `json:"inputs"`
+}
+
+// CreateEmbedding creates embeddings.
+func (c *Client) CreateEmbedding(
+	ctx context.Context,
+	model string,
+	task string,
+	r *EmbeddingRequest,
+) ([][]float64, error) {
+	resp, err := c.createEmbedding(ctx, model, task, &embeddingPayload{
+		Inputs:  r.Inputs,
+		Options: r.Options,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	if len(resp) == 0 {
+		return nil, ErrEmptyResponse
+	}
+
+	return c.convertFloat32ToFloat64(resp), nil
+}
+
+func (c *Client) convertFloat32ToFloat64(input [][]float32) [][]float64 {
+	output := make([][]float64, len(input))
+	for i, row := range input {
+		output[i] = make([]float64, len(row))
+		for j, val := range row {
+			output[i][j] = float64(val)
+		}
+	}
+
+	return output
+}

llms/huggingface/internal/huggingfaceclient/inference.go

@@ -43,15 +43,13 @@ type (
 	}
 )
 
-const hfInferenceAPI = "https://api-inference.huggingface.co/models/"
-
 func (c *Client) runInference(ctx context.Context, payload *inferencePayload) (inferenceResponsePayload, error) {
 	payloadBytes, err := json.Marshal(payload)
 	if err != nil {
 		return nil, err
 	}
 	body := bytes.NewReader(payloadBytes)
-	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("%s%s", c.url, payload.Model), body) //nolint:lll
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("%s/models/%s", c.url, payload.Model), body) //nolint:lll
 	if err != nil {
 		return nil, err
 	}