cybertron

Documentation

Index

• type Cybertron
  • func NewCybertron(opts ...Option) (*Cybertron, error)
  • func (c *Cybertron) CreateEmbedding(ctx context.Context, texts []string) ([][]float32, error)
• type Option
  • func WithEncoder(encoder textencoding.Interface) Option
  • func WithModel(model string) Option
  • func WithModelsDir(dir string) Option
  • func WithPoolingStrategy(strategy bert.PoolingStrategyType) Option

Types

type Cybertron

type Cybertron struct {
	Model           string
	ModelsDir       string
	PoolingStrategy bert.PoolingStrategyType
	// contains filtered or unexported fields
}

Cybertron is the embedder using Cybertron to run embedding models locally.

func NewCybertron

func NewCybertron(opts ...Option) (*Cybertron, error)

NewCybertron returns a new embedding client that uses Cybertron to run embedding models locally (on the CPU). The embedding model will be downloaded and cached automatically. Use `WithModel` and `WithModelsDir` to change which model is used and where it is cached.

func (*Cybertron) CreateEmbedding

func (c *Cybertron) CreateEmbedding(ctx context.Context, texts []string) ([][]float32, error)

CreateEmbedding implements the `embeddings.EmbedderClient` and creates an embedding vector for each of the supplied texts.

type Option

type Option func(c *Cybertron)

Option is a function type that can be used to modify the client.

func WithEncoder

func WithEncoder(encoder textencoding.Interface) Option

WithEncoder is an option for providing the Encoder.

func WithModel

func WithModel(model string) Option

WithModel is an option for providing the model name to use. Default is "sentence-transformers/all-MiniLM-L6-v2". Note that not all embedding models are supported.

func WithModelsDir

func WithModelsDir(dir string) Option

WithModelsDir is an option for setting the directory to store downloaded models. Default is "models".

func WithPoolingStrategy

func WithPoolingStrategy(strategy bert.PoolingStrategyType) Option

WithPoolingStrategy sets the pooling strategy. Default is mean pooling.