Package: chunkdot
ChunkDotMulti-threaded matrix multiplication and cosine similarity calculations for dense and sparse matrices. Appropriate for calculating the K most similar items for a large number of items by chunking the item matrix representation (embeddings) and using Numba to accelerate the calculations.Use for:dense embeddingssparse embeddingssimilarity calculation versus other embeddingsCosineSimilarityTopK scikit-learn transformerRelated blog postsCosine Similarity for 1 Trillion Pairs of VectorsBulk Similarity Calculations for Sparse EmbeddingsUsagepipinstall-UchunkdotDense embeddingsCalculate the 50 most similar and dissimilar items for 100K items.importnumpyasnpfromchunkdotimportcosine_similarity_top_kembeddings=np.random.randn(100000,256)# using all you system's memorycosine_similarity_top_k(embeddings,top_k=50)# most dissimilar items using 20GBcosine_similarity_top_k(embeddings,top_k=-50,max_memory=20E9)<100000x100000 sparse matrix of type '<class 'numpy.float64'>'
 with 5000000 stored elements in Compressed Sparse Row format># with progress barcosine_similarity_top_k(embeddings,top_k=50,show_progress=True)100%|███████████████████████████████████████████████████████████████| 129.0/129 [01:04<00:00,  1.80it/s]
<100000x100000 sparse matrix of type '<class 'numpy.float64'>'
  with 5000000 stored elements in Compressed Sparse Row format>Execution timefromtimeitimporttimeitimportnumpyasnpfromchunkdotimportcosine_similarity_top_kembeddings=np.random.randn(100000,256)timeit(lambda:cosine_similarity_top_k(embeddings,top_k=50,max_memory=20E9),number=1)58.611996899999994Sparse embeddingsCalculate the 50 most similar and dissimilar items for 100K items. Items represented by 10K dimensional vectors and an embeddings matrix of 0.005 density.fromscipyimportsparsefromchunkdotimportcosine_similarity_top_kembeddings=sparse.rand(100000,10000,density=0.005)# using all you system's memorycosine_similarity_top_k(embeddings,top_k=50)# most dissimilar items using 20GBcosine_similarity_top_k(embeddings,top_k=-50,max_memory=20E9)<100000x100000 sparse matrix of type '<class 'numpy.float64'>'
 with 5000000 stored elements in Compressed Sparse Row format>Execution timefromtimeitimporttimeitfromscipyimportsparsefromchunkdotimportcosine_similarity_top_kembeddings=sparse.rand(100000,10000,density=0.005)timeit(lambda:cosine_similarity_top_k(embeddings,top_k=50,max_memory=20E9),number=1)51.87472256699999Similarity calculation versus other embeddingsGiven 20K items, for each item, find the 50 most similar items in a collection of other 10K items.importnumpyasnpfromchunkdotimportcosine_similarity_top_kembeddings=np.random.randn(20000,256)other_embeddings=np.random.randn(10000,256)cosine_similarity_top_k(embeddings,embeddings_right=other_embeddings,top_k=10)<20000x10000 sparse matrix of type '<class 'numpy.float64'>'
 with 200000 stored elements in Compressed Sparse Row format>CosineSimilarityTopK scikit-learn transformerGiven a pandas DataFrame with 100K rows and2 numerical columns2 categorical columns with 500 categories eachuse scikit-learn transformers, the standard scaler for the numerical columns and the one-hot encoder for the categorical columns, to form an embeddings matrix of dimensions 100K x 1002 and then calculate the top 50 most similar rows per each row.importnumpyasnpimportpandasaspdn_rows=100000n_categories=500df=pd.DataFrame({"A_numeric":np.random.rand(n_rows),"B_numeric":np.random.rand(n_rows),"C_categorical":np.random.randint(n_categories,size=n_rows),"D_categorical":np.random.randint(n_categories,size=n_rows),})fromsklearn.composeimportColumnTransformerfromsklearn.pipelineimportPipelinefromsklearn.preprocessingimportOneHotEncoder,StandardScalerfromchunkdotimportCosineSimilarityTopKnumeric_features=["A_numeric","B_numeric"]numeric_transformer=Pipeline(steps=[("scaler",StandardScaler())])categorical_features=["C_categorical","D_categorical"]categorical_transformer=Pipeline(steps=[("encoder",OneHotEncoder())])preprocessor=ColumnTransformer(transformers=[("num",numeric_transformer,numeric_features),("cat",categorical_transformer,categorical_features),])cos_sim=CosineSimilarityTopK(top_k=50)pipe=Pipeline(steps=[("preprocessor",preprocessor),("cos_sim",cos_sim)])pipe.fit_transform(df)<100000x100000 sparse matrix of type '<class 'numpy.float64'>'
	with 5000000 stored elements in Compressed Sparse Row format>Execution timefromtimeitimporttimeittimeit(lambda:pipe.fit_transform(df),number=1)24.45172154181637