def _read_array_header(file: BinaryIO) -> Tuple[int, int]:
"""
Helper method to read the header of an NdArray chunk.
The method reads the shape tuple, verifies the TypeId and seeks the file to the start
of the array. The shape tuple is returned.
Parameters
----------
file : BinaryIO
finalfusion file with a storage at the start of a NdArray chunk.
Returns
-------
shape : Tuple[int, int]
Shape of the storage.
Raises
------
FinalfusionFormatError
If the TypeId does not match TypeId.f32
"""
rows, cols = _read_required_binary(file, "<QI")
type_id = TypeId(_read_required_binary(file, "<I")[0])
if TypeId.f32 != type_id:
raise FinalfusionFormatError(
f"Invalid Type, expected {TypeId.f32}, got {type_id}")
file.seek(_pad_float32(file.tell()), 1)
return rows, cols
def write_chunk(self, file: BinaryIO):
_write_binary(file, "<I", int(self.chunk_identifier()))
padding = _pad_float32(file.tell())
chunk_len = struct.calcsize("QI") + padding + struct.calcsize(
f"<{self.size}f")
_write_binary(file, f"<QQI{padding}x", chunk_len, self.size,
int(TypeId.f32))
_serialize_array_as_le(file, self)
def write_chunk(self, file: BinaryIO):
_write_binary(file, "<I", int(self.chunk_identifier()))
padding = _pad_float32(file.tell())
chunk_len = struct.calcsize("<QII") + padding + struct.calcsize(
f'<{self.size}f')
# pylint: disable=unpacking-non-sequence
rows, cols = self.shape
_write_binary(file, "<QQII", chunk_len, rows, cols, int(TypeId.f32))
_write_binary(file, f"{padding}x")
_serialize_array_as_le(file, self)
def read_chunk(file: BinaryIO) -> 'Norms':
n_norms, dtype = _read_required_binary(file, "<QI")
type_id = TypeId(dtype)
if TypeId.f32 != type_id:
raise FinalfusionFormatError(
f"Invalid Type, expected {TypeId.f32}, got {str(type_id)}")
padding = _pad_float32(file.tell())
file.seek(padding, 1)
array = np.fromfile(file=file, count=n_norms, dtype=np.float32)
if sys.byteorder == "big":
array.byteswap(inplace=True)
return Norms(array)
def _read_quantized_header(
file: BinaryIO
) -> Tuple[PQ, Tuple[int, int], Optional[np.ndarray]]:
"""
Helper method to read the header of a quantized array chunk.
Returns a tuple containing PQ, quantized_shape and optional norms.
"""
projection = _read_required_binary(file, '<I')[0] != 0
read_norms = _read_required_binary(file, '<I')[0] != 0
quantized_len = _read_required_binary(file, '<I')[0]
reconstructed_len = _read_required_binary(file, '<I')[0]
n_centroids = _read_required_binary(file, '<I')[0]
n_embeddings = _read_required_binary(file, '<Q')[0]
assert reconstructed_len % quantized_len == 0
type_id = _read_required_binary(file, '<I')[0]
if int(TypeId.u8) != type_id:
raise FinalfusionFormatError(
f"Invalid Type, expected {str(TypeId.u8)}, got {type_id}")
type_id = _read_required_binary(file, '<I')[0]
if int(TypeId.f32) != type_id:
raise FinalfusionFormatError(
f"Invalid Type, expected {str(TypeId.f32)}, got {type_id}")
file.seek(_pad_float32(file.tell()), 1)
if projection:
projection = _read_array_as_native(file, np.float32,
reconstructed_len**2)
projection_shape = (reconstructed_len, reconstructed_len)
projection = projection.reshape(projection_shape)
else:
projection = None
quantizer_shape = (quantized_len, n_centroids,
reconstructed_len // quantized_len)
quantizers_size = quantized_len * n_centroids * (reconstructed_len //
quantized_len)
quantizers = _read_array_as_native(file, np.float32, quantizers_size)
quantizers = quantizers.reshape(quantizer_shape)
if read_norms:
norms = _read_array_as_native(file, np.float32, n_embeddings)
else:
norms = None
quantizer = PQ(quantizers, projection)
return quantizer, (n_embeddings, quantized_len), norms
def write_chunk(self, file: BinaryIO):
_write_binary(file, "<I", int(self.chunk_identifier()))
padding = _pad_float32(file.tell())
chunk_len = struct.calcsize("<IIIIIQII") + padding
proj = self._quantizer.projection is not None
if proj:
chunk_len += struct.calcsize(
f"<{pow(self._quantizer.reconstructed_len, 2)}f")
chunk_len += struct.calcsize(f"<{self._quantizer.subquantizers.size}f")
norms = self._norms is not None
if self._norms is not None:
chunk_len += struct.calcsize(f"<{self._norms.size}f")
chunk_len += self._quantized_embeddings.size
chunk_header = (chunk_len, proj, norms, self.quantized_len,
self.shape[1], self.quantizer.n_centroids,
self.shape[0], int(TypeId.u8), int(TypeId.f32))
_write_binary(file, "<QIIIIIQII", *chunk_header)
file.write(struct.pack(f"{padding}x"))
if proj:
_serialize_array_as_le(file, self.quantizer.projection)
_serialize_array_as_le(file, self.quantizer.subquantizers)
if norms:
_serialize_array_as_le(file, self._norms)
self._quantized_embeddings.tofile(file)
