task for each process
(iterator, tokenizer, seq_length, eot, parallel_writer=True)
| 143 | |
| 144 | |
| 145 | def task_unit(iterator, tokenizer, seq_length, eot, parallel_writer=True): |
| 146 | """task for each process""" |
| 147 | p = current_process() |
| 148 | index = p.pid if p.pid else 0 |
| 149 | |
| 150 | item_iter = tokenize_openwebtext(tokenizer, iterator, seq_length, eot) |
| 151 | batch_size = 1024 # size of write batch |
| 152 | count = 0 |
| 153 | while True: |
| 154 | data_batch = [] |
| 155 | try: |
| 156 | for _ in range(batch_size): |
| 157 | data_batch.append(next(item_iter)) |
| 158 | count += 1 |
| 159 | writer.write_raw_data(data_batch, parallel_writer=parallel_writer) |
| 160 | print("Process {} transformed {} records.".format(index, count)) |
| 161 | except StopIteration: |
| 162 | if data_batch: |
| 163 | writer.write_raw_data( |
| 164 | data_batch, parallel_writer=parallel_writer |
| 165 | ) |
| 166 | print("Process {} transformed {} records.".format(index, count)) |
| 167 | break |
| 168 | |
| 169 | |
| 170 | if __name__ == "__main__": |
nothing calls this directly
no test coverage detected