Skip to content

Commit f119cc8

Browse files
authored
Refactor gguf scripts to improve metadata handling
Added contents method to ReaderField class Added endianess property to GGUFReader class
1 parent fe163d5 commit f119cc8

File tree

1 file changed

+54
-1
lines changed

1 file changed

+54
-1
lines changed

gguf-py/gguf/gguf_reader.py

Lines changed: 54 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
import logging
88
import os
9+
import sys
910
from collections import OrderedDict
1011
from typing import Any, Literal, NamedTuple, TypeVar, Union
1112

@@ -15,7 +16,6 @@
1516
from .quants import quant_shape_to_byte_shape
1617

1718
if __name__ == "__main__":
18-
import sys
1919
from pathlib import Path
2020

2121
# Allow running file in package as a script.
@@ -28,6 +28,7 @@
2828
GGUF_VERSION,
2929
GGMLQuantizationType,
3030
GGUFValueType,
31+
GGUFEndian,
3132
)
3233

3334
logger = logging.getLogger(__name__)
@@ -53,6 +54,48 @@ class ReaderField(NamedTuple):
5354

5455
types: list[GGUFValueType] = []
5556

57+
def contents(self, index_or_slice: int | slice = slice(None)) -> Any:
58+
if self.types:
59+
to_string = lambda x: str(x.tobytes(), encoding='utf-8')
60+
main_type = self.types[0]
61+
62+
if main_type == GGUFValueType.ARRAY:
63+
sub_type = self.types[-1]
64+
65+
if sub_type == GGUFValueType.STRING:
66+
indices = self.data[index_or_slice]
67+
68+
if isinstance(index_or_slice, int):
69+
return to_string(self.parts[indices])
70+
else:
71+
return [to_string(self.parts[idx]) for idx in indices]
72+
else:
73+
# FIXME: When/if _get_field_parts() support multi-dimensional arrays, this must do so too
74+
75+
# Check if it's unsafe to perform slice optimization on data
76+
# if any(True for idx in self.data if len(self.parts[idx]) != 1):
77+
# optim_slice = slice(None)
78+
# else:
79+
# optim_slice = index_or_slice
80+
# index_or_slice = slice(None)
81+
82+
# if isinstance(optim_slice, int):
83+
# return self.parts[self.data[optim_slice]].tolist()[0]
84+
# else:
85+
# return [pv for idx in self.data[optim_slice] for pv in self.parts[idx].tolist()][index_or_slice]
86+
87+
if isinstance(index_or_slice, int):
88+
return self.parts[self.data[index_or_slice]].tolist()[0]
89+
else:
90+
return [pv for idx in self.data[index_or_slice] for pv in self.parts[idx].tolist()]
91+
92+
if main_type == GGUFValueType.STRING:
93+
return to_string(self.parts[-1])
94+
else:
95+
return self.parts[-1].tolist()[0]
96+
97+
return None
98+
5699

57100
class ReaderTensor(NamedTuple):
58101
name: str
@@ -105,6 +148,15 @@ def __init__(self, path: os.PathLike[str] | str, mode: Literal['r', 'r+', 'c'] =
105148
version = temp_version[0]
106149
if version not in READER_SUPPORTED_VERSIONS:
107150
raise ValueError(f'Sorry, file appears to be version {version} which we cannot handle')
151+
if sys.byteorder == "little":
152+
# Host is little endian
153+
host_endian = GGUFEndian.LITTLE
154+
swapped_endian = GGUFEndian.BIG
155+
else:
156+
# Sorry PDP or other weird systems that don't use BE or LE.
157+
host_endian = GGUFEndian.BIG
158+
swapped_endian = GGUFEndian.LITTLE
159+
self.endianess = swapped_endian if self.byte_order == "S" else host_endian
108160
self.fields: OrderedDict[str, ReaderField] = OrderedDict()
109161
self.tensors: list[ReaderTensor] = []
110162
offs += self._push_field(ReaderField(offs, 'GGUF.version', [temp_version], [0], [GGUFValueType.UINT32]))
@@ -190,6 +242,7 @@ def _get_field_parts(
190242
offs += int(alen.nbytes)
191243
aparts: list[npt.NDArray[Any]] = [raw_itype, alen]
192244
data_idxs: list[int] = []
245+
# FIXME: Handle multi-dimensional arrays properly instead of flattening
193246
for idx in range(alen[0]):
194247
curr_size, curr_parts, curr_idxs, curr_types = self._get_field_parts(offs, raw_itype[0])
195248
if idx == 0:

0 commit comments

Comments
 (0)