reader.go 1.89 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
	tensorKindFP32 uint32 = iota
	tensorKindFP16
	tensorKindBF16  = 30
37
	tensorKindMXFP4 = 39
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
		t.name == "v.positional_embedding_vlm" ||
		t.name == "v.tile_position_embd.weight" ||
		t.name == "v.pre_tile_position_embd.weight" ||
Michael Yang's avatar
Michael Yang committed
47
48
49
50
		t.name == "v.post_tile_position_embd.weight" ||
		t.name == "s.position_embd" ||
		strings.HasSuffix(t.name, "rel_pos_h") ||
		strings.HasSuffix(t.name, "rel_pos_w") {
Michael Yang's avatar
Michael Yang committed
51
		// these tensors are always F32
Michael Yang's avatar
Michael Yang committed
52
		return tensorKindFP32
Michael Yang's avatar
Michael Yang committed
53
54
55
56
57
58
	}

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

Michael Yang's avatar
llama4  
Michael Yang committed
65
func (t *tensorBase) SetRepacker(fn Repacker) {
Michael Yang's avatar
Michael Yang committed
66
67
68
	t.repacker = fn
}

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

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

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

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

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