mColBERT / colbert /indexing /index_manager.py
vjeronymo2's picture
Adding model and checkpoint
828992f
raw
history blame contribute delete
457 Bytes
import torch
import faiss
import numpy as np
from colbert.utils.utils import print_message
class IndexManager():
def __init__(self, dim):
self.dim = dim
def save(self, tensor, path_prefix):
torch.save(tensor, path_prefix)
def load_index_part(filename, verbose=True):
part = torch.load(filename)
if type(part) == list: # for backward compatibility
part = torch.cat(part)
return part