feat: use gRPC for transformers

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

pkg/grpc/llm/transformers/dolly.go

@@ -0,0 +1,42 @@
+package transformers
+
+// This is a wrapper to statisfy the GRPC service interface
+// It is meant to be used by the main executable that is the server for the specific backend type (falcon, gpt3, etc)
+import (
+	"fmt"
+
+	pb "github.com/go-skynet/LocalAI/pkg/grpc/proto"
+
+	transformers "github.com/go-skynet/go-ggml-transformers.cpp"
+)
+
+type Dolly struct {
+	dolly *transformers.Dolly
+}
+
+func (llm *Dolly) Load(opts *pb.ModelOptions) error {
+	model, err := transformers.NewDolly(opts.Model)
+	llm.dolly = model
+	return err
+}
+
+func (llm *Dolly) Embeddings(opts *pb.PredictOptions) ([]float32, error) {
+	return nil, fmt.Errorf("not implemented")
+}
+
+func (llm *Dolly) Predict(opts *pb.PredictOptions) (string, error) {
+	return llm.dolly.Predict(opts.Prompt, buildPredictOptions(opts)...)
+}
+
+// fallback to Predict
+func (llm *Dolly) PredictStream(opts *pb.PredictOptions, results chan string) {
+	go func() {
+		res, err := llm.dolly.Predict(opts.Prompt, buildPredictOptions(opts)...)
+
+		if err != nil {
+			fmt.Println("err: ", err)
+		}
+		results <- res
+		close(results)
+	}()
+}

pkg/grpc/llm/transformers/gpt2.go

@@ -0,0 +1,42 @@
+package transformers
+
+// This is a wrapper to statisfy the GRPC service interface
+// It is meant to be used by the main executable that is the server for the specific backend type (falcon, gpt3, etc)
+import (
+	"fmt"
+
+	pb "github.com/go-skynet/LocalAI/pkg/grpc/proto"
+
+	transformers "github.com/go-skynet/go-ggml-transformers.cpp"
+)
+
+type GPT2 struct {
+	gpt2 *transformers.GPT2
+}
+
+func (llm *GPT2) Load(opts *pb.ModelOptions) error {
+	model, err := transformers.New(opts.Model)
+	llm.gpt2 = model
+	return err
+}
+
+func (llm *GPT2) Embeddings(opts *pb.PredictOptions) ([]float32, error) {
+	return nil, fmt.Errorf("not implemented")
+}
+
+func (llm *GPT2) Predict(opts *pb.PredictOptions) (string, error) {
+	return llm.gpt2.Predict(opts.Prompt, buildPredictOptions(opts)...)
+}
+
+// fallback to Predict
+func (llm *GPT2) PredictStream(opts *pb.PredictOptions, results chan string) {
+	go func() {
+		res, err := llm.gpt2.Predict(opts.Prompt, buildPredictOptions(opts)...)
+
+		if err != nil {
+			fmt.Println("err: ", err)
+		}
+		results <- res
+		close(results)
+	}()
+}

pkg/grpc/llm/transformers/gptj.go

@@ -0,0 +1,42 @@
+package transformers
+
+// This is a wrapper to statisfy the GRPC service interface
+// It is meant to be used by the main executable that is the server for the specific backend type (falcon, gpt3, etc)
+import (
+	"fmt"
+
+	pb "github.com/go-skynet/LocalAI/pkg/grpc/proto"
+
+	transformers "github.com/go-skynet/go-ggml-transformers.cpp"
+)
+
+type GPTJ struct {
+	gptj *transformers.GPTJ
+}
+
+func (llm *GPTJ) Load(opts *pb.ModelOptions) error {
+	model, err := transformers.NewGPTJ(opts.Model)
+	llm.gptj = model
+	return err
+}
+
+func (llm *GPTJ) Embeddings(opts *pb.PredictOptions) ([]float32, error) {
+	return nil, fmt.Errorf("not implemented")
+}
+
+func (llm *GPTJ) Predict(opts *pb.PredictOptions) (string, error) {
+	return llm.gptj.Predict(opts.Prompt, buildPredictOptions(opts)...)
+}
+
+// fallback to Predict
+func (llm *GPTJ) PredictStream(opts *pb.PredictOptions, results chan string) {
+	go func() {
+		res, err := llm.gptj.Predict(opts.Prompt, buildPredictOptions(opts)...)
+
+		if err != nil {
+			fmt.Println("err: ", err)
+		}
+		results <- res
+		close(results)
+	}()
+}

pkg/grpc/llm/transformers/gptneox.go

@@ -0,0 +1,42 @@
+package transformers
+
+// This is a wrapper to statisfy the GRPC service interface
+// It is meant to be used by the main executable that is the server for the specific backend type (falcon, gpt3, etc)
+import (
+	"fmt"
+
+	pb "github.com/go-skynet/LocalAI/pkg/grpc/proto"
+
+	transformers "github.com/go-skynet/go-ggml-transformers.cpp"
+)
+
+type GPTNeoX struct {
+	gptneox *transformers.GPTNeoX
+}
+
+func (llm *GPTNeoX) Load(opts *pb.ModelOptions) error {
+	model, err := transformers.NewGPTNeoX(opts.Model)
+	llm.gptneox = model
+	return err
+}
+
+func (llm *GPTNeoX) Embeddings(opts *pb.PredictOptions) ([]float32, error) {
+	return nil, fmt.Errorf("not implemented")
+}
+
+func (llm *GPTNeoX) Predict(opts *pb.PredictOptions) (string, error) {
+	return llm.gptneox.Predict(opts.Prompt, buildPredictOptions(opts)...)
+}
+
+// fallback to Predict
+func (llm *GPTNeoX) PredictStream(opts *pb.PredictOptions, results chan string) {
+	go func() {
+		res, err := llm.gptneox.Predict(opts.Prompt, buildPredictOptions(opts)...)
+
+		if err != nil {
+			fmt.Println("err: ", err)
+		}
+		results <- res
+		close(results)
+	}()
+}

