ollama/convert/convert_test.go

129 lines
2.5 KiB
Go
Raw Normal View History

2024-05-21 18:28:16 +00:00
package convert
import (
"crypto/sha256"
2024-08-01 21:52:15 +00:00
"encoding/hex"
"encoding/json"
"flag"
"fmt"
"io"
2024-06-29 23:53:59 +00:00
"io/fs"
"log/slog"
"math"
2024-05-21 18:28:16 +00:00
"os"
"path/filepath"
"slices"
2024-05-21 18:28:16 +00:00
"testing"
"golang.org/x/exp/maps"
2024-08-01 21:52:15 +00:00
"github.com/ollama/ollama/llm"
2024-05-21 18:28:16 +00:00
)
2024-06-29 23:53:59 +00:00
func convertFull(t *testing.T, fsys fs.FS) (*os.File, llm.KV, llm.Tensors) {
2024-05-21 18:28:16 +00:00
t.Helper()
f, err := os.CreateTemp(t.TempDir(), "f16")
if err != nil {
t.Fatal(err)
}
defer f.Close()
2024-06-29 23:53:59 +00:00
if err := Convert(fsys, f); err != nil {
2024-05-21 18:28:16 +00:00
t.Fatal(err)
}
r, err := os.Open(f.Name())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { r.Close() })
2024-05-21 18:28:16 +00:00
m, _, err := llm.DecodeGGML(r, math.MaxInt)
2024-05-21 18:28:16 +00:00
if err != nil {
t.Fatal(err)
}
if _, err := r.Seek(0, io.SeekStart); err != nil {
t.Fatal(err)
}
return r, m.KV(), m.Tensors()
}
func TestMain(m *testing.M) {
var level slog.Level
flag.TextVar(&level, "level", slog.LevelInfo, "log level")
flag.Parse()
slog.SetLogLoggerLevel(level)
os.Exit(m.Run())
2024-05-21 18:28:16 +00:00
}
func TestConvertFull(t *testing.T) {
cases := []string{
"Meta-Llama-3-8B-Instruct",
"Mistral-7B-Instruct-v0.2",
"Mixtral-8x7B-Instruct-v0.1",
"gemma-2b-it",
2024-05-21 18:28:16 +00:00
}
for i := range cases {
tt := cases[i]
t.Run(tt, func(t *testing.T) {
t.Parallel()
p := filepath.Join("testdata", tt)
if testing.Short() {
t.Skip("skipping in short mode")
} else if _, err := os.Stat(p); err != nil {
2024-05-21 18:28:16 +00:00
t.Skipf("%s not found", p)
}
2024-06-29 23:53:59 +00:00
f, kv, tensors := convertFull(t, os.DirFS(p))
actual := make(map[string]string)
for k, v := range kv {
if s, ok := v.(json.Marshaler); !ok {
actual[k] = fmt.Sprintf("%v", v)
} else {
bts, err := json.Marshal(s)
if err != nil {
t.Fatal(err)
}
actual[k] = fmt.Sprintf("%x", sha256.Sum256(bts))
}
}
for _, tensor := range tensors.Items {
sha256sum := sha256.New()
sr := io.NewSectionReader(f, int64(tensors.Offset+tensor.Offset), int64(tensor.Size()))
if _, err := io.Copy(sha256sum, sr); err != nil {
t.Fatal(err)
}
2024-05-21 18:28:16 +00:00
2024-08-01 21:52:15 +00:00
actual[tensor.Name] = hex.EncodeToString(sha256sum.Sum(nil))
2024-05-21 18:28:16 +00:00
}
expectFile, err := os.Open(filepath.Join("testdata", fmt.Sprintf("%s.json", tt)))
if err != nil {
t.Fatal(err)
2024-05-21 18:28:16 +00:00
}
var expect map[string]string
if err := json.NewDecoder(expectFile).Decode(&expect); err != nil {
t.Fatal(err)
2024-05-21 18:28:16 +00:00
}
keys := maps.Keys(expect)
slices.Sort(keys)
for _, k := range keys {
if v, ok := actual[k]; !ok {
t.Errorf("missing %s", k)
} else if v != expect[k] {
t.Errorf("unexpected %s: want %s, got %s", k, expect[k], v)
}
2024-05-21 18:28:16 +00:00
}
})
}
}