(self, idx)
| 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] |
| 237 | doc_idx = int(doc_idx) # NumPy int => Python int |
| 238 | |
| 239 | input_ids = self.input_ids_indexed_dataset[doc_idx] |
| 240 | # print_rank_0(f"input_ids={input_ids}") |
| 241 | attention_mask = self.attention_mask_index_dataset[doc_idx] |
| 242 | labels = self.labels_indexed_dataset[doc_idx] |
| 243 | |
| 244 | res = { |
| 245 | "input_ids": np.array(input_ids, dtype=np.int64), |
| 246 | "attention_mask": np.array(attention_mask, dtype=np.int64), |
| 247 | "labels": np.array(labels, dtype=np.int64), |
| 248 | } |
| 249 | |
| 250 | return res |
| 251 | |
| 252 | |
| 253 | def _build_index_mappings( |
nothing calls this directly
no outgoing calls
no test coverage detected