Package: embedding-reader
embedding_readerEmbedding reader is a module to make it easy to read efficiently a large collection of embeddings stored in any file system.400GB of embeddings read in 8min using an nvme drive400GB of embeddings read in 40min using an hdd drive400GB of embeddings read in 1.3h from aws s3Installpip install embedding_readerPython examplesCheckout these examples to call this as a lib:example.pySimple examplefromembedding_readerimportEmbeddingReaderembedding_reader=EmbeddingReader(embeddings_folder="embedding_folder",file_format="npy")print("embedding count",embedding_reader.count)print("dimension",embedding_reader.dimension)print("total size",embedding_reader.total_size)print("byte per item",embedding_reader.byte_per_item)foremb,metainembedding_reader(batch_size=10**6,start=0,end=embedding_reader.count):print(emb.shape)Laion5B exampleInlaion5Byou can find 5B ViT-L/14 image embeddings, you can read them with that code:fromembedding_readerimportEmbeddingReaderembedding_reader=EmbeddingReader(embeddings_folder="https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/img_emb/",file_format="npy")print("embedding count",embedding_reader.count)print("dimension",embedding_reader.dimension)print("total size",embedding_reader.total_size)print("byte per item",embedding_reader.byte_per_item)foremb,metainembedding_reader(batch_size=10**6,start=0,end=embedding_reader.count):print(emb.shape)It takes about 3h to read laion2B-en embeddings at 300MB/sNumpy & Parquet Metadata ExampleThe parquet_npy format supports reading from both a .npy collection and a .parquet collection that are in the same order.
Here is an example of usage:fromembedding_readerimportEmbeddingReaderembedding_reader=EmbeddingReader(embeddings_folder="embedding_folder",metadata_folder="metadata_folder",meta_columns=['image_path','caption'],file_format="parquet_npy")foremb,metainembedding_reader(batch_size=10**6,start=0,end=embedding_reader.count):print(emb.shape)print(meta["image_path"],meta["caption"])embis a numpy array like the previous examples whilemetais a pandas dataframe with the columns requested inmeta_columns.Who is using embedding reader?Some use cases of embedding reader include:building knn indices inautofaisscomputing zero shot attributes using cliprunning training or inferences of linear layer models on top of embeddingsEmbeddings are a powerful concept, they allow turning highly complex data into point in a linearly separable space.
Embeddings are also much smaller and more efficient to manipulate than usual data (images, audio, video, text, interaction items, ...)To learn more about embeddings readSemantic search blogpostFile system supportThanks tofsspec, embedding_reader supports reading and writing files inmany file systems.
To use it, simply use the prefix of your filesystem before the path. For examplehdfs://,s3://,http://, orgcs://.
Some of these file systems require installing an additional package (for example s3fs for s3, gcsfs for gcs).
See fsspec doc for all the details.APIThis module exposes one class:EmbeddingReader(folder, file_format, embedding_column="embedding", meta_columns=None, metadata_folder=None)initialize the reader by listing all files and retrieving their metadatafolderembeddings folder. Can also be a list of folders. (required)file_formatparquet, npy or parquet_npy. (required)embedding_columnembedding column in parquet. (default embedding)meta_columnsmeta columns in parquet. (default None)metadata_foldermetadata folder, used by the parquet_npy reader (default None).embeddings_folderthe embedding folder.counttotal number of embedding in this folder.dimensiondimension of one embedding.byte_per_itemsize of one embedding in bytes.total_sizesize in bytes of the collection.nb_filestotal number of embedding files in this folder.max_file_sizemax size in bytes of the embedding files of the collectioncall(batch_size, start=0, end=None, max_piece_size=None, parallel_pieces=None, show_progress=True, max_ram_usage_in_bytes=2**31)Produces an iterator that yields tuples (data, meta) with the given batch_sizebatch_sizeamount of embeddings in one batch. (required)startstart of the subset of the collection to read. (default0)endend of the subset of the collection to read. (defaultend of collection)max_piece_sizemaximum size of a piece. The default value works for most cases. Increase or decrease based on your file system performances (defaultmax(number of embedding for 50MB, batch size in MB))parallel_piecesNumber of pieces to read in parallel. Increase or decrease depending on your filesystem. (default~min(round(max_ram_usage_in_bytes/max_piece_sizebyte_per_item), 50)*)show_progressDisplay a tqdm bar with the number of pieces done. (defaultTrue)max_ram_usage_in_bytesConstraint the ram usage of embedding reader. The exact max ram usage ismin(max_ram_usage_in_bytes, size of a batch in bytes). (default 4GB)Architecture notes and benchmarkThe main architecture choice of this lib is thebuild_piecesfunction that builds decently sizes pieces of embedding files (typically 50MB) initially.
These pieces metadata can then be used to fetch in parallel these pieces, which are then used to build the embedding batches and provided to the user.
In order to reach the maximal speed, it is better to read files of equal size. The number of threads used is constrained by the maximum size of your embeddings files: the lower the size, the more threads are used (you can also set a custom number of threads, but ram consumption will be higher).In practice, it has been observed speed of up to 100MB/s when fetching embeddings from s3, 1GB/s when fetching from an nvme drive.
That means reading 400GB of embeddings in 8 minutes (400M embeddings in float16 and dimension 512)
The memory usage stays low and flat thanks to the absence of copy. Decreasing the batch size decreases the amount of memory consumed, you can also set max_ram_usage_in_bytes to have a better control on the ram usage.For developmentEither locally, or ingitpod(doexport PIP_USER=falsethere)Setup a virtualenv:python3 -m venv .env
source .env/bin/activate
pip install -e .to run tests:pip install -r requirements-test.txtthenmake lint
make testYou can usemake blackto reformat the codepython -m pytest -x -s -v tests -k "dummy"to run a specific test