Convert to numpy and return a sample consumed by the batch producer.
(ids, types=None, paddings=None, positions=None, masks=None, label=None, unique_id=None, target=None,
logit_mask=None, segment_ids=None, prompt_ids=None)
| 264 | |
| 265 | |
| 266 | def build_sample(ids, types=None, paddings=None, positions=None, masks=None, label=None, unique_id=None, target=None, |
| 267 | logit_mask=None, segment_ids=None, prompt_ids=None): |
| 268 | """Convert to numpy and return a sample consumed by the batch producer.""" |
| 269 | |
| 270 | ids_np = np.array(ids, dtype=np.int64) |
| 271 | sample = {'text': ids_np, 'label': int(label)} |
| 272 | if types is not None: |
| 273 | types_np = np.array(types, dtype=np.int64) |
| 274 | sample['types'] = types_np |
| 275 | if paddings is not None: |
| 276 | paddings_np = np.array(paddings, dtype=np.int64) |
| 277 | sample['padding_mask'] = paddings_np |
| 278 | if positions is not None: |
| 279 | positions_np = np.array(positions, dtype=np.int64) |
| 280 | sample['position'] = positions_np |
| 281 | if masks is not None: |
| 282 | masks_np = np.array(masks, dtype=np.int64) |
| 283 | sample['mask'] = masks_np |
| 284 | if target is not None: |
| 285 | target_np = np.array(target, dtype=np.int64) |
| 286 | sample['target'] = target_np |
| 287 | if logit_mask is not None: |
| 288 | logit_mask_np = np.array(logit_mask, dtype=np.int64) |
| 289 | sample['logit_mask'] = logit_mask_np |
| 290 | if segment_ids is not None: |
| 291 | segment_ids = np.array(segment_ids, dtype=np.int64) |
| 292 | sample['segment_id'] = segment_ids |
| 293 | if prompt_ids is not None: |
| 294 | prompt_ids = np.array(prompt_ids, dtype=np.int64) |
| 295 | sample['prompt_pos'] = prompt_ids |
| 296 | if unique_id is not None: |
| 297 | sample['uid'] = unique_id |
| 298 | return sample |
| 299 | |
| 300 | |
| 301 | def build_decoder_sample(sample, dec_ids, dec_position, dec_masks, dec_target, dec_logit_mask): |