convert.go 3.8 KB
Newer Older
1
2
3
4
package convert

import (
	"encoding/json"
Michael Yang's avatar
Michael Yang committed
5
	"errors"
6
	"fmt"
Michael Yang's avatar
Michael Yang committed
7
	"io"
8
	"io/fs"
9
	"log/slog"
Michael Yang's avatar
Michael Yang committed
10
	"strings"
11

12
	"github.com/ollama/ollama/llm"
13
14
)

Michael Yang's avatar
Michael Yang committed
15
16
17
type Parameters struct {
	Architectures []string `json:"architectures"`
	VocabSize     uint32   `json:"vocab_size"`
18
19
}

Michael Yang's avatar
Michael Yang committed
20
21
22
23
24
25
26
27
28
29
func (Parameters) KV(t *Tokenizer) llm.KV {
	kv := llm.KV{
		"general.file_type":            uint32(1),
		"general.quantization_version": uint32(2),
		"tokenizer.ggml.pre":           t.Pre,
		"tokenizer.ggml.model":         t.Vocabulary.Model,
		"tokenizer.ggml.tokens":        t.Vocabulary.Tokens,
		"tokenizer.ggml.scores":        t.Vocabulary.Scores,
		"tokenizer.ggml.token_type":    t.Vocabulary.Types,
	}
30

31
32
33
34
	if len(t.Merges) > 0 {
		kv["tokenizer.ggml.merges"] = t.Merges
	}

Michael Yang's avatar
Michael Yang committed
35
36
37
	if t.Template != "" {
		kv["tokenizer.chat_template"] = t.Template
	}
38

Michael Yang's avatar
Michael Yang committed
39
40
41
42
	for _, sv := range t.SpecialVocabulary {
		kv[fmt.Sprintf("tokenizer.ggml.%s_token_id", sv.Key())] = uint32(sv.ID)
		kv[fmt.Sprintf("tokenizer.ggml.add_%s_token", sv.Key())] = sv.AddToken
	}
43

Michael Yang's avatar
Michael Yang committed
44
	return kv
45
46
}

Michael Yang's avatar
Michael Yang committed
47
func (Parameters) specialTokenTypes() []string {
Michael Yang's avatar
Michael Yang committed
48
49
	return []string{
		"bos", "eos", "unk", "sep", "pad", "cls", "mask",
50
	}
Michael Yang's avatar
Michael Yang committed
51
}
52

Michael Yang's avatar
Michael Yang committed
53
func (Parameters) writeFile(ws io.WriteSeeker, kv llm.KV, ts []llm.Tensor) error {
Michael Yang's avatar
Michael Yang committed
54
	return llm.WriteGGUF(ws, kv, ts)
55
56
}

Michael Yang's avatar
Michael Yang committed
57
58
59
60
type Converter interface {
	// KV maps parameters to LLM key-values
	KV(*Tokenizer) llm.KV
	// Tensors maps input tensors to LLM tensors. Model specific modifications can be done here.
Michael Yang's avatar
Michael Yang committed
61
	Tensors([]Tensor) []llm.Tensor
Michael Yang's avatar
Michael Yang committed
62
63
64
	// Replacements returns a list of string pairs to replace in tensor names.
	// See [strings.Replacer](https://pkg.go.dev/strings#Replacer) for details
	Replacements() []string
Michael Yang's avatar
Michael Yang committed
65

Michael Yang's avatar
Michael Yang committed
66
67
	// specialTokenTypes returns any special token types the model uses
	specialTokenTypes() []string
Michael Yang's avatar
Michael Yang committed
68
	// writeFile writes the model to the provided io.WriteSeeker
Michael Yang's avatar
Michael Yang committed
69
	writeFile(io.WriteSeeker, llm.KV, []llm.Tensor) error
70
71
}

Michael Yang's avatar
bert  
Michael Yang committed
72
73
74
75
type moreParser interface {
	parseMore(fs.FS) error
}

