Simplify model conversion (#3422)

convert/gemma.go

@@ -0,0 +1,136 @@
+package convert
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+	"log/slog"
+	"os"
+	"strings"
+
+	"github.com/d4l3k/go-bfloat16"
+	"github.com/pdevine/tensor"
+	"github.com/pdevine/tensor/native"
+
+	"github.com/ollama/ollama/llm"
+)
+
+type GemmaModel struct {
+	ModelData
+}
+
+func gemmaLayerHandler(w io.Writer, r safetensorWriterTo, f *os.File) error {
+	slog.Debug(fmt.Sprintf("converting '%s'", r.t.Name))
+
+	data := make([]byte, r.end-r.start)
+	if err := binary.Read(f, r.bo, data); err != nil {
+		return err
+	}
+
+	tDataF32 := bfloat16.DecodeFloat32(data)
+
+	var err error
+	tDataF32, err = addOnes(tDataF32, int(r.t.Shape[0]))
+	if err != nil {
+		return err
+	}
+
+	if err := binary.Write(w, r.bo, tDataF32); err != nil {
+		return err
+	}
+	return nil
+}
+
+func addOnes(data []float32, vectorSize int) ([]float32, error) {
+	n := tensor.New(tensor.WithShape(vectorSize), tensor.WithBacking(data))
+	ones := tensor.Ones(tensor.Float32, vectorSize)
+
+	var err error
+	n, err = n.Add(ones)
+	if err != nil {
+		return []float32{}, err
+	}
+
+	newN, err := native.SelectF32(n, 0)
+	if err != nil {
+		return []float32{}, err
+	}
+
+	var fullTensor []float32
+	for _, v := range newN {
+		fullTensor = append(fullTensor, v...)
+	}
+
+	return fullTensor, nil
+}
+
+func (m *GemmaModel) GetTensors() error {
+	t, err := GetSafeTensors(m.Path, m.Params)
+	if err != nil {
+		return err
+	}
+
+	m.Tensors = []llm.Tensor{}
+
+	for _, l := range t {
+		if strings.HasSuffix(l.Name, "norm.weight") {
+			wt := l.WriterTo.(safetensorWriterTo)
+			wt.handler = gemmaLayerHandler
+			l.WriterTo = wt
+		}
+		m.Tensors = append(m.Tensors, l)
+	}
+
+	return nil
+}
+
+func (m *GemmaModel) LoadVocab() error {
+	v, err := LoadSentencePieceTokens(m.Path, m.Params.VocabSize)
+	if err != nil {
+		return err
+	}
+	m.Vocab = v
+	return nil
+}
+
+func (m *GemmaModel) WriteGGUF() (string, error) {
+	kv := llm.KV{
+		"general.architecture":                   "gemma",
+		"general.name":                           m.Name,
+		"gemma.context_length":                   uint32(m.Params.ContextSize),
+		"gemma.embedding_length":                 uint32(m.Params.HiddenSize),
+		"gemma.block_count":                      uint32(m.Params.HiddenLayers),
+		"gemma.feed_forward_length":              uint32(m.Params.IntermediateSize),
+		"gemma.attention.head_count":             uint32(m.Params.AttentionHeads),
+		"gemma.attention.head_count_kv":          uint32(m.Params.KeyValHeads),
+		"gemma.attention.layer_norm_rms_epsilon": float32(m.Params.NormEPS),
+		"gemma.attention.key_length":             uint32(m.Params.HeadDimension),
+		"gemma.attention.value_length":           uint32(m.Params.HeadDimension),
+		"general.file_type":                      uint32(1),
+		"tokenizer.ggml.model":                   "llama",
+
+		"tokenizer.ggml.tokens":     m.Vocab.Tokens,
+		"tokenizer.ggml.scores":     m.Vocab.Scores,
+		"tokenizer.ggml.token_type": m.Vocab.Types,
+
+		"tokenizer.ggml.bos_token_id":     uint32(m.Params.BoSTokenID),
+		"tokenizer.ggml.eos_token_id":     uint32(m.Params.EoSTokenID),
+		"tokenizer.ggml.padding_token_id": uint32(m.Params.PaddingTokenID),
+		"tokenizer.ggml.unknown_token_id": uint32(3),
+		"tokenizer.ggml.add_bos_token":    true,
+		"tokenizer.ggml.add_eos_token":    false,
+	}
+
+	f, err := os.CreateTemp("", "ollama-gguf")
+	if err != nil {
+		return "", err
+	}
+	defer f.Close()
+
+	mod := llm.NewGGUFV3(m.Params.ByteOrder)
+	if err := mod.Encode(f, kv, m.Tensors); err != nil {
+		return "", err
+	}
+
+	return f.Name(), nil
+}