Select continuous chunks of data from non-overlapping regions of the dataset. Args: dataset: The dataset to sample from total_sample_size: Total approximate number of samples to select num_chunks: Number of continuous chunks to select seed: Random seed f
(dataset, total_sample_size, num_chunks=10, seed=42)
| 63 | return index_path |
| 64 | |
| 65 | def select_continuous_chunks(dataset, total_sample_size, num_chunks=10, seed=42): |
| 66 | """ |
| 67 | Select continuous chunks of data from non-overlapping regions of the dataset. |
| 68 | |
| 69 | Args: |
| 70 | dataset: The dataset to sample from |
| 71 | total_sample_size: Total approximate number of samples to select |
| 72 | num_chunks: Number of continuous chunks to select |
| 73 | seed: Random seed for reproducibility |
| 74 | |
| 75 | Returns: |
| 76 | Numpy array of selected samples |
| 77 | """ |
| 78 | logger.info(f"Selecting ~{total_sample_size} samples in {num_chunks} continuous chunks") |
| 79 | np.random.seed(seed) |
| 80 | dataset_size = len(dataset) |
| 81 | |
| 82 | # Calculate samples per chunk (approximate) |
| 83 | samples_per_chunk = max(1, total_sample_size // num_chunks) |
| 84 | |
| 85 | # Dataset is too small for chunking |
| 86 | if dataset_size <= total_sample_size: |
| 87 | logger.warning(f"Dataset size ({dataset_size}) is smaller than requested sample size, using all data") |
| 88 | return np.array(dataset[:]['keys']).astype(np.float32) |
| 89 | |
| 90 | # Divide dataset into non-overlapping regions |
| 91 | region_size = dataset_size // num_chunks |
| 92 | |
| 93 | all_samples = [] |
| 94 | total_selected = 0 |
| 95 | |
| 96 | logger.info(f"Selecting {num_chunks} chunks with ~{samples_per_chunk} samples each") |
| 97 | for i in range(num_chunks): |
| 98 | # Calculate region boundaries |
| 99 | region_start = i * region_size |
| 100 | region_end = (i + 1) * region_size if i < num_chunks - 1 else dataset_size |
| 101 | |
| 102 | # Calculate valid starting range within this region |
| 103 | # The starting point should allow the chunk to fit within the region |
| 104 | max_start = max(region_start, region_end - samples_per_chunk - 1) |
| 105 | |
| 106 | # If the region is smaller than samples_per_chunk, use the whole region |
| 107 | if max_start <= region_start: |
| 108 | start_idx = region_start |
| 109 | end_idx = region_end |
| 110 | else: |
| 111 | # Randomly select a starting point within valid range |
| 112 | start_idx = np.random.randint(region_start, max_start + 1) |
| 113 | end_idx = min(start_idx + samples_per_chunk, region_end) |
| 114 | |
| 115 | chunk_size = end_idx - start_idx |
| 116 | |
| 117 | logger.info(f"Chunk {i+1}/{num_chunks}: Region [{region_start}:{region_end}], Selected [{start_idx}:{end_idx}] ({chunk_size} samples)") |
| 118 | chunk_data = np.array(dataset.select(range(start_idx, end_idx))['keys']).astype(np.float32) |
| 119 | all_samples.append(chunk_data) |
| 120 | total_selected += chunk_size |
| 121 | |
| 122 | # Concatenate all chunks |