Michael Yang's avatar
Michael Yang committed
76
77
78
79
// Convert writes an Ollama compatible model to the provided io.WriteSeeker based on configurations
// and files it finds in the input path.
// Supported input model formats include safetensors.
// Supported input tokenizers files include tokenizer.json (preferred) and tokenizer.model.
80
81
func Convert(fsys fs.FS, ws io.WriteSeeker) error {
	bts, err := fs.ReadFile(fsys, "config.json")
82
	if err != nil {
Michael Yang's avatar
Michael Yang committed
83
		return err
84
85
	}

Michael Yang's avatar
Michael Yang committed
86
	var p Parameters
Michael Yang's avatar
Michael Yang committed
87
	if err := json.Unmarshal(bts, &p); err != nil {
Michael Yang's avatar
Michael Yang committed
88
		return err
89
90
	}

Michael Yang's avatar
Michael Yang committed
91
92
	if len(p.Architectures) < 1 {
		return errors.New("unknown architecture")
93
94
	}

Michael Yang's avatar
Michael Yang committed
95
	var conv Converter
Michael Yang's avatar
Michael Yang committed
96
97
	switch p.Architectures[0] {
	case "LlamaForCausalLM", "MistralForCausalLM":
Michael Yang's avatar
Michael Yang committed
98
		conv = &llama{}
Michael Yang's avatar
Michael Yang committed
99
	case "MixtralForCausalLM":
Michael Yang's avatar
Michael Yang committed
100
		conv = &mixtral{}
Michael Yang's avatar
Michael Yang committed
101
	case "GemmaForCausalLM":
Michael Yang's avatar
Michael Yang committed
102
		conv = &gemma{}
Michael Yang's avatar
Michael Yang committed
103
104
	case "Gemma2ForCausalLM":
		conv = &gemma2{}
105
106
	case "Phi3ForCausalLM":
		conv = &phi3{}
Michael Yang's avatar
bert  
Michael Yang committed
107
108
	case "BertModel":
		conv = &bert{}
Michael Yang's avatar
Michael Yang committed
109
110
	default:
		return errors.New("unsupported architecture")
111
112
	}

Michael Yang's avatar
Michael Yang committed
113
	if err := json.Unmarshal(bts, conv); err != nil {
Michael Yang's avatar
Michael Yang committed
114
		return err
115
116
	}

Michael Yang's avatar
bert  
Michael Yang committed
117
118
119
120
121
122
	if t, ok := conv.(moreParser); ok {
		if err := t.parseMore(fsys); err != nil {
			return err
		}
	}

123
	t, err := parseTokenizer(fsys, conv.specialTokenTypes())
Michael Yang's avatar
Michael Yang committed
124
125
	if err != nil {
		return err
126
127
	}

Michael Yang's avatar
Michael Yang committed
128
129
130
131
132
133
	if vocabSize := int(p.VocabSize); vocabSize > len(t.Vocabulary.Tokens) {
		slog.Warn("vocabulary is smaller than expected, padding with dummy tokens", "expect", p.VocabSize, "actual", len(t.Vocabulary.Tokens))
		for i := range vocabSize - len(t.Vocabulary.Tokens) {
			t.Vocabulary.Tokens = append(t.Vocabulary.Tokens, fmt.Sprintf("[PAD%d]", i))
			t.Vocabulary.Scores = append(t.Vocabulary.Scores, -1)
			t.Vocabulary.Types = append(t.Vocabulary.Types, tokenTypeUserDefined)
134
		}
Michael Yang's avatar
Michael Yang committed
135
136
	} else {
		slog.Debug("vocabulary", "size", len(t.Vocabulary.Tokens))
137
	}
Michael Yang's avatar
Michael Yang committed
138

Michael Yang's avatar
Michael Yang committed
139
	ts, err := parseTensors(fsys, strings.NewReplacer(conv.Replacements()...))
Michael Yang's avatar
Michael Yang committed
140
141
	if err != nil {
		return err
142
143
	}

Michael Yang's avatar
Michael Yang committed
144
	return conv.writeFile(ws, conv.KV(t), conv.Tensors(ts))
145
}