initial commit

neurips23/ood/vamana/Dockerfile

@@ -0,0 +1,24 @@
+FROM neurips23
+
+RUN apt update
+RUN apt install -y software-properties-common
+RUN add-apt-repository -y ppa:git-core/ppa
+RUN apt update
+RUN DEBIAN_FRONTEND=noninteractive apt install -y git make cmake g++ libaio-dev libgoogle-perftools-dev libunwind-dev clang-format libboost-dev libboost-program-options-dev libmkl-full-dev libcpprest-dev python3.10
+
+RUN echo "docker caching is so painful pt. 7"
+
+RUN git clone -b pybindings https://github.com/cmuparlay/ParlayANN.git && cd ParlayANN && git submodule update --init --recursive && cd python && pip install pybind11 && bash compile.sh 
+# WORKDIR /home/app/ParlayANN
+# RUN git submodule update --init --recursive
+# WORKDIR /home/app/ParlayANN/python
+
+# RUN pip install pybind11
+
+# RUN bash compile.sh
+
+ENV PYTHONPATH=$PYTHONPATH:/home/app/ParlayANN/python
+
+# ENV PARLAY_NUM_THREADS=1
+
+WORKDIR /home/app

neurips23/ood/vamana/config.yaml

@@ -0,0 +1,27 @@
+random-xs:
+    vamana: 
+      docker-tag: neurips23-ood-vamana
+      module: neurips23.ood.vamana.vamana
+      constructor: vamana
+      base-args: ["@metric"]
+      run-groups:
+        base:
+          args: |
+            [{"R":32, "L":50}]
+          query-args: |
+            [{"Ls":50, "T":8}]
+text2image-10M:
+    vamana: 
+      docker-tag: neurips23-ood-vamana
+      module: neurips23.ood.vamana.vamana
+      constructor: vamana
+      base-args: ["@metric"]
+      run-groups:
+        base:
+          args: |
+            [{"R":64, "L":500}]
+          query-args: |
+            [{"Ls":30, "T":8},
+             {"Ls":50, "T":8},
+             {"Ls":70, "T":8},
+             {"Ls":100, "T":8}]

neurips23/ood/vamana/vamana.py

@@ -0,0 +1,101 @@
+from __future__ import absolute_import
+import psutil
+import os
+import time
+import numpy as np
+import wrapper as pann
+
+from neurips23.ood.base import BaseOODANN
+from benchmark.datasets import DATASETS, download_accelerated
+
+class vamana(BaseOODANN):
+    def __init__(self, metric, index_params):
+        self.name = "vamana"
+        if (index_params.get("R")==None):
+            print("Error: missing parameter R")
+            return
+        if (index_params.get("L")==None):
+            print("Error: missing parameter L")
+            return
+        self._index_params = index_params
+        self._metric = self.translate_dist_fn(metric)
+
+        self.R = int(index_params.get("R"))
+        self.L = int(index_params.get("L"))
+
+    def index_name(self):
+        return f"R{self.R}_L{self.L}"
+
+    def create_index_dir(self, dataset):
+        index_dir = os.path.join(os.getcwd(), "data", "indices", "ood")
+        os.makedirs(index_dir, mode=0o777, exist_ok=True)
+        index_dir = os.path.join(index_dir, 'vamana')
+        os.makedirs(index_dir, mode=0o777, exist_ok=True)
+        index_dir = os.path.join(index_dir, dataset.short_name())
+        os.makedirs(index_dir, mode=0o777, exist_ok=True)
+        index_dir = os.path.join(index_dir, self.index_name())
+        os.makedirs(index_dir, mode=0o777, exist_ok=True)
+        return os.path.join(index_dir, self.index_name())
+
+    def translate_dist_fn(self, metric):
+        if metric == 'euclidean':
+            return 'Euclidian'
+        elif metric == 'ip':
+            return 'mips'
+        else:
+            raise Exception('Invalid metric')
+
+    def translate_dtype(self, dtype:str):
+        if dtype == 'float32':
+            return 'float'
+        else:
+            return dtype
+
+    def fit(self, dataset):
+        """
+        Build the index for the data points given in dataset name.
+        """
+        ds = DATASETS[dataset]()
+        d = ds.d
+
+        index_dir = self.create_index_dir(ds)
+
+        if hasattr(self, 'index'):
+            print("Index already exists")
+            return
+        else:
+            start = time.time()
+            # ds.ds_fn is the name of the dataset file but probably needs a prefix
+            # choosing 1.2 for alpha but this is probably provided in index_params
+            pann.build_vamana_index(self._metric, self.translate_dtype(ds.dtype), ds.get_dataset_fn(), index_dir, self.R, self.L, 1.2)
+            end = time.time()
+            print("Indexing time: ", end - start)
+            print(f"Wrote index to {index_dir}")
+
+        self.index = pann.load_vamana_index(self._metric, self.translate_dtype(ds.dtype), ds.get_dataset_fn(), index_dir, ds.nb, d)
+        print("Index loaded")
+
+    def query(self, X, k):
+        nq, d = X.shape
+        self.res, self.query_dists = self.index.batch_search(X, nq, k, self.Ls)
+        print(f"self.res shape: {self.res.shape}")
+        print(f"self.res[:5]: {self.res[:5]}")
+        print(f"self.query_dists shape: {self.query_dists.shape}")
+        print(f"self.query_dists[:5]: {self.query_dists[:5]}")
+
+    def set_query_arguments(self, query_args):
+        self._query_args = query_args
+        self.Ls = 0 if query_args.get("Ls") is None else query_args.get("Ls")
+        self.search_threads = self._query_args.get("T", 16)
+
+    def load_index(self, dataset):
+        ds = DATASETS[dataset]()
+        d = ds.d
+        index_dir = self.create_index_dir(ds)
+        try:
+            self.index = pann.load_vamana_index(self._metric, self.translate_dtype(ds.dtype), ds.get_dataset_fn(), index_dir, ds.nb, d)
+            print("Index loaded")
+            return True
+        except:
+            print("Index not found")
+            return False