chore(refactor): do not expose internal backend Loader

mudler · mudler · commit 97060ae9a533 · 2024-11-08T19:15:30.000+01:00
Signed-off-by: Ettore Di Giacinto &lt;mudler@localai.io&gt;
diff --git a/core/backend/embeddings.go b/core/backend/embeddings.go
@@ -11,17 +11,9 @@ import (
 
 func ModelEmbedding(s string, tokens []int, loader *model.ModelLoader, backendConfig config.BackendConfig, appConfig *config.ApplicationConfig) (func() ([]float32, error), error) {
 
-	var inferenceModel interface{}
-	var err error
-
 	opts := ModelOptions(backendConfig, appConfig)
 
-	if backendConfig.Backend == "" {
-		inferenceModel, err = loader.GreedyLoader(opts...)
-	} else {
-		opts = append(opts, model.WithBackendString(backendConfig.Backend))
-		inferenceModel, err = loader.BackendLoader(opts...)
-	}
+	inferenceModel, err := loader.Load(opts...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/core/backend/image.go b/core/backend/image.go
@@ -10,8 +10,7 @@ import (
 func ImageGeneration(height, width, mode, step, seed int, positive_prompt, negative_prompt, src, dst string, loader *model.ModelLoader, backendConfig config.BackendConfig, appConfig *config.ApplicationConfig) (func() error, error) {
 
 	opts := ModelOptions(backendConfig, appConfig)
-
-	inferenceModel, err := loader.BackendLoader(
+	inferenceModel, err := loader.Load(
 		opts...,
 	)
 	if err != nil {
diff --git a/core/backend/llm.go b/core/backend/llm.go
@@ -16,7 +16,6 @@ import (
 	"github.com/mudler/LocalAI/core/schema"
 
 	"github.com/mudler/LocalAI/core/gallery"
-	"github.com/mudler/LocalAI/pkg/grpc"
 	"github.com/mudler/LocalAI/pkg/grpc/proto"
 	model "github.com/mudler/LocalAI/pkg/model"
 	"github.com/mudler/LocalAI/pkg/utils"
@@ -35,15 +34,6 @@ type TokenUsage struct {
 func ModelInference(ctx context.Context, s string, messages []schema.Message, images, videos, audios []string, loader *model.ModelLoader, c config.BackendConfig, o *config.ApplicationConfig, tokenCallback func(string, TokenUsage) bool) (func() (LLMResponse, error), error) {
 	modelFile := c.Model
 
-	var inferenceModel grpc.Backend
-	var err error
-
-	opts := ModelOptions(c, o)
-
-	if c.Backend != "" {
-		opts = append(opts, model.WithBackendString(c.Backend))
-	}
-
 	// Check if the modelFile exists, if it doesn't try to load it from the gallery
 	if o.AutoloadGalleries { // experimental
 		if _, err := os.Stat(modelFile); os.IsNotExist(err) {
@@ -56,12 +46,8 @@ func ModelInference(ctx context.Context, s string, messages []schema.Message, im
 		}
 	}
 
-	if c.Backend == "" {
-		inferenceModel, err = loader.GreedyLoader(opts...)
-	} else {
-		inferenceModel, err = loader.BackendLoader(opts...)
-	}
-
+	opts := ModelOptions(c, o)
+	inferenceModel, err := loader.Load(opts...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/core/backend/rerank.go b/core/backend/rerank.go
@@ -12,7 +12,7 @@ import (
 func Rerank(modelFile string, request *proto.RerankRequest, loader *model.ModelLoader, appConfig *config.ApplicationConfig, backendConfig config.BackendConfig) (*proto.RerankResult, error) {
 
 	opts := ModelOptions(backendConfig, appConfig, model.WithModel(modelFile))
-	rerankModel, err := loader.BackendLoader(opts...)
+	rerankModel, err := loader.Load(opts...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/core/backend/soundgeneration.go b/core/backend/soundgeneration.go
@@ -26,8 +26,7 @@ func SoundGeneration(
 ) (string, *proto.Result, error) {
 
 	opts := ModelOptions(backendConfig, appConfig, model.WithModel(modelFile))
-
-	soundGenModel, err := loader.BackendLoader(opts...)
+	soundGenModel, err := loader.Load(opts...)
 	if err != nil {
 		return "", nil, err
 	}
diff --git a/core/backend/stores.go b/core/backend/stores.go
@@ -8,16 +8,15 @@ import (
 )
 
 func StoreBackend(sl *model.ModelLoader, appConfig *config.ApplicationConfig, storeName string) (grpc.Backend, error) {
-    if storeName == "" {
-      storeName = "default"
-    }
+	if storeName == "" {
+		storeName = "default"
+	}
 
-    sc := []model.Option{
-      model.WithBackendString(model.LocalStoreBackend),
-      model.WithAssetDir(appConfig.AssetsDestination),
-      model.WithModel(storeName),
-    }
+	sc := []model.Option{
+		model.WithBackendString(model.LocalStoreBackend),
+		model.WithAssetDir(appConfig.AssetsDestination),
+		model.WithModel(storeName),
+	}
 
-    return sl.BackendLoader(sc...)
+	return sl.Load(sc...)
 }
-
diff --git a/core/backend/token_metrics.go b/core/backend/token_metrics.go
@@ -16,7 +16,7 @@ func TokenMetrics(
 	backendConfig config.BackendConfig) (*proto.MetricsResponse, error) {
 
 	opts := ModelOptions(backendConfig, appConfig, model.WithModel(modelFile))
-	model, err := loader.BackendLoader(opts...)
+	model, err := loader.Load(opts...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/core/backend/tokenize.go b/core/backend/tokenize.go
@@ -17,10 +17,10 @@ func ModelTokenize(s string, loader *model.ModelLoader, backendConfig config.Bac
 	opts := ModelOptions(backendConfig, appConfig, model.WithModel(modelFile))
 
 	if backendConfig.Backend == "" {
-		inferenceModel, err = loader.GreedyLoader(opts...)
+		inferenceModel, err = loader.Load(opts...)
 	} else {
 		opts = append(opts, model.WithBackendString(backendConfig.Backend))
-		inferenceModel, err = loader.BackendLoader(opts...)
+		inferenceModel, err = loader.Load(opts...)
 	}
 	if err != nil {
 		return schema.TokenizeResponse{}, err
diff --git a/core/backend/transcript.go b/core/backend/transcript.go
@@ -20,7 +20,7 @@ func ModelTranscription(audio, language string, translate bool, ml *model.ModelL
 
 	opts := ModelOptions(backendConfig, appConfig)
 
-	transcriptionModel, err := ml.BackendLoader(opts...)
+	transcriptionModel, err := ml.Load(opts...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/core/backend/tts.go b/core/backend/tts.go
@@ -29,7 +29,7 @@ func ModelTTS(
 	}
 
 	opts := ModelOptions(backendConfig, appConfig, model.WithBackendString(bb), model.WithModel(modelFile))
-	ttsModel, err := loader.BackendLoader(opts...)
+	ttsModel, err := loader.Load(opts...)
 	if err != nil {
 		return "", nil, err
 	}
diff --git a/core/startup/startup.go b/core/startup/startup.go
@@ -163,12 +163,7 @@ func Startup(opts ...config.AppOption) (*config.BackendConfigLoader, *model.Mode
 			o := backend.ModelOptions(*cfg, options)
 
 			var backendErr error
-			if cfg.Backend != "" {
-				o = append(o, model.WithBackendString(cfg.Backend))
-				_, backendErr = ml.BackendLoader(o...)
-			} else {
-				_, backendErr = ml.GreedyLoader(o...)
-			}
+			_, backendErr = ml.Load(o...)
 			if backendErr != nil {
 				return nil, nil, nil, err
 			}
diff --git a/pkg/model/initializers.go b/pkg/model/initializers.go
@@ -455,7 +455,7 @@ func (ml *ModelLoader) ListAvailableBackends(assetdir string) ([]string, error)
 	return orderBackends(backends)
 }
 
-func (ml *ModelLoader) BackendLoader(opts ...Option) (client grpc.Backend, err error) {
+func (ml *ModelLoader) backendLoader(opts ...Option) (client grpc.Backend, err error) {
 	o := NewOptions(opts...)
 
 	log.Info().Msgf("Loading model '%s' with backend %s", o.modelID, o.backendString)
@@ -500,7 +500,7 @@ func (ml *ModelLoader) stopActiveBackends(modelID string, singleActiveBackend bo
 	}
 }
 
-func (ml *ModelLoader) GreedyLoader(opts ...Option) (grpc.Backend, error) {
+func (ml *ModelLoader) Load(opts ...Option) (grpc.Backend, error) {
 	o := NewOptions(opts...)
 
 	// Return earlier if we have a model already loaded
@@ -513,6 +513,10 @@ func (ml *ModelLoader) GreedyLoader(opts ...Option) (grpc.Backend, error) {
 
 	ml.stopActiveBackends(o.modelID, o.singleActiveBackend)
 
+	if o.backendString != "" {
+		return ml.backendLoader(opts...)
+	}
+
 	var err error
 
 	// get backends embedded in the binary
@@ -536,7 +540,7 @@ func (ml *ModelLoader) GreedyLoader(opts ...Option) (grpc.Backend, error) {
 			WithBackendString(key),
 		}...)
 
-		model, modelerr := ml.BackendLoader(options...)
+		model, modelerr := ml.backendLoader(options...)
 		if modelerr == nil && model != nil {
 			log.Info().Msgf("[%s] Loads OK", key)
 			return model, nil
diff --git a/tests/integration/stores_test.go b/tests/integration/stores_test.go
@@ -57,7 +57,7 @@ var _ = Describe("Integration tests for the stores backend(s) and internal APIs"
 			}
 
 			sl = model.NewModelLoader("")
-			sc, err = sl.BackendLoader(storeOpts...)
+			sc, err = sl.Load(storeOpts...)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(sc).ToNot(BeNil())
 		})

Original file line number	Diff line number	Diff line change
`@@ -10,8 +10,7 @@ import (`
`10`	`10`	`func ImageGeneration(height, width, mode, step, seed int, positive_prompt, negative_prompt, src, dst string, loader model.ModelLoader, backendConfig config.BackendConfig, appConfig config.ApplicationConfig) (func() error, error) {`
`11`	`11`
`12`	`12`	`opts := ModelOptions(backendConfig, appConfig)`
`13`		`-`
`14`		`- inferenceModel, err := loader.BackendLoader(`
	`13`	`+ inferenceModel, err := loader.Load(`
`15`	`14`	`opts...,`
`16`	`15`	`)`
`17`	`16`	`if err != nil {`
Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,7 @@ import (`
`12`	`12`	`func Rerank(modelFile string, request proto.RerankRequest, loader model.ModelLoader, appConfig config.ApplicationConfig, backendConfig config.BackendConfig) (proto.RerankResult, error) {`
`13`	`13`
`14`	`14`	`opts := ModelOptions(backendConfig, appConfig, model.WithModel(modelFile))`
`15`		`- rerankModel, err := loader.BackendLoader(opts...)`
	`15`	`+ rerankModel, err := loader.Load(opts...)`
`16`	`16`	`if err != nil {`
`17`	`17`	`return nil, err`
`18`	`18`	`}`
Original file line number	Diff line number	Diff line change
`@@ -26,8 +26,7 @@ func SoundGeneration(`
`26`	`26`	`) (string, *proto.Result, error) {`
`27`	`27`
`28`	`28`	`opts := ModelOptions(backendConfig, appConfig, model.WithModel(modelFile))`
`29`		`-`
`30`		`- soundGenModel, err := loader.BackendLoader(opts...)`
	`29`	`+ soundGenModel, err := loader.Load(opts...)`
`31`	`30`	`if err != nil {`
`32`	`31`	`return "", nil, err`
`33`	`32`	`}`
Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,7 @@ func TokenMetrics(`
`16`	`16`	`backendConfig config.BackendConfig) (*proto.MetricsResponse, error) {`
`17`	`17`
`18`	`18`	`opts := ModelOptions(backendConfig, appConfig, model.WithModel(modelFile))`
`19`		`- model, err := loader.BackendLoader(opts...)`
	`19`	`+ model, err := loader.Load(opts...)`
`20`	`20`	`if err != nil {`
`21`	`21`	`return nil, err`
`22`	`22`	`}`
Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@ func ModelTranscription(audio, language string, translate bool, ml *model.ModelL`
`20`	`20`
`21`	`21`	`opts := ModelOptions(backendConfig, appConfig)`
`22`	`22`
`23`		`- transcriptionModel, err := ml.BackendLoader(opts...)`
	`23`	`+ transcriptionModel, err := ml.Load(opts...)`
`24`	`24`	`if err != nil {`
`25`	`25`	`return nil, err`
`26`	`26`	`}`
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ func ModelTTS(`
`29`	`29`	`}`
`30`	`30`
`31`	`31`	`opts := ModelOptions(backendConfig, appConfig, model.WithBackendString(bb), model.WithModel(modelFile))`
`32`		`- ttsModel, err := loader.BackendLoader(opts...)`
	`32`	`+ ttsModel, err := loader.Load(opts...)`
`33`	`33`	`if err != nil {`
`34`	`34`	`return "", nil, err`
`35`	`35`	`}`