| 5 | |
| 6 | |
| 7 | class IndexedDataset: |
| 8 | def __init__(self, path, num_cache=1): |
| 9 | super().__init__() |
| 10 | self.path = path |
| 11 | self.data_file = None |
| 12 | self.data_offsets = np.load(f"{path}.idx", allow_pickle=True).item()['offsets'] |
| 13 | self.data_file = open(f"{path}.data", 'rb', buffering=-1) |
| 14 | self.cache = [] |
| 15 | self.num_cache = num_cache |
| 16 | |
| 17 | def check_index(self, i): |
| 18 | if i < 0 or i >= len(self.data_offsets) - 1: |
| 19 | raise IndexError('index out of range') |
| 20 | |
| 21 | def __del__(self): |
| 22 | if self.data_file: |
| 23 | self.data_file.close() |
| 24 | |
| 25 | def __getitem__(self, i): |
| 26 | self.check_index(i) |
| 27 | if self.num_cache > 0: |
| 28 | for c in self.cache: |
| 29 | if c[0] == i: |
| 30 | return c[1] |
| 31 | self.data_file.seek(self.data_offsets[i]) |
| 32 | b = self.data_file.read(self.data_offsets[i + 1] - self.data_offsets[i]) |
| 33 | item = pickle.loads(b) |
| 34 | if self.num_cache > 0: |
| 35 | self.cache = [(i, deepcopy(item))] + self.cache[:-1] |
| 36 | return item |
| 37 | |
| 38 | def __len__(self): |
| 39 | return len(self.data_offsets) - 1 |
| 40 | |
| 41 | class IndexedDatasetBuilder: |
| 42 | def __init__(self, path): |
no outgoing calls
no test coverage detected