diff --git a/edge0-infer/tools/assets/format_gguf.py b/edge0-infer/tools/assets/format_gguf.py index 51517e2e..7eb2b1af 100644 --- a/edge0-infer/tools/assets/format_gguf.py +++ b/edge0-infer/tools/assets/format_gguf.py @@ -79,6 +79,16 @@ def _value(f: BinaryIO, type: int) -> Value: n = _unpack(f, 'Q') if et == 9 or n > 10_000_000: raise ValueError('unsupported or oversized GGUF array') + if et in SCALARS: + # Numeric arrays can contain a vocabulary-sized score/type table. + # Read bounded blocks instead of dispatching once per scalar. + values = [] + fmt = SCALARS[et] + width = struct.calcsize('<' + fmt) + for start in range(0, n, 8192): + count = min(8192, n - start) + values.extend(struct.unpack('<' + str(count) + fmt, _read(f, count * width))) + return Value(type, values, et) return Value(type, [_value(f, et).data for _ in range(n)], et) raise ValueError(f'unsupported metadata type {type}') diff --git a/edge0-infer/tools/assets/tests/test_gguf_lite.py b/edge0-infer/tools/assets/tests/test_gguf_lite.py index 7a45a016..fe865279 100644 --- a/edge0-infer/tools/assets/tests/test_gguf_lite.py +++ b/edge0-infer/tools/assets/tests/test_gguf_lite.py @@ -58,3 +58,41 @@ def test_invalid_payloads(tmp_path): pack_q4_1(np.full(32, np.nan)) with pytest.raises(ValueError): pack_q4_1(np.zeros(31)) + + +def test_numeric_metadata_arrays_use_bounded_block_reads(tmp_path, monkeypatch): + import format_gguf + + path = tmp_path / 'scores.gguf' + values = list(range(25000)) + write(path, {'scores': Value(9, values, 4)}, []) + original = format_gguf._read + requests = [] + + def counted(stream, size): + requests.append(size) + return original(stream, size) + + monkeypatch.setattr(format_gguf, '_read', counted) + reader = Reader(path) + assert reader.metadata['scores'].data == values + print(f'read_calls={len(requests)}, max_read={max(requests)}') + assert len(requests) < 32 + assert max(requests) <= 65536 + + +@pytest.mark.parametrize('type,value', [(0, 255), (1, -128), (2, 65535), (3, -32768), + (4, 2**32-1), (5, -2**31), (6, 0.5), (7, True), + (10, 2**64-1), (11, -2**63), (12, 0.125)]) +def test_numeric_array_block_boundaries(tmp_path, type, value): + path = tmp_path / 'array.gguf' + values = [value] * 8193 + write(path, {'values': Value(9, values, type), 'empty': Value(9, [], type)}, []) + reader = Reader(path) + assert reader.metadata['values'].data == values + assert reader.metadata['empty'].data == [] + # Cut inside the array, not in the writer's final padding. + size = struct.calcsize('<' + __import__('format_gguf').SCALARS[type]) + path.write_bytes(path.read_bytes()[:24 + 8 + len('values') + 4 + 4 + 8 + len(values)*size - 1]) + with pytest.raises(ValueError, match='truncated'): + Reader(path)