feat: various refactorings

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

api/backend/embeddings.go

@@ -0,0 +1,107 @@
+package backend
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	config "github.com/go-skynet/LocalAI/api/config"
+	"github.com/go-skynet/LocalAI/api/options"
+	"github.com/go-skynet/LocalAI/pkg/grpc"
+	model "github.com/go-skynet/LocalAI/pkg/model"
+	bert "github.com/go-skynet/go-bert.cpp"
+)
+
+func ModelEmbedding(s string, tokens []int, loader *model.ModelLoader, c config.Config, o *options.Option) (func() ([]float32, error), error) {
+	if !c.Embeddings {
+		return nil, fmt.Errorf("endpoint disabled for this model by API configuration")
+	}
+
+	modelFile := c.Model
+
+	grpcOpts := gRPCModelOpts(c)
+
+	var inferenceModel interface{}
+	var err error
+
+	opts := []model.Option{
+		model.WithLoadGRPCOpts(grpcOpts),
+		model.WithThreads(uint32(c.Threads)),
+		model.WithAssetDir(o.AssetsDestination),
+		model.WithModelFile(modelFile),
+	}
+
+	if c.Backend == "" {
+		inferenceModel, err = loader.GreedyLoader(opts...)
+	} else {
+		opts = append(opts, model.WithBackendString(c.Backend))
+		inferenceModel, err = loader.BackendLoader(opts...)
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	var fn func() ([]float32, error)
+	switch model := inferenceModel.(type) {
+	case *grpc.Client:
+		fn = func() ([]float32, error) {
+			predictOptions := gRPCPredictOpts(c, loader.ModelPath)
+			if len(tokens) > 0 {
+				embeds := []int32{}
+
+				for _, t := range tokens {
+					embeds = append(embeds, int32(t))
+				}
+				predictOptions.EmbeddingTokens = embeds
+
+				res, err := model.Embeddings(context.TODO(), predictOptions)
+				if err != nil {
+					return nil, err
+				}
+
+				return res.Embeddings, nil
+			}
+			predictOptions.Embeddings = s
+
+			res, err := model.Embeddings(context.TODO(), predictOptions)
+			if err != nil {
+				return nil, err
+			}
+
+			return res.Embeddings, nil
+		}
+
+	// bert embeddings
+	case *bert.Bert:
+		fn = func() ([]float32, error) {
+			if len(tokens) > 0 {
+				return model.TokenEmbeddings(tokens, bert.SetThreads(c.Threads))
+			}
+			return model.Embeddings(s, bert.SetThreads(c.Threads))
+		}
+	default:
+		fn = func() ([]float32, error) {
+			return nil, fmt.Errorf("embeddings not supported by the backend")
+		}
+	}
+
+	return func() ([]float32, error) {
+		// This is still needed, see: https://github.com/ggerganov/llama.cpp/discussions/784
+		l := Lock(modelFile)
+		defer l.Unlock()
+
+		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
+}

api/backend/image.go

@@ -0,0 +1,56 @@
+package backend
+
+import (
+	"fmt"
+	"sync"
+
+	config "github.com/go-skynet/LocalAI/api/config"
+	"github.com/go-skynet/LocalAI/api/options"
+	model "github.com/go-skynet/LocalAI/pkg/model"
+	"github.com/go-skynet/LocalAI/pkg/stablediffusion"
+)
+
+func ImageGeneration(height, width, mode, step, seed int, positive_prompt, negative_prompt, dst string, loader *model.ModelLoader, c config.Config, o *options.Option) (func() error, error) {
+	if c.Backend != model.StableDiffusionBackend {
+		return nil, fmt.Errorf("endpoint only working with stablediffusion models")
+	}
+
+	inferenceModel, err := loader.BackendLoader(
+		model.WithBackendString(c.Backend),
+		model.WithAssetDir(o.AssetsDestination),
+		model.WithThreads(uint32(c.Threads)),
+		model.WithModelFile(c.ImageGenerationAssets),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	var fn func() error
+	switch model := inferenceModel.(type) {
+	case *stablediffusion.StableDiffusion:
+		fn = func() error {
+			return model.GenerateImage(height, width, mode, step, seed, positive_prompt, negative_prompt, dst)
+		}
+
+	default:
+		fn = func() error {
+			return fmt.Errorf("creation of images not supported by the backend")
+		}
+	}
+
+	return func() error {
+		// This is still needed, see: https://github.com/ggerganov/llama.cpp/discussions/784
+		mutexMap.Lock()
+		l, ok := mutexes[c.Backend]
+		if !ok {
+			m := &sync.Mutex{}
+			mutexes[c.Backend] = m
+			l = m
+		}
+		mutexMap.Unlock()
+		l.Lock()
+		defer l.Unlock()
+
+		return fn()
+	}, nil
+}

api/backend/llm.go

@@ -0,0 +1,160 @@
+package backend
+
+import (
+	"context"
+	"regexp"
+	"strings"
+	"sync"
+
+	"github.com/donomii/go-rwkv.cpp"
+	config "github.com/go-skynet/LocalAI/api/config"
+	"github.com/go-skynet/LocalAI/api/options"
+	"github.com/go-skynet/LocalAI/pkg/grpc"
+	"github.com/go-skynet/LocalAI/pkg/langchain"
+	model "github.com/go-skynet/LocalAI/pkg/model"
+	"github.com/go-skynet/bloomz.cpp"
+)
+
+func ModelInference(s string, loader *model.ModelLoader, c config.Config, o *options.Option, tokenCallback func(string) bool) (func() (string, error), error) {
+	supportStreams := false
+	modelFile := c.Model
+
+	grpcOpts := gRPCModelOpts(c)
+
+	var inferenceModel interface{}
+	var err error
+
+	opts := []model.Option{
+		model.WithLoadGRPCOpts(grpcOpts),
+		model.WithThreads(uint32(c.Threads)), // GPT4all uses this
+		model.WithAssetDir(o.AssetsDestination),
+		model.WithModelFile(modelFile),
+	}
+
+	if c.Backend == "" {
+		inferenceModel, err = loader.GreedyLoader(opts...)
+	} else {
+		opts = append(opts, model.WithBackendString(c.Backend))
+		inferenceModel, err = loader.BackendLoader(opts...)
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	var fn func() (string, error)
+
+	switch model := inferenceModel.(type) {
+	case *rwkv.RwkvState:
+		supportStreams = true
+
+		fn = func() (string, error) {
+			stopWord := "\n"
+			if len(c.StopWords) > 0 {
+				stopWord = c.StopWords[0]
+			}
+
+			if err := model.ProcessInput(s); err != nil {
+				return "", err
+			}
+
+			response := model.GenerateResponse(c.Maxtokens, stopWord, float32(c.Temperature), float32(c.TopP), tokenCallback)
+
+			return response, nil
+		}
+	case *bloomz.Bloomz:
+		fn = func() (string, error) {
+			// Generate the prediction using the language model
+			predictOptions := []bloomz.PredictOption{
+				bloomz.SetTemperature(c.Temperature),
+				bloomz.SetTopP(c.TopP),
+				bloomz.SetTopK(c.TopK),
+				bloomz.SetTokens(c.Maxtokens),
+				bloomz.SetThreads(c.Threads),
+			}
+
+			if c.Seed != 0 {
+				predictOptions = append(predictOptions, bloomz.SetSeed(c.Seed))
+			}
+
+			return model.Predict(
+				s,
+				predictOptions...,
+			)
+		}
+
+	case *grpc.Client:
+		// in GRPC, the backend is supposed to answer to 1 single token if stream is not supported
+		supportStreams = true
+		fn = func() (string, error) {
+
+			opts := gRPCPredictOpts(c, loader.ModelPath)
+			opts.Prompt = s
+			if tokenCallback != nil {
+				ss := ""
+				err := model.PredictStream(context.TODO(), opts, func(s string) {
+					tokenCallback(s)
+					ss += s
+				})
+				return ss, err
+			} else {
+				reply, err := model.Predict(context.TODO(), opts)
+				return reply.Message, err
+			}
+		}
+	case *langchain.HuggingFace:
+		fn = func() (string, error) {
+
+			// Generate the prediction using the language model
+			predictOptions := []langchain.PredictOption{
+				langchain.SetModel(c.Model),
+				langchain.SetMaxTokens(c.Maxtokens),
+				langchain.SetTemperature(c.Temperature),
+				langchain.SetStopWords(c.StopWords),
+			}
+
+			pred, er := model.PredictHuggingFace(s, predictOptions...)
+			if er != nil {
+				return "", er
+			}
+			return pred.Completion, nil
+		}
+	}
+
+	return func() (string, error) {
+		// This is still needed, see: https://github.com/ggerganov/llama.cpp/discussions/784
+		l := Lock(modelFile)
+		defer l.Unlock()
+
+		res, err := fn()
+		if tokenCallback != nil && !supportStreams {
+			tokenCallback(res)
+		}
+		return res, err
+	}, nil
+}
+
+var cutstrings map[string]*regexp.Regexp = make(map[string]*regexp.Regexp)
+var mu sync.Mutex = sync.Mutex{}
+
+func Finetune(config config.Config, input, prediction string) string {
+	if config.Echo {
+		prediction = input + prediction
+	}
+
+	for _, c := range config.Cutstrings {
+		mu.Lock()
+		reg, ok := cutstrings[c]
+		if !ok {
+			cutstrings[c] = regexp.MustCompile(c)
+			reg = cutstrings[c]
+		}
+		mu.Unlock()
+		prediction = reg.ReplaceAllString(prediction, "")
+	}
+
+	for _, c := range config.TrimSpace {
+		prediction = strings.TrimSpace(strings.TrimPrefix(prediction, c))
+	}
+	return prediction
+
+}

api/backend/lock.go

@@ -0,0 +1,22 @@
+package backend
+
+import "sync"
+
+// mutex still needed, see: https://github.com/ggerganov/llama.cpp/discussions/784
+var mutexMap sync.Mutex
+var mutexes map[string]*sync.Mutex = make(map[string]*sync.Mutex)
+
+func Lock(s string) *sync.Mutex {
+	// This is still needed, see: https://github.com/ggerganov/llama.cpp/discussions/784
+	mutexMap.Lock()
+	l, ok := mutexes[s]
+	if !ok {
+		m := &sync.Mutex{}
+		mutexes[s] = m
+		l = m
+	}
+	mutexMap.Unlock()
+	l.Lock()
+
+	return l
+}

api/backend/options.go

@@ -0,0 +1,98 @@
+package backend
+
+import (
+	"os"
+	"path/filepath"
+
+	pb "github.com/go-skynet/LocalAI/pkg/grpc/proto"
+
+	config "github.com/go-skynet/LocalAI/api/config"
+	"github.com/go-skynet/LocalAI/pkg/langchain"
+	"github.com/go-skynet/bloomz.cpp"
+)
+
+func langchainOptions(c config.Config) []langchain.PredictOption {
+	return []langchain.PredictOption{
+		langchain.SetModel(c.Model),
+		langchain.SetMaxTokens(c.Maxtokens),
+		langchain.SetTemperature(c.Temperature),
+		langchain.SetStopWords(c.StopWords),
+	}
+}
+
+func bloomzOptions(c config.Config) []bloomz.PredictOption {
+	// Generate the prediction using the language model
+	predictOptions := []bloomz.PredictOption{
+		bloomz.SetTemperature(c.Temperature),
+		bloomz.SetTopP(c.TopP),
+		bloomz.SetTopK(c.TopK),
+		bloomz.SetTokens(c.Maxtokens),
+		bloomz.SetThreads(c.Threads),
+	}
+
+	if c.Seed != 0 {
+		predictOptions = append(predictOptions, bloomz.SetSeed(c.Seed))
+	}
+	return predictOptions
+}
+func gRPCModelOpts(c config.Config) *pb.ModelOptions {
+	b := 512
+	if c.Batch != 0 {
+		b = c.Batch
+	}
+	return &pb.ModelOptions{
+		ContextSize: int32(c.ContextSize),
+		Seed:        int32(c.Seed),
+		NBatch:      int32(b),
+		F16Memory:   c.F16,
+		MLock:       c.MMlock,
+		NUMA:        c.NUMA,
+		Embeddings:  c.Embeddings,
+		LowVRAM:     c.LowVRAM,
+		NGPULayers:  int32(c.NGPULayers),
+		MMap:        c.MMap,
+		MainGPU:     c.MainGPU,
+		Threads:     int32(c.Threads),
+		TensorSplit: c.TensorSplit,
+	}
+}
+
+func gRPCPredictOpts(c config.Config, modelPath string) *pb.PredictOptions {
+	promptCachePath := ""
+	if c.PromptCachePath != "" {
+		p := filepath.Join(modelPath, c.PromptCachePath)
+		os.MkdirAll(filepath.Dir(p), 0755)
+		promptCachePath = p
+	}
+	return &pb.PredictOptions{
+		Temperature:     float32(c.Temperature),
+		TopP:            float32(c.TopP),
+		TopK:            int32(c.TopK),
+		Tokens:          int32(c.Maxtokens),
+		Threads:         int32(c.Threads),
+		PromptCacheAll:  c.PromptCacheAll,
+		PromptCacheRO:   c.PromptCacheRO,
+		PromptCachePath: promptCachePath,
+		F16KV:           c.F16,
+		DebugMode:       c.Debug,
+		Grammar:         c.Grammar,
+
+		Mirostat:          int32(c.Mirostat),
+		MirostatETA:       float32(c.MirostatETA),
+		MirostatTAU:       float32(c.MirostatTAU),
+		Debug:             c.Debug,
+		StopPrompts:       c.StopWords,
+		Repeat:            int32(c.RepeatPenalty),
+		NKeep:             int32(c.Keep),
+		Batch:             int32(c.Batch),
+		IgnoreEOS:         c.IgnoreEOS,
+		Seed:              int32(c.Seed),
+		FrequencyPenalty:  float32(c.FrequencyPenalty),
+		MLock:             c.MMlock,
+		MMap:              c.MMap,
+		MainGPU:           c.MainGPU,
+		TensorSplit:       c.TensorSplit,
+		TailFreeSamplingZ: float32(c.TFZ),
+		TypicalP:          float32(c.TypicalP),
+	}
+}