@@ -13,9 +13,11 @@ def ivecs_read(filename: str | Path) -> np.ndarray:
1313 d = a [0 ]
1414 return a .reshape (- 1 , d + 1 )[:, 1 :].copy ()
1515
16+
1617def fvecs_read (filename : str | Path ) -> np .ndarray :
1718 return ivecs_read (filename ).view (np .float32 )
1819
20+
1921DATASET_METADATA : Dict [str , Dict [str , Any ]] = {
2022 "sift1m" : {
2123 "name" : "SIFT1M" ,
@@ -84,9 +86,11 @@ def fvecs_read(filename: str | Path) -> np.ndarray:
8486 },
8587}
8688
89+
8790def resolve_dataset_key (key : str ) -> str :
8891 return key .lower ()
8992
93+
9094def get_default_cache_dir () -> Path :
9195 """Returns the default dataset cache directory (~/.cache/deg_datasets or DEG_CACHE_DIR)."""
9296 env_dir = os .environ .get ("DEG_CACHE_DIR" )
@@ -97,6 +101,7 @@ def get_default_cache_dir() -> Path:
97101 path .mkdir (parents = True , exist_ok = True )
98102 return path
99103
104+
100105def download_file (url : str , dest_path : Path ):
101106 """Downloads a file from url to dest_path with progress indication."""
102107 print (f"Downloading { url } to { dest_path } ..." )
@@ -127,6 +132,7 @@ def download_file(url: str, dest_path: Path):
127132 temp_path .rename (dest_path )
128133 print (f"Download complete: { dest_path } " )
129134
135+
130136def ivecs_write (filename : str | Path , data : np .ndarray ):
131137 """Writes a 2D int32 numpy array to ivecs format."""
132138 filename = Path (filename )
@@ -137,6 +143,7 @@ def ivecs_write(filename: str | Path, data: np.ndarray):
137143 with_d = np .hstack ([d_col , data .astype (np .int32 )])
138144 with_d .tofile (filename )
139145
146+
140147def fvecs_write (filename : str | Path , data : np .ndarray ):
141148 """Writes a 2D float32 numpy array to fvecs format."""
142149 filename = Path (filename )
@@ -148,6 +155,7 @@ def fvecs_write(filename: str | Path, data: np.ndarray):
148155 with_d = np .hstack ([d_col , data .astype (np .float32 )])
149156 with_d .tofile (filename )
150157
158+
151159def cleanup_legacy_gt (folder : Path ):
152160 """Deletes legacy/broken groundtruth files extracted from old dataset archives."""
153161 legacy_files = [
@@ -172,15 +180,16 @@ def cleanup_legacy_gt(folder: Path):
172180 except Exception as e :
173181 print (f"Warning: Failed to remove legacy GT file { p } : { e } " )
174182
183+
175184def ensure_dataset (dataset_key : str , cache_dir : Path ) -> Path :
176185 """Ensures the dataset is downloaded and extracted in cache_dir, returning folder path."""
177186 key = resolve_dataset_key (dataset_key )
178187 if key not in DATASET_METADATA :
179188 raise ValueError (f"Unknown dataset '{ dataset_key } '. Choose from: { list (DATASET_METADATA .keys ())} " )
180-
189+
181190 meta = DATASET_METADATA [key ]
182191 archive_path = cache_dir / meta ["archive" ]
183-
192+
184193 # 1. Check if direct folder (e.g., D:\Data\DEG\sift1m or D:\Data\DEG\sift) exists
185194 extracted_folder = cache_dir / meta ["folder" ]
186195 if not extracted_folder .is_dir ():
@@ -198,7 +207,7 @@ def ensure_dataset(dataset_key: str, cache_dir: Path) -> Path:
198207 archive_path = archive_matches [0 ]
199208 else :
200209 download_file (meta ["url" ], archive_path )
201-
210+
202211 print (f"Extracting { archive_path } into { cache_dir } ..." )
203212 with tarfile .open (archive_path , "r:gz" ) as tar :
204213 tar .extractall (path = cache_dir )
@@ -215,50 +224,63 @@ def ensure_dataset(dataset_key: str, cache_dir: Path) -> Path:
215224 return extracted_folder
216225
217226
218-
219- def compute_and_save_anns_gt (base_vecs : np .ndarray , query_vecs : np .ndarray , float_space : FloatSpace , k : int , out_path : Path ):
227+ def compute_and_save_anns_gt (
228+ base_vecs : np .ndarray , query_vecs : np .ndarray , float_space : FloatSpace , k : int , out_path : Path
229+ ):
220230 """Computes ANNS ground truth against the full base dataset."""
221231 n_queries = len (query_vecs )
222232 gt_matrix = np .zeros ((n_queries , k ), dtype = np .int32 )
223233 print (f"Generating ANNS Ground Truth (Full Dataset, top-{ k } )..." )
224234 start_t = time .perf_counter ()
225-
235+
226236 for q in range (n_queries ):
227237 dists = float_space .compute_distances (query_vecs [q ], base_vecs )
228238 top_k = np .argpartition (dists , k )[:k ]
229239 top_k = top_k [np .argsort (dists [top_k ])]
230240 gt_matrix [q , :] = top_k
231241 if (q + 1 ) % 100 == 0 or (q + 1 ) == n_queries :
232242 elapsed_ms = int ((time .perf_counter () - start_t ) * 1000 )
233- print (f"\r Ground truth progress: { q + 1 } /{ n_queries } ({ 100.0 * (q + 1 ) / n_queries :.1f} %) after { elapsed_ms } ms" , end = "" , flush = True )
243+ print (
244+ f"\r Ground truth progress: { q + 1 } /{ n_queries } ({ 100.0 * (q + 1 ) / n_queries :.1f} %) after { elapsed_ms } ms" ,
245+ end = "" ,
246+ flush = True ,
247+ )
234248 print ()
235249 ivecs_write (out_path , gt_matrix )
236250 print (f"ANNS Ground Truth generated and saved to { out_path } " )
237251
238- def compute_and_save_anns_gt_half (base_vecs : np .ndarray , query_vecs : np .ndarray , float_space : FloatSpace , k : int , out_path : Path ):
252+
253+ def compute_and_save_anns_gt_half (
254+ base_vecs : np .ndarray , query_vecs : np .ndarray , float_space : FloatSpace , k : int , out_path : Path
255+ ):
239256 """Computes ANNS ground truth against only the first half of the base dataset.
240257
241258 Mirrors the C++ generate_anns_groundtruth_files() with include_half=true,
242259 used for dynamic benchmarks where only half the base vectors are active.
243260 """
244- half_base = base_vecs [:len (base_vecs ) // 2 ]
261+ half_base = base_vecs [: len (base_vecs ) // 2 ]
245262 n_queries = len (query_vecs )
246263 gt_matrix = np .zeros ((n_queries , k ), dtype = np .int32 )
247264 print (f"Generating ANNS Ground Truth (Half Dataset = { len (half_base )} vectors, top-{ k } )..." )
248265 start_t = time .perf_counter ()
249-
266+
250267 for q in range (n_queries ):
251268 dists = float_space .compute_distances (query_vecs [q ], half_base )
252269 top_k = np .argpartition (dists , k )[:k ]
253270 top_k = top_k [np .argsort (dists [top_k ])]
254271 gt_matrix [q , :] = top_k
255272 if (q + 1 ) % 100 == 0 or (q + 1 ) == n_queries :
256273 elapsed_ms = int ((time .perf_counter () - start_t ) * 1000 )
257- print (f"\r Half GT progress: { q + 1 } /{ n_queries } ({ 100.0 * (q + 1 ) / n_queries :.1f} %) after { elapsed_ms } ms" , end = "" , flush = True )
274+ print (
275+ f"\r Half GT progress: { q + 1 } /{ n_queries } ({ 100.0 * (q + 1 ) / n_queries :.1f} %) after { elapsed_ms } ms" ,
276+ end = "" ,
277+ flush = True ,
278+ )
258279 print ()
259280 ivecs_write (out_path , gt_matrix )
260281 print (f"ANNS Half Ground Truth generated and saved to { out_path } " )
261282
283+
262284def load_dataset_for_dynamic (
263285 dataset_key : str , cache_dir : Path
264286) -> Tuple [np .ndarray , np .ndarray , np .ndarray , np .ndarray , Dict [str , Any ]]:
0 commit comments