Fix absolute path names + gguf detection (#8428)

This commit is contained in:
Patrick Devine 2025-01-14 19:01:24 -08:00 committed by GitHub
parent 61676fb506
commit 2539f2dbf9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 217 additions and 28 deletions

View file

@ -3,6 +3,7 @@ package server
import (
"bytes"
"cmp"
"crypto/sha256"
"encoding/json"
"fmt"
"io"
@ -710,3 +711,100 @@ func TestCreateDetectTemplate(t *testing.T) {
})
})
}
func TestDetectModelTypeFromFiles(t *testing.T) {
t.Run("gguf file", func(t *testing.T) {
_, digest := createBinFile(t, nil, nil)
files := map[string]string{
"model.gguf": digest,
}
modelType := detectModelTypeFromFiles(files)
if modelType != "gguf" {
t.Fatalf("expected model type 'gguf', got %q", modelType)
}
})
t.Run("gguf file w/o extension", func(t *testing.T) {
_, digest := createBinFile(t, nil, nil)
files := map[string]string{
fmt.Sprintf("%x", digest): digest,
}
modelType := detectModelTypeFromFiles(files)
if modelType != "gguf" {
t.Fatalf("expected model type 'gguf', got %q", modelType)
}
})
t.Run("safetensors file", func(t *testing.T) {
files := map[string]string{
"model.safetensors": "sha256:abc123",
}
modelType := detectModelTypeFromFiles(files)
if modelType != "safetensors" {
t.Fatalf("expected model type 'safetensors', got %q", modelType)
}
})
t.Run("unsupported file type", func(t *testing.T) {
p := t.TempDir()
t.Setenv("OLLAMA_MODELS", p)
data := []byte("12345678")
digest := fmt.Sprintf("sha256:%x", sha256.Sum256(data))
if err := os.MkdirAll(filepath.Join(p, "blobs"), 0o755); err != nil {
t.Fatal(err)
}
f, err := os.Create(filepath.Join(p, "blobs", fmt.Sprintf("sha256-%s", strings.TrimPrefix(digest, "sha256:"))))
if err != nil {
t.Fatal(err)
}
defer f.Close()
if _, err := f.Write(data); err != nil {
t.Fatal(err)
}
files := map[string]string{
"model.bin": digest,
}
modelType := detectModelTypeFromFiles(files)
if modelType != "" {
t.Fatalf("expected empty model type for unsupported file, got %q", modelType)
}
})
t.Run("file with less than 4 bytes", func(t *testing.T) {
p := t.TempDir()
t.Setenv("OLLAMA_MODELS", p)
data := []byte("123")
digest := fmt.Sprintf("sha256:%x", sha256.Sum256(data))
if err := os.MkdirAll(filepath.Join(p, "blobs"), 0o755); err != nil {
t.Fatal(err)
}
f, err := os.Create(filepath.Join(p, "blobs", fmt.Sprintf("sha256-%s", strings.TrimPrefix(digest, "sha256:"))))
if err != nil {
t.Fatal(err)
}
defer f.Close()
if _, err := f.Write(data); err != nil {
t.Fatal(err)
}
files := map[string]string{
"noext": digest,
}
modelType := detectModelTypeFromFiles(files)
if modelType != "" {
t.Fatalf("expected empty model type for small file, got %q", modelType)
}
})
}