forked from tmc/langchaingo
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
embeddings: add cybertron local embeddings
Added a new package 'cybertron' under 'embeddings' to run embedding models locally. This includes the main 'cybertron.go' file, a test file 'cybertron_test.go', and an 'options.go' file for providing various options for the client. Resolves tmc#494
- Loading branch information
Showing
6 changed files
with
189 additions
and
7 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,38 @@ | ||
| package cybertron | ||
|
|
||
| import ( | ||
| "context" | ||
|
|
||
| "github.com/nlpodyssey/cybertron/pkg/models/bert" | ||
| "github.com/nlpodyssey/cybertron/pkg/tasks/textencoding" | ||
| "github.com/tmc/langchaingo/embeddings" | ||
| ) | ||
|
|
||
| // Cybertron is the embedder using Cybertron to run embedding models locally. | ||
| type Cybertron struct { | ||
| encoder textencoding.Interface | ||
| Model string | ||
| ModelsDir string | ||
| PoolingStrategy bert.PoolingStrategyType | ||
| } | ||
|
|
||
| var _ embeddings.EmbedderClient = (*Cybertron)(nil) | ||
|
|
||
| func NewCybertron(opts ...Option) (*Cybertron, error) { | ||
| return applyOptions(opts...) | ||
| } | ||
|
|
||
| func (c *Cybertron) CreateEmbedding(ctx context.Context, texts []string) ([][]float32, error) { | ||
| result := make([][]float32, 0, len(texts)) | ||
|
|
||
| for _, text := range texts { | ||
| embedding, err := c.encoder.Encode(ctx, text, int(c.PoolingStrategy)) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| result = append(result, embedding.Vector.Normalize2().Data().F32()) | ||
| } | ||
|
|
||
| return result, nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,34 @@ | ||
| package cybertron | ||
|
|
||
| import ( | ||
| "context" | ||
| "os" | ||
| "testing" | ||
|
|
||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| func TestCybertronEmbeddings(t *testing.T) { | ||
| t.Parallel() | ||
|
|
||
| _, err := os.Stat(_defaultModelsDir) | ||
| if os.IsNotExist(err) && os.Getenv("CYBERTRON_DO_DOWNLOAD") == "" { | ||
| // Cybertron downloads the embedding model and caches it in ModelsDir. Doing this as | ||
| // part of the tests would be costly in terms of time and bandwidth and likely make | ||
| // the test flaky. | ||
| t.Skipf("ModelsDir %q doesn't exist", _defaultModelsDir) | ||
| } | ||
|
|
||
| emb, err := NewCybertron( | ||
| WithModelsDir(_defaultModelsDir), | ||
| WithModel(_defaultModel), | ||
| WithPoolingStrategy(_defaultPoolingStrategy), | ||
| ) | ||
| require.NoError(t, err) | ||
|
|
||
| res, err := emb.CreateEmbedding(context.Background(), []string{ | ||
| "Hello world", "The world is ending", "good bye", | ||
| }) | ||
| require.NoError(t, err) | ||
| require.Len(t, res, 3) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| package cybertron | ||
|
|
||
| import ( | ||
| "github.com/nlpodyssey/cybertron/pkg/models/bert" | ||
| "github.com/nlpodyssey/cybertron/pkg/tasks" | ||
| "github.com/nlpodyssey/cybertron/pkg/tasks/textencoding" | ||
| ) | ||
|
|
||
| const ( | ||
| _defaultModel = "sentence-transformers/all-MiniLM-L6-v2" | ||
| _defaultModelsDir = "models" | ||
| _defaultPoolingStrategy = bert.MeanPooling | ||
| ) | ||
|
|
||
| // Option is a function type that can be used to modify the client. | ||
| type Option func(c *Cybertron) | ||
|
|
||
| // apply the option to the instance. | ||
| func (o Option) apply(c *Cybertron) { | ||
| o(c) | ||
| } | ||
|
|
||
| // WithModel is an option for providing the model name to use. | ||
| func WithModel(model string) Option { | ||
| return func(c *Cybertron) { | ||
| c.Model = model | ||
| } | ||
| } | ||
|
|
||
| // WithModelsDir is an option for setting the directory to store downloaded models. | ||
| func WithModelsDir(dir string) Option { | ||
| return func(c *Cybertron) { | ||
| c.ModelsDir = dir | ||
| } | ||
| } | ||
|
|
||
| func WithPoolingStrategy(strategy bert.PoolingStrategyType) Option { | ||
| return func(c *Cybertron) { | ||
| c.PoolingStrategy = strategy | ||
| } | ||
| } | ||
|
|
||
| // WithEncoder is an option for providing the Encoder. | ||
| func WithEncoder(encoder textencoding.Interface) Option { | ||
| return func(c *Cybertron) { | ||
| c.encoder = encoder | ||
| } | ||
| } | ||
|
|
||
| func applyOptions(opts ...Option) (*Cybertron, error) { | ||
| c := &Cybertron{ | ||
| Model: _defaultModel, | ||
| ModelsDir: _defaultModelsDir, | ||
| PoolingStrategy: _defaultPoolingStrategy, | ||
| encoder: nil, | ||
| } | ||
|
|
||
| for _, opt := range opts { | ||
| opt.apply(c) | ||
| } | ||
|
|
||
| if c.encoder == nil { | ||
| encoder, err := tasks.Load[textencoding.Interface](&tasks.Config{ | ||
| ModelsDir: c.ModelsDir, | ||
| ModelName: c.Model, | ||
| }) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| c.encoder = encoder | ||
| } | ||
|
|
||
| return c, nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters