354 lines
8.0 KiB
Go
354 lines
8.0 KiB
Go
package main
|
|
|
|
// rawkv.go implements a minimal read-only GGUF metadata (KV) section walker.
|
|
//
|
|
// Why not go-gguf for this: the library's MetadataEntry.Value() deliberately
|
|
// returns an empty Value for BTypeArray (see readLazyValue in lazy.go), which is
|
|
// fine for scalar/string consumers but loses the tokenizer arrays
|
|
// (tokenizer.ggml.tokens / scores / token_type / merges). To reproduce the
|
|
// reference file's metadata byte-for-byte we need the raw array element bytes.
|
|
// go-gguf is still used for all tensor reading/writing and for scalar metadata;
|
|
// this walker only fills the array gap. It never writes GGUF.
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"fmt"
|
|
"hash/crc32"
|
|
"io"
|
|
"math"
|
|
"os"
|
|
)
|
|
|
|
// ggufBType wire enum (mirrors gguf.BType).
|
|
type ggufBType uint32
|
|
|
|
const (
|
|
btUint8 ggufBType = 0
|
|
btInt8 ggufBType = 1
|
|
btUint16 ggufBType = 2
|
|
btInt16 ggufBType = 3
|
|
btUint32 ggufBType = 4
|
|
btInt32 ggufBType = 5
|
|
btFloat32 ggufBType = 6
|
|
btBool ggufBType = 7
|
|
btString ggufBType = 8
|
|
btArray ggufBType = 9
|
|
btUint64 ggufBType = 10
|
|
btInt64 ggufBType = 11
|
|
btFloat64 ggufBType = 12
|
|
)
|
|
|
|
func (t ggufBType) name() string {
|
|
switch t {
|
|
case btUint8:
|
|
return "UINT8"
|
|
case btInt8:
|
|
return "INT8"
|
|
case btUint16:
|
|
return "UINT16"
|
|
case btInt16:
|
|
return "INT16"
|
|
case btUint32:
|
|
return "UINT32"
|
|
case btInt32:
|
|
return "INT32"
|
|
case btFloat32:
|
|
return "FLOAT32"
|
|
case btBool:
|
|
return "BOOL"
|
|
case btString:
|
|
return "STRING"
|
|
case btArray:
|
|
return "ARRAY"
|
|
case btUint64:
|
|
return "UINT64"
|
|
case btInt64:
|
|
return "INT64"
|
|
case btFloat64:
|
|
return "FLOAT64"
|
|
default:
|
|
return fmt.Sprintf("BTYPE(%d)", uint32(t))
|
|
}
|
|
}
|
|
|
|
func (t ggufBType) size() int {
|
|
switch t {
|
|
case btUint8, btInt8, btBool:
|
|
return 1
|
|
case btUint16, btInt16:
|
|
return 2
|
|
case btUint32, btInt32, btFloat32:
|
|
return 4
|
|
case btUint64, btInt64, btFloat64:
|
|
return 8
|
|
default:
|
|
return 0
|
|
}
|
|
}
|
|
|
|
// rawKV is one metadata entry with its raw wire value preserved.
|
|
type rawKV struct {
|
|
Key string
|
|
BType ggufBType
|
|
Int int64 // integer/bool scalars
|
|
Float float64 // float scalars
|
|
Str string // string scalars
|
|
ElemType ggufBType // arrays: element type
|
|
Count uint64 // arrays: element count
|
|
Raw []byte // arrays: concatenated element bytes (excludes elem_type+count)
|
|
}
|
|
|
|
// rawKVHeader holds the GGUF header fields.
|
|
type rawKVHeader struct {
|
|
Version uint32
|
|
NTensors uint64
|
|
NKV uint64
|
|
}
|
|
|
|
// readRawKV opens path and returns the header plus every KV entry in file order.
|
|
func readRawKV(path string) (rawKVHeader, []rawKV, error) {
|
|
f, err := os.Open(path)
|
|
if err != nil {
|
|
return rawKVHeader{}, nil, err
|
|
}
|
|
defer f.Close()
|
|
|
|
var hdr rawKVHeader
|
|
var b [24]byte
|
|
if _, err := io.ReadFull(f, b[:]); err != nil {
|
|
return hdr, nil, fmt.Errorf("read header: %w", err)
|
|
}
|
|
if string(b[0:4]) != "GGUF" {
|
|
return hdr, nil, fmt.Errorf("bad magic %q", b[0:4])
|
|
}
|
|
hdr.Version = binary.LittleEndian.Uint32(b[4:8])
|
|
hdr.NTensors = binary.LittleEndian.Uint64(b[8:16])
|
|
hdr.NKV = binary.LittleEndian.Uint64(b[16:24])
|
|
|
|
br := newBufReader(f)
|
|
out := make([]rawKV, 0, hdr.NKV)
|
|
for i := uint64(0); i < hdr.NKV; i++ {
|
|
kv, err := readOneKV(br)
|
|
if err != nil {
|
|
return hdr, nil, fmt.Errorf("kv[%d]: %w", i, err)
|
|
}
|
|
out = append(out, kv)
|
|
}
|
|
return hdr, out, nil
|
|
}
|
|
|
|
func readOneKV(br *bufReader) (rawKV, error) {
|
|
var kv rawKV
|
|
key, err := br.str()
|
|
if err != nil {
|
|
return kv, fmt.Errorf("key: %w", err)
|
|
}
|
|
kv.Key = key
|
|
bt, err := br.u32()
|
|
if err != nil {
|
|
return kv, fmt.Errorf("btype: %w", err)
|
|
}
|
|
kv.BType = ggufBType(bt)
|
|
|
|
switch kv.BType {
|
|
case btBool, btUint8, btInt8:
|
|
v, err := br.u8()
|
|
if err != nil {
|
|
return kv, err
|
|
}
|
|
kv.Int = int64(v)
|
|
case btUint16, btInt16:
|
|
v, err := br.u16()
|
|
if err != nil {
|
|
return kv, err
|
|
}
|
|
if kv.BType == btInt16 {
|
|
kv.Int = int64(int16(v))
|
|
} else {
|
|
kv.Int = int64(v)
|
|
}
|
|
case btUint32, btInt32:
|
|
v, err := br.u32()
|
|
if err != nil {
|
|
return kv, err
|
|
}
|
|
if kv.BType == btInt32 {
|
|
kv.Int = int64(int32(v))
|
|
} else {
|
|
kv.Int = int64(v)
|
|
}
|
|
case btFloat32:
|
|
v, err := br.u32()
|
|
if err != nil {
|
|
return kv, err
|
|
}
|
|
kv.Float = float64(math.Float32frombits(v))
|
|
case btUint64, btInt64:
|
|
v, err := br.u64()
|
|
if err != nil {
|
|
return kv, err
|
|
}
|
|
kv.Int = int64(v)
|
|
case btFloat64:
|
|
v, err := br.u64()
|
|
if err != nil {
|
|
return kv, err
|
|
}
|
|
kv.Float = math.Float64frombits(v)
|
|
case btString:
|
|
s, err := br.str()
|
|
if err != nil {
|
|
return kv, err
|
|
}
|
|
kv.Str = s
|
|
case btArray:
|
|
et, err := br.u32()
|
|
if err != nil {
|
|
return kv, fmt.Errorf("array elem_type: %w", err)
|
|
}
|
|
kv.ElemType = ggufBType(et)
|
|
cnt, err := br.u64()
|
|
if err != nil {
|
|
return kv, fmt.Errorf("array count: %w", err)
|
|
}
|
|
kv.Count = cnt
|
|
raw, err := readArrayElems(br, kv.ElemType, cnt)
|
|
if err != nil {
|
|
return kv, fmt.Errorf("array elems: %w", err)
|
|
}
|
|
kv.Raw = raw
|
|
default:
|
|
return kv, fmt.Errorf("unsupported btype %d", kv.BType)
|
|
}
|
|
return kv, nil
|
|
}
|
|
|
|
// readArrayElems reads `count` elements of elemType, returning their concatenated
|
|
// wire bytes (exactly the bytes go-gguf's writer expects in Value.Raw).
|
|
func readArrayElems(br *bufReader, elemType ggufBType, count uint64) ([]byte, error) {
|
|
if elemType == btArray || elemType == btString {
|
|
// strings (and nested arrays, unsupported) are variable length
|
|
if elemType != btString {
|
|
return nil, fmt.Errorf("nested array element type unsupported")
|
|
}
|
|
var out []byte
|
|
for i := uint64(0); i < count; i++ {
|
|
slen, err := br.u64()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, leU64(slen)...)
|
|
buf := make([]byte, slen)
|
|
if _, err := io.ReadFull(br.r, buf); err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, buf...)
|
|
}
|
|
return out, nil
|
|
}
|
|
sz := elemType.size()
|
|
if sz == 0 {
|
|
return nil, fmt.Errorf("unknown array element type %d", elemType)
|
|
}
|
|
total := int(count) * sz
|
|
// sanity cap: 256 MiB
|
|
if uint64(total) > 256<<20 {
|
|
return nil, fmt.Errorf("array too large: %d bytes", total)
|
|
}
|
|
buf := make([]byte, total)
|
|
if _, err := io.ReadFull(br.r, buf); err != nil {
|
|
return nil, err
|
|
}
|
|
return buf, nil
|
|
}
|
|
|
|
// rawKVHash is a stable fingerprint of a KV entry's full wire value.
|
|
func (kv rawKV) rawKVHash() string {
|
|
h := crc32.NewIEEE()
|
|
h.Write([]byte(kv.Key))
|
|
var b [8]byte
|
|
binary.LittleEndian.PutUint32(b[:4], uint32(kv.BType))
|
|
h.Write(b[:4])
|
|
switch kv.BType {
|
|
case btBool, btUint8, btInt8, btUint16, btInt16, btUint32, btInt32, btUint64, btInt64:
|
|
binary.LittleEndian.PutUint64(b[:], uint64(kv.Int))
|
|
h.Write(b[:])
|
|
case btFloat32:
|
|
binary.LittleEndian.PutUint32(b[:], math.Float32bits(float32(kv.Float)))
|
|
h.Write(b[:4])
|
|
case btFloat64:
|
|
binary.LittleEndian.PutUint64(b[:], math.Float64bits(kv.Float))
|
|
h.Write(b[:])
|
|
case btString:
|
|
h.Write([]byte(kv.Str))
|
|
case btArray:
|
|
binary.LittleEndian.PutUint32(b[:4], uint32(kv.ElemType))
|
|
h.Write(b[:4])
|
|
binary.LittleEndian.PutUint64(b[:], kv.Count)
|
|
h.Write(b[:])
|
|
h.Write(kv.Raw)
|
|
}
|
|
return fmt.Sprintf("%08x", h.Sum32())
|
|
}
|
|
|
|
// --- tiny buffered reader helpers ---
|
|
|
|
type bufReader struct {
|
|
r io.Reader
|
|
}
|
|
|
|
func newBufReader(r io.Reader) *bufReader { return &bufReader{r: r} }
|
|
|
|
func (b *bufReader) u8() (uint8, error) {
|
|
var x [1]byte
|
|
if _, err := io.ReadFull(b.r, x[:]); err != nil {
|
|
return 0, err
|
|
}
|
|
return x[0], nil
|
|
}
|
|
|
|
func (b *bufReader) u16() (uint16, error) {
|
|
var x [2]byte
|
|
if _, err := io.ReadFull(b.r, x[:]); err != nil {
|
|
return 0, err
|
|
}
|
|
return binary.LittleEndian.Uint16(x[:]), nil
|
|
}
|
|
|
|
func (b *bufReader) u32() (uint32, error) {
|
|
var x [4]byte
|
|
if _, err := io.ReadFull(b.r, x[:]); err != nil {
|
|
return 0, err
|
|
}
|
|
return binary.LittleEndian.Uint32(x[:]), nil
|
|
}
|
|
|
|
func (b *bufReader) u64() (uint64, error) {
|
|
var x [8]byte
|
|
if _, err := io.ReadFull(b.r, x[:]); err != nil {
|
|
return 0, err
|
|
}
|
|
return binary.LittleEndian.Uint64(x[:]), nil
|
|
}
|
|
|
|
func (b *bufReader) str() (string, error) {
|
|
n, err := b.u64()
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if n > 1<<28 {
|
|
return "", fmt.Errorf("string too long: %d", n)
|
|
}
|
|
buf := make([]byte, n)
|
|
if _, err := io.ReadFull(b.r, buf); err != nil {
|
|
return "", err
|
|
}
|
|
return string(buf), nil
|
|
}
|
|
|
|
func leU64(v uint64) []byte {
|
|
var b [8]byte
|
|
binary.LittleEndian.PutUint64(b[:], v)
|
|
return b[:]
|
|
}
|