pkg/grpc/llm/transformers/mpt.go

@@ -0,0 +1,42 @@
+package transformers
+
+// This is a wrapper to statisfy the GRPC service interface
+// It is meant to be used by the main executable that is the server for the specific backend type (falcon, gpt3, etc)
+import (
+	"fmt"
+
+	pb "github.com/go-skynet/LocalAI/pkg/grpc/proto"
+
+	transformers "github.com/go-skynet/go-ggml-transformers.cpp"
+)
+
+type MPT struct {
+	mpt *transformers.MPT
+}
+
+func (llm *MPT) Load(opts *pb.ModelOptions) error {
+	model, err := transformers.NewMPT(opts.Model)
+	llm.mpt = model
+	return err
+}
+
+func (llm *MPT) Embeddings(opts *pb.PredictOptions) ([]float32, error) {
+	return nil, fmt.Errorf("not implemented")
+}
+
+func (llm *MPT) Predict(opts *pb.PredictOptions) (string, error) {
+	return llm.mpt.Predict(opts.Prompt, buildPredictOptions(opts)...)
+}
+
+// fallback to Predict
+func (llm *MPT) PredictStream(opts *pb.PredictOptions, results chan string) {
+	go func() {
+		res, err := llm.mpt.Predict(opts.Prompt, buildPredictOptions(opts)...)
+
+		if err != nil {
+			fmt.Println("err: ", err)
+		}
+		results <- res
+		close(results)
+	}()
+}

pkg/grpc/llm/transformers/predict.go

@@ -0,0 +1,26 @@
+package transformers
+
+import (
+	pb "github.com/go-skynet/LocalAI/pkg/grpc/proto"
+	transformers "github.com/go-skynet/go-ggml-transformers.cpp"
+)
+
+func buildPredictOptions(opts *pb.PredictOptions) []transformers.PredictOption {
+	predictOptions := []transformers.PredictOption{
+		transformers.SetTemperature(float64(opts.Temperature)),
+		transformers.SetTopP(float64(opts.TopP)),
+		transformers.SetTopK(int(opts.TopK)),
+		transformers.SetTokens(int(opts.Tokens)),
+		transformers.SetThreads(int(opts.Threads)),
+	}
+
+	if opts.Batch != 0 {
+		predictOptions = append(predictOptions, transformers.SetBatch(int(opts.Batch)))
+	}
+
+	if opts.Seed != 0 {
+		predictOptions = append(predictOptions, transformers.SetSeed(int(opts.Seed)))
+	}
+
+	return predictOptions
+}

pkg/grpc/llm/transformers/replit.go

@@ -0,0 +1,42 @@
+package transformers
+
+// This is a wrapper to statisfy the GRPC service interface
+// It is meant to be used by the main executable that is the server for the specific backend type (falcon, gpt3, etc)
+import (
+	"fmt"
+
+	pb "github.com/go-skynet/LocalAI/pkg/grpc/proto"
+
+	transformers "github.com/go-skynet/go-ggml-transformers.cpp"
+)
+
+type Replit struct {
+	replit *transformers.Replit
+}
+
+func (llm *Replit) Load(opts *pb.ModelOptions) error {
+	model, err := transformers.NewReplit(opts.Model)
+	llm.replit = model
+	return err
+}
+
+func (llm *Replit) Embeddings(opts *pb.PredictOptions) ([]float32, error) {
+	return nil, fmt.Errorf("not implemented")
+}
+
+func (llm *Replit) Predict(opts *pb.PredictOptions) (string, error) {
+	return llm.replit.Predict(opts.Prompt, buildPredictOptions(opts)...)
+}
+
+// fallback to Predict
+func (llm *Replit) PredictStream(opts *pb.PredictOptions, results chan string) {
+	go func() {
+		res, err := llm.replit.Predict(opts.Prompt, buildPredictOptions(opts)...)
+
+		if err != nil {
+			fmt.Println("err: ", err)
+		}
+		results <- res
+		close(results)
+	}()
+}

pkg/grpc/llm/transformers/starcoder.go

@@ -0,0 +1,42 @@
+package transformers
+
+// This is a wrapper to statisfy the GRPC service interface
+// It is meant to be used by the main executable that is the server for the specific backend type (falcon, gpt3, etc)
+import (
+	"fmt"
+
+	pb "github.com/go-skynet/LocalAI/pkg/grpc/proto"
+
+	transformers "github.com/go-skynet/go-ggml-transformers.cpp"
+)
+
+type Starcoder struct {
+	starcoder *transformers.Starcoder
+}
+
+func (llm *Starcoder) Load(opts *pb.ModelOptions) error {
+	model, err := transformers.NewStarcoder(opts.Model)
+	llm.starcoder = model
+	return err
+}
+
+func (llm *Starcoder) Embeddings(opts *pb.PredictOptions) ([]float32, error) {
+	return nil, fmt.Errorf("not implemented")
+}
+
+func (llm *Starcoder) Predict(opts *pb.PredictOptions) (string, error) {
+	return llm.starcoder.Predict(opts.Prompt, buildPredictOptions(opts)...)
+}
+
+// fallback to Predict
+func (llm *Starcoder) PredictStream(opts *pb.PredictOptions, results chan string) {
+	go func() {
+		res, err := llm.starcoder.Predict(opts.Prompt, buildPredictOptions(opts)...)
+
+		if err != nil {
+			fmt.Println("err: ", err)
+		}
+		results <- res
+		close(results)
+	}()
+}