浏览代码

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()
 	embedding := f.KV().EmbeddingLength()
 	heads := f.KV().HeadCount()
 	heads := f.KV().HeadCount()
 	headsKV := f.KV().HeadCountKV()
 	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()
 	embeddingHeads := f.KV().EmbeddingHeadCount()
 	embeddingHeadsK := f.KV().EmbeddingHeadCountK()
 	embeddingHeadsK := f.KV().EmbeddingHeadCountK()

+ 26 - 0
fs/ggml/ggml_test.go

@@ -9,6 +9,32 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"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) {
 func TestTensorLayers(t *testing.T) {
 	tensors := make(map[string]*Tensor)
 	tensors := make(map[string]*Tensor)
 	for _, name := range []string{
 	for _, name := range []string{