File size: 2,555 Bytes
76a12b2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
import struct
from enum import IntEnum
class GGUFValueType(IntEnum):
UINT8 = 0
INT8 = 1
UINT16 = 2
INT16 = 3
UINT32 = 4
INT32 = 5
FLOAT32 = 6
BOOL = 7
STRING = 8
ARRAY = 9
UINT64 = 10
INT64 = 11
FLOAT64 = 12
_simple_value_packing = {
GGUFValueType.UINT8: "<B",
GGUFValueType.INT8: "<b",
GGUFValueType.UINT16: "<H",
GGUFValueType.INT16: "<h",
GGUFValueType.UINT32: "<I",
GGUFValueType.INT32: "<i",
GGUFValueType.FLOAT32: "<f",
GGUFValueType.UINT64: "<Q",
GGUFValueType.INT64: "<q",
GGUFValueType.FLOAT64: "<d",
GGUFValueType.BOOL: "?",
}
value_type_info = {
GGUFValueType.UINT8: 1,
GGUFValueType.INT8: 1,
GGUFValueType.UINT16: 2,
GGUFValueType.INT16: 2,
GGUFValueType.UINT32: 4,
GGUFValueType.INT32: 4,
GGUFValueType.FLOAT32: 4,
GGUFValueType.UINT64: 8,
GGUFValueType.INT64: 8,
GGUFValueType.FLOAT64: 8,
GGUFValueType.BOOL: 1,
}
def get_single(value_type, file):
if value_type == GGUFValueType.STRING:
value_length = struct.unpack("<Q", file.read(8))[0]
value = file.read(value_length)
try:
value = value.decode('utf-8')
except:
pass
else:
type_str = _simple_value_packing.get(value_type)
bytes_length = value_type_info.get(value_type)
value = struct.unpack(type_str, file.read(bytes_length))[0]
return value
def load_metadata(fname):
metadata = {}
with open(fname, 'rb') as file:
GGUF_MAGIC = struct.unpack("<I", file.read(4))[0]
GGUF_VERSION = struct.unpack("<I", file.read(4))[0]
ti_data_count = struct.unpack("<Q", file.read(8))[0]
kv_data_count = struct.unpack("<Q", file.read(8))[0]
if GGUF_VERSION == 1:
raise Exception('You are using an outdated GGUF, please download a new one.')
for i in range(kv_data_count):
key_length = struct.unpack("<Q", file.read(8))[0]
key = file.read(key_length)
value_type = GGUFValueType(struct.unpack("<I", file.read(4))[0])
if value_type == GGUFValueType.ARRAY:
ltype = GGUFValueType(struct.unpack("<I", file.read(4))[0])
length = struct.unpack("<Q", file.read(8))[0]
arr = [get_single(ltype, file) for _ in range(length)]
metadata[key.decode()] = arr
else:
value = get_single(value_type, file)
metadata[key.decode()] = value
return metadata
|