129 lines
3.9 KiB
Go
129 lines
3.9 KiB
Go
package rocm
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// writeDiscoverTestGGUF creates a minimal GGUF v3 file in dir with the given
|
|
// filename and metadata KV pairs. Returns the full path to the created file.
|
|
func writeDiscoverTestGGUF(t *testing.T, dir, filename string, kvs [][2]any) string {
|
|
t.Helper()
|
|
|
|
path := filepath.Join(dir, filename)
|
|
|
|
f, err := os.Create(path)
|
|
require.NoError(t, err)
|
|
defer f.Close()
|
|
|
|
// Magic: "GGUF" in little-endian
|
|
require.NoError(t, binary.Write(f, binary.LittleEndian, uint32(0x46554747)))
|
|
// Version 3
|
|
require.NoError(t, binary.Write(f, binary.LittleEndian, uint32(3)))
|
|
// Tensor count (uint64): 0
|
|
require.NoError(t, binary.Write(f, binary.LittleEndian, uint64(0)))
|
|
// KV count (uint64)
|
|
require.NoError(t, binary.Write(f, binary.LittleEndian, uint64(len(kvs))))
|
|
|
|
for _, kv := range kvs {
|
|
key := kv[0].(string)
|
|
writeDiscoverKV(t, f, key, kv[1])
|
|
}
|
|
|
|
return path
|
|
}
|
|
|
|
func writeDiscoverKV(t *testing.T, f *os.File, key string, val any) {
|
|
t.Helper()
|
|
|
|
// Key: uint64 length + bytes
|
|
require.NoError(t, binary.Write(f, binary.LittleEndian, uint64(len(key))))
|
|
_, err := f.Write([]byte(key))
|
|
require.NoError(t, err)
|
|
|
|
switch v := val.(type) {
|
|
case string:
|
|
// Type: 8 (string)
|
|
require.NoError(t, binary.Write(f, binary.LittleEndian, uint32(8)))
|
|
// String value: uint64 length + bytes
|
|
require.NoError(t, binary.Write(f, binary.LittleEndian, uint64(len(v))))
|
|
_, err := f.Write([]byte(v))
|
|
require.NoError(t, err)
|
|
case uint32:
|
|
// Type: 4 (uint32)
|
|
require.NoError(t, binary.Write(f, binary.LittleEndian, uint32(4)))
|
|
require.NoError(t, binary.Write(f, binary.LittleEndian, v))
|
|
default:
|
|
t.Fatalf("writeDiscoverKV: unsupported value type %T", val)
|
|
}
|
|
}
|
|
|
|
func TestDiscoverModels(t *testing.T) {
|
|
dir := t.TempDir()
|
|
|
|
// Create two valid GGUF model files.
|
|
writeDiscoverTestGGUF(t, dir, "gemma3-4b-q4km.gguf", [][2]any{
|
|
{"general.architecture", "gemma3"},
|
|
{"general.name", "Gemma 3 4B Instruct"},
|
|
{"general.file_type", uint32(15)},
|
|
{"general.size_label", "4B"},
|
|
{"gemma3.context_length", uint32(32768)},
|
|
{"gemma3.block_count", uint32(34)},
|
|
})
|
|
|
|
writeDiscoverTestGGUF(t, dir, "llama-3.1-8b-q4km.gguf", [][2]any{
|
|
{"general.architecture", "llama"},
|
|
{"general.name", "Llama 3.1 8B Instruct"},
|
|
{"general.file_type", uint32(15)},
|
|
{"general.size_label", "8B"},
|
|
{"llama.context_length", uint32(131072)},
|
|
{"llama.block_count", uint32(32)},
|
|
})
|
|
|
|
// Create a non-GGUF file that should be ignored (no .gguf extension).
|
|
require.NoError(t, os.WriteFile(filepath.Join(dir, "README.txt"), []byte("not a model"), 0644))
|
|
|
|
models, err := DiscoverModels(dir)
|
|
require.NoError(t, err)
|
|
require.Len(t, models, 2)
|
|
|
|
// Sort order from Glob is lexicographic, so gemma3 comes first.
|
|
gemma := models[0]
|
|
assert.Equal(t, filepath.Join(dir, "gemma3-4b-q4km.gguf"), gemma.Path)
|
|
assert.Equal(t, "gemma3", gemma.Architecture)
|
|
assert.Equal(t, "Gemma 3 4B Instruct", gemma.Name)
|
|
assert.Equal(t, "Q4_K_M", gemma.Quantisation)
|
|
assert.Equal(t, "4B", gemma.Parameters)
|
|
assert.Equal(t, uint32(32768), gemma.ContextLen)
|
|
assert.Greater(t, gemma.FileSize, int64(0))
|
|
|
|
llama := models[1]
|
|
assert.Equal(t, filepath.Join(dir, "llama-3.1-8b-q4km.gguf"), llama.Path)
|
|
assert.Equal(t, "llama", llama.Architecture)
|
|
assert.Equal(t, "Llama 3.1 8B Instruct", llama.Name)
|
|
assert.Equal(t, "Q4_K_M", llama.Quantisation)
|
|
assert.Equal(t, "8B", llama.Parameters)
|
|
assert.Equal(t, uint32(131072), llama.ContextLen)
|
|
assert.Greater(t, llama.FileSize, int64(0))
|
|
}
|
|
|
|
func TestDiscoverModels_EmptyDir(t *testing.T) {
|
|
dir := t.TempDir()
|
|
|
|
models, err := DiscoverModels(dir)
|
|
require.NoError(t, err)
|
|
assert.Empty(t, models)
|
|
}
|
|
|
|
func TestDiscoverModels_NotFound(t *testing.T) {
|
|
// filepath.Glob returns nil, nil for a pattern matching no files,
|
|
// even when the directory does not exist.
|
|
models, err := DiscoverModels("/nonexistent/dir")
|
|
require.NoError(t, err)
|
|
assert.Empty(t, models)
|
|
}
|