Skip to content

Commit

Permalink
FIX apply patch to use requests to download dataset
Browse files Browse the repository at this point in the history
  • Loading branch information
dantegd committed Dec 17, 2024
1 parent a0762ec commit 117d98e
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 0 deletions.
27 changes: 27 additions & 0 deletions context/cuvs-bench/cuvs_bench_get_dataset.patch
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
diff --git a/python/cuvs_bench/cuvs_bench/get_dataset/__main__.py b/python/cuvs_bench/cuvs_bench/get_dataset/__main__.py
index a6b154ef..b023fcbd 100644
--- a/python/cuvs_bench/cuvs_bench/get_dataset/__main__.py
+++ b/python/cuvs_bench/cuvs_bench/get_dataset/__main__.py
@@ -17,7 +17,7 @@ import argparse
import os
import subprocess
import sys
-from urllib.request import urlretrieve
+import requests


def get_dataset_path(name, ann_bench_data_path):
@@ -29,7 +29,12 @@ def get_dataset_path(name, ann_bench_data_path):
def download_dataset(url, path):
if not os.path.exists(path):
print(f"downloading {url} -> {path}...")
- urlretrieve(url, path)
+ with requests.get(url, stream=True) as r:
+ r.raise_for_status()
+ with open(path, "wb") as f:
+ for chunk in r.iter_content(chunk_size=8192):
+ if chunk:
+ f.write(chunk)


def convert_hdf5_to_fbin(path, normalize):
7 changes: 7 additions & 0 deletions context/cuvs-bench/get_datasets.sh
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,13 @@

set -eo pipefail

# find cuVS in the environment
PACKAGE_FILE_PATH=$(python -c "import cuvs-bench; print(package_name.__file__)")
PACKAGE_DIR=$(dirname "$PACKAGE_FILE_PATH")

# Apply the patch
patch "$PACKAGE_DIR/get_dataset/__main__.py" < cuvs_bench_get_dataset.patch

python -m cuvs_bench.get_dataset --dataset deep-image-96-angular --normalize --dataset-path /home/rapids/preloaded_datasets
python -m cuvs_bench.get_dataset --dataset fashion-mnist-784-euclidean --dataset-path /home/rapids/preloaded_datasets
python -m cuvs_bench.get_dataset --dataset glove-50-angular --normalize --dataset-path /home/rapids/preloaded_datasets
Expand Down

0 comments on commit 117d98e

Please sign in to comment.