reader.go 1.71 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
bf16  
Michael Yang committed
34
35
36
	tensorKindFP32 uint32 = iota
	tensorKindFP16
	tensorKindBF16 = 30
Michael Yang's avatar
Michael Yang committed
37
38
)

Michael Yang's avatar
Michael Yang committed
39
func (t tensorBase) Kind() uint32 {
Michael Yang's avatar
Michael Yang committed
40
	if strings.HasSuffix(t.name, ".ffn_gate_inp.weight") ||
Michael Yang's avatar
llama4  
Michael Yang committed
41
		t.name == "token_types.weight" ||
42
43
44
45
		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
46
		// these tensors are always F32
Michael Yang's avatar
bf16  
Michael Yang committed
47
		return tensorKindFP32
Michael Yang's avatar
Michael Yang committed
48
49
50
51
52
53
	}

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

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

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

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

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

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

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