mirror of
https://github.com/mudler/LocalAI.git
synced 2026-03-31 21:25:59 -04:00
76 lines
2.9 KiB
Go
76 lines
2.9 KiB
Go
package vram_test
|
|
|
|
import (
|
|
. "github.com/mudler/LocalAI/pkg/vram"
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
)
|
|
|
|
var _ = Describe("ParseSizeString", func() {
|
|
DescribeTable("valid sizes",
|
|
func(input string, expected uint64) {
|
|
got, err := ParseSizeString(input)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(got).To(BeNumerically("~", expected, float64(expected)*0.0001+1))
|
|
},
|
|
Entry("500MB", "500MB", uint64(500_000_000)),
|
|
Entry("14.5GB", "14.5GB", uint64(14_500_000_000)),
|
|
Entry("2TB", "2TB", uint64(2_000_000_000_000)),
|
|
Entry("1024KB", "1024KB", uint64(1_024_000)),
|
|
Entry("100B", "100B", uint64(100)),
|
|
Entry("75 MB with space", "75 MB", uint64(75_000_000)),
|
|
Entry("1.5 gb lowercase", "1.5 gb", uint64(1_500_000_000)),
|
|
Entry("0.5GB", "0.5GB", uint64(500_000_000)),
|
|
Entry("3PB", "3PB", uint64(3_000_000_000_000_000)),
|
|
Entry("short suffix 100M", "100M", uint64(100_000_000)),
|
|
Entry("short suffix 2G", "2G", uint64(2_000_000_000)),
|
|
Entry("short suffix 1K", "1K", uint64(1_000)),
|
|
)
|
|
|
|
DescribeTable("invalid sizes",
|
|
func(input string) {
|
|
_, err := ParseSizeString(input)
|
|
Expect(err).To(HaveOccurred())
|
|
},
|
|
Entry("empty", ""),
|
|
Entry("suffix only", "MB"),
|
|
Entry("letters only", "abc"),
|
|
Entry("negative", "-5GB"),
|
|
Entry("unknown suffix", "5XB"),
|
|
)
|
|
})
|
|
|
|
var _ = Describe("ExtractHFRepoID", func() {
|
|
DescribeTable("valid repo IDs",
|
|
func(input, expectedID string) {
|
|
gotID, gotOK := ExtractHFRepoID(input)
|
|
Expect(gotOK).To(BeTrue())
|
|
Expect(gotID).To(Equal(expectedID))
|
|
},
|
|
Entry("short form", "Wan-AI/Wan2.2-I2V-A14B-Diffusers", "Wan-AI/Wan2.2-I2V-A14B-Diffusers"),
|
|
Entry("short form 2", "meta-llama/Llama-3-8B", "meta-llama/Llama-3-8B"),
|
|
Entry("https URL", "https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers", "Wan-AI/Wan2.2-I2V-A14B-Diffusers"),
|
|
Entry("http URL", "http://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers", "Wan-AI/Wan2.2-I2V-A14B-Diffusers"),
|
|
Entry("no scheme", "huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers", "Wan-AI/Wan2.2-I2V-A14B-Diffusers"),
|
|
Entry("trailing slash", "https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers/", "Wan-AI/Wan2.2-I2V-A14B-Diffusers"),
|
|
Entry("extra path", "https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers/tree/main", "Wan-AI/Wan2.2-I2V-A14B-Diffusers"),
|
|
Entry("uppercase URL", "HTTPS://HUGGINGFACE.CO/org/model", "org/model"),
|
|
)
|
|
|
|
DescribeTable("invalid inputs",
|
|
func(input string) {
|
|
_, gotOK := ExtractHFRepoID(input)
|
|
Expect(gotOK).To(BeFalse())
|
|
},
|
|
Entry("empty", ""),
|
|
Entry("single word", "single-word"),
|
|
Entry("three parts", "llama-cpp/models/file.gguf"),
|
|
Entry("non-HF URL", "https://example.com/org/model"),
|
|
Entry("wrong scheme", "ftp://huggingface.co/org/model"),
|
|
Entry("has space", "has spaces/model"),
|
|
Entry("incomplete URL", "huggingface.co/"),
|
|
Entry("org only", "huggingface.co/org"),
|
|
Entry("empty org", "huggingface.co//model"),
|
|
)
|
|
})
|