使用滑动窗口提取目录,并对冲突页进行单页投票
(pdf_path: str, client: AsyncOpenAI, model: str, initial_data_dir: str)
| 230 | write_log(f"开始提取目录信息,总页数:{total_pages}") |
| 231 | |
| 232 | # 控制并发度为 8 |
| 233 | semaphore = asyncio.Semaphore(8) |
| 234 | # 记录每一页被判定为目录和非目录的次数 |
| 235 | page_votes = {i: {"is_toc": 0, "not_toc": 0} for i in range(1, total_pages + 1)} |
| 236 | |
| 237 | async def process_window(start_p, end_p): |
| 238 | async with semaphore: |
| 239 | img_filename = f"concat_pages_{start_p}_{end_p}.jpg" |
| 240 | img_save_path = os.path.join(initial_data_dir, img_filename) |
| 241 | raw_filename = f"toc_response_{start_p}_{end_p}.json" |
| 242 | raw_save_path = os.path.join(initial_data_dir, raw_filename) |
| 243 | |
| 244 | b64_img = create_concat_image_b64(doc, start_p, end_p, save_path=img_save_path) |
| 245 | if not b64_img: |
| 246 | return start_p, end_p, None, None, None |
| 247 | |
| 248 | raw_res = await fetch_toc_from_image(client, model, b64_img, start_p, end_p, raw_save_path) |
| 249 | toc_start, toc_end = parse_toc_json(raw_res) |
| 250 | |
| 251 | detail_filename = f"toc_detail_{start_p}_{end_p}.json" |
| 252 | detail_save_path = os.path.join(initial_data_dir, detail_filename) |
| 253 | detail_data = { |
| 254 | "page_range": f"{start_p}-{end_p}", |
| 255 | "parsed_result": {"toc_start": toc_start, "toc_end": toc_end}, |
| 256 | "raw_response": raw_res, |
| 257 | "image_saved_as": img_filename |
| 258 | } |
| 259 | with open(detail_save_path, 'w', encoding='utf-8') as f: |
| 260 | json.dump(detail_data, f, ensure_ascii=False, indent=2) |
| 261 | |
| 262 | return start_p, end_p, toc_start, toc_end, raw_res |
| 263 | |
| 264 | async def run_batch(start_page, end_page_limit): |
| 265 | """执行一个批次的滑动窗口扫描""" |
| 266 | windows = [] |
| 267 | for i in range(start_page, end_page_limit, 2): |
| 268 | if i > total_pages: |
| 269 | break |
| 270 | windows.append((i, min(i + 3, total_pages))) |
| 271 | |
| 272 | if not windows: |
| 273 | return |
| 274 | |
| 275 | tasks = [process_window(s, e) for s, e in windows] |
| 276 | results = await asyncio.gather(*tasks) |
| 277 | |
| 278 | for s, e, t_start, t_end, _ in results: |
| 279 | if t_start is None: continue |
| 280 | for p in range(s, e + 1): |
| 281 | if t_start <= p <= t_end: |
| 282 | page_votes[p]["is_toc"] += 1 |
| 283 | else: |
| 284 | page_votes[p]["not_toc"] += 1 |
| 285 | |
| 286 | def has_toc_in_range(votes_dict, start_p, end_p): |
| 287 | """检查指定范围内是否有被投票为目录的页""" |
| 288 | for p in range(start_p, end_p + 1): |
| 289 | if votes_dict.get(p, {}).get("is_toc", 0) > 0: |
no test coverage detected