wandb

1
package server
2

3
import (
4
	"bytes"
5
	"encoding/binary"
6
	"log/slog"
7
)
8

9
type Header struct {
10
	Magic      uint8
11
	DataLength uint32
12
}
13

14
type Tokenizer struct {
15
	header       Header
16
	headerLength int
17
	headerValid  bool
18
}
19

20
func (x *Tokenizer) Split(data []byte, _ bool) (advance int, token []byte, err error) {
21
	if x.headerLength == 0 {
22
		x.headerLength = binary.Size(x.header)
23
	}
24

25
	advance = 0
26

27
	if !x.headerValid {
28
		if len(data) < x.headerLength {
29
			return
30
		}
31
		buf := bytes.NewReader(data)
32
		err := binary.Read(buf, binary.LittleEndian, &x.header)
33
		if err != nil {
34
			slog.Error("can't read token", "err", err)
35
			return 0, nil, err
36
		}
37
		if x.header.Magic != uint8('W') {
38
			slog.Error("Invalid magic byte in header")
39
		}
40
		x.headerValid = true
41
		advance += x.headerLength
42
		data = data[advance:]
43
	}
44

45
	if len(data) < int(x.header.DataLength) {
46
		return
47
	}
48

49
	advance += int(x.header.DataLength)
50
	token = data[:x.header.DataLength]
51
	x.headerValid = false
52
	return
53
}
54