test: fix

Signed-off-by: thxCode <[email protected]>
gpustack · Nov 20, 2024 · a1e91fe · a1e91fe
1 parent 397b53a
commit a1e91fe
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 10 deletions.
diff --git a/file.go b/file.go
@@ -950,9 +950,6 @@ func (tis GGUFTensorInfos) GetFileType() GGUFFileType {
 
 	cm := make(map[GGMLType]int)
 	for i := range tis {
-		if !strings.HasSuffix(tis[i].Name, ".weight") {
-			continue
-		}
 		cm[tis[i].Type]++
 	}
 
@@ -1210,9 +1207,6 @@ func (ltis GGUFLayerTensorInfos) GetFileType() GGUFFileType {
 	for i := range ltis {
 		switch v := ltis[i].(type) {
 		case GGUFTensorInfo:
-			if !strings.HasSuffix(v.Name, ".weight") {
-				continue
-			}
 			cm[v.Type]++
 		case *GGUFNamedTensorInfos:
 			cm[v.GetFileType().GGMLType()]++

diff --git a/file_metadata.go b/file_metadata.go
@@ -196,7 +196,7 @@ func (gf *GGUFFile) Metadata() (gm GGUFMetadata) {
 	}
 
 	if gm.FileType >= _GGUFFileTypeCount {
-		gm.FileType = gf.guessFileType()
+		gm.FileType = gf.guessFileType(gm.Architecture)
 	}
 
 	gm.LittleEndian = gf.Header.Version < GGUFVersionV3 || gf.Header.Magic == GGUFMagicGGUFLe
@@ -280,15 +280,18 @@ func (t GGUFFileType) GGMLType() GGMLType {
 // statistically analyzing the tensor types,
 // which is inspired by
 // https://huggingface.co/TheBloke/Llama-2-13B-chat-GGML#provided-files.
-func (gf *GGUFFile) guessFileType() GGUFFileType {
+func (gf *GGUFFile) guessFileType(arch string) GGUFFileType {
 	if len(gf.TensorInfos) == 0 {
 		return _GGUFFileTypeCount
 	}
 
 	// Count.
 	cm := make(map[GGMLType]int)
 	for i := range gf.TensorInfos {
-		if !strings.HasSuffix(gf.TensorInfos[i].Name, ".weight") {
+		switch {
+		case arch != "diffusion" && !strings.HasPrefix(gf.TensorInfos[i].Name, "blk."):
+			continue
+		case arch == "diffusion" && !strings.HasSuffix(gf.TensorInfos[i].Name, ".weight"):
 			continue
 		}
 		cm[gf.TensorInfos[i].Type]++

diff --git a/file_metadata_test.go b/file_metadata_test.go
@@ -74,7 +74,8 @@ func TestGGUFFile_guessFileType(t *testing.T) {
 				t.Fatal(err)
 				return
 			}
-			assert.Equal(t, gf.Metadata().FileType.String(), gf.guessFileType().String(), tc+" file type should be equal")
+			md := gf.Metadata()
+			assert.Equal(t, md.FileType.String(), gf.guessFileType(md.Architecture).String(), tc+" file type should be equal")
 		})
 	}
 }