From 4270b7190c99d08749b991085d47fe739537303a Mon Sep 17 00:00:00 2001 From: Siddhant Chaudhary Date: Thu, 25 Jul 2024 17:03:52 +0530 Subject: [PATCH] Qualifying all `JLD2.save` and `JLD2.load` functions explicitly. --- src/indexing/codecs/residual.jl | 2 +- src/indexing/collection_indexer.jl | 6 +++--- src/indexing/index_saver.jl | 18 +++++++++--------- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/src/indexing/codecs/residual.jl b/src/indexing/codecs/residual.jl index d73cc24..aa73b5a 100644 --- a/src/indexing/codecs/residual.jl +++ b/src/indexing/codecs/residual.jl @@ -147,6 +147,6 @@ A vector of codes for the specified chunk. """ function load_codes(codec::ResidualCodec, chunk_idx::Int) codes_path = joinpath(codec.config.indexing_settings.index_path, "$(chunk_idx).codes.jld2") - codes = load(codes_path, "codes") + codes = JLD2.load(codes_path, "codes") codes end diff --git a/src/indexing/collection_indexer.jl b/src/indexing/collection_indexer.jl index 3ad980a..d09de19 100644 --- a/src/indexing/collection_indexer.jl +++ b/src/indexing/collection_indexer.jl @@ -104,7 +104,7 @@ function _sample_embeddings(indexer::CollectionIndexer, sampled_pids::Set{Int}) sample_path = joinpath(indexer.config.indexing_settings.index_path, "sample.jld2") @info "avg_doclen_est = $(indexer.avg_doclen_est) \t length(local_sample) = $(length(local_sample))" @info "Saving sampled embeddings to $(sample_path)." - save(sample_path, Dict("local_sample_embs" => local_sample_embs)) + JLD2.save(sample_path, Dict("local_sample_embs" => local_sample_embs)) indexer.avg_doclen_est end @@ -182,7 +182,7 @@ The tuple `sample, sample_heldout`. function _concatenate_and_split_sample(indexer::CollectionIndexer) # load the sample embeddings sample_path = joinpath(indexer.config.indexing_settings.index_path, "sample.jld2") - sample = load(sample_path, "local_sample_embs") + sample = JLD2.load(sample_path, "local_sample_embs") @debug "Original sample shape: $(size(sample))" # randomly shuffle embeddings @@ -357,7 +357,7 @@ function _build_ivf(indexer::CollectionIndexer) @info "Saving the IVF." ivf_path = joinpath(indexer.config.indexing_settings.index_path, "ivf.jld2") - save(ivf_path, Dict( + JLD2.save(ivf_path, Dict( "ivf" => ivf, "ivf_lengths" => ivf_lengths )) diff --git a/src/indexing/index_saver.jl b/src/indexing/index_saver.jl index d994839..37af938 100644 --- a/src/indexing/index_saver.jl +++ b/src/indexing/index_saver.jl @@ -28,9 +28,9 @@ The path of of the codec is inferred from the config stored in `saver`. """ function load_codec!(saver::IndexSaver) index_path = saver.config.indexing_settings.index_path - centroids = load(joinpath(index_path, "centroids.jld2"), "centroids") - avg_residual = load(joinpath(index_path, "avg_residual.jld2"), "avg_residual") - buckets = load(joinpath(index_path, "buckets.jld2")) + centroids = JLD2.load(joinpath(index_path, "centroids.jld2"), "centroids") + avg_residual = JLD2.load(joinpath(index_path, "avg_residual.jld2"), "avg_residual") + buckets = JLD2.load(joinpath(index_path, "buckets.jld2")) saver.codec = ResidualCodec(saver.config, centroids, avg_residual, buckets["bucket_cutoffs"], buckets["bucket_weights"]) end @@ -57,9 +57,9 @@ function save_codec(saver::IndexSaver) buckets_path = joinpath(index_path, "buckets.jld2") @info "Saving codec to $(centroids_path), $(avg_residual_path) and $(buckets_path)" - save(centroids_path, Dict("centroids" => saver.codec.centroids)) - save(avg_residual_path, Dict("avg_residual" => saver.codec.avg_residual)) - save( + JLD2.save(centroids_path, Dict("centroids" => saver.codec.centroids)) + JLD2.save(avg_residual_path, Dict("avg_residual" => saver.codec.avg_residual)) + JLD2.save( buckets_path, Dict( "bucket_cutoffs" => saver.codec.bucket_cutoffs, @@ -92,13 +92,13 @@ function save_chunk(saver::IndexSaver, chunk_idx::Int, offset::Int, embs::Matrix codes_path = "$(path_prefix).codes.jld2" residuals_path = "$(path_prefix).residuals.jld2" @info "Saving compressed codes to $(codes_path) and residuals to $(residuals_path)" - save(codes_path, Dict("codes" => codes)) - save(residuals_path, Dict("residuals" => residuals)) + JLD2.save(codes_path, Dict("codes" => codes)) + JLD2.save(residuals_path, Dict("residuals" => residuals)) # saving doclens doclens_path = joinpath(saver.config.indexing_settings.index_path, "doclens.$(chunk_idx).jld2") @info "Saving doclens to $(doclens_path)" - save(doclens_path, Dict("doclens" => doclens)) + JLD2.save(doclens_path, Dict("doclens" => doclens)) # the metadata metadata_path = joinpath(saver.config.indexing_settings.index_path, "$(chunk_idx).metadata.json")