Dim reduction support in OIVFBBS (#3290)

Summary:
This PR adds support for dimensionality reduction in OIVFBBS. I tested the code with an index `OPQ64_128,IVF4096,PQ64` using the ssnpp embeddings - this index string is added to the config_ssnpp.yaml to showcase this functionality.

Pull Request resolved: #3290

Reviewed By: junjieqi

Differential Revision: D54878345

Pulled By: mlomeli1

fbshipit-source-id: 98ecdeb2224ce0325e37720cc113d82f9c6c75d6

Author: mlomeli1

demos/offline_ivf/config_ssnpp.yaml

@@ -6,6 +6,7 @@ index:
   non-prod:
   - 'IVF16384,PQ128'
   - 'IVF32768,PQ128'
+  - 'OPQ64_128,IVF4096,PQ64'
 nprobe:
   prod:
     - 512

demos/offline_ivf/offline_ivf.py

@@ -178,7 +178,7 @@ def dedupe(self):
             idxs.append(np.empty((0,), dtype=np.uint32))
         bs = 1_000_000
         i = 0
-        for buffer in tqdm(self.xb_ds.iterate(0, bs, np.float32)):
+        for buffer in tqdm(self._iterate_transformed(self.xb_ds, 0, bs, np.float32)):
             for j in range(len(codecs)):
                 codec, codeset, idx = codecs[j], codesets[j], idxs[j]
                 uniq = codeset.insert(codec.sa_encode(buffer))
@@ -267,11 +267,18 @@ def index_shard(self):
                         ),
                         file=sys.stdout,
                     ):
-                        assert xb_j.shape[1] == index.d
-                        index.add_with_ids(
-                            xb_j,
-                            np.arange(start + jj, start + jj + xb_j.shape[0]),
-                        )
+                        if is_pretransform_index(index):
+                            assert xb_j.shape[1] == index.chain.at(0).d_out
+                            index_ivf.add_with_ids(
+                                xb_j,
+                                np.arange(start + jj, start + jj + xb_j.shape[0]),
+                            )
+                        else:
+                            assert xb_j.shape[1] == index.d
+                            index.add_with_ids(
+                                xb_j,
+                                np.arange(start + jj, start + jj + xb_j.shape[0]),
+                            )
                         jj += xb_j.shape[0]
                         logging.info(jj)
                         assert (
@@ -670,10 +677,14 @@ def search(self):
                             os.remove(Ifn)
                             os.remove(Dfn)
 
-            try:  # TODO: modify shape for pretransform case
+            try:
+                if is_pretransform_index(index):
+                    d = index.chain.at(0).d_out
+                else:
+                    d = self.input_d
                 with open(Ifn, "xb") as f, open(Dfn, "xb") as g:
                     xq_i = np.empty(
-                        shape=(self.xq_bs, self.input_d), dtype=np.float16
+                        shape=(self.xq_bs, d), dtype=np.float16
                     )
                     q_assign = np.empty(
                         (self.xq_bs, self.nprobe), dtype=np.int32
@@ -835,8 +846,7 @@ def consistency_check(self):
             for j in range(SMALL_DATA_SAMPLE):
                 assert np.where(I[j] == j + r)[0].size > 0, (
                     f"I[j]: {I[j]}, j: {j}, i: {i}, shard_size:"
-                    f" {self.shard_size}"
-                )
+                    f" {self.shard_size}")
 
         logging.info("search results...")
         index_ivf.nprobe = self.nprobe