Distance array now stores distances in float32

Fixes `/opt/conda/envs/lmi/lib/python3.11/site-packages/torch/_tensor.py:1089: RuntimeWarning: overflow encountered in cast`

Author: ProchazkaDavid

task1.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+
 os.environ['MKL_NUM_THREADS'] = '27'
 os.environ['OMP_NUM_THREADS'] = '27'
 os.environ['OMP_DYNAMIC'] = 'FALSE'
@@ -126,13 +127,13 @@ def _visit_buckets(
     def search(self, queries: Tensor, k: int, nprobe: int = 100) -> tuple[np.ndarray, np.ndarray]:
         predicted_bucket_ids = self._predict(queries, nprobe)
         n_queries = queries.shape[0]
-        D = np.empty((n_queries, k), dtype=np.float16)
+        D = np.empty((n_queries, k), dtype=np.float32)
         I = np.empty((n_queries, k), dtype=np.int32)
 
         torch.set_num_threads(3)
         faiss.omp_set_num_threads(3)
 
-        with ThreadPoolExecutor(max_workers=9) as executor: 
+        with ThreadPoolExecutor(max_workers=9) as executor:
             results = executor.map(
                 lambda i: self._visit_buckets(k, predicted_bucket_ids[i], queries[i : i + 1], i, nprobe),
                 range(n_queries),

task2.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+
 os.environ['MKL_NUM_THREADS'] = '4'
 os.environ['OMP_NUM_THREADS'] = '4'
 os.environ['OMP_DYNAMIC'] = 'FALSE'
@@ -148,7 +149,7 @@ def search(
     ) -> tuple[np.ndarray, np.ndarray]:
         predicted_bucket_ids = self._predict(queries, nprobe)
         n_queries = queries.shape[0]
-        D = np.empty((n_queries, k), dtype=np.float16)
+        D = np.empty((n_queries, k), dtype=np.float32)
         I = np.empty((n_queries, k), dtype=np.int32)
 
         torch.set_num_threads(2)

task3.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+
 os.environ['MKL_NUM_THREADS'] = '27'
 os.environ['OMP_NUM_THREADS'] = '27'
 os.environ['OMP_DYNAMIC'] = 'FALSE'
@@ -135,13 +136,13 @@ def search(
     ) -> tuple[np.ndarray, np.ndarray]:
         predicted_bucket_ids = self._predict(queries, nprobe)
         n_queries = queries.shape[0]
-        D = np.empty((n_queries, k), dtype=np.float16)
+        D = np.empty((n_queries, k), dtype=np.float32)
         I = np.empty((n_queries, k), dtype=np.int32)
 
         torch.set_num_threads(3)
         faiss.omp_set_num_threads(3)
 
-        with ThreadPoolExecutor(max_workers=9) as executor: # max_workers=9
+        with ThreadPoolExecutor(max_workers=9) as executor:  # max_workers=9
             results = executor.map(
                 lambda i: self._visit_buckets(k, predicted_bucket_ids[i], decomposed_queries[i : i + 1], i, nprobe),
                 range(n_queries),