Implement QuantizedArray storage.

Implement the QuantizedArray storage type and the PQ quantizer.

Author: sebpuetz

src/finalfusion/embeddings.py

@@ -12,7 +12,7 @@
     FinalfusionFormatError, _read_required_chunk_header
 from finalfusion.metadata import Metadata
 from finalfusion.norms import Norms
-from finalfusion.storage import Storage, NdArray
+from finalfusion.storage import Storage, NdArray, QuantizedArray
 from finalfusion.vocab import Vocab, SimpleVocab, FinalfusionBucketVocab, FastTextVocab, \
     ExplicitVocab
 
@@ -39,6 +39,7 @@ class Embeddings:  # pylint: disable=too-many-instance-attributes
 
     1. :class:`~finalfusion.storage.Storage` *(required)*:
         * :class:`~finalfusion.storage.ndarray.NdArray`
+        * :class:`~finalfusion.storage.ndarray.QuantizedArray`
     2. :class:`~finalfusion.vocab.Vocab` *(required)*:
         * :class:`~finalfusion.vocab.simple_vocab.SimpleVocab`,
         * :class:`~finalfusion.vocab.subword.FinalfusionBucketVocab`
@@ -645,7 +646,9 @@ def load_finalfusion(file: Union[str, bytes, int, PathLike],
 
         chunk_id, _ = _read_required_chunk_header(inf)
         if chunk_id == ChunkIdentifier.NdArray:
-            storage = NdArray.load(inf, mmap)
+            storage = NdArray.load(inf, mmap)  # type: Storage
+        elif chunk_id == ChunkIdentifier.QuantizedArray:
+            storage = QuantizedArray.load(inf, mmap)
         else:
             raise FinalfusionFormatError(
                 f'Expected storage chunk, not {str(chunk_id)}')

src/finalfusion/io.py

@@ -324,3 +324,11 @@ def _serialize_array_as_le(file: BinaryIO, array: np.ndarray):
                 row.byteswap(inplace=False).tofile(file)
         else:
             array.byteswap(inplace=False).tofile(file)
+
+
+def _read_array_as_native(file: BinaryIO, dtype: np.dtype,
+                          count: int) -> np.array:
+    array = np.fromfile(file=file, count=count, dtype=dtype)
+    if sys.byteorder == "big":
+        array.byteswap(inplace=True)
+    return array

src/finalfusion/storage/__init__.py

@@ -7,6 +7,7 @@
 from finalfusion.io import ChunkIdentifier, find_chunk
 from finalfusion.storage.storage import Storage
 from finalfusion.storage.ndarray import NdArray, load_ndarray
+from finalfusion.storage.quantized import QuantizedArray, load_quantized_array
 
 
 def load_storage(file: Union[str, bytes, int, PathLike],
@@ -42,4 +43,8 @@ def load_storage(file: Union[str, bytes, int, PathLike],
             if mmap:
                 return NdArray.mmap_chunk(inf)
             return NdArray.read_chunk(inf)
+        if chunk == ChunkIdentifier.QuantizedArray:
+            if mmap:
+                return QuantizedArray.mmap_chunk(inf)
+            return QuantizedArray.read_chunk(inf)
         raise NotImplementedError('Storage type is not yet supported.')

src/finalfusion/storage/ndarray.py

@@ -10,7 +10,8 @@
 import numpy as np
 
 from finalfusion.io import ChunkIdentifier, TypeId, FinalfusionFormatError, find_chunk, \
-    _pad_float32, _read_required_binary, _write_binary, _serialize_array_as_le
+    _pad_float32, _read_required_binary, _write_binary, _serialize_array_as_le, \
+    _read_array_as_native
 from finalfusion.storage.storage import Storage
 
 
@@ -51,9 +52,7 @@ def chunk_identifier() -> ChunkIdentifier:
     @staticmethod
     def read_chunk(file: BinaryIO) -> 'NdArray':
         rows, cols = NdArray._read_array_header(file)
-        array = np.fromfile(file=file, count=rows * cols, dtype=np.float32)
-        if sys.byteorder == "big":
-            array.byteswap(inplace=True)
+        array = _read_array_as_native(file, np.float32, rows * cols)
         array = np.reshape(array, (rows, cols))
         return NdArray(array)