| 15 | |
| 16 | class IndexDictOfArray: |
| 17 | def __init__(self, index_path=None, force_new=False, filename="array_index.h5py", dim_voc=None): |
| 18 | # index_path = None # for debug |
| 19 | if index_path is not None: |
| 20 | self.index_path = index_path |
| 21 | if not os.path.exists(index_path): |
| 22 | os.makedirs(index_path) |
| 23 | self.filename = os.path.join(self.index_path, filename) |
| 24 | if os.path.exists(self.filename) and not force_new: |
| 25 | print("index already exists, loading...") |
| 26 | |
| 27 | self.file = h5py.File(self.filename, "r") |
| 28 | if dim_voc is not None: |
| 29 | dim = dim_voc |
| 30 | else: |
| 31 | dim = self.file["dim"][()] |
| 32 | self.index_doc_id = dict() |
| 33 | self.index_doc_value = dict() |
| 34 | for key in tqdm(range(dim)): |
| 35 | try: |
| 36 | self.index_doc_id[key] = np.array(self.file["index_doc_id_{}".format(key)], |
| 37 | dtype=np.int32) |
| 38 | # ideally we would not convert to np.array() but we cannot give pool an object with hdf5 |
| 39 | self.index_doc_value[key] = np.array(self.file["index_doc_value_{}".format(key)], |
| 40 | dtype=np.float32) |
| 41 | except: |
| 42 | self.index_doc_id[key] = np.array([], dtype=np.int32) |
| 43 | self.index_doc_value[key] = np.array([], dtype=np.float32) |
| 44 | self.file.close() |
| 45 | del self.file |
| 46 | print("done loading index...") |
| 47 | doc_ids = pickle.load(open(os.path.join(self.index_path, "doc_ids.pkl"), "rb")) |
| 48 | self.n = len(doc_ids) |
| 49 | else: |
| 50 | self.n = 0 |
| 51 | print("initializing new index...") |
| 52 | self.index_doc_id = defaultdict(lambda: array.array("I")) |
| 53 | self.index_doc_value = defaultdict(lambda: array.array("f")) |
| 54 | else: |
| 55 | self.n = 0 |
| 56 | print("initializing new index...") |
| 57 | self.index_doc_id = defaultdict(lambda: array.array("I")) |
| 58 | self.index_doc_value = defaultdict(lambda: array.array("f")) |
| 59 | |
| 60 | def add_batch_document(self, row, col, data, n_docs=-1): |
| 61 | """add a batch of documents to the index |