Add Cohere dataset and configuration (#30)

* Add Cohere dataset and configuration

The Cohere HDF5 does not have a "distances" dataset. Since this is not
used, we can make it optional.

* Add SQ8

* Fix distance

* Remove NUM_THREADS and SQ8 from cohere-cal.json file.

* Add vector_count and description to cohere dataset

---------

Co-authored-by: Martin Dimitrov <[email protected]>
Co-authored-by: Marcin Poźniak <[email protected]>

Author: 3 people

dataset_reader/ann_h5_reader.py

@@ -1,3 +1,4 @@
+import itertools
 from typing import Iterator
 
 import h5py
@@ -14,17 +15,18 @@ def __init__(self, path, normalize=False):
 
     def read_queries(self) -> Iterator[Query]:
         data = h5py.File(self.path)
+        distances = data["distances"] if "distances" in data else itertools.repeat(None)
 
         for vector, expected_result, expected_scores in zip(
-            data["test"], data["neighbors"], data["distances"]
+            data["test"], data["neighbors"], distances
         ):
             if self.normalize:
                 vector /= np.linalg.norm(vector)
             yield Query(
                 vector=vector.tolist(),
                 meta_conditions=None,
                 expected_result=expected_result.tolist(),
-                expected_scores=expected_scores.tolist(),
+                expected_scores=expected_scores.tolist() if expected_scores is not None else None,
             )
 
     def read_data(self, *args, **kwargs) -> Iterator[Record]:

datasets/datasets.json

@@ -1296,5 +1296,25 @@
     "path": "random-100-match-kw-small-vocab/random_keywords_1m_vocab_10_no_filters",
     "vector_count": 100,
     "description": "Synthetic data"
+  },
+  {
+    "name": "cohere-768-1M",
+    "vector_size": 768,
+    "distance": "dot",
+    "type": "h5",
+    "path": "cohere-768-1M/cohere-768-1M.hdf5",
+    "link": "https://dbyiw3u3rf9yr.cloudfront.net/corpora/vectorsearch/cohere-wikipedia-22-12-en-embeddings/documents-1m.hdf5.bz2",
+    "vector_count": 1000000,
+    "description": "Wikipedia embeddings"
+  },
+  {
+    "name": "cohere-768-10M",
+    "vector_size": 768,
+    "distance": "dot",
+    "type": "h5",
+    "path": "cohere-768-10M/cohere-768-10M.hdf5",
+    "link": "https://dbyiw3u3rf9yr.cloudfront.net/corpora/vectorsearch/cohere-wikipedia-22-12-en-embeddings/documents-10m.hdf5.bz2",
+    "vector_count": 10000000,
+    "description": "Wikipedia embeddings"
   }
-]
+]