working baseline for both task 1 and task 2

Author: maumueller

datasets.py

@@ -8,14 +8,13 @@ def download(src, dst):
         print('downloading %s -> %s...' % (src, dst))
         urlretrieve(src, dst)
 
-def get_fn(kind):
-    version = "ccnews-small"
-    return os.path.join("data", kind, f"{version}.h5"), os.path.join('data', kind, 'gt', f'{version}.h5')
+def get_fn(dataset, task):
+    return os.path.join("data", dataset, task, f"{dataset}.h5"), os.path.join('data', dataset, task, 'gt', f'gt_{dataset}.h5')
 
-def prepare(kind):
-    url = DATASETS['ccnews-small'][kind]['url']
-    gt_url = DATASETS['ccnews-small'][kind]['gt_url']
-    fn, gt_fn = get_fn(kind)
+def prepare(dataset, task):
+    url = DATASETS[dataset][task]['url']
+    gt_url = DATASETS[dataset][task]['gt_url']
+    fn, gt_fn = get_fn(dataset, task)
 
     download(url, fn)
     download(gt_url, gt_fn)
@@ -28,13 +27,15 @@ def prepare(kind):
             'data': lambda x: x['train'],
             'gt_url': 'https://huggingface.co/datasets/sadit/SISAP2025/resolve/main/benchmark-dev-ccnews-fp16.h5?download=true',
             'gt_I': lambda x: x['itest']['knns'],
+            'k': 30,
         },
         'task2': {
             'url': 'https://huggingface.co/datasets/sadit/SISAP2025/resolve/main/benchmark-dev-ccnews-fp16.h5?download=true',
             'queries': lambda x: x['train'],
             'data': lambda x: x['train'],
             'gt_url': 'https://huggingface.co/datasets/sadit/SISAP2025/resolve/main/allknn-benchmark-dev-ccnews.h5?download=true',
-            'gt_I': lambda x: x['knns']
+            'gt_I': lambda x: x['knns'],
+            'k': 15,
         }
     }
 }

eval.py

@@ -5,24 +5,17 @@
 import csv
 import glob
 from pathlib import Path
-from datasets import DATASETS
-
-def get_groundtruth(size="100K", private=False):
-    # test
-    gt_f = h5py.File(out_fn, "r")
-    true_I = np.array(gt_f['knns'])
-    gt_f.close()
-    return true_I
+from datasets import DATASETS, get_fn, prepare
 
 def get_all_results(dirname):
-    mask = [dirname + "/**/*.h5"]
-    print("search for results matching:")
+    mask = [dirname + "/**/*.h5", dirname + "/**/*/*.h5"]
+    print("Searching for results matching:")
     print("\n".join(mask))
     for m in mask:
         for fn in glob.iglob(m):
             print(fn)
             f = h5py.File(fn, "r")
-            if "knns" not in f or not ("data" in f or "data" in f.attrs):
+            if "knns" not in f or not ("dataset" in f or "dataset" in f.attrs):
                 print("Ignoring " + fn)
                 f.close()
                 continue
@@ -45,36 +38,38 @@ def get_recall(I, gt, k):
     parser.add_argument(
         "--results",
         help='directory in which results are stored',
-        default="result"
+        default="results"
     )
     parser.add_argument(
         '--private',
         help="private queries held out for evaluation",
         action='store_true',
         default=False
     )
-    parser.add_argument(
-        '--dataset',
-        choices = ['ccnews-small'],
-        default='ccnews-small',
-    )
 
     parser.add_argument("csvfile")
     args = parser.parse_args()
     true_I_cache = {}
 
 
-    columns = ["data", "kind", "algo", "buildtime", "querytime", "params", "recall"]
+    columns = ["dataset", "task", "algo", "buildtime", "querytime", "params", "recall"]
 
     with open(args.csvfile, 'w', newline='') as csvfile:
         writer = csv.DictWriter(csvfile, fieldnames=columns)
         writer.writeheader()
         for res in get_all_results(args.results):
-            data = res.attrs["data"]
+            dataset = res.attrs["dataset"]
+            task = res.attrs["task"]
+            assert dataset in DATASETS and task in DATASETS[dataset]
+            prepare(dataset, task)
             d = dict(res.attrs)
