embedding

type BGEProvider ¶ added in v0.1.2

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

BGEProvider creates a provider for BGE models

func NewBGEProvider(modelPath string) (*BGEProvider, error)

func (p *BGEProvider) Dimension() int

func (p *BGEProvider) Embed(ctx context.Context, texts []string) ([][]float32, error)

type BatchOptions ¶

type BatchOptions struct {
	MaxBatchSize  int // Maximum batch size for embedding generation
	MaxConcurrent int // Maximum number of concurrent batches
}

BatchOptions contains configuration options for batch processing.

type BatchProcessor ¶

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

BatchProcessor handles batch processing of embeddings with optimization features including caching, concurrent processing, and progress tracking.

func NewBatchProcessor(provider Provider, options BatchOptions) *BatchProcessor

func (bp *BatchProcessor) Process(ctx context.Context, texts []string) ([][]float32, error)

func (bp *BatchProcessor) ProcessWithProgress(
	ctx context.Context,
	texts []string,
	callback ProgressCallback,
) ([][]float32, error)

type Config ¶ added in v0.1.2

type Config struct {
	Model        EmbeddingModel // Local embedding model implementation
	Dimension    int            // Embedding dimension
	MaxBatchSize int            // Maximum batch size for embedding generation
}

Config contains configuration parameters for the local embedding provider.

type DownloadModelInfo ¶ added in v0.1.2

type DownloadModelInfo struct {
	Name        string   // Model name
	Type        string   // Model type (e.g., "bge", "sentence-bert")
	URLs        []string // URLs of files to download
	Size        string   // Approximate total size
	Description string   // Model description
}

DownloadModelInfo contains information about a model available for download.

type DownloadProgressCallback ¶ added in v0.1.2

type DownloadProgressCallback func(modelName, fileName string, downloaded, total int64)

DownloadProgressCallback is a callback function for tracking download progress.

Parameters: - modelName: Name of the model being downloaded - fileName: Name of the file being downloaded - downloaded: Number of bytes downloaded so far - total: Total number of bytes to download (0 if unknown)

type Downloader ¶ added in v0.1.2

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

Downloader handles the downloading and caching of embedding models.

func NewDownloader(cacheDir string) *Downloader

func (d *Downloader) DownloadModel(modelName string, callback DownloadProgressCallback) (string, error)

func (d *Downloader) GetModelInfo() []DownloadModelInfo

type EmbeddingModel ¶ added in v0.1.2

type EmbeddingModel interface {
	// Run performs inference on the given inputs and returns the model outputs.
	Run(inputs map[string]interface{}) (map[string]interface{}, error)

	// Close releases any resources associated with the model.
	Close() error
}

EmbeddingModel defines the interface for local embedding models. Implementations should handle model loading, inference, and resource cleanup.

type LocalProvider ¶ added in v0.1.2

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

LocalProvider implements the embedding.Provider interface for local models. It handles tokenization, batch processing, and model inference.

func New(config Config) (*LocalProvider, error)

func (p *LocalProvider) Dimension() int

func (p *LocalProvider) Embed(ctx context.Context, texts []string) ([][]float32, error)

type Model ¶ added in v0.1.2

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

Model is a generic model implementation for embedding generation

func NewModel(dimension int, modelPath string) (*Model, error)

func (m *Model) Close() error

func (m *Model) Run(inputs map[string]interface{}) (map[string]interface{}, error)

type ModelInfo ¶ added in v0.1.2

type ModelInfo struct {
	Type      ModelType
	Name      string
	Dimension int
	ModelPath string
}

ModelInfo contains information about a model including its type, name, and dimensions.

func NewModelInfo(modelPath string) (*ModelInfo, error)

type ModelType ¶ added in v0.1.2

type ModelType string

ModelType defines the type of embedding model.

const (
	// ModelTypeBERT represents BERT models
	ModelTypeBERT ModelType = "bert"
	// ModelTypeSentenceBERT represents Sentence-BERT models
	ModelTypeSentenceBERT ModelType = "sentence-bert"
	// ModelTypeBGE represents BGE models
	ModelTypeBGE ModelType = "bge"
	// ModelTypeGPT represents GPT models
	ModelTypeGPT ModelType = "gpt"
	// ModelTypeFastText represents FastText models
	ModelTypeFastText ModelType = "fasttext"
	// ModelTypeGloVe represents GloVe models
	ModelTypeGloVe ModelType = "glove"
)

type ProgressCallback ¶

type ProgressCallback func(current, total int, err error) bool

ProgressCallback is a callback function for tracking batch processing progress.

Parameters: - current: Number of texts processed so far - total: Total number of texts to process - error: Error if any occurred during processing

Returns: - bool: True to continue processing, false to cancel the operation

type Provider ¶

type Provider interface {
	// Embed generates embeddings for the given texts.
	//
	// Parameters:
	// - ctx: Context for cancellation and timeout
	// - texts: Slice of texts to embed
	//
	// Returns:
	// - [][]float32: Slice of embeddings, one for each text
	// - error: Error if embedding generation fails
	Embed(ctx context.Context, texts []string) ([][]float32, error)

	// Dimension returns the dimension of the embeddings generated by this provider.
	//
	// Returns:
	// - int: Embedding dimension
	Dimension() int
}

Provider defines the interface for embedding providers.

This interface is implemented by all embedding model providers and allows the application to generate embeddings for text.

Example implementation:

type LocalProvider struct {
    model  *LocalModel
    dimension int
}

func (p *LocalProvider) Embed(ctx context.Context, texts []string) ([][]float32, error) {
    // Generate embeddings using local model
}

func (p *LocalProvider) Dimension() int {
    return p.dimension
}

func NewProvider(modelPath string) (Provider, error)

type SentenceBERTProvider ¶ added in v0.1.2

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

SentenceBERTProvider creates a provider for Sentence-BERT models

func NewSentenceBERTProvider(modelPath string) (*SentenceBERTProvider, error)

func (p *SentenceBERTProvider) Dimension() int

func (p *SentenceBERTProvider) Embed(ctx context.Context, texts []string) ([][]float32, error)

type Tokenizer ¶ added in v0.1.2

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

Tokenizer handles text tokenization for embedding models. It converts raw text into token IDs that can be processed by embedding models.

func NewTokenizer() (*Tokenizer, error)

func (t *Tokenizer) TokenizeBatch(texts []string) ([][]int64, [][]int64, error)