Преглед на файлове

ggml: check if vocab key is present to use in size estimate

Vocab is expected to be present when estimating graph size,
but we should not panic if its not found.
Bruce MacDonald преди 2 месеца
родител
ревизия
4a8c539c99
променени са 2 файла, в които са добавени 33 реда и са изтрити 1 реда
  1. 7 1
      fs/ggml/ggml.go
  2. 26 0
      fs/ggml/ggml_test.go

+ 7 - 1
fs/ggml/ggml.go

@@ -379,7 +379,13 @@ func (f GGML) GraphSize(context, batch uint64, kvCacheType string) (kv, partialO
 	embedding := f.KV().EmbeddingLength()
 	heads := f.KV().HeadCount()
 	headsKV := f.KV().HeadCountKV()
-	vocab := uint64(f.KV()["tokenizer.ggml.tokens"].(*array).size)
+	var vocab uint64
+	if tokensArray, ok := f.KV()["tokenizer.ggml.tokens"].(*array); ok && tokensArray != nil {
+		vocab = uint64(tokensArray.size)
+	} else {
+		vocab = 1 // ignore vocab size if not found
+		slog.Warn("tokenizer.ggml.tokens not found or nil, using 0 for vocab size in size calculation")
+	}
 
 	embeddingHeads := f.KV().EmbeddingHeadCount()
 	embeddingHeadsK := f.KV().EmbeddingHeadCountK()

+ 26 - 0
fs/ggml/ggml_test.go

@@ -9,6 +9,32 @@ import (
 	"github.com/google/go-cmp/cmp"
 )
 
+type ggufModel struct {
+	kv      KV
+	tensors Tensors
+}
+
+func (m *ggufModel) KV() KV           { return m.kv }
+func (m *ggufModel) Tensors() Tensors { return m.tensors }
+
+func TestGraphNoVocab(t *testing.T) {
+	g := &GGML{
+		container: &containerGGUF{},
+		model: &ggufModel{
+			kv: KV{
+				"general.architecture": "llama",
+				"block_count":          uint32(1),
+			},
+			tensors: Tensors{
+				items: []*Tensor{},
+			},
+		},
+	}
+
+	// This should not panic
+	_, _, _ = g.GraphSize(1, 1, "f16")
+}
+
 func TestTensorLayers(t *testing.T) {
 	tensors := make(map[string]*Tensor)
 	for _, name := range []string{