embed

Documentation

Overview

Package embed provides embedding interfaces and implementations.

Index

• func CosineSimilarity(a, b []float32) float64
• func DefaultModelPath() string
• func DefaultVocabPath() string
• func EnsureModel() (string, string, error)
• func ModelDir() string
• func ModelExists() bool
• func TopK(query []float32, vectors [][]float32, k int) []int
• type Embedder
  • func NewEmbedder(provider, model string) (Embedder, error)
• type OnnxEmbedder
  • func NewOnnxEmbedder(modelPath, vocabPath string) (*OnnxEmbedder, error)
  • func (e *OnnxEmbedder) Close() error
  • func (e *OnnxEmbedder) Dimensions() int
  • func (e *OnnxEmbedder) Embed(text string) ([]float32, error)
  • func (e *OnnxEmbedder) EmbedBatch(texts []string) ([][]float32, error)
• type OpenAIEmbedder
  • func NewOpenAIEmbedder(apiKey string, model string) (*OpenAIEmbedder, error)
  • func (e *OpenAIEmbedder) Close() error
  • func (e *OpenAIEmbedder) Dimensions() int
  • func (e *OpenAIEmbedder) Embed(text string) ([]float32, error)
  • func (e *OpenAIEmbedder) EmbedBatch(texts []string) ([][]float32, error)

Functions

func CosineSimilarity

func CosineSimilarity(a, b []float32) float64

CosineSimilarity computes the cosine similarity between two vectors.

func DefaultModelPath

func DefaultModelPath() string

DefaultModelPath returns the path to the default ONNX model file.

func DefaultVocabPath

func DefaultVocabPath() string

DefaultVocabPath returns the path to the default vocab file.

func EnsureModel

func EnsureModel() (string, string, error)

EnsureModel downloads the ONNX model, vocab, and runtime library if not already present.

func ModelDir

func ModelDir() string

ModelDir returns the directory where ONNX models are stored.

func ModelExists

func ModelExists() bool

ModelExists checks if the ONNX model files are present.

func TopK

func TopK(query []float32, vectors [][]float32, k int) []int

TopK returns indices of the top-k most similar vectors to query.

Types

type Embedder

type Embedder interface {
	Embed(text string) ([]float32, error)
	EmbedBatch(texts []string) ([][]float32, error)
	Dimensions() int
	Close() error
}

Embedder generates vector embeddings from text.

func NewEmbedder

func NewEmbedder(provider, model string) (Embedder, error)

NewEmbedder creates an embedder based on the provider name. provider: "openai", "local"/"onnx", or "" (returns nil).

type OnnxEmbedder

type OnnxEmbedder struct {
	// contains filtered or unexported fields
}

OnnxEmbedder uses a local ONNX model for embeddings.

func NewOnnxEmbedder

func NewOnnxEmbedder(modelPath, vocabPath string) (*OnnxEmbedder, error)

NewOnnxEmbedder creates a new ONNX embedder with persistent session.

func (*OnnxEmbedder) Close

func (e *OnnxEmbedder) Close() error

Close releases ONNX session resources.

func (*OnnxEmbedder) Dimensions

func (e *OnnxEmbedder) Dimensions() int

Dimensions returns the embedding vector dimensions (384 for MiniLM).

func (*OnnxEmbedder) Embed

func (e *OnnxEmbedder) Embed(text string) ([]float32, error)

Embed generates an embedding for a single text.

func (*OnnxEmbedder) EmbedBatch

func (e *OnnxEmbedder) EmbedBatch(texts []string) ([][]float32, error)

EmbedBatch generates embeddings for multiple texts.

type OpenAIEmbedder

type OpenAIEmbedder struct {
	// contains filtered or unexported fields
}

OpenAIEmbedder uses the OpenAI API for embeddings.

func NewOpenAIEmbedder

func NewOpenAIEmbedder(apiKey string, model string) (*OpenAIEmbedder, error)

NewOpenAIEmbedder creates a new OpenAI embedder. apiKey must be a valid OpenAI API key. model can be empty to use the default (text-embedding-3-small).

func (*OpenAIEmbedder) Close

func (e *OpenAIEmbedder) Close() error

Close is a no-op for OpenAI (no local resources).

func (*OpenAIEmbedder) Dimensions

func (e *OpenAIEmbedder) Dimensions() int

Dimensions returns the embedding vector dimensions.

func (*OpenAIEmbedder) Embed

func (e *OpenAIEmbedder) Embed(text string) ([]float32, error)

Embed generates an embedding for a single text.

func (*OpenAIEmbedder) EmbedBatch

func (e *OpenAIEmbedder) EmbedBatch(texts []string) ([][]float32, error)

EmbedBatch generates embeddings for multiple texts in one API call.