| 231 | |
| 232 | |
| 233 | class IndexedCachedDataset(IndexedDataset): |
| 234 | def __init__(self, path): |
| 235 | super().__init__(path) |
| 236 | self.cache = None |
| 237 | self.cache_index = {} |
| 238 | |
| 239 | @property |
| 240 | def supports_prefetch(self): |
| 241 | return True |
| 242 | |
| 243 | def prefetch(self, indices): |
| 244 | if all(i in self.cache_index for i in indices): |
| 245 | return |
| 246 | if not self.data_file: |
| 247 | self.read_data(self.path) |
| 248 | indices = sorted(set(indices)) |
| 249 | total_size = 0 |
| 250 | for i in indices: |
| 251 | total_size += self.data_offsets[i + 1] - self.data_offsets[i] |
| 252 | self.cache = np.empty(total_size, dtype=self.dtype) |
| 253 | ptx = 0 |
| 254 | self.cache_index.clear() |
| 255 | for i in indices: |
| 256 | self.cache_index[i] = ptx |
| 257 | size = self.data_offsets[i + 1] - self.data_offsets[i] |
| 258 | a = self.cache[ptx : ptx + size] |
| 259 | self.data_file.seek(self.data_offsets[i] * self.element_size) |
| 260 | self.data_file.readinto(a) |
| 261 | ptx += size |
| 262 | if self.data_file: |
| 263 | # close and delete data file after prefetch so we can pickle |
| 264 | self.data_file.close() |
| 265 | self.data_file = None |
| 266 | |
| 267 | # @lru_cache(maxsize=8) |
| 268 | def __getitem__(self, idx): |
| 269 | if isinstance(idx, int): |
| 270 | i = idx |
| 271 | self.check_index(i) |
| 272 | tensor_size = self.sizes[self.dim_offsets[i] : self.dim_offsets[i + 1]] |
| 273 | a = np.empty(tensor_size, dtype=self.dtype) |
| 274 | ptx = self.cache_index[i] |
| 275 | np.copyto(a, self.cache[ptx : ptx + a.size]) |
| 276 | return a |
| 277 | elif isinstance(idx, slice): |
| 278 | # Hack just to make this work, can optimizer later if necessary |
| 279 | sents = [] |
| 280 | for i in range(*idx.indices(len(self))): |
| 281 | sents.append(self[i]) |
| 282 | return sents |
| 283 | |
| 284 | |
| 285 | class IndexedDatasetBuilder(object): |