diff --git a/fs/ggml/gguf_test.go b/fs/ggml/gguf_test.go index e56bab8d26..091e5feea5 100644 --- a/fs/ggml/gguf_test.go +++ b/fs/ggml/gguf_test.go @@ -4,6 +4,8 @@ import ( "bytes" "math/rand/v2" "os" + "slices" + "strconv" "strings" "testing" @@ -81,3 +83,47 @@ func TestWriteGGUF(t *testing.T) { }) } } + +func BenchmarkReadArray(b *testing.B) { + b.ReportAllocs() + + create := func(tb testing.TB, kv KV) string { + tb.Helper() + f, err := os.CreateTemp(b.TempDir(), "") + if err != nil { + b.Fatal(err) + } + defer f.Close() + + if err := WriteGGUF(f, kv, nil); err != nil { + b.Fatal(err) + } + + return f.Name() + } + + cases := map[string]any{ + "int32": slices.Repeat([]int32{42}, 1_000_000), + "uint32": slices.Repeat([]uint32{42}, 1_000_000), + "float32": slices.Repeat([]float32{42.}, 1_000_000), + "string": slices.Repeat([]string{"42"}, 1_000_000), + } + + for name, bb := range cases { + for _, maxArraySize := range []int{-1, 0, 1024} { + b.Run(name+"-maxArraySize="+strconv.Itoa(maxArraySize), func(b *testing.B) { + p := create(b, KV{"array": bb}) + for b.Loop() { + f, err := os.Open(p) + if err != nil { + b.Fatal(err) + } + if _, err := Decode(f, maxArraySize); err != nil { + b.Fatal(err) + } + f.Close() + } + }) + } + } +} diff --git a/fs/gguf/gguf_test.go b/fs/gguf/gguf_test.go index eea28a4808..b119c5cf12 100644 --- a/fs/gguf/gguf_test.go +++ b/fs/gguf/gguf_test.go @@ -3,6 +3,7 @@ package gguf_test import ( "bytes" "os" + "slices" "strconv" "strings" "testing" @@ -247,3 +248,43 @@ func BenchmarkRead(b *testing.B) { f.Close() } } + +func BenchmarkReadArray(b *testing.B) { + b.ReportAllocs() + + create := func(tb testing.TB, kv ggml.KV) string { + tb.Helper() + f, err := os.CreateTemp(b.TempDir(), "") + if err != nil { + b.Fatal(err) + } + defer f.Close() + + if err := ggml.WriteGGUF(f, kv, nil); err != nil { + b.Fatal(err) + } + + return f.Name() + } + + cases := map[string]any{ + "int32": slices.Repeat([]int32{42}, 1_000_000), + "uint32": slices.Repeat([]uint32{42}, 1_000_000), + "float32": slices.Repeat([]float32{42.}, 1_000_000), + "string": slices.Repeat([]string{"42"}, 1_000_000), + } + + for name, bb := range cases { + b.Run(name, func(b *testing.B) { + p := create(b, ggml.KV{"array": bb}) + for b.Loop() { + f, err := gguf.Open(p) + if err != nil { + b.Fatal(err) + } + _ = f.KeyValue("array") + f.Close() + } + }) + } +}