adding realm_index.py

megatron/data/realm_index.py

+from collections import defaultdict
+import itertools
+import os
+import pickle
+import shutil
+
+import faiss
+import numpy as np
+import torch
+
+from megatron import get_args, mpu
+
+
+def detach(tensor):
+    return tensor.detach().cpu().numpy()
+
+
+class BlockData(object):
+    """Serializable data structure for holding data for blocks -- embeddings and necessary metadata for REALM"""
+    def __init__(self, block_data_path=None, rank=None):
+        self.embed_data = dict()
+        self.meta_data = dict()
+        if block_data_path is None:
+            args = get_args()
+            block_data_path = args.block_data_path
+            rank = args.rank
+        self.block_data_path = block_data_path
+        self.rank = rank
+
+        block_data_name = os.path.splitext(self.block_data_path)[0]
+        self.temp_dir_name = block_data_name + '_tmp'
+
+    def state(self):
+        return {
+            'embed_data': self.embed_data,
+            'meta_data': self.meta_data,
+        }
+
+    def clear(self):
+        """Clear the embedding data structures to save memory.
+        The metadata ends up getting used, and is also much smaller in dimensionality
+        so it isn't really worth clearing.
+        """
+        self.embed_data = dict()
+
+    @classmethod
+    def load_from_file(cls, fname):
+        print("\n> Unpickling BlockData", flush=True)
+        state_dict = pickle.load(open(fname, 'rb'))
+        print(">> Finished unpickling BlockData\n", flush=True)
+
+        new_index = cls()
+        new_index.embed_data = state_dict['embed_data']
+        new_index.meta_data = state_dict['meta_data']
+        return new_index
+
+    def add_block_data(self, block_indices, block_embeds, block_metas, allow_overwrite=False):
+        for idx, embed, meta in zip(block_indices, block_embeds, block_metas):
+            if not allow_overwrite and idx in self.embed_data:
+                raise ValueError("Unexpectedly tried to overwrite block data")
+
+            self.embed_data[idx] = np.float16(embed)
+            self.meta_data[idx] = meta
+
+    def save_shard(self):
+        if not os.path.isdir(self.temp_dir_name):
+            os.makedirs(self.temp_dir_name, exist_ok=True)
+
+        # save the data for each shard
+        with open('{}/{}.pkl'.format(self.temp_dir_name, self.rank), 'wb') as data_file:
+            pickle.dump(self.state(), data_file)
+
+    def merge_shards_and_save(self):
+        """Combine all the shards made using self.save_shard()"""
+        shard_names = os.listdir(self.temp_dir_name)
+        seen_own_shard = False
+
+        for fname in os.listdir(self.temp_dir_name):
+            shard_rank = int(os.path.splitext(fname)[0])
+            if shard_rank == self.rank:
+                seen_own_shard = True
+                continue
+
+            with open('{}/{}'.format(self.temp_dir_name, fname), 'rb') as f:
+                data = pickle.load(f)
+                old_size = len(self.embed_data)
+                shard_size = len(data['embed_data'])
+
+                # add the shard's data and check to make sure there is no overlap
+                self.embed_data.update(data['embed_data'])
+                self.meta_data.update(data['meta_data'])
+                assert len(self.embed_data) == old_size + shard_size
+
+        assert seen_own_shard
+
+        # save the consolidated shards and remove temporary directory
+        with open(self.block_data_path, 'wb') as final_file:
+            pickle.dump(self.state(), final_file)
+        shutil.rmtree(self.temp_dir_name, ignore_errors=True)
+
+        print("Finished merging {} shards for a total of {} embeds".format(
+            len(shard_names), len(self.embed_data)), flush=True)
+
+
+class FaissMIPSIndex(object):
+    """Wrapper object for a BlockData which similarity search via FAISS under the hood"""
+    def __init__(self, index_type, embed_size, use_gpu=False):
+        self.index_type = index_type
+        self.embed_size = embed_size
+        self.use_gpu = use_gpu
+        self.id_map = dict()
+
+        self.block_mips_index = None
+        self._set_block_index()
+
+    def _set_block_index(self):
+        INDEX_TYPES = ['flat_ip']
+        if self.index_type not in INDEX_TYPES:
+            raise ValueError("Invalid index type specified")
+
+        print("\n> Building index", flush=True)
+        self.block_mips_index = faiss.index_factory(self.embed_size, 'Flat', faiss.METRIC_INNER_PRODUCT)
+
+        if self.use_gpu:
+            # create resources and config for GpuIndex
+            res = faiss.StandardGpuResources()
+            config = faiss.GpuIndexFlatConfig()
+            config.device = torch.cuda.current_device()
+            config.useFloat16 = True
+
+            self.block_mips_index = faiss.GpuIndexFlat(res, self.block_mips_index, config)
+            print(">>> Finished building index on GPU {}\n".format(self.block_mips_index.getDevice()), flush=True)
+        else:
+            # CPU index supports IDs so wrap with IDMap
+            self.block_mips_index = faiss.IndexIDMap(self.block_mips_index)
+            print(">> Finished building index\n", flush=True)
+
+    def reset_index(self):
+        """Delete existing index and create anew"""
+        del self.block_mips_index
+        self._set_block_index()
+
+    def add_block_embed_data(self, all_block_data):
+        """Add the embedding of each block to the underlying FAISS index"""
+        block_indices, block_embeds = zip(*all_block_data.embed_data.items())
+        if self.use_gpu:
+            for i, idx in enumerate(block_indices):
+                self.id_map[i] = idx
+
+        all_block_data.clear()
+        if self.use_gpu:
+            self.block_mips_index.add(np.float32(np.array(block_embeds)))
+        else:
+            self.block_mips_index.add_with_ids(np.float32(np.array(block_embeds)), np.array(block_indices))
+
+    def search_mips_index(self, query_embeds, top_k, reconstruct=True):
+        """Get the top-k blocks by the index distance metric.
+
+        :param reconstruct: if True: return a [num_queries x k x embed_dim] array of blocks
+                            if False: return [num_queries x k] array of distances, and another for indices
+        """
+        query_embeds = np.float32(detach(query_embeds))
+
+        with torch.no_grad():
+            if reconstruct:
+                top_k_block_embeds = self.block_mips_index.search_and_reconstruct(query_embeds, top_k)
+                return top_k_block_embeds
+            else:
+                distances, block_indices = self.block_mips_index.search(query_embeds, top_k)
+                if self.use_gpu:
+                    fresh_indices = np.zeros(block_indices.shape)
+                    for i, j in itertools.product(block_indices.shape):
+                        fresh_indices[i, j] = self.id_map[block_indices[i, j]]
+                    block_indices = fresh_indices
+                return distances, block_indices