| 249 | return self.cumulative_lengths[-1] |
| 250 | |
| 251 | def cache_dataset(self, dataset_index): |
| 252 | if self.current_cache_index >= 0: |
| 253 | # Remove the least recently used cache entry |
| 254 | self.cache_data[self.current_cache_index] = None |
| 255 | self.cache_indices[self.current_cache_index] = None |
| 256 | |
| 257 | # Select a random cache slot for the new dataset |
| 258 | self.current_cache_index = random.randint(0, self.cache_size - 1) |
| 259 | |
| 260 | # Cache the data from the CSV file |
| 261 | df = pl.read_csv(self.csv_files[dataset_index]).to_pandas() |
| 262 | |
| 263 | self.cache_data[self.current_cache_index] = df.values[:, 1:].astype(np.float32) |
| 264 | self.cache_indices[self.current_cache_index] = dataset_index |
| 265 | |
| 266 | def __getitem__(self, index): |
| 267 | try: |