Fix gemma2 2b conversion (#6645)

convert/convert_test.go

@@ -15,6 +15,7 @@ import (
 	"os"
 	"path/filepath"
 	"slices"
+	"strings"
 	"testing"
 
 	"golang.org/x/exp/maps"
@@ -22,6 +23,12 @@ import (
 	"github.com/ollama/ollama/llm"
 )
 
+type tensorData struct {
+	Offsets []int  `json:"data_offsets"`
+	Type    string `json:"dtype"`
+	Shape   []int  `json:"shape"`
+}
+
 func convertFull(t *testing.T, fsys fs.FS) (*os.File, llm.KV, llm.Tensors) {
 	t.Helper()
 
@@ -96,6 +103,7 @@ func TestConvertModel(t *testing.T) {
 		"Mistral-7B-Instruct-v0.2",
 		"Mixtral-8x7B-Instruct-v0.1",
 		"gemma-2b-it",
+		"gemma-2-2b-it",
 		// microsoft/Phi-3-mini-128-instruct@d548c233192db00165d842bf8edff054bb3212f8
 		"Phi-3-mini-128k-instruct",
 		"all-MiniLM-L6-v2",
@@ -140,6 +148,36 @@ func TestConvertModel(t *testing.T) {
 	}
 }
 
+func TestConvertInvalidTensorNames(t *testing.T) {
+	f, err := os.CreateTemp(t.TempDir(), "testmodel")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f.Close()
+
+	tempDir := t.TempDir()
+
+	td := map[string]*tensorData{}
+	offset := 4096
+
+	td["model.layers.0.self_attn.q_proj.weight"] = &tensorData{
+		Offsets: []int{0, offset},
+		Type:    "F32",
+		Shape:   []int{4096, 4096},
+	}
+	td["blk.0.attn_q.weight"] = &tensorData{
+		Offsets: []int{offset, offset * 2},
+		Type:    "F32",
+		Shape:   []int{4096, 4096},
+	}
+	generateSafetensorTestData(t, tempDir, td)
+
+	err = ConvertModel(os.DirFS(tempDir), f)
+	if err == nil || !strings.HasPrefix(err.Error(), "duplicate tensor name") {
+		t.Errorf("expected error but didn't get one")
+	}
+}
+
 func TestConvertInvalidDatatype(t *testing.T) {
 	f, err := os.CreateTemp(t.TempDir(), "testmodel")
 	if err != nil {
@@ -148,23 +186,10 @@ func TestConvertInvalidDatatype(t *testing.T) {
 	defer f.Close()
 
 	tempDir := t.TempDir()
-	generateSafetensorTestData(t, tempDir)
-
-	err = ConvertModel(os.DirFS(tempDir), f)
-	if err == nil || err.Error() != "unsupported safetensors model" {
-		t.Errorf("expected error but didn't get one")
-	}
-}
-
-func generateSafetensorTestData(t *testing.T, tempDir string) {
-	type tensorData struct {
-		Offsets []int  `json:"data_offsets"`
-		Type    string `json:"dtype"`
-		Shape   []int  `json:"shape"`
-	}
-	offset := 4096 * 14336
 
 	td := map[string]*tensorData{}
+	offset := 4096 * 14336
+
 	td["model.layers.0.mlp.down_proj.weight"] = &tensorData{
 		Offsets: []int{0, offset},
 		Type:    "I8",
@@ -175,8 +200,16 @@ func generateSafetensorTestData(t *testing.T, tempDir string) {
 		Type:    "U8",
 		Shape:   []int{},
 	}
+	generateSafetensorTestData(t, tempDir, td)
 
-	data, err := json.Marshal(td)
+	err = ConvertModel(os.DirFS(tempDir), f)
+	if err == nil || err.Error() != "unsupported safetensors model" {
+		t.Errorf("expected error but didn't get one")
+	}
+}
+
+func generateSafetensorTestData(t *testing.T, tempDir string, tensorData map[string]*tensorData) {
+	data, err := json.Marshal(tensorData)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -322,11 +355,6 @@ func TestConvertAdapter(t *testing.T) {
 }
 
 func generateLoraTestData(t *testing.T, tempDir string) {
-	type tensorData struct {
-		Offsets []int  `json:"data_offsets"`
-		Type    string `json:"dtype"`
-		Shape   []int  `json:"shape"`
-	}
 	offset := 4096 * 8 * 4
 
 	td := map[string]*tensorData{"__metadata__": nil}