Update embeddings.Embedder interface to require []byte for Image meth…

…ods; fix type value for llamafile API image uploads (#4) Co-authored-by: sfomuseumbot <sfomuseumbot@localhost>
whosonfirst · Aug 23, 2024 · eb1db3a · eb1db3a
1 parent 90d83d3
commit eb1db3a
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 17 deletions.
diff --git a/embeddings/chromem_ollama.go b/embeddings/chromem_ollama.go
@@ -53,10 +53,10 @@ func (e *ChromemOllamaEmbedder) Embeddings32(ctx context.Context, content string
 	return e.embeddings_func(ctx, content)
 }
 
-func (e *ChromemOllamaEmbedder) ImageEmbeddings(ctx context.Context, content string) ([]float64, error) {
+func (e *ChromemOllamaEmbedder) ImageEmbeddings(ctx context.Context, data []byte) ([]float64, error) {
 	return nil, dedupe.NotImplemented()
 }
 
-func (e *ChromemOllamaEmbedder) ImageEmbeddings32(ctx context.Context, content string) ([]float32, error) {
+func (e *ChromemOllamaEmbedder) ImageEmbeddings32(ctx context.Context, data []byte) ([]float32, error) {
 	return nil, dedupe.NotImplemented()
 }
diff --git a/embeddings/embedder.go b/embeddings/embedder.go
@@ -19,9 +19,9 @@ type Embedder interface {
 	// Embeddings32 returns the embeddings for a string as a list of float32 values.
 	Embeddings32(context.Context, string) ([]float32, error)
 	// ImageEmbeddings returns the embeddings for a base64-encoded image as a list of float64 values.
-	ImageEmbeddings(context.Context, string) ([]float64, error)
+	ImageEmbeddings(context.Context, []byte) ([]float64, error)
 	// ImageEmbeddings32 returns the embeddings for a base64-encoded image as a list of float32 values.
-	ImageEmbeddings32(context.Context, string) ([]float32, error)
+	ImageEmbeddings32(context.Context, []byte) ([]float32, error)
 }
 
 // EmbedderInitializationFunc is a function defined by individual embedder package and used to create

diff --git a/embeddings/embeddings.go b/embeddings/embeddings.go
@@ -5,8 +5,20 @@ func asFloat32(data []float64) []float32 {
 	e32 := make([]float32, len(data))
 
 	for idx, v := range data {
+		// Really, check for max float32here...
 		e32[idx] = float32(v)
 	}
 
 	return e32
 }
+
+func asFloat64(data []float32) []float64 {
+
+	e64 := make([]float64, len(data))
+
+	for idx, v := range data {
+		e64[idx] = float64(v)
+	}
+
+	return e64
+}
diff --git a/embeddings/llamafile.go b/embeddings/llamafile.go
@@ -3,21 +3,25 @@
 package embeddings
 
 // https://github.com/Mozilla-Ocho/llamafile/blob/main/llama.cpp/server/README.md#api-endpoints
+// https://github.com/Mozilla-Ocho/llamafile?tab=readme-ov-file#other-example-llamafiles
+//
 // curl --request POST --url http://localhost:8080/embedding --header "Content-Type: application/json" --data '{"content": "Hello world" }'
 
 import (
 	"bytes"
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	_ "io"
 	"net/http"
 	"net/url"
 	"strconv"
+	"time"
 )
 
 type LlamafileImageDataEmbeddingRequest struct {
-	Id   string `json:"id"`
+	Id   int64  `json:"id"`
 	Data string `json:"data"`
 }
 
@@ -119,11 +123,16 @@ func (e *LlamafileEmbedder) Embeddings32(ctx context.Context, content string) ([
 	return asFloat32(e64), nil
 }
 
-func (e *LlamafileEmbedder) ImageEmbeddings(ctx context.Context, content string) ([]float64, error) {
+func (e *LlamafileEmbedder) ImageEmbeddings(ctx context.Context, data []byte) ([]float64, error) {
+
+	data_b64 := base64.StdEncoding.EncodeToString(data)
+
+	now := time.Now()
+	ts := now.Unix()
 
 	image_req := &LlamafileImageDataEmbeddingRequest{
-		Data: content,
-		Id:   "1",
+		Data: data_b64,
+		Id:   ts,
 	}
 
 	req := &LlamafileEmbeddingRequest{
@@ -141,9 +150,9 @@ func (e *LlamafileEmbedder) ImageEmbeddings(ctx context.Context, content string)
 	return rsp.Embeddings, nil
 }
 
-func (e *LlamafileEmbedder) ImageEmbeddings32(ctx context.Context, content string) ([]float32, error) {
+func (e *LlamafileEmbedder) ImageEmbeddings32(ctx context.Context, data []byte) ([]float32, error) {
 
-	e64, err := e.ImageEmbeddings(ctx, content)
+	e64, err := e.ImageEmbeddings(ctx, data)
 
 	if err != nil {
 		return nil, err

diff --git a/embeddings/llamafile_test.go b/embeddings/llamafile_test.go
@@ -4,7 +4,6 @@ package embeddings
 
 import (
 	"context"
-	"encoding/base64"
 	"io"
 	"os"
 	"testing"
@@ -57,9 +56,7 @@ func TestLlamafileImageEmbeddings(t *testing.T) {
 		t.Fatalf("Failed to read data from %s, %v", im_path, err)
 	}
 
-	im_b64 := base64.StdEncoding.EncodeToString(im_body)
-
-	rsp, err := emb.ImageEmbeddings(ctx, im_b64)
+	rsp, err := emb.ImageEmbeddings(ctx, im_body)
 
 	if err != nil {
 		t.Fatalf("Failed to derive embeddings, %v", err)

diff --git a/embeddings/ollama.go b/embeddings/ollama.go
@@ -112,10 +112,10 @@ func (e *OllamaEmbedder) Embeddings32(ctx context.Context, content string) ([]fl
 	return asFloat32(e64), nil
 }
 
-func (e *OllamaEmbedder) ImageEmbeddings(ctx context.Context, content string) ([]float64, error) {
+func (e *OllamaEmbedder) ImageEmbeddings(ctx context.Context, data []byte) ([]float64, error) {
 	return nil, dedupe.NotImplemented()
 }
 
-func (e *OllamaEmbedder) ImageEmbeddings32(ctx context.Context, content string) ([]float32, error) {
+func (e *OllamaEmbedder) ImageEmbeddings32(ctx context.Context, data []byte) ([]float32, error) {
 	return nil, dedupe.NotImplemented()
 }
diff --git a/embeddings/ollama_test.go b/embeddings/ollama_test.go
@@ -23,7 +23,9 @@ func TestOllamafileImageEmbeddings(t *testing.T) {
 		t.Fatalf("Failed to create embedder, %v", err)
 	}
 
-	_, err = emb.ImageEmbeddings(ctx, "")
+	var data []byte
+
+	_, err = emb.ImageEmbeddings(ctx, data)
 
 	if !dedupe.IsNotImplementedError(err) {
 		t.Fatalf("Unexpected error, %v", err)