Spaces:
Runtime error
Runtime error
File size: 1,826 Bytes
58627fa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
import os
import torch
import tqdm
from colbert.indexing.loaders import load_doclens
from colbert.utils.utils import print_message, flatten
def optimize_ivf(orig_ivf, orig_ivf_lengths, index_path):
print_message("#> Optimizing IVF to store map from centroids to list of pids..")
print_message("#> Building the emb2pid mapping..")
all_doclens = load_doclens(index_path, flatten=False)
# assert self.num_embeddings == sum(flatten(all_doclens))
all_doclens = flatten(all_doclens)
total_num_embeddings = sum(all_doclens)
emb2pid = torch.zeros(total_num_embeddings, dtype=torch.int)
"""
EVENTUALLY: Use two tensors. emb2pid_offsets will have every 256th element.
emb2pid_delta will have the delta from the corresponding offset,
"""
offset_doclens = 0
for pid, dlength in enumerate(all_doclens):
emb2pid[offset_doclens: offset_doclens + dlength] = pid
offset_doclens += dlength
print_message("len(emb2pid) =", len(emb2pid))
ivf = emb2pid[orig_ivf]
unique_pids_per_centroid = []
ivf_lengths = []
offset = 0
for length in tqdm.tqdm(orig_ivf_lengths.tolist()):
pids = torch.unique(ivf[offset:offset+length])
unique_pids_per_centroid.append(pids)
ivf_lengths.append(pids.shape[0])
offset += length
ivf = torch.cat(unique_pids_per_centroid)
ivf_lengths = torch.tensor(ivf_lengths)
original_ivf_path = os.path.join(index_path, 'ivf.pt')
optimized_ivf_path = os.path.join(index_path, 'ivf.pid.pt')
torch.save((ivf, ivf_lengths), optimized_ivf_path)
print_message(f"#> Saved optimized IVF to {optimized_ivf_path}")
if os.path.exists(original_ivf_path):
print_message(f"#> Original IVF at path \"{original_ivf_path}\" can now be removed")
return ivf, ivf_lengths
|