reader.go 1.77 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
4
5
package convert

import (
	"errors"
	"io"
6
	"io/fs"
Michael Yang's avatar
Michael Yang committed
7
8
9
10
11
12
13
	"strings"
)

type Tensor interface {
	Name() string
	Shape() []uint64
	Kind() uint32
Michael Yang's avatar
llama4  
Michael Yang committed
14
	SetRepacker(Repacker)
Michael Yang's avatar
Michael Yang committed
15
	WriteTo(io.Writer) (int64, error)
Michael Yang's avatar
llama4  
Michael Yang committed
16
	Clone() Tensor
Michael Yang's avatar
Michael Yang committed
17
18
19
}

type tensorBase struct {
Michael Yang's avatar
llama4  
Michael Yang committed
20
21
22
	name     string
	shape    []uint64
	repacker Repacker
Michael Yang's avatar
Michael Yang committed
23
24
25
26
27
28
29
30
31
32
}

func (t tensorBase) Name() string {
	return t.name
}

func (t tensorBase) Shape() []uint64 {
	return t.shape
}

Michael Yang's avatar
Michael Yang committed
33
const (
Michael Yang's avatar
Michael Yang committed
34
35
36
37
	tensorKindFP32 uint32 = iota
	tensorKindFP16
	tensorKindMXFP4 = 4
	tensorKindBF16  = 30
Michael Yang's avatar
Michael Yang committed
38
39
)

Michael Yang's avatar
Michael Yang committed
40
func (t tensorBase) Kind() uint32 {
Michael Yang's avatar
Michael Yang committed
41
	if strings.HasSuffix(t.name, ".ffn_gate_inp.weight") ||
42
		strings.HasSuffix(t.name, ".bias") ||
Michael Yang's avatar
llama4  
Michael Yang committed
43
		t.name == "token_types.weight" ||
44
45
46
47
		t.name == "v.positional_embedding_vlm" ||
		t.name == "v.tile_position_embd.weight" ||
		t.name == "v.pre_tile_position_embd.weight" ||
		t.name == "v.post_tile_position_embd.weight" {
Michael Yang's avatar
Michael Yang committed
48
		// these tensors are always F32
Michael Yang's avatar
Michael Yang committed
49
		return tensorKindFP32
Michael Yang's avatar
Michael Yang committed
50
51
52
53
54
55
	}

	switch len(t.shape) {
	case 0:
		panic("invalid tensor shape")
	case 1:
Michael Yang's avatar
Michael Yang committed
56
		return tensorKindFP32
Michael Yang's avatar
Michael Yang committed
57
	default:
Michael Yang's avatar
Michael Yang committed
58
		return tensorKindFP16
Michael Yang's avatar
Michael Yang committed
59
60
61
	}
}

Michael Yang's avatar
llama4  
Michael Yang committed
62
func (t *tensorBase) SetRepacker(fn Repacker) {
Michael Yang's avatar
Michael Yang committed
63
64
65
	t.repacker = fn
}

Michael Yang's avatar
llama4  
Michael Yang committed
66
type Repacker func(string, []float32, []uint64) ([]float32, error)
Michael Yang's avatar
Michael Yang committed
67

Michael Yang's avatar
Michael Yang committed
68
func parseTensors(fsys fs.FS, replacer *strings.Replacer) ([]Tensor, error) {
Michael Yang's avatar
Michael Yang committed
69
70
	patterns := []struct {
		Pattern string
Michael Yang's avatar
Michael Yang committed
71
		Func    func(fs.FS, *strings.Replacer, ...string) ([]Tensor, error)
Michael Yang's avatar
Michael Yang committed
72
	}{
73
		{"*.safetensors", parseSafetensors},
Michael Yang's avatar
Michael Yang committed
74
75
76
		{"pytorch_model-*-of-*.bin", parseTorch},
		{"pytorch_model.bin", parseTorch},
		{"consolidated.*.pth", parseTorch},
Michael Yang's avatar
Michael Yang committed
77
78
	}

Michael Yang's avatar
Michael Yang committed
79
80
	for _, pattern := range patterns {
		matches, err := fs.Glob(fsys, pattern.Pattern)
Michael Yang's avatar
Michael Yang committed
81
82
83
84
85
		if err != nil {
			return nil, err
		}

		if len(matches) > 0 {
Michael Yang's avatar
Michael Yang committed
86
			return pattern.Func(fsys, replacer, matches...)
Michael Yang's avatar
Michael Yang committed
87
88
89
90
91
		}
	}

	return nil, errors.New("unknown tensor format")
}