Process a batch of data
(args)
| 67 | return {(item.get('source', ''), item.get('refined_statement', '')) for item in results} |
| 68 | |
| 69 | def process_batch(args): |
| 70 | """Process a batch of data""" |
| 71 | start_idx, end_idx, data, sampling_params, process_id, num_batches, checkpoint_dir = args |
| 72 | |
| 73 | # Create a unique checkpoint file for each process |
| 74 | checkpoint_file = os.path.join(checkpoint_dir, f'checkpoint_process_{process_id}.json') |
| 75 | batch_results = [] |
| 76 | |
| 77 | # Load this process's checkpoint |
| 78 | existing_results = load_checkpoint(checkpoint_file) |
| 79 | processed_items = get_processed_items(existing_results) |
| 80 | |
| 81 | for i in tqdm(range(start_idx, end_idx), desc=f"Process {os.getpid()} progress"): |
| 82 | item = data[i] |
| 83 | # Check if already processed |
| 84 | if (item.get('source', ''), item.get('refined_statement', '')) in processed_items: |
| 85 | continue |
| 86 | |
| 87 | result = process_single_item(item, sampling_params, num_batches) |
| 88 | if result: |
| 89 | batch_results.append(result) |
| 90 | |
| 91 | # Periodically save checkpoint |
| 92 | if len(batch_results) % 10 == 0: # Save every 10 items |
| 93 | existing_results.extend(batch_results) |
| 94 | with open(checkpoint_file, 'w') as f: |
| 95 | json.dump(existing_results, f, ensure_ascii=False, indent=2) |
| 96 | batch_results = [] # Clear saved results |
| 97 | |
| 98 | # Save remaining results |
| 99 | if batch_results: |
| 100 | existing_results.extend(batch_results) |
| 101 | with open(checkpoint_file, 'w') as f: |
| 102 | json.dump(existing_results, f, ensure_ascii=False, indent=2) |
| 103 | |
| 104 | return checkpoint_file |
| 105 | |
| 106 | def merge_checkpoints(checkpoint_files, output_file): |
| 107 | """Merge results from all checkpoint files""" |
nothing calls this directly
no test coverage detected