| 65 | from sat.data_utils import load_hf_dataset |
| 66 | def create_dataset_function(path, args): |
| 67 | def process_fn(row): |
| 68 | pack, label = _encode(row['passage'], row['question']), int(row['label']) |
| 69 | return { |
| 70 | 'input_ids': np.array(pack['input_ids'], dtype=np.int64), |
| 71 | 'position_ids': np.array(pack['position_ids'], dtype=np.int64), |
| 72 | 'attention_mask': np.array(pack['attention_mask'], dtype=np.int64), |
| 73 | 'token_type_ids': np.array(pack['token_type_ids'], dtype=np.int64), |
| 74 | 'label': label |
| 75 | } |
| 76 | return load_hf_dataset(path, process_fn, columns = ["input_ids", "position_ids", "token_type_ids", "attention_mask", "label"], cache_dir=args.data_root, offline=False, transformer_name="boolq_transformer") |
| 77 | |
| 78 | if __name__ == '__main__': |