diff --git a/api/config/config.go b/api/config/config.go index 7ed7061af91..4c28f49b33a 100644 --- a/api/config/config.go +++ b/api/config/config.go @@ -1,6 +1,7 @@ package api_config import ( + "errors" "fmt" "io/fs" "os" @@ -270,15 +271,19 @@ func (cm *ConfigLoader) Preload(modelPath string) error { cm.Lock() defer cm.Unlock() + log.Info().Msgf("Preloading models from %s", modelPath) + for i, config := range cm.configs { + modelURL := config.PredictionOptions.Model modelURL = utils.ConvertURL(modelURL) - if strings.HasPrefix(modelURL, "http://") || strings.HasPrefix(modelURL, "https://") { + + if utils.LooksLikeURL(modelURL) { // md5 of model name md5Name := utils.MD5(modelURL) // check if file exists - if _, err := os.Stat(filepath.Join(modelPath, md5Name)); err == os.ErrNotExist { + if _, err := os.Stat(filepath.Join(modelPath, md5Name)); errors.Is(err, os.ErrNotExist) { err := utils.DownloadFile(modelURL, filepath.Join(modelPath, md5Name), "", func(fileName, current, total string, percent float64) { log.Info().Msgf("Downloading %s: %s/%s (%.2f%%)", fileName, current, total, percent) }) diff --git a/pkg/utils/uri.go b/pkg/utils/uri.go index 260e65a9a1a..5dde047decb 100644 --- a/pkg/utils/uri.go +++ b/pkg/utils/uri.go @@ -71,10 +71,18 @@ func GetURI(url string, f func(url string, i []byte) error) error { return f(url, body) } +const ( + HuggingFacePrefix = "huggingface://" +) + +func LooksLikeURL(s string) bool { + return strings.HasPrefix(s, "http://") || strings.HasPrefix(s, "https://") || strings.HasPrefix(s, HuggingFacePrefix) +} + func ConvertURL(s string) string { switch { - case strings.HasPrefix(s, "huggingface://"): - repository := strings.Replace(s, "huggingface://", "", 1) + case strings.HasPrefix(s, HuggingFacePrefix): + repository := strings.Replace(s, HuggingFacePrefix, "", 1) // convert repository to a full URL. // e.g. TheBloke/Mixtral-8x7B-v0.1-GGUF/mixtral-8x7b-v0.1.Q2_K.gguf@main -> https://huggingface.co/TheBloke/Mixtral-8x7B-v0.1-GGUF/resolve/main/mixtral-8x7b-v0.1.Q2_K.gguf owner := strings.Split(repository, "/")[0]