feat: add nomic embedder

embedder/nomic/api.go

@@ -0,0 +1,85 @@
+package nomicembedder
+
+import (
+ "bytes"
+ "encoding/json"
+ "io"
+
+ "github.com/henomis/lingoose/embedder"
+ "github.com/henomis/restclientgo"
+)
+
+type Model string
+
+const (
+ ModelNomicEmbedTextV1 Model = "nomic-embed-text-v1"
+ ModelAllMiniLML6V2 Model = "all-MiniLM-L6-v2"
+)
+
+type TaskType string
+
+const (
+ TaskTypeSearchQuery TaskType = "search_query"
+ TaskTypeSearchDocument TaskType = "search_document"
+ TaskTypeClustering TaskType = "clustering"
+ TaskTypeClassification TaskType = "classification"
+)
+
+type request struct {
+ Model string `json:"model"`
+ Texts []string `json:"texts"`
+ TaskType TaskType `json:"task_type,omitempty"`
+}
+
+func (r *request) Path() (string, error) {
+ return "/embedding/text", nil
+}
+
+func (r *request) Encode() (io.Reader, error) {
+ jsonBytes, err := json.Marshal(r)
+ if err != nil {
+ return nil, err
+ }
+
+ return bytes.NewReader(jsonBytes), nil
+}
+
+func (r *request) ContentType() string {
+ return "application/json"
+}
+
+type response struct {
+ HTTPStatusCode int `json:"-"`
+ Embeddings []embedder.Embedding `json:"embeddings"`
+ Usage Usage `json:"usage"`
+ RawBody string `json:"-"`
+}
+
+type Usage struct {
+ TotalTokens int `json:"total_tokens"`
+}
+
+func (r *response) Decode(body io.Reader) error {
+ return json.NewDecoder(body).Decode(r)
+}
+
+func (r *response) SetBody(body io.Reader) error {
+ b, err := io.ReadAll(body)
+ if err != nil {
+ return err
+ }
+
+ r.RawBody = string(b)
+ return nil
+}
+
+func (r *response) AcceptContentType() string {
+ return "application/json"
+}
+
+func (r *response) SetStatusCode(code int) error {
+ r.HTTPStatusCode = code
+ return nil
+}
+
+func (r *response) SetHeaders(_ restclientgo.Headers) error { return nil }

embedder/nomic/nomic.go

@@ -0,0 +1,74 @@
+package nomicembedder
+
+import (
+ "context"
+ "net/http"
+ "os"
+
+ "github.com/henomis/lingoose/embedder"
+ "github.com/henomis/restclientgo"
+)
+
+const (
+ defaultEndpoint = "https://api-atlas.nomic.ai/v1"
+ defaultModel = ModelNomicEmbedTextV1
+)
+
+type Embedder struct {
+ taskType TaskType
+ model Model
+ restClient *restclientgo.RestClient
+}
+
+func New() *Embedder {
+ apiKey := os.Getenv("NOMIC_API_KEY")
+
+ return &Embedder{
+ restClient: restclientgo.New(defaultEndpoint).WithRequestModifier(
+ func(req *http.Request) *http.Request {
+ req.Header.Set("Authorization", "Bearer "+apiKey)
+ return req
+ },
+ ),
+ model: defaultModel,
+ }
+}
+
+func (e *Embedder) WithAPIKey(apiKey string) *Embedder {
+ e.restClient = restclientgo.New(defaultEndpoint).WithRequestModifier(
+ func(req *http.Request) *http.Request {
+ req.Header.Set("Authorization", "Bearer "+apiKey)
+ return req
+ },
+ )
+ return e
+}
+
+func (e *Embedder) WithTaskType(taskType TaskType) *Embedder {
+ e.taskType = taskType
+ return e
+}
+
+func (e *Embedder) WithModel(model Model) *Embedder {
+ e.model = model
+ return e
+}
+
+// Embed returns the embeddings for the given texts
+func (e *Embedder) Embed(ctx context.Context, texts []string) ([]embedder.Embedding, error) {
+ var resp response
+ err := e.restClient.Post(
+ ctx,
+ &request{
+ Texts: texts,
+ Model: string(e.model),
+ TaskType: e.taskType,
+ },
+ &resp,
+ )
+ if err != nil {
+ return nil, err
+ }
+
+ return resp.Embeddings, nil
+}

examples/embeddings/nomic/main.go

@@ -0,0 +1,103 @@
+package main
+
+import (
+ "context"
+ "fmt"
+
+ nomicembedder "github.com/henomis/lingoose/embedder/nomic"
+ "github.com/henomis/lingoose/index"
+ indexoption "github.com/henomis/lingoose/index/option"
+ "github.com/henomis/lingoose/index/vectordb/jsondb"
+ "github.com/henomis/lingoose/llm/openai"
+ "github.com/henomis/lingoose/loader"
+ "github.com/henomis/lingoose/prompt"
+ "github.com/henomis/lingoose/textsplitter"
+)
+
+// download https://raw.githubusercontent.com/hwchase17/chat-your-data/master/state_of_the_union.txt
+
+func main() {
+
+ index := index.New(
+ jsondb.New().WithPersist("db.json"),
+ nomicembedder.New(),
+ ).WithIncludeContents(true).WithAddDataCallback(func(data *index.Data) error {
+ data.Metadata["contentLen"] = len(data.Metadata["content"].(string))
+ return nil
+ })
+
+ indexIsEmpty, _ := index.IsEmpty(context.Background())
+
+ if indexIsEmpty {
+ err := ingestData(index)
+ if err != nil {
+ panic(err)
+ }
+ }
+
+ query := "What is the purpose of the NATO Alliance?"
+ similarities, err := index.Query(
+ context.Background(),
+ query,
+ indexoption.WithTopK(3),
+ )
+ if err != nil {
+ panic(err)
+ }
+
+ for _, similarity := range similarities {
+ fmt.Printf("Similarity: %f\n", similarity.Score)
+ fmt.Printf("Document: %s\n", similarity.Content())
+ fmt.Println("Metadata: ", similarity.Metadata)
+ fmt.Println("----------")
+ }
+
+ documentContext := ""
+ for _, similarity := range similarities {
+ documentContext += similarity.Content() + "\n\n"
+ }
+
+ llmOpenAI := openai.NewCompletion().WithVerbose(true)
+ prompt1 := prompt.NewPromptTemplate(
+ "Based on the following context answer to the question.\n\nContext:\n{{.context}}\n\nQuestion: {{.query}}").WithInputs(
+ map[string]string{
+ "query": query,
+ "context": documentContext,
+ },
+ )
+
+ err = prompt1.Format(nil)
+ if err != nil {
+ panic(err)
+ }
+
+ output, err := llmOpenAI.Completion(context.Background(), prompt1.String())
+ if err != nil {
+ panic(err)
+ }
+
+ fmt.Println(output)
+}
+
+func ingestData(index *index.Index) error {
+
+ fmt.Printf("Ingesting data...")
+
+ documents, err := loader.NewDirectoryLoader(".", ".txt").Load(context.Background())
+ if err != nil {
+ return err
+ }
+
+ textSplitter := textsplitter.NewRecursiveCharacterTextSplitter(1000, 20)
+
+ documentChunks := textSplitter.SplitDocuments(documents)
+
+ err = index.LoadFromDocuments(context.Background(), documentChunks)
+ if err != nil {
+ return err
+ }
+
+ fmt.Printf("Done!\n")
+
+ return nil
+}