批量处理三元组(多线程加速) :return: 处理后的三元组列表
(self)
| 146 | return triple_score |
| 147 | |
| 148 | def run(self) -> dict: |
| 149 | """ |
| 150 | 批量处理三元组(多线程加速) |
| 151 | :return: 处理后的三元组列表 |
| 152 | """ |
| 153 | results = [] |
| 154 | scores = [] |
| 155 | |
| 156 | # 使用线程池并发处理 |
| 157 | # max_workers = multiprocessing.cpu_count() - 1 |
| 158 | max_workers = 64 |
| 159 | logger.info(f"Using {max_workers} threads for processing.") |
| 160 | with ThreadPoolExecutor(max_workers=max_workers) as executor: |
| 161 | # 创建future到data的映射 |
| 162 | future_to_data = {} |
| 163 | for data in self.triples: |
| 164 | # 提前处理source_text |
| 165 | data["source_text"] = self.triple_sources[(data["page_idx"], data["paragraph_idx"])] |
| 166 | future = executor.submit(self.run_one, data) |
| 167 | future_to_data[future] = data |
| 168 | |
| 169 | # 进度条跟踪 |
| 170 | for future in tqdm(as_completed(future_to_data), |
| 171 | total=len(future_to_data), |
| 172 | desc="Scoring triples..."): |
| 173 | data = future_to_data[future] |
| 174 | try: |
| 175 | result = future.result() |
| 176 | data["scores"] = result |
| 177 | results.append(data) |
| 178 | scores.append(result["score"]) |
| 179 | except Exception as e: |
| 180 | logger.error(f"Processing error: {str(e)}") |
| 181 | continue |
| 182 | |
| 183 | logger.info(f"Scoring complete. Results num: {len(results)}") |
| 184 | |
| 185 | io_file.write(self.output_path, results, mode="w") |
| 186 | |
| 187 | logger.info(f"Results saved to: {self.output_path}") |
| 188 | |
| 189 | if scores: |
| 190 | valid_scores = [score for score in scores if score != -1.0] |
| 191 | logger.info(f"Valid scores count: {len(valid_scores)} | Invalid score count: {len(scores) - len(valid_scores)} | Total scores count: {len(scores)}") |
| 192 | if valid_scores: |
| 193 | logger.info(f"Average score: {sum(valid_scores) / len(valid_scores):.2f}") |
| 194 | else: |
| 195 | logger.info("No valid scores found.") |
| 196 | |
| 197 | return results |
| 198 | |
| 199 | |
| 200 | if __name__ == "__main__": |