(self)
| 111 | self.writers = writers |
| 112 | |
| 113 | def process(self): |
| 114 | if os.path.isdir(self.PATH): |
| 115 | paths = [os.path.join(top, name) for top, _, names in os.walk(self.PATH) for name in names] |
| 116 | # paths = [entry.path for entry in os.scandir(self.PATH) if |
| 117 | # not entry.is_dir() and not entry.name.endswith("bz2")] |
| 118 | else: |
| 119 | paths = [self.PATH] |
| 120 | task_queue, done_queue, info_queue = Queue(maxsize=self.TASK_QUEUE_LIMIT), Queue( |
| 121 | maxsize=self.DONE_QUEUE_LIMIT), Queue() |
| 122 | processes = [] |
| 123 | for i in range(NUM_PROCESSES): |
| 124 | process = Process(target=self.tokenize_worker, |
| 125 | args=(task_queue, done_queue, info_queue, self.tokenizer, self.tokenize)) |
| 126 | process.start() |
| 127 | processes.append(process) |
| 128 | |
| 129 | def read_input_to_queue(): |
| 130 | for path in paths: |
| 131 | print_rank_0(f"Start reading {path}") |
| 132 | with open(path) as file: |
| 133 | if self.split_row: |
| 134 | for row in file: |
| 135 | task_queue.put(row) |
| 136 | else: |
| 137 | items = json.load(file) |
| 138 | for item in items["RECORDS"]: |
| 139 | task_queue.put(item) |
| 140 | print_rank_0("Read input complete") |
| 141 | for i in range(len(processes)): |
| 142 | task_queue.put('STOP') |
| 143 | |
| 144 | process = Process(target=read_input_to_queue) |
| 145 | process.start() |
| 146 | count = len(processes) |
| 147 | progress_bar = tqdm.tqdm() |
| 148 | while True: |
| 149 | data = done_queue.get() |
| 150 | if data == 'COMPLETE': |
| 151 | count -= 1 |
| 152 | if count == 0: |
| 153 | break |
| 154 | else: |
| 155 | self.write_result(data, self.writers) |
| 156 | progress_bar.update() |
| 157 | progress_bar.close() |
| 158 | self.print_info(info_queue) |
| 159 | |
| 160 | @staticmethod |
| 161 | def write_result(data, writers): |
no test coverage detected