LocalAI/pkg/model/initializers.go

208 lines
6.1 KiB
Go
Raw Normal View History

2023-05-11 14:34:16 +00:00
package model
import (
"fmt"
"path/filepath"
2023-05-11 14:34:16 +00:00
"strings"
rwkv "github.com/donomii/go-rwkv.cpp"
whisper "github.com/ggerganov/whisper.cpp/bindings/go/pkg/whisper"
"github.com/go-skynet/LocalAI/pkg/langchain"
"github.com/go-skynet/LocalAI/pkg/stablediffusion"
2023-05-11 14:34:16 +00:00
bloomz "github.com/go-skynet/bloomz.cpp"
bert "github.com/go-skynet/go-bert.cpp"
transformers "github.com/go-skynet/go-ggml-transformers.cpp"
2023-05-11 14:34:16 +00:00
llama "github.com/go-skynet/go-llama.cpp"
"github.com/hashicorp/go-multierror"
gpt4all "github.com/nomic-ai/gpt4all/gpt4all-bindings/golang"
2023-05-11 14:34:16 +00:00
"github.com/rs/zerolog/log"
)
const tokenizerSuffix = ".tokenizer.json"
const (
LlamaBackend = "llama"
BloomzBackend = "bloomz"
StarcoderBackend = "starcoder"
GPTJBackend = "gptj"
DollyBackend = "dolly"
MPTBackend = "mpt"
GPTNeoXBackend = "gptneox"
ReplitBackend = "replit"
Gpt2Backend = "gpt2"
Gpt4AllLlamaBackend = "gpt4all-llama"
Gpt4AllMptBackend = "gpt4all-mpt"
Gpt4AllJBackend = "gpt4all-j"
Gpt4All = "gpt4all"
FalconBackend = "falcon"
BertEmbeddingsBackend = "bert-embeddings"
RwkvBackend = "rwkv"
WhisperBackend = "whisper"
StableDiffusionBackend = "stablediffusion"
LCHuggingFaceBackend = "langchain-huggingface"
2023-05-11 14:34:16 +00:00
)
var autoLoadBackends []string = []string{
2023-05-11 14:34:16 +00:00
LlamaBackend,
Gpt4All,
2023-05-11 14:34:16 +00:00
RwkvBackend,
GPTNeoXBackend,
WhisperBackend,
2023-05-11 14:34:16 +00:00
BertEmbeddingsBackend,
GPTJBackend,
Gpt2Backend,
DollyBackend,
FalconBackend,
MPTBackend,
ReplitBackend,
2023-05-11 18:20:07 +00:00
StarcoderBackend,
BloomzBackend,
2023-05-11 18:20:07 +00:00
}
var starCoder = func(modelFile string) (interface{}, error) {
return transformers.NewStarcoder(modelFile)
2023-05-11 14:34:16 +00:00
}
var mpt = func(modelFile string) (interface{}, error) {
return transformers.NewMPT(modelFile)
2023-05-11 14:34:16 +00:00
}
var dolly = func(modelFile string) (interface{}, error) {
return transformers.NewDolly(modelFile)
2023-05-11 14:34:16 +00:00
}
2023-05-12 09:36:35 +00:00
var gptNeoX = func(modelFile string) (interface{}, error) {
return transformers.NewGPTNeoX(modelFile)
2023-05-12 09:36:35 +00:00
}
var replit = func(modelFile string) (interface{}, error) {
return transformers.NewReplit(modelFile)
2023-05-12 09:36:35 +00:00
}
var gptJ = func(modelFile string) (interface{}, error) {
return transformers.NewGPTJ(modelFile)
2023-05-11 14:34:16 +00:00
}
var falcon = func(modelFile string) (interface{}, error) {
return transformers.NewFalcon(modelFile)
}
2023-05-11 14:34:16 +00:00
var bertEmbeddings = func(modelFile string) (interface{}, error) {
return bert.New(modelFile)
}
var bloomzLM = func(modelFile string) (interface{}, error) {
return bloomz.New(modelFile)
}
var transformersLM = func(modelFile string) (interface{}, error) {
return transformers.New(modelFile)
2023-05-11 14:34:16 +00:00
}
var stableDiffusion = func(assetDir string) (interface{}, error) {
return stablediffusion.New(assetDir)
}
2023-05-11 14:34:16 +00:00
var whisperModel = func(modelFile string) (interface{}, error) {
return whisper.New(modelFile)
}
var lcHuggingFace = func(repoId string) (interface{}, error) {
return langchain.NewHuggingFace(repoId)
}
2023-05-11 14:34:16 +00:00
func llamaLM(opts ...llama.ModelOption) func(string) (interface{}, error) {
return func(s string) (interface{}, error) {
return llama.New(s, opts...)
}
}
func gpt4allLM(opts ...gpt4all.ModelOption) func(string) (interface{}, error) {
return func(s string) (interface{}, error) {
return gpt4all.New(s, opts...)
}
}
func rwkvLM(tokenFile string, threads uint32) func(string) (interface{}, error) {
return func(s string) (interface{}, error) {
log.Debug().Msgf("Loading RWKV", s, tokenFile)
2023-05-11 14:34:16 +00:00
model := rwkv.LoadFiles(s, tokenFile, threads)
if model == nil {
return nil, fmt.Errorf("could not load model")
}
return model, nil
}
}
func (ml *ModelLoader) BackendLoader(backendString string, modelFile string, llamaOpts []llama.ModelOption, threads uint32, assetDir string) (model interface{}, err error) {
log.Debug().Msgf("Loading model %s from %s", backendString, modelFile)
2023-05-11 14:34:16 +00:00
switch strings.ToLower(backendString) {
case LlamaBackend:
return ml.LoadModel(modelFile, llamaLM(llamaOpts...))
case BloomzBackend:
return ml.LoadModel(modelFile, bloomzLM)
case GPTJBackend:
return ml.LoadModel(modelFile, gptJ)
2023-05-11 14:34:16 +00:00
case DollyBackend:
return ml.LoadModel(modelFile, dolly)
case MPTBackend:
return ml.LoadModel(modelFile, mpt)
2023-05-11 14:34:16 +00:00
case Gpt2Backend:
return ml.LoadModel(modelFile, transformersLM)
case FalconBackend:
return ml.LoadModel(modelFile, falcon)
2023-05-12 09:36:35 +00:00
case GPTNeoXBackend:
return ml.LoadModel(modelFile, gptNeoX)
case ReplitBackend:
return ml.LoadModel(modelFile, replit)
case StableDiffusionBackend:
return ml.LoadModel(modelFile, stableDiffusion)
2023-05-11 18:20:07 +00:00
case StarcoderBackend:
return ml.LoadModel(modelFile, starCoder)
case Gpt4AllLlamaBackend, Gpt4AllMptBackend, Gpt4AllJBackend, Gpt4All:
return ml.LoadModel(modelFile, gpt4allLM(gpt4all.SetThreads(int(threads)), gpt4all.SetLibrarySearchPath(filepath.Join(assetDir, "backend-assets", "gpt4all"))))
2023-05-11 14:34:16 +00:00
case BertEmbeddingsBackend:
return ml.LoadModel(modelFile, bertEmbeddings)
case RwkvBackend:
return ml.LoadModel(modelFile, rwkvLM(filepath.Join(ml.ModelPath, modelFile+tokenizerSuffix), threads))
2023-05-11 14:34:16 +00:00
case WhisperBackend:
return ml.LoadModel(modelFile, whisperModel)
case LCHuggingFaceBackend:
return ml.LoadModel(modelFile, lcHuggingFace)
2023-05-11 14:34:16 +00:00
default:
return nil, fmt.Errorf("backend unsupported: %s", backendString)
}
}
func (ml *ModelLoader) GreedyLoader(modelFile string, llamaOpts []llama.ModelOption, threads uint32, assetDir string) (interface{}, error) {
log.Debug().Msgf("Loading model '%s' greedly", modelFile)
2023-05-11 14:34:16 +00:00
ml.mu.Lock()
m, exists := ml.models[modelFile]
if exists {
log.Debug().Msgf("Model '%s' already loaded", modelFile)
2023-05-11 14:34:16 +00:00
ml.mu.Unlock()
return m, nil
}
ml.mu.Unlock()
var err error
for _, b := range autoLoadBackends {
2023-05-11 14:34:16 +00:00
if b == BloomzBackend || b == WhisperBackend || b == RwkvBackend { // do not autoload bloomz/whisper/rwkv
continue
}
log.Debug().Msgf("[%s] Attempting to load", b)
model, modelerr := ml.BackendLoader(b, modelFile, llamaOpts, threads, assetDir)
2023-05-11 14:34:16 +00:00
if modelerr == nil && model != nil {
log.Debug().Msgf("[%s] Loads OK", b)
return model, nil
} else if modelerr != nil {
err = multierror.Append(err, modelerr)
log.Debug().Msgf("[%s] Fails: %s", b, modelerr.Error())
}
}
return nil, fmt.Errorf("could not load model - all backends returned error: %s", err.Error())
}