Quantized

class finalfusion.storage.quantized.QuantizedArray(pq: finalfusion.storage.quantized.PQ, quantized_embeddings: numpy.ndarray, norms: Optional[numpy.ndarray])[source]

Bases: finalfusion.storage.storage.Storage

QuantizedArray storage.

QuantizedArrays support slicing, indexing with integers, lists of integers and arbitrary dimensional integer arrays. Slicing a QuantizedArray returns a new QuantizedArray but does not copy any buffers.

QuantizedArrays offer two ways of indexing:

  1. QuantizedArray.__getitem__():

    • passing a slice returns a new view of the QuantizedArray.

    • passing an integer returns a single embedding, lists and arrays return ndims + 1 dimensional embeddings.

  2. QuantizedArray.embedding():

    • embeddings can be written to an output buffer.

    • passing a slice returns a matrix holding reconstructed embeddings.

    • otherwise, this method behaves like __getitem__()

A QuantizedArray can be treated as numpy.ndarray through numpy.asarray(). This restores the original matrix and copies into a new buffer.

Using common numpy functions on a QuantizedArray will produce a regular ndarray in the process and is therefore an expensive operation.

__init__(pq: finalfusion.storage.quantized.PQ, quantized_embeddings: numpy.ndarray, norms: Optional[numpy.ndarray])[source]

Initialize a QuantizedArray.

Parameters

  • pq (PQ) – A product quantizer

  • quantized_embeddings (numpy.ndarray) – The quantized embeddings

  • norms (numpy.ndarray, optional) – Optional norms corresponding to the quantized embeddings. Reconstructed embeddings are scaled by their norm.

property shape

Get the shape of the storage.

Returns

shape – Tuple containing (rows, columns)

Return type

Tuple[int, int]

embedding(key, out: numpy.ndarray = None) → numpy.ndarray[source]

Get embeddings.

  • if key is an integer, a single reconstructed embedding is returned.

  • if key is a list of integers or a slice, a matrix of reconstructed embeddings is returned.

  • if key is an n-dimensional array, a tensor with reconstructed embeddings is returned. This tensor has one new axis in the last dimension containing the embeddings.

If out is passed, the reconstruction is written to this buffer. out.shape needs to match the dimensions described above.

Parameters

  • key (int, list, numpy.ndarray, slice) – Key specifying which embeddings to retrieve.

  • out (numpy.ndarray) – Array to reconstruct the embeddings into.

Returns

reconstruction – The reconstructed embedding or embeddings.

Return type

numpy.ndarray

property quantized_len

Length of the quantized embeddings.

Returns

quantized_len – Length of quantized embeddings.

Return type

int

property quantizer

Get the quantizer.

Returns

pq – The Product Quantizer.

Return type

PQ

classmethod load(file: BinaryIO, mmap=False) → finalfusion.storage.quantized.QuantizedArray[source]

Load Storage from the given finalfusion file.

Parameters

  • file (IO[bytes]) – Finalfusion file with a storage chunk

  • mmap (bool)

Returns

  • storage (Storage) – The first storage in the input file

  • mmap (bool) – Toggles memory mapping the storage buffer as read-only.

Raises

ValueError – If the file did not contain a storage.

static read_chunk(file: BinaryIO) → finalfusion.storage.quantized.QuantizedArray[source]

Read the Chunk and return it.

The file must be positioned before the contents of the Chunk but after its header.

Parameters

file (BinaryIO) – a finalfusion file containing the given Chunk

Returns

chunk – The chunk read from the file.

Return type

Chunk

static mmap_chunk(file: BinaryIO) → finalfusion.storage.quantized.QuantizedArray[source]

Memory maps the storage as a read-only buffer.

Parameters

file (IO[bytes]) – Finalfusion file with a storage chunk

Returns

storage – The first storage in the input file

Return type

Storage

Raises

ValueError – If the file did not contain a storage.

write_chunk(file: BinaryIO)[source]

Write the Chunk to a file.

Parameters

file (BinaryIO) – Output file for the Chunk

static chunk_identifier() → finalfusion.io.ChunkIdentifier[source]

Get the ChunkIdentifier for this Chunk.

Returns

chunk_identifier

Return type

ChunkIdentifier

write(file: Union[str, bytes, int, os.PathLike])

Write the Chunk as a standalone finalfusion file.

Parameters

file (Union[str, bytes, int, PathLike]) – Output path

Raises

TypeError – If the Chunk is a Header.

class finalfusion.storage.quantized.PQ(quantizers: numpy.ndarray, projection: Optional[numpy.ndarray])[source]

Product Quantizer

Product Quantizers are vector quantizers which decompose high dimensional vector spaces into subspaces. Each of these subspaces is a slice of the the original vector space. Embeddings are quantized by assigning their ith slice to the closest centroid.

Product Quantizers can reconstruct vectors by concatenating the slices of the quantized vector.

__init__(quantizers: numpy.ndarray, projection: Optional[numpy.ndarray])[source]

Initializes a Product Quantizer.

Parameters

  • quantizers (np.ndarray) – 3-d ndarray with dtype uint8

  • projection (np.ndarray, optional) – Projection matrix, must be a square matrix with shape [reconstructed_len, reconstructed_len]

Raises

AssertionError – If the projection shape does not match the reconstructed_len

property n_centroids

Number of centroids per quantizer.

Returns

n_centroids – The number of centroids per quantizer.

Return type

int

property projection

Projection matrix.

Returns

projection – Projection Matrix (2-d numpy array with datatype float32) or None.

Return type

np.ndarray, optional

property reconstructed_len

Reconstructed length.

Returns

reconstructed_len – Length of the reconstructed vectors.

Return type

int

property subquantizers

Get the quantizers.

Returns a 3-d array with shape quantizers * n_centroids * reconstructed_len / quantizers

Returns

  • quantizers (np.ndarray) – 3-d np.ndarray with dtype=np.uint8

  • @return (3d tensor of quantizers)

reconstruct(quantized: numpy.ndarray, out: numpy.ndarray = None) → numpy.ndarray[source]

Reconstruct vectors.

Input

Parameters

  • quantized (np.ndarray) – Batch of quantized vectors. 2-d np.ndarray with integers required.

  • out (np.ndarray, optional) – 2-d np.ndarray to write the output into.

Returns

out – Batch of reconstructed vectors.

Return type

np.ndarray

Raises

AssertionError – If out is passed and its last dimension does not match reconstructed_len or its first n-1 dimensions do not match the first n-1 dimensions of quantized.

finalfusion.storage.quantized.load_quantized_array(file: Union[str, bytes, int, os.PathLike], mmap: bool = False) → finalfusion.storage.quantized.QuantizedArray[source]

Load a quantized array chunk from the given file.

Parameters

  • file (str, bytes, int, PathLike) – Finalfusion file with a quantized array chunk.

  • mmap (bool) – Toggles memory mapping the array buffer as read only.

Returns

storage – The QuantizedArray storage from the file.

Return type

QuantizedArray

Raises

ValueError – If the file did not contain a QuantizedArray chunk.