| 177 | |
| 178 | |
| 179 | class PromptDataset(torch.utils.data.Dataset): |
| 180 | def __init__( |
| 181 | self, |
| 182 | name, |
| 183 | data_prefix, |
| 184 | documents, |
| 185 | input_ids_indexed_dataset, |
| 186 | attention_mask_index_dataset, |
| 187 | labels_indexed_dataset, |
| 188 | num_samples, |
| 189 | seq_length, |
| 190 | seed, |
| 191 | ): |
| 192 | """ |
| 193 | Args: |
| 194 | name: name of the dataset. |
| 195 | data_prefix: prefix of the data. |
| 196 | documents: list of document indices. |
| 197 | input_ids_indexed_dataset: indexed dataset for prompts. |
| 198 | attention_mask_index_dataset: indexed dataset for text. |
| 199 | labels_indexed_dataset: indexed dataset for labels. |
| 200 | num_samples: number of samples to draw from the indexed dataset. |
| 201 | seq_length: sequence length. |
| 202 | seed: seed for random number generator. |
| 203 | """ |
| 204 | |
| 205 | self.name = name |
| 206 | self.input_ids_indexed_dataset = input_ids_indexed_dataset |
| 207 | self.attention_mask_index_dataset = attention_mask_index_dataset |
| 208 | self.labels_indexed_dataset = labels_indexed_dataset |
| 209 | self.seq_length = seq_length |
| 210 | self.eod_token = get_tokenizer().eod |
| 211 | |
| 212 | # Checks |
| 213 | assert np.min(documents) >= 0 |
| 214 | assert np.max(documents) < input_ids_indexed_dataset.sizes.shape[0] |
| 215 | assert input_ids_indexed_dataset.sizes.shape[0] == attention_mask_index_dataset.sizes.shape[0] |
| 216 | assert attention_mask_index_dataset.sizes.shape[0] == labels_indexed_dataset.sizes.shape[0] |
| 217 | |
| 218 | # Build index mappings. |
| 219 | self.doc_idx = _build_index_mappings( |
| 220 | self.name, |
| 221 | data_prefix, |
| 222 | documents, |
| 223 | self.input_ids_indexed_dataset.sizes, |
| 224 | num_samples, |
| 225 | seq_length, |
| 226 | seed, |
| 227 | ) |
| 228 | |
| 229 | def __len__(self): |
| 230 | # -1 is due to data structure used to retieve the index: |
| 231 | # sample i --> [sample_idx[i], sample_idx[i+1]) |
| 232 | return self.doc_idx.shape[0] |
| 233 | |
| 234 | def __getitem__(self, idx): |
| 235 | # get the doc index |
| 236 | doc_idx = self.doc_idx[idx] |