mirror of
https://github.com/mudler/LocalAI.git
synced 2025-05-30 23:44:59 +00:00
feat: HF /scan endpoint (#2566)
* start by checking /scan during the checksum update Signed-off-by: Dave Lee <dave@gray101.com> * add back in golang side features: downloader/uri gets struct and scan function, gallery uses it, and secscan/models calls it. Signed-off-by: Dave Lee <dave@gray101.com> * add a param to scan specific urls - useful for debugging Signed-off-by: Dave Lee <dave@gray101.com> * helpful printouts Signed-off-by: Dave Lee <dave@gray101.com> * fix offsets Signed-off-by: Dave Lee <dave@gray101.com> * fix error and naming Signed-off-by: Dave Lee <dave@gray101.com> * expose error Signed-off-by: Dave Lee <dave@gray101.com> * fix json tags Signed-off-by: Dave Lee <dave@gray101.com> * slight wording change Signed-off-by: Dave Lee <dave@gray101.com> * go mod tidy - getting warnings Signed-off-by: Dave Lee <dave@gray101.com> * split out python to make editing easier, add some simple code to delete contaminated entries from gallery Signed-off-by: Dave Lee <dave@gray101.com> * o7 to my favorite part of our old name, go-skynet Signed-off-by: Dave Lee <dave@gray101.com> * merge fix Signed-off-by: Dave Lee <dave@gray101.com> * merge fix Signed-off-by: Dave Lee <dave@gray101.com> * merge fix Signed-off-by: Dave Lee <dave@gray101.com> * address review comments Signed-off-by: Dave Lee <dave@gray101.com> * forgot secscan could accept multiple URL at once Signed-off-by: Dave Lee <dave@gray101.com> * invert naming and actually use it Signed-off-by: Dave Lee <dave@gray101.com> * missed cli/models.go Signed-off-by: Dave Lee <dave@gray101.com> * Update .github/check_and_update.py Co-authored-by: Ettore Di Giacinto <mudler@users.noreply.github.com> Signed-off-by: Dave <dave@gray101.com> --------- Signed-off-by: Dave Lee <dave@gray101.com> Signed-off-by: Dave <dave@gray101.com> Co-authored-by: Ettore Di Giacinto <mudler@users.noreply.github.com>
This commit is contained in:
Dave
GitHub
parent
cbb93bd8ec
commit
133987b1fb
15 changed files with 282 additions and 125 deletions
|
|
@ -3,6 +3,8 @@ package downloader
|
|||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
|
|
@ -129,6 +131,7 @@ func ConvertURL(s string) string {
|
|||
// 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]
|
||||
repo := strings.Split(repository, "/")[1]
|
||||
|
||||
branch := "main"
|
||||
if strings.Contains(repo, "@") {
|
||||
branch = strings.Split(repository, "@")[1]
|
||||
|
|
@ -353,3 +356,42 @@ func calculateSHA(filePath string) (string, error) {
|
|||
|
||||
return fmt.Sprintf("%x", hash.Sum(nil)), nil
|
||||
}
|
||||
|
||||
type HuggingFaceScanResult struct {
|
||||
RepositoryId string `json:"repositoryId"`
|
||||
Revision string `json:"revision"`
|
||||
HasUnsafeFiles bool `json:"hasUnsafeFile"`
|
||||
ClamAVInfectedFiles []string `json:"clamAVInfectedFiles"`
|
||||
DangerousPickles []string `json:"dangerousPickles"`
|
||||
ScansDone bool `json:"scansDone"`
|
||||
}
|
||||
|
||||
var ErrNonHuggingFaceFile = errors.New("not a huggingface repo")
|
||||
var ErrUnsafeFilesFound = errors.New("unsafe files found")
|
||||
|
||||
func HuggingFaceScan(uri string) (*HuggingFaceScanResult, error) {
|
||||
cleanParts := strings.Split(ConvertURL(uri), "/")
|
||||
if len(cleanParts) <= 4 || cleanParts[2] != "huggingface.co" {
|
||||
return nil, ErrNonHuggingFaceFile
|
||||
}
|
||||
results, err := http.Get(fmt.Sprintf("https://huggingface.co/api/models/%s/%s/scan", cleanParts[3], cleanParts[4]))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if results.StatusCode != 200 {
|
||||
return nil, fmt.Errorf("unexpected status code during HuggingFaceScan: %d", results.StatusCode)
|
||||
}
|
||||
scanResult := &HuggingFaceScanResult{}
|
||||
bodyBytes, err := io.ReadAll(results.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
err = json.Unmarshal(bodyBytes, scanResult)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if scanResult.HasUnsafeFiles {
|
||||
return scanResult, ErrUnsafeFilesFound
|
||||
}
|
||||
return scanResult, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ import (
|
|||
// InstallModels will preload models from the given list of URLs and galleries
|
||||
// It will download the model if it is not already present in the model path
|
||||
// It will also try to resolve if the model is an embedded model YAML configuration
|
||||
func InstallModels(galleries []config.Gallery, modelLibraryURL string, modelPath string, downloadStatus func(string, string, string, float64), models ...string) error {
|
||||
func InstallModels(galleries []config.Gallery, modelLibraryURL string, modelPath string, enforceScan bool, downloadStatus func(string, string, string, float64), models ...string) error {
|
||||
// create an error that groups all errors
|
||||
var err error
|
||||
|
||||
|
|
@ -113,7 +113,7 @@ func InstallModels(galleries []config.Gallery, modelLibraryURL string, modelPath
|
|||
}
|
||||
} else {
|
||||
// Check if it's a model gallery, or print a warning
|
||||
e, found := installModel(galleries, url, modelPath, downloadStatus)
|
||||
e, found := installModel(galleries, url, modelPath, downloadStatus, enforceScan)
|
||||
if e != nil && found {
|
||||
log.Error().Err(err).Msgf("[startup] failed installing model '%s'", url)
|
||||
err = errors.Join(err, e)
|
||||
|
|
@ -127,7 +127,7 @@ func InstallModels(galleries []config.Gallery, modelLibraryURL string, modelPath
|
|||
return err
|
||||
}
|
||||
|
||||
func installModel(galleries []config.Gallery, modelName, modelPath string, downloadStatus func(string, string, string, float64)) (error, bool) {
|
||||
func installModel(galleries []config.Gallery, modelName, modelPath string, downloadStatus func(string, string, string, float64), enforceScan bool) (error, bool) {
|
||||
models, err := gallery.AvailableGalleryModels(galleries, modelPath)
|
||||
if err != nil {
|
||||
return err, false
|
||||
|
|
@ -143,7 +143,7 @@ func installModel(galleries []config.Gallery, modelName, modelPath string, downl
|
|||
}
|
||||
|
||||
log.Info().Str("model", modelName).Str("license", model.License).Msg("installing model")
|
||||
err = gallery.InstallModelFromGallery(galleries, modelName, modelPath, gallery.GalleryModel{}, downloadStatus)
|
||||
err = gallery.InstallModelFromGallery(galleries, modelName, modelPath, gallery.GalleryModel{}, downloadStatus, enforceScan)
|
||||
if err != nil {
|
||||
return err, true
|
||||
}
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ var _ = Describe("Preload test", func() {
|
|||
libraryURL := "https://raw.githubusercontent.com/mudler/LocalAI/master/embedded/model_library.yaml"
|
||||
fileName := fmt.Sprintf("%s.yaml", "1701d57f28d47552516c2b6ecc3cc719")
|
||||
|
||||
InstallModels([]config.Gallery{}, libraryURL, tmpdir, nil, "phi-2")
|
||||
InstallModels([]config.Gallery{}, libraryURL, tmpdir, true, nil, "phi-2")
|
||||
|
||||
resultFile := filepath.Join(tmpdir, fileName)
|
||||
|
||||
|
|
@ -38,7 +38,7 @@ var _ = Describe("Preload test", func() {
|
|||
url := "https://raw.githubusercontent.com/mudler/LocalAI/master/examples/configurations/phi-2.yaml"
|
||||
fileName := fmt.Sprintf("%s.yaml", utils.MD5(url))
|
||||
|
||||
InstallModels([]config.Gallery{}, "", tmpdir, nil, url)
|
||||
InstallModels([]config.Gallery{}, "", tmpdir, true, nil, url)
|
||||
|
||||
resultFile := filepath.Join(tmpdir, fileName)
|
||||
|
||||
|
|
@ -52,7 +52,7 @@ var _ = Describe("Preload test", func() {
|
|||
Expect(err).ToNot(HaveOccurred())
|
||||
url := "phi-2"
|
||||
|
||||
InstallModels([]config.Gallery{}, "", tmpdir, nil, url)
|
||||
InstallModels([]config.Gallery{}, "", tmpdir, true, nil, url)
|
||||
|
||||
entry, err := os.ReadDir(tmpdir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
|
@ -70,7 +70,7 @@ var _ = Describe("Preload test", func() {
|
|||
url := "mistral-openorca"
|
||||
fileName := fmt.Sprintf("%s.yaml", utils.MD5(url))
|
||||
|
||||
InstallModels([]config.Gallery{}, "", tmpdir, nil, url)
|
||||
InstallModels([]config.Gallery{}, "", tmpdir, true, nil, url)
|
||||
|
||||
resultFile := filepath.Join(tmpdir, fileName)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue