baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

comments

Browse source
This commit is contained in:
Michael Yang 2024-07-08 16:59:48 -07:00
parent 5e9db9fb0b
commit df993fa37b
12 changed files with 63 additions and 61 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
convert/convert.go
View file

@ -40,13 +40,13 @@ func (Parameters) KV(t *Tokenizer) llm.KV {
return kv return kv
} }
func (Parameters) specialTypes() []string { func (Parameters) specialTokenTypes() []string {
return []string{ return []string{
"bos", "eos", "unk", "sep", "pad", "cls", "mask", "bos", "eos", "unk", "sep", "pad", "cls", "mask",
} }
} }
func (Parameters) writeFile(ws io.WriteSeeker, kv llm.KV, ts []*llm.Tensor) error { func (Parameters) writeFile(ws io.WriteSeeker, kv llm.KV, ts []llm.Tensor) error {
return llm.WriteGGUF(ws, kv, ts) return llm.WriteGGUF(ws, kv, ts)
} }
@ -54,24 +54,27 @@ type Converter interface {
// KV maps parameters to LLM key-values // KV maps parameters to LLM key-values
KV(*Tokenizer) llm.KV KV(*Tokenizer) llm.KV
// Tensors maps input tensors to LLM tensors. Model specific modifications can be done here. // Tensors maps input tensors to LLM tensors. Model specific modifications can be done here.
Tensors([]Tensor) []*llm.Tensor Tensors([]Tensor) []llm.Tensor
// tensorName returns the LLM tensor name for a specific input name // tensorName returns the LLM tensor name for a specific input name
tensorName(string) string tensorName(string) string
// specialTypes returns any special token types the model uses // specialTokenTypes returns any special token types the model uses
specialTypes() []string specialTokenTypes() []string
writeFile(io.WriteSeeker, llm.KV, []*llm.Tensor) error writeFile(io.WriteSeeker, llm.KV, []llm.Tensor) error
} }
func Convert(d string, ws io.WriteSeeker) error { // Convert writes an Ollama compatible model to the provided io.WriteSeeker based on configurations
f, err := os.Open(filepath.Join(d, "config.json")) // 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.
func Convert(path string, ws io.WriteSeeker) error {
bts, err := os.ReadFile(filepath.Join(path, "config.json"))
if err != nil { if err != nil {
return err return err
} }
defer f.Close()
var p Parameters var p Parameters
if err := json.NewDecoder(f).Decode(&p); err != nil { if err := json.Unmarshal(bts, &p); err != nil {
return err return err
} }
@ -79,28 +82,23 @@ func Convert(d string, ws io.WriteSeeker) error {
return errors.New("unknown architecture") return errors.New("unknown architecture")
} }
var c Converter var conv Converter
switch p.Architectures[0] { switch p.Architectures[0] {
case "LlamaForCausalLM", "MistralForCausalLM": case "LlamaForCausalLM", "MistralForCausalLM":
c = &llama{} conv = &llama{}
case "MixtralForCausalLM": case "MixtralForCausalLM":
c = &mixtral{} conv = &mixtral{}
case "GemmaForCausalLM": case "GemmaForCausalLM":
c = &gemma{} conv = &gemma{}
default: default:
return errors.New("unsupported architecture") return errors.New("unsupported architecture")
} }
bts, err := os.ReadFile(filepath.Join(d, "config.json")) if err := json.Unmarshal(bts, conv); err != nil {
if err != nil {
return err return err
} }
if err := json.Unmarshal(bts, c); err != nil { t, err := parseTokenizer(path, conv.specialTokenTypes())
return err
}
t, err := parseTokenizer(d, c.specialTypes())
if err != nil { if err != nil {
return err return err
} }
@ -112,12 +110,14 @@ func Convert(d string, ws io.WriteSeeker) error {
t.Vocabulary.Scores = append(t.Vocabulary.Scores, -1) t.Vocabulary.Scores = append(t.Vocabulary.Scores, -1)
t.Vocabulary.Types = append(t.Vocabulary.Types, tokenTypeUserDefined) t.Vocabulary.Types = append(t.Vocabulary.Types, tokenTypeUserDefined)
} }
} else {
slog.Debug("vocabulary", "size", len(t.Vocabulary.Tokens))
} }
ts, err := parseTensors(d) ts, err := parseTensors(path)
if err != nil { if err != nil {
return err return err
} }
return c.writeFile(ws, c.KV(t), c.Tensors(ts)) return conv.writeFile(ws, conv.KV(t), conv.Tensors(ts))
} }

6
convert/convert_gemma.go
View file

@ -43,15 +43,15 @@ func (p *gemma) KV(t *Tokenizer) llm.KV {
return kv return kv
} }
func (p *gemma) Tensors(ts []Tensor) []*llm.Tensor { func (p *gemma) Tensors(ts []Tensor) []llm.Tensor {
var out []*llm.Tensor var out []llm.Tensor
for _, t := range ts { for _, t := range ts {
name := p.tensorName(t.Name()) name := p.tensorName(t.Name())
if strings.HasSuffix(name, "_norm.weight") { if strings.HasSuffix(name, "_norm.weight") {
t.SetRepacker(p.addOne) t.SetRepacker(p.addOne)
} }
out = append(out, &llm.Tensor{ out = append(out, llm.Tensor{
Name: name, Name: name,
Kind: t.Kind(), Kind: t.Kind(),
Shape: t.Shape(), Shape: t.Shape(),

6
convert/convert_llama.go
View file

@ -96,8 +96,8 @@ func (p *llama) KV(t *Tokenizer) llm.KV {
return kv return kv
} }
func (p *llama) Tensors(ts []Tensor) []*llm.Tensor { func (p *llama) Tensors(ts []Tensor) []llm.Tensor {
var out []*llm.Tensor var out []llm.Tensor
for _, t := range ts { for _, t := range ts {
name := p.tensorName(t.Name()) name := p.tensorName(t.Name())
if strings.HasSuffix(name, "attn_q.weight") || if strings.HasSuffix(name, "attn_q.weight") ||
@ -105,7 +105,7 @@ func (p *llama) Tensors(ts []Tensor) []*llm.Tensor {
t.SetRepacker(p.repack) t.SetRepacker(p.repack)
} }
out = append(out, &llm.Tensor{ out = append(out, llm.Tensor{
Name: name, Name: name,
Kind: t.Kind(), Kind: t.Kind(),
Shape: t.Shape(), Shape: t.Shape(),

6
convert/convert_mixtral.go
View file

@ -31,7 +31,7 @@ func (p *mixtral) KV(t *Tokenizer) llm.KV {
return kv return kv
} }
func (p *mixtral) Tensors(ts []Tensor) []*llm.Tensor { func (p *mixtral) Tensors(ts []Tensor) []llm.Tensor {
oldnew := []string{ oldnew := []string{
"model.layers", "blk", "model.layers", "blk",
"w1", "ffn_gate_exps", "w1", "ffn_gate_exps",
@ -58,10 +58,10 @@ func (p *mixtral) Tensors(ts []Tensor) []*llm.Tensor {
return true return true
}) })
var out []*llm.Tensor var out []llm.Tensor
for n, e := range experts { for n, e := range experts {
// TODO(mxyng): sanity check experts // TODO(mxyng): sanity check experts
out = append(out, &llm.Tensor{ out = append(out, llm.Tensor{
Name: n, Name: n,
Kind: e[0].Kind(), Kind: e[0].Kind(),
Shape: append([]uint64{uint64(len(e))}, e[0].Shape()...), Shape: append([]uint64{uint64(len(e))}, e[0].Shape()...),

9
convert/reader.go
View file

@ -29,6 +29,11 @@ func (t tensorBase) Shape() []uint64 {
return t.shape return t.shape
} }
const (
tensorKindF32 uint32 = iota
tensorKindF16
)
func (t tensorBase) Kind() uint32 { func (t tensorBase) Kind() uint32 {
if strings.HasSuffix(t.name, ".block_sparse_moe.gate.weight") { if strings.HasSuffix(t.name, ".block_sparse_moe.gate.weight") {
return 0 return 0
@ -38,9 +43,9 @@ func (t tensorBase) Kind() uint32 {
case 0: case 0:
panic("invalid tensor shape") panic("invalid tensor shape")
case 1: case 1:
return 0 return tensorKindF32
default: default:
return 1 return tensorKindF16
} }
} }

5
convert/reader_safetensors.go
View file

@ -66,6 +66,7 @@ func parseSafetensors(ps ...string) ([]Tensor, error) {
return ts, nil return ts, nil
} }
// safetensorsPad returns the padded size of the safetensors file given a length n and offset s
func safetensorsPad(n, s int64) int64 { func safetensorsPad(n, s int64) int64 {
return 8 + n + s return 8 + n + s
} }
@ -125,9 +126,9 @@ func (st safetensor) WriteTo(w io.Writer) (int64, error) {
} }
switch st.Kind() { switch st.Kind() {
case 0: case tensorKindF32:
return 0, binary.Write(w, binary.LittleEndian, f32s) return 0, binary.Write(w, binary.LittleEndian, f32s)
case 1: case tensorKindF16:
f16s := make([]uint16, len(f32s)) f16s := make([]uint16, len(f32s))
for i := range f32s { for i := range f32s {
f16s[i] = float16.Fromfloat32(f32s[i]).Bits() f16s[i] = float16.Fromfloat32(f32s[i]).Bits()

13
convert/tokenizer.go
View file

@ -32,7 +32,7 @@ type Tokenizer struct {
Template string Template string
} }
func parseTokenizer(d string, specialTypes []string) (*Tokenizer, error) { func parseTokenizer(d string, specialTokenTypes []string) (*Tokenizer, error) {
v, err := parseVocabulary(d) v, err := parseVocabulary(d)
if err != nil { if err != nil {
return nil, err return nil, err
@ -66,6 +66,8 @@ func parseTokenizer(d string, specialTypes []string) (*Tokenizer, error) {
switch pt.Type { switch pt.Type {
case "Split": case "Split":
if pt.Pattern.Regex != "" { if pt.Pattern.Regex != "" {
// create a checksum of all Split pretokenizers which should be sufficient
// to identify the pretokenizer
sha256sum.Write([]byte(pt.Pattern.Regex)) sha256sum.Write([]byte(pt.Pattern.Regex))
} }
} }
@ -102,7 +104,7 @@ func parseTokenizer(d string, specialTypes []string) (*Tokenizer, error) {
} }
} }
for _, st := range specialTypes { for _, st := range specialTokenTypes {
sv := SpecialVocabulary{Type: st} sv := SpecialVocabulary{Type: st}
if bts, ok := p[fmt.Sprintf("add_%s_token", st)]; ok { if bts, ok := p[fmt.Sprintf("add_%s_token", st)]; ok {
if err := json.Unmarshal(bts, &sv.AddToken); err != nil { if err := json.Unmarshal(bts, &sv.AddToken); err != nil {
@ -224,15 +226,14 @@ func parseVocabulary(d string) (*Vocabulary, error) {
} }
for pattern, parseFn := range patterns { for pattern, parseFn := range patterns {
matches, err := filepath.Glob(filepath.Join(d, pattern)) if _, err := os.Stat(filepath.Join(d, pattern)); errors.Is(err, os.ErrNotExist) {
if err != nil { continue
} else if err != nil {
return nil, err return nil, err
} }
if len(matches) > 0 {
return parseFn(d) return parseFn(d)
} }
}
return nil, errors.New("unknown tensor format") return nil, errors.New("unknown tensor format")
} }

17
llm/gguf.go
View file

@ -489,6 +489,7 @@ func readGGUFArray(llm *gguf, r io.Reader) (*array, error) {
return a, nil return a, nil
} }
// writeGGUFArray writes a slice s of type E to the write with a gguf type of t
func writeGGUFArray[S ~[]E, E any](w io.Writer, t uint32, s S) error { func writeGGUFArray[S ~[]E, E any](w io.Writer, t uint32, s S) error {
if err := binary.Write(w, binary.LittleEndian, ggufTypeArray); err != nil { if err := binary.Write(w, binary.LittleEndian, ggufTypeArray); err != nil {
return err return err
@ -502,16 +503,10 @@ func writeGGUFArray[S ~[]E, E any](w io.Writer, t uint32, s S) error {
return err return err
} }
for _, e := range s { return binary.Write(w, binary.LittleEndian, s)
if err := binary.Write(w, binary.LittleEndian, e); err != nil {
return err
}
} }
return nil func WriteGGUF(ws io.WriteSeeker, kv KV, ts []Tensor) error {
}
func WriteGGUF(ws io.WriteSeeker, kv KV, ts []*Tensor) error {
if err := binary.Write(ws, binary.LittleEndian, []byte("GGUF")); err != nil { if err := binary.Write(ws, binary.LittleEndian, []byte("GGUF")); err != nil {
return err return err
} }
@ -537,7 +532,7 @@ func WriteGGUF(ws io.WriteSeeker, kv KV, ts []*Tensor) error {
} }
} }
slices.SortFunc(ts, func(a, b *Tensor) int { slices.SortFunc(ts, func(a, b Tensor) int {
var i, j int var i, j int
if n, err := fmt.Sscanf(a.Name, "blk.%d", &i); err != nil || n != 1 { if n, err := fmt.Sscanf(a.Name, "blk.%d", &i); err != nil || n != 1 {
return cmp.Compare(a.Name, b.Name) return cmp.Compare(a.Name, b.Name)
@ -622,7 +617,7 @@ func ggufWriteKV(ws io.WriteSeeker, k string, v any) error {
return err return err
} }
func ggufWriteTensorInfo(ws io.WriteSeeker, t *Tensor) error { func ggufWriteTensorInfo(ws io.WriteSeeker, t Tensor) error {
slog.Debug(t.Name, "kind", t.Kind, "shape", t.Shape, "offset", t.Offset) slog.Debug(t.Name, "kind", t.Kind, "shape", t.Shape, "offset", t.Offset)
if err := binary.Write(ws, binary.LittleEndian, uint64(len(t.Name))); err != nil { if err := binary.Write(ws, binary.LittleEndian, uint64(len(t.Name))); err != nil {
return err return err
@ -649,7 +644,7 @@ func ggufWriteTensorInfo(ws io.WriteSeeker, t *Tensor) error {
return binary.Write(ws, binary.LittleEndian, t.Offset) return binary.Write(ws, binary.LittleEndian, t.Offset)
} }
func ggufWriteTensor(ws io.WriteSeeker, t *Tensor, alignment int64) error { func ggufWriteTensor(ws io.WriteSeeker, t Tensor, alignment int64) error {
offset, err := ws.Seek(0, io.SeekCurrent) offset, err := ws.Seek(0, io.SeekCurrent)
if err != nil { if err != nil {
return err return err

2
llm/memory_test.go
View file

@ -21,7 +21,7 @@ func TestEstimateGPULayers(t *testing.T) {
defer f.Close() defer f.Close()
inputLayerCount := 5 inputLayerCount := 5
tensors := []*Tensor{ tensors := []Tensor{
{Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, {Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},
{Name: "blk.1.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, {Name: "blk.1.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},
{Name: "blk.2.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, {Name: "blk.2.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},

2
server/routes_create_test.go
View file

@ -19,7 +19,7 @@ import (
var stream bool = false var stream bool = false
func createBinFile(t *testing.T, kv map[string]any, ti []*llm.Tensor) string { func createBinFile(t *testing.T, kv map[string]any, ti []llm.Tensor) string {
t.Helper() t.Helper()
f, err := os.CreateTemp(t.TempDir(), "") f, err := os.CreateTemp(t.TempDir(), "")

8
server/routes_generate_test.go
View file

@ -101,7 +101,7 @@ func TestGenerateChat(t *testing.T) {
"tokenizer.ggml.tokens": []string{""}, "tokenizer.ggml.tokens": []string{""},
"tokenizer.ggml.scores": []float32{0}, "tokenizer.ggml.scores": []float32{0},
"tokenizer.ggml.token_type": []int32{0}, "tokenizer.ggml.token_type": []int32{0},
}, []*llm.Tensor{ }, []llm.Tensor{
{Name: "token_embd.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, {Name: "token_embd.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))},
{Name: "blk.0.attn_norm.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, {Name: "blk.0.attn_norm.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))},
{Name: "blk.0.ffn_down.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, {Name: "blk.0.ffn_down.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))},
@ -149,7 +149,7 @@ func TestGenerateChat(t *testing.T) {
Modelfile: fmt.Sprintf("FROM %s", createBinFile(t, llm.KV{ Modelfile: fmt.Sprintf("FROM %s", createBinFile(t, llm.KV{
"general.architecture": "bert", "general.architecture": "bert",
"bert.pooling_type": uint32(0), "bert.pooling_type": uint32(0),
}, []*llm.Tensor{})), }, []llm.Tensor{})),
Stream: &stream, Stream: &stream,
}) })
@ -399,7 +399,7 @@ func TestGenerate(t *testing.T) {
"tokenizer.ggml.tokens": []string{""}, "tokenizer.ggml.tokens": []string{""},
"tokenizer.ggml.scores": []float32{0}, "tokenizer.ggml.scores": []float32{0},
"tokenizer.ggml.token_type": []int32{0}, "tokenizer.ggml.token_type": []int32{0},
}, []*llm.Tensor{ }, []llm.Tensor{
{Name: "token_embd.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, {Name: "token_embd.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))},
{Name: "blk.0.attn_norm.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, {Name: "blk.0.attn_norm.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))},
{Name: "blk.0.ffn_down.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, {Name: "blk.0.ffn_down.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))},
@ -447,7 +447,7 @@ func TestGenerate(t *testing.T) {
Modelfile: fmt.Sprintf("FROM %s", createBinFile(t, llm.KV{ Modelfile: fmt.Sprintf("FROM %s", createBinFile(t, llm.KV{
"general.architecture": "bert", "general.architecture": "bert",
"bert.pooling_type": uint32(0), "bert.pooling_type": uint32(0),
}, []*llm.Tensor{})), }, []llm.Tensor{})),
Stream: &stream, Stream: &stream,
}) })

2
server/sched_test.go
View file

@ -124,7 +124,7 @@ func newScenarioRequest(t *testing.T, ctx context.Context, modelName string, est
"tokenizer.ggml.tokens": []string{" "}, "tokenizer.ggml.tokens": []string{" "},
"tokenizer.ggml.scores": []float32{0}, "tokenizer.ggml.scores": []float32{0},
"tokenizer.ggml.token_type": []int32{0}, "tokenizer.ggml.token_type": []int32{0},
}, []*llm.Tensor{ }, []llm.Tensor{
{Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, {Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},
{Name: "output.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, {Name: "output.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},
})) }))