fix(download): correctly check for not found error (#1514)

mudler · Dec 30, 2023 · a95bb05 · a95bb05
1 parent e2311a1
commit a95bb05
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git 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
@@ -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]