Quantized

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

Bases: ffp.io.Chunk, ffp.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: ffp.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

The storage shape

Returns

(rows, cols) – Tuple with storage dimensions

Return type

Tuple[int, int]

embedding(key, out: numpy.ndarray = None)[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) → ffp.storage.quantized.QuantizedArray[source]

Load Storage from the given finalfusion file.

Parameters

  • file (BinaryIO) – File at the beginning of a finalfusion storage

  • mmap (bool) – Toggles memory mapping the buffer.

Returns

storage – The storage from the file.

Return type

Storage

static read_chunk(file: BinaryIO) → ffp.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_storage(file: BinaryIO) → ffp.storage.quantized.QuantizedArray[source]

Memory map the storage.

Parallel method to ffp.io.Chunk.read_chunk(). Instead of storing the Storage in-memory, it memory maps the embeddings.

Parameters

file (BinaryIO) – File at the beginning of a finalfusion storage

Returns

storage – The memory mapped storage.

Return type

Storage

write_chunk(file: BinaryIO)[source]

Write the Chunk to a file.

Parameters

file (BinaryIO) – Output file for the Chunk

static chunk_identifier() → ffp.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 (str, bytes, int, PathLike) – Output file

Raises

TypeError – If the Chunk is a Header.

class ffp.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.

ffp.storage.quantized.load_quantized_array(file: Union[str, bytes, int, os.PathLike], mmap: bool = False) → ffp.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.