Add gemma safetensors conversion (#3250)

Co-authored-by: Michael Yang <mxyng@pm.me>

.golangci.yaml

@@ -15,13 +15,3 @@ linters:
     - misspell
     - nilerr
     - unused
-linters-settings:
-  errcheck:
-    # exclude the following functions since we don't generally
-    # need to be concerned with the returned errors
-    exclude-functions:
-      - encoding/binary.Read
-      - (*os.File).Seek
-      - (*bufio.Writer).WriteString
-      - (*github.com/spf13/pflag.FlagSet).Set
-      - (*github.com/ollama/ollama/llm.readSeekOffset).Seek

cmd/cmd.go

@@ -213,7 +213,10 @@ func createBlob(cmd *cobra.Command, client *api.Client, path string) (string, er
 	if _, err := io.Copy(hash, bin); err != nil {
 		return "", err
 	}
-	bin.Seek(0, io.SeekStart)
+
+	if _, err := bin.Seek(0, io.SeekStart); err != nil {
+		return "", err
+	}
 
 	digest := fmt.Sprintf("sha256:%x", hash.Sum(nil))
 	if err = client.CreateBlob(cmd.Context(), digest, bin); err != nil {

cmd/interactive.go

@@ -295,10 +295,14 @@ func generateInteractive(cmd *cobra.Command, opts runOptions) error {
 					opts.WordWrap = false
 					fmt.Println("Set 'nowordwrap' mode.")
 				case "verbose":
-					cmd.Flags().Set("verbose", "true")
+					if err := cmd.Flags().Set("verbose", "true"); err != nil {
+						return err
+					}
 					fmt.Println("Set 'verbose' mode.")
 				case "quiet":
-					cmd.Flags().Set("verbose", "false")
+					if err := cmd.Flags().Set("verbose", "false"); err != nil {
+						return err
+					}
 					fmt.Println("Set 'quiet' mode.")
 				case "format":
 					if len(args) < 3 || args[2] != "json" {

convert/convert.go

@@ -12,8 +12,13 @@ import (
 	"path/filepath"
 	"regexp"
 	"slices"
+	"strings"
 
+	"github.com/d4l3k/go-bfloat16"
 	"github.com/mitchellh/mapstructure"
+	"github.com/pdevine/tensor"
+	"github.com/pdevine/tensor/native"
+	"github.com/x448/float16"
 	"google.golang.org/protobuf/proto"
 
 	"github.com/ollama/ollama/convert/sentencepiece"
@@ -33,6 +38,15 @@ type Params struct {
 	RopeFreqBase     float64  `json:"rope_theta"`
 	BoSTokenID       int      `json:"bos_token_id"`
 	EoSTokenID       int      `json:"eos_token_id"`
+	HeadDimension    int      `json:"head_dim"`
+	PaddingTokenID   int      `json:"pad_token_id"`
+
+	ByteOrder
+}
+
+type ByteOrder interface {
+	binary.ByteOrder
+	binary.AppendByteOrder
 }
 
 type MetaData struct {
@@ -41,27 +55,29 @@ type MetaData struct {
 	Offsets []int  `mapstructure:"data_offsets"`
 }
 
-func ReadSafeTensors(fn string, offset uint64) ([]llm.Tensor, uint64, error) {
+func ReadSafeTensors(fn string, offset uint64, params *Params) ([]llm.Tensor, uint64, error) {
 	f, err := os.Open(fn)
 	if err != nil {
-		return []llm.Tensor{}, 0, err
+		return nil, 0, err
 	}
 	defer f.Close()
 
 	var jsonSize uint64
-	binary.Read(f, binary.LittleEndian, &jsonSize)
+	if err := binary.Read(f, binary.LittleEndian, &jsonSize); err != nil {
+		return nil, 0, err
+	}
 
 	buf := make([]byte, jsonSize)
 	_, err = io.ReadFull(f, buf)
 	if err != nil {
-		return []llm.Tensor{}, 0, err
+		return nil, 0, err
 	}
 
 	d := json.NewDecoder(bytes.NewBuffer(buf))
 	d.UseNumber()
 	var parsed map[string]interface{}
 	if err = d.Decode(&parsed); err != nil {
-		return []llm.Tensor{}, 0, err
+		return nil, 0, err
 	}
 
 	var keys []string
@@ -78,7 +94,7 @@ func ReadSafeTensors(fn string, offset uint64) ([]llm.Tensor, uint64, error) {
 		vals := parsed[k].(map[string]interface{})
 		var data MetaData
 		if err = mapstructure.Decode(vals, &data); err != nil {
-			return []llm.Tensor{}, 0, err
+			return nil, 0, err
 		}
 
 		var size uint64
@@ -100,7 +116,7 @@ func ReadSafeTensors(fn string, offset uint64) ([]llm.Tensor, uint64, error) {
 		ggufName, err := GetTensorName(k)
 		if err != nil {
 			slog.Error("%v", err)
-			return []llm.Tensor{}, 0, err
+			return nil, 0, err
 		}
 
 		shape := []uint64{0, 0, 0, 0}
@@ -109,14 +125,24 @@ func ReadSafeTensors(fn string, offset uint64) ([]llm.Tensor, uint64, error) {
 		}
 
 		t := llm.Tensor{
-			Name:          ggufName,
-			Kind:          kind,
-			Offset:        offset,
-			Shape:         shape[:],
-			FileName:      fn,
-			OffsetPadding: 8 + jsonSize,
-			FileOffsets:   []uint64{uint64(data.Offsets[0]), uint64(data.Offsets[1])},
+			Name:   ggufName,
+			Kind:   kind,
+			Offset: offset,
+			Shape:  shape[:],
 		}
+
+		t.WriterTo = safetensorWriterTo{
+			t:           &t,
+			params:      params,
+			bo:          params.ByteOrder,
+			headCount:   uint32(params.AttentionHeads),
+			headCountKV: uint32(params.KeyValHeads),
+			filename:    fn,
+			start:       uint64(data.Offsets[0]),
+			end:         uint64(data.Offsets[1]),
+			padding:     8 + jsonSize,
+		}
+
 		slog.Debug(fmt.Sprintf("%v", t))
 		tensors = append(tensors, t)
 		offset += size
@@ -124,21 +150,21 @@ func ReadSafeTensors(fn string, offset uint64) ([]llm.Tensor, uint64, error) {
 	return tensors, offset, nil
 }
 
-func GetSafeTensors(dirpath string) ([]llm.Tensor, error) {
+func GetSafeTensors(dirpath string, params *Params) ([]llm.Tensor, error) {
 	var tensors []llm.Tensor
 	files, err := filepath.Glob(filepath.Join(dirpath, "/model-*.safetensors"))
 	if err != nil {
-		return []llm.Tensor{}, err
+		return nil, err
 	}
 
 	var offset uint64
 	for _, f := range files {
 		var t []llm.Tensor
 		var err error
-		t, offset, err = ReadSafeTensors(f, offset)
+		t, offset, err = ReadSafeTensors(f, offset, params)
 		if err != nil {
 			slog.Error("%v", err)
-			return []llm.Tensor{}, err
+			return nil, err
 		}
 		tensors = append(tensors, t...)
 	}
@@ -160,6 +186,7 @@ func GetParams(dirpath string) (*Params, error) {
 		return nil, err
 	}
 
+	params.ByteOrder = binary.LittleEndian
 	return &params, nil
 }
 
@@ -171,7 +198,7 @@ type Vocab struct {
 	Types  []int32
 }
 
-func LoadTokens(dirpath string) (*Vocab, error) {
+func LoadTokens(dirpath string, params *Params) (*Vocab, error) {
 	slog.Info(fmt.Sprintf("reading vocab from %s", filepath.Join(dirpath, "tokenizer.model")))
 	in, err := os.ReadFile(filepath.Join(dirpath, "tokenizer.model"))
 	if err != nil {
@@ -196,6 +223,14 @@ func LoadTokens(dirpath string) (*Vocab, error) {
 		v.Tokens = append(v.Tokens, p.GetPiece())
 		v.Scores = append(v.Scores, p.GetScore())
 		t := p.GetType()
+		switch t {
+		case sentencepiece.ModelProto_SentencePiece_UNKNOWN:
+		case sentencepiece.ModelProto_SentencePiece_CONTROL:
+		case sentencepiece.ModelProto_SentencePiece_UNUSED:
+		case sentencepiece.ModelProto_SentencePiece_BYTE:
+		default:
+			t = sentencepiece.ModelProto_SentencePiece_NORMAL
+		}
 		v.Types = append(v.Types, int32(t))
 	}
 
@@ -243,6 +278,16 @@ func LoadTokens(dirpath string) (*Vocab, error) {
 	}
 	slog.Info(fmt.Sprintf("vocab size w/ extra tokens: %d", len(v.Tokens)))
 
+	if params.VocabSize > len(v.Tokens) {
+		missingTokens := params.VocabSize - len(v.Tokens)
+		slog.Warn(fmt.Sprintf("vocab is missing %d tokens", missingTokens))
+		for cnt := 0; cnt < missingTokens; cnt++ {
+			v.Tokens = append(v.Tokens, fmt.Sprintf("<dummy%05d>", cnt+1))
+			v.Scores = append(v.Scores, -1)
+			v.Types = append(v.Types, int32(llm.GGUFTokenUserDefined))
+		}
+	}
+
 	return v, nil
 }
 
@@ -279,42 +324,287 @@ func GetTensorName(n string) (string, error) {
 	return "", fmt.Errorf("couldn't find a layer name for '%s'", n)
 }
 
+type safetensorWriterTo struct {
+	t *llm.Tensor
+
+	params      *Params
+	bo          ByteOrder
+	headCount   uint32
+	headCountKV uint32
+
+	filename string
+
+	start, end, padding uint64
+}
+
+func (r safetensorWriterTo) addOnes(data []float32) ([]float32, error) {
+	n := tensor.New(tensor.WithShape(int(r.t.Shape[0])), tensor.WithBacking(data))
+	ones := tensor.Ones(tensor.Float32, int(r.t.Shape[0]))
+
+	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 (r safetensorWriterTo) repack(data []uint16, heads int) ([]uint16, error) {
+	n := tensor.New(tensor.WithShape(int(r.t.Shape[0]), int(r.t.Shape[1])), tensor.WithBacking(data))
+	origShape := n.Shape().Clone()
+
+	// reshape the tensor and swap axes 1 and 2 to unpack the layer for gguf
+	if err := n.Reshape(heads, 2, origShape[0]/heads/2, origShape[1]); err != nil {
+		return nil, err
+	}
+
+	if err := n.T(0, 2, 1, 3); err != nil {
+		return nil, err
+	}
+
+	if err := n.Reshape(origShape...); err != nil {
+		return nil, err
+	}
+
+	if err := n.Transpose(); err != nil {
+		return nil, err
+	}
+	newN, err := native.SelectU16(n, 1)
+	if err != nil {
+		return nil, err
+	}
+
+	var fullTensor []uint16
+	for _, v := range newN {
+		fullTensor = append(fullTensor, v...)
+	}
+	return fullTensor, nil
+}
+
+func (r safetensorWriterTo) WriteTo(w io.Writer) (n int64, err error) {
+	arch, err := getArchFromParams(r.params)
+	if err != nil {
+		return 0, err
+	}
+
+	f, err := os.Open(r.filename)
+	if err != nil {
+		return 0, err
+	}
+	defer f.Close()
+
+	if _, err = f.Seek(int64(r.padding+r.start), 0); err != nil {
+		return 0, err
+	}
+
+	switch arch {
+	case "llama":
+
+		pattern := `^blk\.[0-9]+\.attn_(?P<layer>q|k)\.weight$`
+		re, err := regexp.Compile(pattern)
+		if err != nil {
+			return 0, err
+		}
+
+		matches := re.FindAllStringSubmatch(r.t.Name, -1)
+		if len(matches) > 0 {
+			layerSize := r.end - r.start
+
+			var err error
+			tData := make([]uint16, layerSize/2)
+			if err = binary.Read(f, r.bo, tData); err != nil {
+				return 0, err
+			}
+
+			layerType := matches[0][re.SubexpIndex("layer")]
+			var heads uint32
+			switch layerType {
+			case "q":
+				heads = r.headCount
+			case "k":
+				heads = r.headCountKV
+				if heads == 0 {
+					heads = r.headCount
+				}
+			}
+
+			tData, err = r.repack(tData, int(heads))
+			if err != nil {
+				return 0, err
+			}
+
+			var buf []byte
+			for _, n := range tData {
+				buf = r.bo.AppendUint16(buf, n)
+			}
+
+			tempBuf := make([]uint16, len(tData))
+			tDataF32 := bfloat16.DecodeFloat32(buf)
+			for cnt, v := range tDataF32 {
+				tDataF16 := float16.Fromfloat32(v)
+				tempBuf[cnt] = uint16(tDataF16)
+			}
+
+			if err = binary.Write(w, r.bo, tempBuf); err != nil {
+				return 0, err
+			}
+
+			return 0, nil
+		}
+
+	case "gemma":
+		if strings.HasSuffix(r.t.Name, "norm.weight") {
+			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 0, err
+			}
+
+			tDataF32 := bfloat16.DecodeFloat32(data)
+
+			var err error
+			tDataF32, err = r.addOnes(tDataF32)
+			if err != nil {
+				return 0, err
+			}
+
+			if err := binary.Write(w, r.bo, tDataF32); err != nil {
+				return 0, err
+			}
+			return 0, nil
+		}
+	}
+
+	remaining := r.end - r.start
+
+	bufSize := uint64(10240)
+	var finished bool
+	for {
+		data := make([]byte, min(bufSize, remaining))
+
+		b, err := io.ReadFull(f, data)
+		remaining -= uint64(b)
+
+		if err == io.EOF || remaining <= 0 {
+			finished = true
+		} else if err != nil {
+			return 0, err
+		}
+
+		// convert bfloat16 -> ieee float32
+		tDataF32 := bfloat16.DecodeFloat32(data)
+
+		switch r.t.Kind {
+		case 0:
+			if err := binary.Write(w, r.bo, tDataF32); err != nil {
+				return 0, err
+			}
+		case 1:
+			// convert float32 -> float16
+			tempBuf := make([]uint16, len(data)/2)
+			for cnt, v := range tDataF32 {
+				tDataF16 := float16.Fromfloat32(v)
+				tempBuf[cnt] = uint16(tDataF16)
+			}
+			if err := binary.Write(w, binary.LittleEndian, tempBuf); err != nil {
+				return 0, err
+			}
+		}
+		if finished {
+			break
+		}
+	}
+	return 0, nil
+}
+
+func getArchFromParams(params *Params) (string, error) {
+	var arch string
+	switch len(params.Architectures) {
+	case 0:
+		return "", fmt.Errorf("No architecture specified to convert")
+	case 1:
+		switch params.Architectures[0] {
+		case "MistralForCausalLM":
+			arch = "llama"
+		case "GemmaForCausalLM":
+			arch = "gemma"
+		default:
+			return "", fmt.Errorf("Models based on '%s' are not yet supported", params.Architectures[0])
+		}
+	default:
+		return "", fmt.Errorf("Multimodal models are not yet supported")
+	}
+
+	return arch, nil
+}
+
 func WriteGGUF(name string, tensors []llm.Tensor, params *Params, vocab *Vocab) (string, error) {
-	c := llm.ContainerGGUF{
-		ByteOrder: binary.LittleEndian,
-	}
-
-	m := llm.NewGGUFModel(&c)
-	m.Tensors = tensors
-	m.KV["general.architecture"] = "llama"
-	m.KV["general.name"] = name
-	m.KV["llama.context_length"] = uint32(params.ContextSize)
-	m.KV["llama.embedding_length"] = uint32(params.HiddenSize)
-	m.KV["llama.block_count"] = uint32(params.HiddenLayers)
-	m.KV["llama.feed_forward_length"] = uint32(params.IntermediateSize)
-	m.KV["llama.rope.dimension_count"] = uint32(128)
-	m.KV["llama.attention.head_count"] = uint32(params.AttentionHeads)
-	m.KV["llama.attention.head_count_kv"] = uint32(params.KeyValHeads)
-	m.KV["llama.attention.layer_norm_rms_epsilon"] = float32(params.NormEPS)
-	m.KV["llama.rope.freq_base"] = float32(params.RopeFreqBase)
-	m.KV["general.file_type"] = uint32(1)
-	m.KV["tokenizer.ggml.model"] = "llama"
-
-	m.KV["tokenizer.ggml.tokens"] = vocab.Tokens
-	m.KV["tokenizer.ggml.scores"] = vocab.Scores
-	m.KV["tokenizer.ggml.token_type"] = vocab.Types
-
-	m.KV["tokenizer.ggml.bos_token_id"] = uint32(params.BoSTokenID)
-	m.KV["tokenizer.ggml.eos_token_id"] = uint32(params.EoSTokenID)
-	m.KV["tokenizer.ggml.unknown_token_id"] = uint32(0)
-	m.KV["tokenizer.ggml.add_bos_token"] = true
-	m.KV["tokenizer.ggml.add_eos_token"] = false
-
-	// llamacpp sets the chat template, however we don't need to set it since we pass it in through a layer
-	// m.KV["tokenizer.chat_template"] = "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}" // XXX removeme
-
-	c.V3.NumTensor = uint64(len(tensors))
-	c.V3.NumKV = uint64(len(m.KV))
+	arch, err := getArchFromParams(params)
+	if err != nil {
+		return "", err
+	}
+
+	kv := llm.KV{
+		"general.architecture": arch,
+		"general.name":         name,
+	}
+
+	switch arch {
+	case "llama":
+		kv["llama.context_length"] = uint32(params.ContextSize)
+		kv["llama.embedding_length"] = uint32(params.HiddenSize)
+		kv["llama.block_count"] = uint32(params.HiddenLayers)
+		kv["llama.feed_forward_length"] = uint32(params.IntermediateSize)
+		kv["llama.rope.dimension_count"] = uint32(params.HiddenSize / params.AttentionHeads)
+		slog.Debug(fmt.Sprintf("rope dim count = %d", kv["llama.rope.dimension_count"]))
+		kv["llama.attention.head_count"] = uint32(params.AttentionHeads)
+		kv["llama.attention.head_count_kv"] = uint32(params.KeyValHeads)
+		kv["llama.attention.layer_norm_rms_epsilon"] = float32(params.NormEPS)
+		kv["llama.rope.freq_base"] = float32(params.RopeFreqBase)
+	case "gemma":
+		kv["gemma.context_length"] = uint32(params.ContextSize)
+		kv["gemma.embedding_length"] = uint32(params.HiddenSize)
+		kv["gemma.block_count"] = uint32(params.HiddenLayers)
+		kv["gemma.feed_forward_length"] = uint32(params.IntermediateSize)
+		kv["gemma.attention.head_count"] = uint32(params.AttentionHeads)
+		kv["gemma.attention.head_count_kv"] = uint32(params.KeyValHeads)
+		kv["gemma.attention.layer_norm_rms_epsilon"] = float32(params.NormEPS)
+		kv["gemma.attention.key_length"] = uint32(params.HeadDimension)
+		kv["gemma.attention.value_length"] = uint32(params.HeadDimension)
+	}
+
+	kv["general.file_type"] = uint32(1)
+	kv["tokenizer.ggml.model"] = "llama"
+
+	kv["tokenizer.ggml.tokens"] = vocab.Tokens
+	kv["tokenizer.ggml.scores"] = vocab.Scores
+	kv["tokenizer.ggml.token_type"] = vocab.Types
+
+	kv["tokenizer.ggml.bos_token_id"] = uint32(params.BoSTokenID)
+	kv["tokenizer.ggml.eos_token_id"] = uint32(params.EoSTokenID)
+
+	switch arch {
+	case "llama":
+		kv["tokenizer.ggml.unknown_token_id"] = uint32(0)
+	case "gemma":
+		kv["tokenizer.ggml.padding_token_id"] = uint32(params.PaddingTokenID)
+		kv["tokenizer.ggml.unknown_token_id"] = uint32(3)
+	}
+
+	kv["tokenizer.ggml.add_bos_token"] = true
+	kv["tokenizer.ggml.add_eos_token"] = false
 
 	f, err := os.CreateTemp("", "ollama-gguf")
 	if err != nil {
@@ -322,8 +612,8 @@ func WriteGGUF(name string, tensors []llm.Tensor, params *Params, vocab *Vocab)
 	}
 	defer f.Close()
 
-	err = m.Encode(f)
-	if err != nil {
+	m := llm.NewGGUFV3(params.ByteOrder)
+	if err := m.Encode(f, kv, tensors); err != nil {
 		return "", err
 	}
 

go.mod

@@ -9,7 +9,7 @@ require (
 	github.com/d4l3k/go-bfloat16 v0.0.0-20211005043715-690c3bdd05f1
 	github.com/emirpasic/gods v1.18.1
 	github.com/gin-gonic/gin v1.9.1
-	github.com/golang/protobuf v1.5.0
+	github.com/golang/protobuf v1.5.0 // indirect
 	github.com/google/uuid v1.0.0
 	github.com/mitchellh/mapstructure v1.5.0
 	github.com/olekukonko/tablewriter v0.0.5

llm/ggla.go

@@ -7,16 +7,18 @@ import (
 	"slices"
 )
 
-type ContainerGGLA struct {
+type containerGGLA struct {
 	version uint32
 }
 
-func (c *ContainerGGLA) Name() string {
+func (c *containerGGLA) Name() string {
 	return "ggla"
 }
 
-func (c *ContainerGGLA) Decode(rs io.ReadSeeker) (model, error) {
-	binary.Read(rs, binary.LittleEndian, &c.version)
+func (c *containerGGLA) Decode(rs io.ReadSeeker) (model, error) {
+	if err := binary.Read(rs, binary.LittleEndian, &c.version); err != nil {
+		return nil, err
+	}
 
 	switch c.version {
 	case 1:
@@ -24,26 +26,26 @@ func (c *ContainerGGLA) Decode(rs io.ReadSeeker) (model, error) {
 		return nil, errors.New("invalid version")
 	}
 
-	model := newModelGGLA(c)
+	model := newGGLA(c)
 	err := model.decode(rs)
 	return model, err
 }
 
-type ModelGGLA struct {
-	*ContainerGGLA
+type ggla struct {
+	*containerGGLA
 
 	kv      KV
 	tensors []Tensor
 }
 
-func newModelGGLA(container *ContainerGGLA) *ModelGGLA {
-	return &ModelGGLA{
-		ContainerGGLA: container,
+func newGGLA(container *containerGGLA) *ggla {
+	return &ggla{
+		containerGGLA: container,
 		kv:            make(KV),
 	}
 }
 
-func (m *ModelGGLA) decode(rs io.ReadSeeker) error {
+func (m *ggla) decode(rs io.ReadSeeker) error {
 	var r uint32
 	if err := binary.Read(rs, binary.LittleEndian, &r); err != nil {
 		return err
@@ -109,7 +111,7 @@ func (m *ModelGGLA) decode(rs io.ReadSeeker) error {
 
 		t.Offset = uint64(offset)
 
-		if _, err := rs.Seek(int64(t.Size()), io.SeekCurrent); err != nil {
+		if _, err := rs.Seek(int64(t.size()), io.SeekCurrent); err != nil {
 			return err
 		}
 
@@ -117,46 +119,46 @@ func (m *ModelGGLA) decode(rs io.ReadSeeker) error {
 	}
 }
 
-func (m *ModelGGLA) KV() KV {
+func (m *ggla) KV() KV {
 	return m.kv
 }
 
-func (m *ModelGGLA) Tensor() []Tensor {
+func (m *ggla) Tensor() []Tensor {
 	return m.tensors
 }
 
-func (*ModelGGLA) ModelFamily() string {
+func (*ggla) ModelFamily() string {
 	return "ggla"
 }
 
-func (*ModelGGLA) ModelType() string {
+func (*ggla) ModelType() string {
 	panic("not implemented")
 }
 
-func (*ModelGGLA) FileType() string {
+func (*ggla) FileType() string {
 	panic("not implemented")
 }
 
-func (*ModelGGLA) NumLayers() uint32 {
+func (*ggla) NumLayers() uint32 {
 	panic("not implemented")
 }
 
-func (*ModelGGLA) NumGQA() uint32 {
+func (*ggla) NumGQA() uint32 {
 	panic("not implemented")
 }
 
-func (*ModelGGLA) NumEmbed() uint32 {
+func (*ggla) NumEmbed() uint32 {
 	panic("not implemented")
 }
 
-func (*ModelGGLA) NumHead() uint32 {
+func (*ggla) NumHead() uint32 {
 	panic("not implemented")
 }
 
-func (*ModelGGLA) NumHeadKv() uint32 {
+func (*ggla) NumHeadKv() uint32 {
 	panic("not implemented")
 }
 
-func (*ModelGGLA) NumCtx() uint32 {
+func (*ggla) NumCtx() uint32 {
 	panic("not implemented")
 }

llm/ggml.go

@@ -101,6 +101,85 @@ type model interface {
 	NumCtx() uint32
 }
 
+type KV map[string]any
+
+type Tensor struct {
+	Name   string
+	Kind   uint32
+	Offset uint64
+
+	// Shape is the number of elements in each dimension
+	Shape []uint64
+
+	io.WriterTo
+}
+
+func (t Tensor) blockSize() uint64 {
+	switch {
+	case t.Kind < 2:
+		return 1
+	case t.Kind < 10:
+		return 32
+	default:
+		return 256
+	}
+}
+
+func (t Tensor) typeSize() uint64 {
+	blockSize := t.blockSize()
+
+	switch t.Kind {
+	case 0: // FP32
+		return 4
+	case 1: // FP16
+		return 2
+	case 2: // Q4_0
+		return 2 + blockSize/2
+	case 3: // Q4_1
+		return 2 + 2 + blockSize/2
+	case 6: // Q5_0
+		return 2 + 4 + blockSize/2
+	case 7: // Q5_1
+		return 2 + 2 + 4 + blockSize/2
+	case 8: // Q8_0
+		return 2 + blockSize
+	case 9: // Q8_1
+		return 4 + 4 + blockSize
+	case 10: // Q2_K
+		return blockSize/16 + blockSize/4 + 2 + 2
+	case 11: // Q3_K
+		return blockSize/8 + blockSize/4 + 12 + 2
+	case 12: // Q4_K
+		return 2 + 2 + 12 + blockSize/2
+	case 13: // Q5_K
+		return 2 + 2 + 12 + blockSize/8 + blockSize/2
+	case 14: // Q6_K
+		return blockSize/2 + blockSize/4 + blockSize/16 + 2
+	case 15: // Q8_K
+		return 2 + blockSize + 2*blockSize/16
+	case 16: // IQ2_XXS
+		return 2 + 2*blockSize/8
+	case 17: // IQ2_XS
+		return 2 + 2*blockSize/8 + blockSize/32
+	case 18: // IQ3_XXS
+		return 2 + 3*blockSize/8
+	default:
+		return 0
+	}
+}
+
+func (t Tensor) parameters() uint64 {
+	var count uint64 = 1
+	for _, n := range t.Shape {
+		count *= n
+	}
+	return count
+}
+
+func (t Tensor) size() uint64 {
+	return t.parameters() * t.typeSize() / t.blockSize()
+}
+
 type container interface {
 	Name() string
 	Decode(io.ReadSeeker) (model, error)
@@ -133,11 +212,11 @@ func DecodeGGML(rs io.ReadSeeker) (*GGML, error) {
 	case FILE_MAGIC_GGML, FILE_MAGIC_GGMF, FILE_MAGIC_GGJT:
 		return nil, ErrUnsupportedFormat
 	case FILE_MAGIC_GGLA:
-		c = &ContainerGGLA{}
+		c = &containerGGLA{}
 	case FILE_MAGIC_GGUF_LE:
-		c = &ContainerGGUF{ByteOrder: binary.LittleEndian}
+		c = &containerGGUF{ByteOrder: binary.LittleEndian}
 	case FILE_MAGIC_GGUF_BE:
-		c = &ContainerGGUF{ByteOrder: binary.BigEndian}
+		c = &containerGGUF{ByteOrder: binary.BigEndian}
 	default:
 		return nil, errors.New("invalid file magic")
 	}

readline/history.go

@@ -142,7 +142,9 @@ func (h *History) Save() error {
 	for cnt := 0; cnt < h.Size(); cnt++ {
 		v, _ := h.Buf.Get(cnt)
 		line, _ := v.([]rune)
-		buf.WriteString(string(line) + "\n")
+		if _, err := buf.WriteString(string(line) + "\n"); err != nil {
+			return err
+		}
 	}
 	buf.Flush()
 	f.Close()

server/images.go

@@ -321,7 +321,7 @@ func CreateModel(ctx context.Context, name, modelFileDir string, commands []pars
 
 			pathName := realpath(modelFileDir, c.Args)
 
-			ggufName, err := convertSafetensors(name, pathName)
+			ggufName, err := convertSafetensors(name, pathName, fn)
 			if err != nil {
 				var pathErr *fs.PathError
 				switch {
@@ -336,6 +336,7 @@ func CreateModel(ctx context.Context, name, modelFileDir string, commands []pars
 
 			if ggufName != "" {
 				pathName = ggufName
+				slog.Debug(fmt.Sprintf("new image layer path: %s", pathName))
 				defer os.RemoveAll(ggufName)
 			}
 
@@ -422,10 +423,13 @@ func CreateModel(ctx context.Context, name, modelFileDir string, commands []pars
 		CREATE:
 			for {
 				fn(api.ProgressResponse{Status: "creating model layer"})
+				if _, err := bin.Seek(offset, io.SeekStart); err != nil {
+					return err
+				}
 
-				bin.Seek(offset, io.SeekStart)
 				ggml, err := llm.DecodeGGML(bin)
 				if err != nil {
+					slog.Error(fmt.Sprintf("error decoding gguf file: %q", err))
 					switch {
 					case errors.Is(err, io.EOF):
 						break CREATE
@@ -621,8 +625,8 @@ func CreateModel(ctx context.Context, name, modelFileDir string, commands []pars
 	return nil
 }
 
-func convertSafetensors(name, fn string) (string, error) {
-	r, err := zip.OpenReader(fn)
+func convertSafetensors(name, path string, fn func(resp api.ProgressResponse)) (string, error) {
+	r, err := zip.OpenReader(path)
 	if err != nil {
 		return "", err
 	}
@@ -634,6 +638,7 @@ func convertSafetensors(name, fn string) (string, error) {
 	}
 	defer os.RemoveAll(tempDir)
 
+	fn(api.ProgressResponse{Status: "unpacking model metadata"})
 	for _, f := range r.File {
 		fpath := filepath.Join(tempDir, f.Name)
 		outFile, err := os.OpenFile(fpath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
@@ -662,6 +667,7 @@ func convertSafetensors(name, fn string) (string, error) {
 
 	SupportedArchs := []string{
 		"MistralForCausalLM",
+		"GemmaForCausalLM",
 	}
 
 	for _, arch := range params.Architectures {
@@ -670,22 +676,24 @@ func convertSafetensors(name, fn string) (string, error) {
 		}
 	}
 
-	t, err := convert.GetSafeTensors(tempDir)
+	fn(api.ProgressResponse{Status: "processing safetensors"})
+	t, err := convert.GetSafeTensors(tempDir, params)
 	if err != nil {
 		return "", err
 	}
 
-	vocab, err := convert.LoadTokens(tempDir)
+	vocab, err := convert.LoadTokens(tempDir, params)
 	if err != nil {
 		return "", err
 	}
 
-	fn, err = convert.WriteGGUF(name, t, params, vocab)
+	fn(api.ProgressResponse{Status: "converting model"})
+	path, err = convert.WriteGGUF(name, t, params, vocab)
 	if err != nil {
 		return "", err
 	}
 
-	return fn, nil
+	return path, nil
 }
 
 func CopyModel(src, dest string) error {

server/routes_test.go

@@ -3,6 +3,7 @@ package server
 import (
 	"bytes"
 	"context"
+	"encoding/binary"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -31,13 +32,22 @@ func Test_Routes(t *testing.T) {
 	}
 
 	createTestFile := func(t *testing.T, name string) string {
+		t.Helper()
+
 		f, err := os.CreateTemp(t.TempDir(), name)
 		assert.Nil(t, err)
 		defer f.Close()
 
-		_, err = f.Write([]byte("GGUF"))
+		err = binary.Write(f, binary.LittleEndian, []byte("GGUF"))
+		assert.Nil(t, err)
+
+		err = binary.Write(f, binary.LittleEndian, uint32(3))
 		assert.Nil(t, err)
-		_, err = f.Write([]byte{0x2, 0})
+
+		err = binary.Write(f, binary.LittleEndian, uint64(0))
+		assert.Nil(t, err)
+
+		err = binary.Write(f, binary.LittleEndian, uint64(0))
 		assert.Nil(t, err)
 
 		return f.Name()