forestcalled/text-generation-webui

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