mirror of
https://github.com/mudler/LocalAI.git
synced 2024-06-07 19:40:48 +00:00
89351f1a7d
Certain engines requires to know during model loading if the embedding feature has to be enabled, however, it is impractical to have to set it to ALL the backends that supports embeddings. There are transformers and sentencentransformers that seamelessly handle both cases, without having this settings to be explicitly enabled. The case sussist only for ggml-based models that needs to enable featuresets during model loading (and thus settings `embedding` is required), however most of the other engines does not require this. This change disables the check done at code side, making easier to use embeddings by not having to specify explicitly `embeddings: true`. Part of: https://github.com/mudler/LocalAI/issues/1373
89 lines
2.1 KiB
Go
89 lines
2.1 KiB
Go
package backend
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"github.com/go-skynet/LocalAI/core/config"
|
|
|
|
"github.com/go-skynet/LocalAI/pkg/grpc"
|
|
model "github.com/go-skynet/LocalAI/pkg/model"
|
|
)
|
|
|
|
func ModelEmbedding(s string, tokens []int, loader *model.ModelLoader, backendConfig config.BackendConfig, appConfig *config.ApplicationConfig) (func() ([]float32, error), error) {
|
|
modelFile := backendConfig.Model
|
|
|
|
grpcOpts := gRPCModelOpts(backendConfig)
|
|
|
|
var inferenceModel interface{}
|
|
var err error
|
|
|
|
opts := modelOpts(backendConfig, appConfig, []model.Option{
|
|
model.WithLoadGRPCLoadModelOpts(grpcOpts),
|
|
model.WithThreads(uint32(*backendConfig.Threads)),
|
|
model.WithAssetDir(appConfig.AssetsDestination),
|
|
model.WithModel(modelFile),
|
|
model.WithContext(appConfig.Context),
|
|
})
|
|
|
|
if backendConfig.Backend == "" {
|
|
inferenceModel, err = loader.GreedyLoader(opts...)
|
|
} else {
|
|
opts = append(opts, model.WithBackendString(backendConfig.Backend))
|
|
inferenceModel, err = loader.BackendLoader(opts...)
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var fn func() ([]float32, error)
|
|
switch model := inferenceModel.(type) {
|
|
case grpc.Backend:
|
|
fn = func() ([]float32, error) {
|
|
predictOptions := gRPCPredictOpts(backendConfig, loader.ModelPath)
|
|
if len(tokens) > 0 {
|
|
embeds := []int32{}
|
|
|
|
for _, t := range tokens {
|
|
embeds = append(embeds, int32(t))
|
|
}
|
|
predictOptions.EmbeddingTokens = embeds
|
|
|
|
res, err := model.Embeddings(appConfig.Context, predictOptions)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return res.Embeddings, nil
|
|
}
|
|
predictOptions.Embeddings = s
|
|
|
|
res, err := model.Embeddings(appConfig.Context, predictOptions)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return res.Embeddings, nil
|
|
}
|
|
default:
|
|
fn = func() ([]float32, error) {
|
|
return nil, fmt.Errorf("embeddings not supported by the backend")
|
|
}
|
|
}
|
|
|
|
return func() ([]float32, error) {
|
|
embeds, err := fn()
|
|
if err != nil {
|
|
return embeds, err
|
|
}
|
|
// Remove trailing 0s
|
|
for i := len(embeds) - 1; i >= 0; i-- {
|
|
if embeds[i] == 0.0 {
|
|
embeds = embeds[:i]
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
return embeds, nil
|
|
}, nil
|
|
}
|