-            print(d)
-            gt_I = np.array(DATASETS['ccnews-small'][data]['gt_I'](res))
-            recall = get_recall(np.array(res["knns"]), gt_I, 10)
+            # print(d)
+            _, gt_f = get_fn(dataset, task)
+            print(f"Using groundtruth in {gt_f}")
+            f = h5py.File(gt_f)
+            gt_I = np.array(DATASETS[dataset][task]['gt_I'](f))
+            f.close()
+            recall = get_recall(np.array(res["knns"]), gt_I, DATASETS[dataset][task]['k'])
             d['recall'] = recall
-            print(d["data"], d["algo"], d["params"], "=>", recall)
+            print(d["dataset"], d["task"], d["algo"], d["params"], "=>", recall)
             writer.writerow(d)

search.py

@@ -7,37 +7,36 @@
 import time
 from datasets import DATASETS, prepare, get_fn
 
-def store_results(dst, algo, kind, D, I, buildtime, querytime, params):
+def store_results(dst, algo, dataset, task, D, I, buildtime, querytime, params):
     os.makedirs(Path(dst).parent, exist_ok=True)
     f = h5py.File(dst, 'w')
     f.attrs['algo'] = algo
-    f.attrs['data'] = kind
+    f.attrs['dataset'] = dataset
+    f.attrs['task'] = task
     f.attrs['buildtime'] = buildtime
     f.attrs['querytime'] = querytime
     f.attrs['params'] = params
     f.create_dataset('knns', I.shape, dtype=I.dtype)[:] = I
     f.create_dataset('dists', D.shape, dtype=D.dtype)[:] = D
     f.close()
 
-def run(kind, params):
-    print("Running", kind)
+def run(dataset, task, k):
+    print(f'Running {task} on {dataset}')
 
-    prepare(kind)
+    prepare(dataset, task)
 
-    fn, _ = get_fn(kind)
+    fn, _ = get_fn(dataset, task)
     f = h5py.File(fn)
-    data = np.array(DATASETS['ccnews-small'][kind]['data'](f))
-    queries = np.array(DATASETS['ccnews-small'][kind]['queries'](f))
+    data = np.array(DATASETS[dataset][task]['data'](f))
+    queries = np.array(DATASETS[dataset][task]['queries'](f))
     f.close()
 
     n, d = data.shape
-    k = params['k']
+    if task == 'task2':
+        k = k + 1 # need to search for one more NN since we cannot remove self-loop
 
     nlist = 1024 # number of clusters/centroids to build the IVF from
-    if kind == 'task1':
-        index_identifier = f"IVF{nlist},SQfp16"
-    elif kind == 'task2':
-        index_identifier = f"IVF{nlist},PQ{d//2}x4fs"
+    index_identifier = f"IVF{nlist},SQfp16"
 
     index = faiss.index_factory(d, index_identifier)
 
@@ -49,10 +48,6 @@ def run(kind, params):
     print(f"Done training in {elapsed_build}s.")
     assert index.is_trained
 
-    if kind == "task2":
-        index = faiss.IndexRefineFlat(index, faiss.swig_ptr(data.astype('float32')))
-        index.k_factor = 200
-
     for nprobe in [1, 2, 5, 10, 20, 50, 100]:
         print(f"Starting search on {queries.shape} with nprobe={nprobe}")
         start = time.time()
@@ -65,7 +60,7 @@ def run(kind, params):
 
         identifier = f"index=({index_identifier}),query=(nprobe={nprobe})"
 
-        store_results(os.path.join("result/", kind, f"{identifier}.h5"), "faissIVF", kind, D, I, elapsed_build, elapsed_search, identifier)
+        store_results(os.path.join("results/", dataset, task, f"{identifier}.h5"), "faissIVF", dataset, task, D, I, elapsed_build, elapsed_search, identifier)
 
 if __name__ == "__main__":
 
@@ -78,22 +73,11 @@ def run(kind, params):
 
     parser.add_argument(
         '--dataset',
-        choices=[
-            'ccnews-small',
-        ],
+        choices=DATASETS.keys(),
         default='ccnews-small'
     )
 
-    params = {
-        'task1': {
-            "k": 30,
-        },
-        'task2': {
-            "k": 15,
-        }
-    }
 
     args = parser.parse_args()
-
-    run(args.task, params[args.task])
+    run(args.dataset, args.task, DATASETS[args.dataset][args.task]['k'])