diff --git a/docs/COMMANDS.md b/docs/COMMANDS.md index 4f2e8d5e..e4990d2e 100644 --- a/docs/COMMANDS.md +++ b/docs/COMMANDS.md @@ -10,89 +10,285 @@ |-----------|---------|------|---------| | `/pf-sync` | `paperforge sync` | 同步 Zotero，生成正式笔记 | BBT JSON 导出 | | `/pf-ocr` | `paperforge ocr` | PDF OCR 文本与图表提取 | `do_ocr: true` | -| `/pf-status` | `paperforge status` | 查看系统状态 | 配置完成 | -| `/pf-deep ` | `paperforge deep-reading` | 三阶段精读 | OCR done + `analyze: true` | -| `/pf-paper ` | — | 文献问答 | 正式笔记存在 | +| `/pf-status` | `paperforge status` / `paperforge runtime-health` | 查看系统状态 | 配置完成 | +| `/pf-deep ` | `paperforge deep-reading` (队列) | 三阶段精读 (Agent 层) | OCR done + `analyze: true` | +| `/pf-paper ` | — | 文献问答 (Agent 层) | 正式笔记存在 | +| `/pf-end ` | — | 结束精读/问答 Session (Agent 层) | 精读或问答进行中 | +| `/pf-log-reading` | `paperforge reading-log --write ` | 记录阅读笔记 | 精读进行中 | +| `/pf-log-session` | `paperforge project-log --write` | 总结会话决策 | 会话结束 | --- ## CLI 命令 -### `paperforge sync` +### `paperforge paths` ```bash -paperforge sync # 完整同步 -paperforge sync --dry-run # 预览 -paperforge sync --rebuild-index # 强制重建索引 -paperforge sync --json # JSON 输出 -``` - -### `paperforge ocr` - -```bash -paperforge ocr # 处理队列 -paperforge ocr --key ABCDEFG # 处理指定文献 -paperforge ocr --diagnose # 诊断模式 -paperforge ocr --json # JSON 输出 +paperforge paths # 打印 vault 路径清单 +paperforge paths --json # JSON 格式输出 ``` ### `paperforge status` ```bash -paperforge status # 完整状态 -paperforge status --json # JSON 输出 +paperforge status # 完整状态检查 +paperforge status --json # JSON 输出 ``` -### `paperforge doctor` +### `paperforge sync` ```bash -paperforge doctor # 验证安装配置 -paperforge doctor --json # JSON 输出 +paperforge sync # 完整同步 +paperforge sync --dry-run # 预览 +paperforge sync --selection # 仅 selection-sync +paperforge sync --index # 仅 index-refresh +paperforge sync --rebuild-index # 强制重建索引 +paperforge sync --prune # 预览清理孤儿笔记 (dry-run) +paperforge sync --prune-force # 执行清理 +paperforge sync --json # JSON 输出 ``` -### `paperforge repair` +### `paperforge selection-sync` ```bash -paperforge repair # 扫描分歧（dry-run） -paperforge repair --fix # 修复 -paperforge repair --fix-paths # 修复 PDF 路径 -paperforge repair --json # JSON 输出 +paperforge selection-sync # 同步 Zotero 选择到文献记录 (向后兼容) +``` + +### `paperforge index-refresh` + +```bash +paperforge index-refresh # 刷新正式文献笔记 (向后兼容) ``` ### `paperforge deep-reading` ```bash -paperforge deep-reading # 查看精读队列 -paperforge deep-reading --verbose # 含修复指令 +paperforge deep-reading # 查精读队列状态 +paperforge deep-reading --json # JSON 输出 +``` + +### `paperforge deep-finalize` + +```bash +paperforge deep-finalize # 标记精读完成并通知 Dashboard +paperforge deep-finalize --json +``` + +### `paperforge repair` + +```bash +paperforge repair # 扫描分歧 (dry-run) +paperforge repair --fix # 修复 +paperforge repair --fix-paths # 修复 PDF 路径 +paperforge repair --json # JSON 输出 +``` + +### `paperforge ocr` + +```bash +paperforge ocr # 处理队列 (旧式) +paperforge ocr --key # 处理指定文献 +paperforge ocr --diagnose # 诊断模式 +paperforge ocr --json # JSON 输出 +paperforge ocr run # 运行 OCR 队列 (新式) +paperforge ocr doctor # 诊断 OCR 配置和连接 +paperforge ocr doctor --live # 运行真实 PDF 测试 (L4) +paperforge ocr redo # 重新运行所有标记 ocr_redo 的文献 +paperforge ocr redo [KEY...] # 重新运行指定文献 +paperforge ocr redo --dry-run # 预览待重做的文献 +paperforge ocr list # 列出所有 OCR 维护状态 +paperforge ocr list --json # JSON 输出 +paperforge ocr list --manifest # 输出 key→sha256 清单 +paperforge ocr list --keys ... # 仅指定 key 的信息 +paperforge ocr rebuild # 从已有 raw blocks 重建 OCR 产物 +paperforge ocr rebuild ... # 重建指定文献 +paperforge ocr rebuild --all # 重建全部 +paperforge ocr rebuild --status done # 按状态过滤 +paperforge ocr rebuild --dry-run # 预览 +paperforge ocr rebuild --resume # 跳过已有检查点的文献 +``` + +### `paperforge context` + +```bash +paperforge context # 单篇文献的 AI context 包 (JSON) +paperforge context --domain 骨科 # 按 domain 过滤 (JSON 数组) +paperforge context --collection "path/to/coll" # 按 collection 过滤 +paperforge context --all # 输出全部 canonical index +``` + +### `paperforge dashboard` + +```bash +paperforge dashboard # 聚合统计和 Dashboard 权限 +paperforge dashboard --json # JSON 输出 ``` ### `paperforge embed` ```bash -paperforge embed build # 构建向量索引 -paperforge embed build --resume # 续建 -paperforge embed status # 查看状态 -paperforge embed stop # 停止构建 +paperforge embed build # 构建向量索引 +paperforge embed build --resume # 续建 +paperforge embed build --force # 强制重建 +paperforge embed build --json # JSON 输出 +paperforge embed status # 查看向量 DB 状态 +paperforge embed status --json # JSON 输出 +paperforge embed stop # 停止构建 +paperforge embed stop --json # JSON 输出 +``` + +### `paperforge retrieve` + +```bash +paperforge retrieve "" # 跨 OCR 全文语义检索 +paperforge retrieve "PEMF" --limit 10 --json +paperforge retrieve "75 Hz" --no-expand +``` + +### `paperforge query-plan` + +```bash +paperforge query-plan "" --intent discover # 文献发现 +paperforge query-plan "" --intent content # 内容检索 +paperforge query-plan "" --intent known-paper # 已知论文定位 +paperforge query-plan "" --intent discover --json +``` + +### `paperforge prune` + +```bash +paperforge prune # 预览删除孤儿产物 (dry-run) +paperforge prune --force # 实际删除 +paperforge prune [KEY...] # 仅处理指定 key +paperforge prune --json # JSON 输出 ``` ### `paperforge memory` ```bash -paperforge memory build # 构建 memory DB -paperforge memory status # 查看状态 +paperforge memory build # 构建 memory DB +paperforge memory build --json # JSON 输出 +paperforge memory status # 查看 memory DB 状态 +paperforge memory status --json # JSON 输出 ``` ### `paperforge search` ```bash +paperforge search "" # 元数据全文搜索 paperforge search "" --json paperforge search "PEMF" --domain 骨科 --ocr done --year-from 2020 +paperforge search "PEMF" --deep done --lifecycle fulltext_ready +paperforge search "PEMF" --next-step ocr --limit 10 +``` + +### `paperforge paper-status` + +```bash +paperforge paper-status # 查找文献状态 (key/DOI/title/alias) +paperforge paper-status "XGT9Z257" --json +``` + +### `paperforge paper-context` + +```bash +paperforge paper-context # 文献完整上下文 +paperforge paper-context "XGT9Z257" --json +``` + +### `paperforge reading-log` + +```bash +paperforge reading-log --write --section "Discussion P12" --excerpt "..." --usage "..." # 写阅读笔记 +paperforge reading-log --render # 渲染 reading-log.md +paperforge reading-log --lookup # 查某文献的阅读笔记 +paperforge reading-log --since 2026-01-01 --limit 100 --output notes.md # 导出 +paperforge reading-log --validate reading-log.md +paperforge reading-log --import reading-log.md +paperforge reading-log --correct --correction "..." --reason "..." # 修正笔记 +paperforge reading-log --json +``` + +### `paperforge project-log` + +```bash +paperforge project-log --write --project "ProjectX" --payload '{"hours":2,"status":"progress"}' # 写工作日志 +paperforge project-log --list --project "ProjectX" +paperforge project-log --render --project "ProjectX" +paperforge project-log --project "ProjectX" --json +``` + +### `paperforge agent-context` + +```bash +paperforge agent-context # 生成 Agent 引导上下文 +paperforge agent-context --json # JSON 输出 ``` ### `paperforge runtime-health` ```bash -paperforge runtime-health --json +paperforge runtime-health # 检查 memory layer 运行时健康 +paperforge runtime-health --json # JSON 输出 +``` + +### `paperforge base-refresh` + +```bash +paperforge base-refresh # 刷新 Obsidian Base 视图 +paperforge base-refresh --force # 强制全量重建 +``` + +### `paperforge doctor` + +```bash +paperforge doctor # 验证安装配置 +paperforge doctor --json # JSON 输出 +``` + +### `paperforge update` + +```bash +paperforge update # 更新到最新版本 +``` + +### `paperforge setup` + +```bash +paperforge setup # 交互式设置 +paperforge setup --headless # 非交互式 (AI 代理) +paperforge setup --agent opencode # 指定 AI Agent 平台 +paperforge setup --paddleocr-key "KEY" --paddleocr-url "URL" +paperforge setup --skip-checks # 跳过环境检查 +paperforge setup --modular # 使用模块化设置 (v2.1+) +``` + +### `paperforge paper-lookup` + +```bash +paperforge paper-lookup "" # L4 网关: 定位文献 +paperforge paper-lookup "PEMF" --json --limit 10 +``` + +### `paperforge content-discovery` + +```bash +paperforge content-discovery "" # L4 网关: 内容发现 +paperforge content-discovery "骨科 生物力学" --json --limit 10 +``` + +### `paperforge paper-navigation` + +```bash +paperforge paper-navigation "" # L4 网关: 文献结构导航 +paperforge paper-navigation "XGT9Z257" --json +``` + +### `paperforge scoped-fetch` + +```bash +paperforge scoped-fetch "" # L4 网关: 范围受限获取 +paperforge scoped-fetch "PEMF 75 Hz treatment protocol" --json --limit 10 ``` --- @@ -130,3 +326,4 @@ $env:PAPERFORGE_REAL_OCR_VAULT="D:\path\to\Literature-hub"; $env:PAPERFORGE_REAL ``` 不设置环境变量时，真实文献测试自动 SKIP。 + diff --git a/docs/superpowers/plans/2026-07-05-rebuild-speed-and-batch-improvements-plan.md b/docs/superpowers/plans/2026-07-05-rebuild-speed-and-batch-improvements-plan.md new file mode 100644 index 00000000..14dc5e3a --- /dev/null +++ b/docs/superpowers/plans/2026-07-05-rebuild-speed-and-batch-improvements-plan.md @@ -0,0 +1,464 @@ +# Rebuild 加速与批量改进 — 实现计划 + +> Branch: `feat/rebuild-speed` +> Base: `master` +> Spec: `docs/superpowers/specs/2026-07-05-rebuild-speed-and-batch-improvements.md` +> Review: 8.8/10 — 可以进入实现 + +## 总览 + +5 个 Wave，文件不重叠，可并行。 + +| Wave | 文件 | 改动 | 复杂度 | +|------|------|------|--------| +| 1 | `ocr_rebuild.py` + `commands/ocr.py` + `cli.py` | 阶段化 + 跨论文并行 + checkpoint 重构 | 高 | +| 2 | `ocr_objects.py` | 资产裁剪并行（复用现有坐标逻辑，不重写） | 中 | +| 3 | `ocr_maintenance.py` | Display 封装（用真实字段） | 低 | +| 4 | `commands/embed.py` | Embed 人类进度（保留机器协议） | 低 | +| 5 | `docs/COMMANDS.md` | CLI 文档补齐（从 parser 核对） | 低 | + +--- + +## Wave 1: 阶段化 + 跨论文并行 + +### 1.1 前置 + +```python +import logging +logger = logging.getLogger(__name__) +``` + +### 1.2 改写 `run_derived_rebuild_for_keys()` + +**签名：** + +```python +def run_derived_rebuild_for_keys( + vault: Path, + keys: list[str], + progress_bar=None, + checkpoint_dir: Path | None = None, + parallel: int = 4, +) -> dict: +``` + +**函数体：checkpoint 过滤在入口统一执行，串行/并行路径共享一处过滤逻辑。** + +```python +def run_derived_rebuild_for_keys(vault, keys, progress_bar=None, + checkpoint_dir=None, parallel=4): + # 入口统一过滤 checkpoint，串行/并行路径共享 + if checkpoint_dir: + checkpoint_dir.mkdir(parents=True, exist_ok=True) + keys = _filter_completed_keys(keys, checkpoint_dir) + + if parallel and len(keys) > 1: + return _run_parallel_rebuild(vault, keys, parallel, checkpoint_dir) + + # 串行路径 + rebuilt_count = 0 + keys_iter = progress_bar(keys, desc="OCR rebuild") if progress_bar else keys + for key in keys_iter: + result = _rebuild_one_paper(vault, key) + if result["ok"]: + rebuilt_count += 1 + if checkpoint_dir: + _write_done_marker(checkpoint_dir, key) + return {"rebuild_count": rebuilt_count} +``` + +### 1.3 提取 `_rebuild_one_paper()`（module-level） + +**约束：module-level，Windows/spawn pickle-safe。不写成内部闭包。** + +**import 策略：保持重 imports 在阶段函数内部。** 只把 `logging`, `json`, `Path`, `datetime`, `concurrent.futures` 等 stdlib 提到模块顶部。fitz/figure/render/health 等不动。 + +```python +def _rebuild_one_paper(vault: Path, key: str) -> dict: + try: + ocr_root = pipeline_paths(vault)["ocr"] + artifacts = artifact_paths_for_root(ocr_root, key) + paper_root = artifacts.paper_root + if not paper_root.exists(): + return {"key": key, "ok": False, "error": "paper root missing", "rebuild_count": 0} + if not artifacts.blocks_raw.exists(): + return {"key": key, "ok": False, "error": "raw blocks missing", "rebuild_count": 0} + + raw_blocks = list(read_jsonl(artifacts.blocks_raw)) + meta = read_json(artifacts.meta_json) if artifacts.meta_json.exists() else {} + + # Phase 1–5 + ... + return {"key": key, "ok": True, "rebuild_count": 1} + except Exception as e: + logger.exception(f"Rebuild failed for {key}: {e}") + return {"key": key, "ok": False, "error": str(e), "rebuild_count": 0} +``` + +#### Phase 1: `_phase_read_and_span()` + +读 raw blocks → backfill 去重 → span backfill（有效跳过）→ PDF 行提取 → source metadata enrich + +Returns: `(raw_blocks, meta, source_pdf_path, source_meta, page_pdf_lines, span_meta_patch)` + +#### Phase 2: `_phase_structure_and_metadata()` + +`build_structured_blocks()` → `write_role_span_profiles()` → metadata resolve → write + +Returns: `(structured, doc_structure, resolved)` + +#### Phase 3: `_phase_figures_tables_objects()` + +**必须保护以下顺序（`apply_object_writebacks` 必须在 `write_figure_inventory` 前）：** + +```text +1. build_figure_inventory() +2. write_back_figure_roles() +3. residual_author_bio_pass() + post_ref_bio_cleanup() +4. synthesize_reader_figures() +5. build_table_inventory() +6. resolve_media_asset_conflicts() +7. attach_ownership_conflicts() +8. apply_object_writebacks() ← 这个顺序不能动 +9. write_figure_inventory() +10. write_back_table_roles() +11. write_table_inventory() +12. write_structured_blocks_jsonl() +13. extract_and_write_objects() +``` + +#### Phase 4: `_phase_render_and_health()` + +`render_fulltext_markdown()` → health → decision log → write + +#### Phase 5: `_phase_index_and_meta()` + +role index → structure tree → version flags → render outputs → meta.json + +### 1.4 `_run_parallel_rebuild()` + +```python +def _run_parallel_rebuild(vault, keys, workers, checkpoint_dir=None): + from concurrent.futures import ProcessPoolExecutor, as_completed + import sys, time, logging + logger = logging.getLogger(__name__) + + with ProcessPoolExecutor(max_workers=workers) as pool: + futures = {pool.submit(_rebuild_one_paper, vault, k): k for k in keys} + rebuilt_count = 0 + errors = [] + start = time.monotonic() + for i, f in enumerate(as_completed(futures), 1): + key = futures[f] + try: + result = f.result() + except Exception as exc: + logger.error(f"Parallel worker failed for {key}: {exc}") + errors.append({"key": key, "ok": False, "error": repr(exc), "rebuild_count": 0}) + continue + if result["ok"]: + rebuilt_count += 1 + if checkpoint_dir: + _write_done_marker(checkpoint_dir, key) + else: + errors.append(result) + print(f"[{i}/{len(keys)}] {key} — total {time.monotonic()-start:.1f}s", + file=sys.stderr, flush=True) + return {"rebuild_count": rebuilt_count, "errors": errors} +``` + +### 1.5 Checkpoint 辅助函数 + +```python +def _write_done_marker(checkpoint_dir: Path, key: str): + (checkpoint_dir / f".done.{key}").write_text("", encoding="utf-8") + +def _filter_completed_keys(keys: list[str], checkpoint_dir: Path) -> list[str]: + done = set() + if checkpoint_dir.exists(): + for f in checkpoint_dir.iterdir(): + if f.name.startswith(".done."): + done.add(f.name[len(".done."):]) + return [k for k in keys if k not in done] +``` + +### 1.6 CLI 接口 + +**`cli.py` parser：** + +```python +group = parser.add_mutually_exclusive_group() +group.add_argument("--parallel", type=int, nargs="?", const=4, default=4) +group.add_argument("--no-parallel", action="store_true") +``` + +**normalize：** + +```python +parallel_workers = 0 if args.no_parallel else max(1, int(args.parallel or 4)) +``` + +**`commands/ocr.py` dispatch：** + +```python +result = run_derived_rebuild_for_keys( + vault, keys, + checkpoint_dir=checkpoint_dir, + parallel=parallel_workers, + progress_bar=progress_bar, +) +``` + +--- + +## Wave 2: 资产裁剪并行 + +### 2.1 设计约束 + +**不能重写坐标转换。** 当前 `_crop_asset_from_pdf()` 已经正确处理了 OCR bbox → PDF rect 缩放和 rotation。Wave 2 只负责并发调度，不重写裁剪算法。 + +```text +Object tasks should call existing _crop_asset_from_pdf() without shared pdf_doc_provider. +Preserve OCR bbox -> PDF rect conversion and rotation handling. +``` + +### 2.2 重构方案 + +```text +1. Serial: _build_object_tasks() — 创建目录、清理旧文件、构建 task list、预分配稳定 ID +2. Task executors: + _write_figure_object_task(task) → _crop_asset_from_pdf(pdf_doc=None, pdf_doc_provider=None) + _write_table_object_task(task) → _crop_asset_from_pdf(pdf_doc=None, pdf_doc_provider=None) + _write_orphan_object_task(task) → _crop_asset_from_pdf(pdf_doc=None, pdf_doc_provider=None) +3. Each task opens its own fitz.Document (not shared) +4. Parallel: ThreadPoolExecutor over tasks +``` + +### 2.3 调用方式 + +```python +was_cropped = _crop_asset_from_pdf( + pdf_path=pdf_path, + page_num=page, + bbox=crop_bbox, + dst=asset_path, + page_width=page_width, + page_height=page_height, + page_cache_dir=threadsafe_page_cache_dir, + pdf_doc=None, + pdf_doc_provider=None, + rotation_deg=rotation_deg, +) +``` + +### 2.4 Worker 数 + +```python +max_workers = min(2, os.cpu_count() or 4) +``` + +### 2.5 Thread-safe page cache + +**接入点：** thread-safe helper 必须接入 `_crop_asset_from_pdf()` 内部，替换当前 page cache 分支。不能只写一个新函数但不被调用。 + +**Recommended：** 在 `_crop_asset_from_pdf()` 内部，将 page cache 写入替换为以下线程安全版本： + +```python +_RENDER_LOCKS: dict[tuple[str, int], threading.Lock] = {} +_RENDER_LOCKS_GUARD = threading.Lock() + +def _get_page_lock(cache_dir: Path, page_num: int) -> threading.Lock: + key = (str(cache_dir.resolve()), page_num) + with _RENDER_LOCKS_GUARD: + if key not in _RENDER_LOCKS: + _RENDER_LOCKS[key] = threading.Lock() + return _RENDER_LOCKS[key] +``` + +**key 用 `(cache_dir, page_num)` 而不仅是 `page_num`**，防止同一进程中不同 cache 目录的 page 冲突。 + +**备选：** 如果接入成本高，并行 task 传 `page_cache_dir=None` 禁用共享缓存，代价是每次从头渲染。 + +### 2.6 Orphan 处理 + +串行预分配 ID，再并行裁剪： + +```python +orphan_tasks = [] +for i, orphan in enumerate(unmatched_assets): + orphan_id = f"orphan_{page:03d}_{i:03d}" + orphan_tasks.append((orphan, orphans_asset_dir / f"{orphan_id}.jpg", + figures_render_dir / f"{orphan_id}.md")) +``` + +--- + +## Wave 3: Maintenance 显示封装 + +### 3.1 当前真实字段 + +当前 `OCRMaintenanceRow` 字段：`status`, `health`, `version`, `can_redo`, `can_rebuild`, `error_stage`, `error_summary`, `degraded_reasons`, ... + +当前 `_compute_display_fields()` 签名： + +```python +def _compute_display_fields( + status: str, health_overall: str, version: str, + can_redo: bool, can_rebuild: bool, + error_stage: str = "", error_summary: str = "", + degraded_reasons: list[str] | None = None, +) -> dict: +``` + +### 3.2 改写 `OCRMaintenanceRow` + +```python +@dataclass +class OCRMaintenanceRow: + ... # 现有字段不变 + + def __post_init__(self): + df = self.compute_display_fields( + status=self.status, health_overall=self.health, + version=self.version, can_redo=self.can_redo, + can_rebuild=self.can_rebuild, error_stage=self.error_stage, + error_summary=self.error_summary, + degraded_reasons=self.degraded_reasons, + ) + self.display_action = df["display_action"] + self.display_label = df["display_label"] + self.display_label_key = df.get("display_label_key", "") + self.display_reason = df["display_reason"] + self.display_reason_key = df.get("display_reason_key", "") + self.display_group = df["display_group"] + self.display_severity = df["display_severity"] + self.visible_in_maintenance = df["visible_in_maintenance"] + self.show_in_base = df["show_in_base"] + + @staticmethod + def compute_display_fields(...) -> dict: + """原 _compute_display_fields() 逻辑移入此处。""" + ... +``` + +### 3.3 修改 `compute_maintenance_manifest()` + +```python +df = OCRMaintenanceRow.compute_display_fields( + status=status, health_overall=health_overall, version=version, + can_redo=can_redo, can_rebuild=can_rebuild, + error_stage=_error_stage(meta), error_summary=_error_summary(meta), + degraded_reasons=health.get("degraded_reasons", []) or [], +) +``` + +### 3.4 删除或保留外部函数 + +- **推荐：** 保留兼容 wrapper `_compute_display_fields = OCRMaintenanceRow.compute_display_fields` +- 或：直接删除，同步更新 `tests/test_ocr_maintenance.py` 中所有 import 和调用 + +### 3.5 测试更新 + +```python +# tests/test_ocr_maintenance.py +# 原: from paperforge.worker.ocr_maintenance import _compute_display_fields +# 改: from paperforge.worker.ocr_maintenance import OCRMaintenanceRow +# OCRMaintenanceRow.compute_display_fields(...) +``` + +--- + +## Wave 4: Embed 人类进度 + +### 4.1 当前状态 + +已有 stdout 机器协议： + +```python +print(f"EMBED_START:{total}", flush=True) +for i, entry in enumerate(papers): + ... + print(f"EMBED_PROGRESS:{i}:{total}:{key}:{count}", flush=True) +print("EMBED_DONE", flush=True) +``` + +### 4.2 叠加 tqdm + +**只包 iterable，不改变现有 `i += 1` 语义。** 当前 `i` 只在真正 embed 的 paper 递增（跳过 resume 已存在的）。 + +```python +from paperforge.worker._progress import progress_bar + +use_tqdm = not getattr(args, "json", False) +papers_iter = progress_bar(papers, desc="Embedding", disable=not use_tqdm) + +print(f"EMBED_START:{total}", flush=True) +i = 0 +for entry in papers_iter: + key = entry["key"] + if resume and _already_embedded(...): + continue + chunks = chunker.chunk(entry["fulltext"]) + count = embed_paper(vault, key, chunks) + i += 1 + print(f"EMBED_PROGRESS:{i}:{total}:{key}:{count}", flush=True) +print("EMBED_DONE", flush=True) +``` + +### 4.3 不变 + +- `embedding/builder.py` 不动 +- `EMBED_*` stdout 协议完全保留 +- 插件端不受影响 + +--- + +## Wave 5: CLI 文档补齐 + +### 5.1 从 `build_parser()` 逐项核对 + +运行 `python paperforge/cli.py --help` 获取完整命令树，逐项填入。 + +| 命令 | 用途 | 示例 | +|------|------|------| +| `ocr` | 运行 OCR | `paperforge ocr --key KEY` | +| `ocr rebuild` | 重新生成 derived 产物 | `paperforge ocr rebuild [KEY...] [--all] [--status done_degraded] [--parallel 4]` | +| `sync` | 同步 Zotero → Obsidian | `paperforge sync` | +| `content-discovery` | 全文内容检索 | `paperforge content-discovery ` | +| `scoped-fetch` | 按查询获取全文块 | `paperforge scoped-fetch [--limit 5]` | +| `paper-navigation` | 论文结构导航 | `paperforge paper-navigation ` | +| `paper-lookup` | 模糊查找论文 | `paperforge paper-lookup ` | +| `reading-log` | 阅读日志操作 | `paperforge reading-log --lookup/--write/--render/--validate/--import` | +| `project-log` | 项目日志操作 | `paperforge project-log --write/--list/--render/--project/--payload` | +| `paper-status` | 论文状态摘要 | `paperforge paper-status ` | +| `agent-context` | Agent 上下文 | `paperforge agent-context ` | +| `runtime-health` | 运行时健康检查 | `paperforge runtime-health` | +| `embed build` | 构建向量索引 | `paperforge embed build [--force] [--resume]` | +| `embed status` | 向量索引状态 | `paperforge embed status` | +| `embed stop` | 停止向量构建 | `paperforge embed stop` | +| `status` | 全局状态 | `paperforge status` | +| `search` | 检索 | `paperforge search ` | + +**所有示例必须从 `build_parser()` 实际参数核对，不要手写猜测。** + +--- + +## 测试计划 + +| Wave | 测试 | 覆盖 | +|------|------|------| +| 1 | `_rebuild_one_paper()` | 成功、paper root 缺失、raw blocks 缺失 | +| 1 | `_run_parallel_rebuild()` | 3 papers → count=3；1/3 fail → count=2+1 error | +| 1 | `_filter_completed_keys()` | 有标记 → 过滤 | +| 1 | 已有 `test_ocr_rebuild` | 21 个全部通过（串行不退化） | +| 2 | `_write_figure_object_task()` | 坐标正确性（复用 `_crop_asset_from_pdf`） | +| 2 | `_get_page_lock()` 线程安全 | 并发同一页 → 不竞态 | +| 2 | 整体并行 | 3 figures → 全部产出 | +| 3 | 已有 `test_ocr_maintenance` | 28 个全部通过（行为不退化） | +| 3 | `OCRMaintenanceRow.compute_display_fields()` | 替换旧 `_compute_display_fields` import | +| 4 | stdout 抓取 | `EMBED_START / EMBED_PROGRESS / EMBED_DONE` 仍然存在 | +| 5 | `test_command_docs.py` | 更新（如果存在） | + +Wave 1+2 并行 → review → Wave 3+4+5 并行 → review → 全量测试 → 合入 master。 + +重点审核：Wave 2 坐标安全、Wave 3 字段匹配、Wave 4 机器协议保留。 diff --git a/docs/superpowers/specs/2026-07-05-rebuild-speed-and-batch-improvements.md b/docs/superpowers/specs/2026-07-05-rebuild-speed-and-batch-improvements.md new file mode 100644 index 00000000..364a256d --- /dev/null +++ b/docs/superpowers/specs/2026-07-05-rebuild-speed-and-batch-improvements.md @@ -0,0 +1,453 @@ +# Rebuild 加速与批量改进设计 + +> 设计讨论：2026-07-05 +> 状态：设计稿（已审核，待实现） +> 审核结论：7.6/10 — A1/P0#1 先做，P0#2 和 embed 进度已按审核意见重写 + +## 设计目标 + +本设计解决一个核心问题和三个附带问题： + +**核心问题**：rebuild 全串行，单线程逐个跑，CPU 空闲、I/O 空闲。大规模批量 rebuild（50+ 篇）耗时过长。 + +**附带问题**： +1. `run_derived_rebuild_for_keys()` 是 300 行内联编排，不可测试、不可插桩 +2. rebuild 缺少并行完成进度；embed 已有 stdout 机器进度协议（`EMBED_START/EMBED_PROGRESS/EMBED_DONE`），但缺少面向终端用户的 human-readable tqdm +3. CLI 文档覆盖不全，用户不知道有哪些命令可用 +4. Maintenance 显示逻辑在数据类外散布 + +## 范围 + +### 目标内 + +1. **跨论文并行**（P0#1）：`ProcessPoolExecutor` 驱动，默认 4 进程 +2. **资产裁剪并行**（P0#2）：`ThreadPoolExecutor` 驱动，默认 2 线程/进程 +3. **Rebuild 阶段化**（A1）：把 300 行内联编排拆成 `_rebuild_one_paper()` + 5 个阶段函数 +4. **Rebuild 进度**：并行模式 `print [i/N] key — t秒` 到 stderr；串行保持 tqdm +5. **Embed 进度**：`commands/embed.py` 中在 stderr 叠加 tqdm，保留 `EMBED_START/EMBED_PROGRESS/EMBED_DONE` 机器协议 +6. **Maintenance 显示封装**（A4）：`compute_display_fields()` 转为静态方法，`__post_init__` 和 `compute_maintenance_manifest()` 共享 +7. **CLI 文档补齐**（B1）：`docs/COMMANDS.md` 覆盖全部 30+ 命令，名称从 `build_parser()` 逐项核对 + +### 非目标 + +- Figure passes 内部并行（pass 间依赖，收益不高，scope 太大） +- JSONL 增量写入（ponytail 标记但实际不是瓶颈） +- Section-aware chunking（B3，shipping 后做） +- Gateway 连接管理器（A2，代码味而非性能问题） +- CLI 导入体操（A3，值不值得做不确定） +- LanceBackend 实现/删除（A5，你说了放放） + +## 设计决策 + +### 1. 跨论文并行（P0#1） + +#### 1.1 Worker 数 + +默认为 4，不随 `os.cpu_count()` 变化。用户通过 `--parallel N` 自定义。 + +理由：rebuild 是 CPU + I/O 混合负载，4 进程 SSD I/O 争用可控。 + +#### 1.2 CLI 接口 + +```python +group = parser.add_mutually_exclusive_group() +group.add_argument("--parallel", type=int, nargs="?", const=4, default=4) +group.add_argument("--no-parallel", action="store_true") +``` + +**必须 normalize**：argparse 不会自动把 `--no-parallel` 和 `--parallel` 关联。如果只看 `args.parallel`，`--no-parallel` 时它仍然是默认值 `4`。 + +```python +parallel_workers = 0 if args.no_parallel else max(1, int(args.parallel or 4)) +run_derived_rebuild_for_keys(..., parallel=parallel_workers) +``` + +| CLI | `args.parallel` | `args.no_parallel` | 实际 workers | +|-----|-----------------|-------------------|-------------| +| 默认 | 4 | False | 4 | +| `--parallel` | 4 | False | 4 | +| `--parallel 2` | 2 | False | 2 | +| `--no-parallel` | 4 | True | 0 / serial | + +`parallel=0` 或 `parallel=False` 走串行路径。 + except Exception as e: + # 注意：ocr_rebuild.py 当前没有 logger，需要补 + import logging + logger = logging.getLogger(__name__) + logger.exception(f"Rebuild failed for {key}: {e}") + return {"key": key, "ok": False, "error": str(e), "rebuild_count": 0} + + +def _run_parallel_rebuild(vault, keys, workers, ...) -> dict: + from concurrent.futures import ProcessPoolExecutor, as_completed + import sys + import time + import logging + + logger = logging.getLogger(__name__) + # checkpoint：主进程先扫描 .done.* 过滤已完成 keys + cp_dir = checkpoint_dir + if cp_dir: + cp_dir.mkdir(parents=True, exist_ok=True) + keys = _filter_completed_keys(keys, cp_dir) + + with ProcessPoolExecutor(max_workers=workers) as pool: + futures = {pool.submit(_rebuild_one_paper, vault, k): k for k in keys} + rebuilt_count = 0 + errors = [] + start = time.monotonic() + for i, f in enumerate(as_completed(futures), 1): + key = futures[f] + try: + result = f.result() # 业务异常已被 worker 内部捕获 + except Exception as exc: + # 来自子进程 crash / unpicklable / OOM / BrokenProcessPool / KeyboardInterrupt + logger.error(f"Parallel rebuild worker failed for {key}: {exc}") + errors.append({"key": key, "ok": False, "error": repr(exc), "rebuild_count": 0}) + continue + + if result["ok"]: + rebuilt_count += 1 + if cp_dir: + _write_done_marker(cp_dir, key) + else: + errors.append(result) + + total_elapsed = time.monotonic() - start + print(f"[{i}/{len(keys)}] {key} — total {total_elapsed:.1f}s", + file=sys.stderr, flush=True) + + return {"rebuild_count": rebuilt_count, "errors": errors} +``` + +#### 1.4 Checkpoint 适配 + +并行时用 **独立目录**，不重用旧 JSON checkpoint 文件路径。 + +```python +# 旧：vault / "System" / "PaperForge" / ".ocr_rebuild_checkpoint.json" +# 新：vault / "System" / "PaperForge" / ".ocr_rebuild_checkpoint" / (目录) + +# 主进程写标记：worker 只返回 result，主进程确认 ok 后写 .done. +def _write_done_marker(checkpoint_dir: Path, key: str): + (checkpoint_dir / f".done.{key}").write_text("", encoding="utf-8") + +# 主进程侧：扫描过滤 +def _filter_completed_keys(keys: list[str], checkpoint_dir: Path) -> list[str]: + done = set() + if checkpoint_dir.exists(): + for f in checkpoint_dir.iterdir(): + if f.name.startswith(".done."): + done.add(f.name[len(".done."):]) + return [k for k in keys if k not in done] +``` + +CLI resume 逻辑同步改为读 `.done.*` 文件，不再读旧 JSON。 + +#### 1.5 日志 + +补 `logging.getLogger(__name__)`。 + +```python +import logging +logger = logging.getLogger(__name__) +``` + +并行模式下子进程 stderr 通过 ProcessPoolExecutor 管道自动合并到主进程，无需特殊处理。 + +#### 1.6 进度输出到 stderr + +```python +print(f"[{i}/{len(keys)}] {key} — total {total_elapsed:.1f}s", + file=sys.stderr, flush=True) +``` + +避免污染 `--json` 输出。 + +### 2. 资产裁剪并行（P0#2） + +审核意见：**不能共享 fitz.Document 并宣称线程安全。** 本方案已重写。 + +#### 2.1 Worker 数 + +`min(2, os.cpu_count())`。不暴露 CLI。 + +#### 2.2 核心约束 + +1. **不能共享 fitz.Document**：避免多线程并发访问同一个 PDF 文件的竞态 +2. **页面缓存写入用临时文件 + atomic replace**：避免读到半成品缓存 +3. **Orphan ID 串行预分配**：在提交 futures 之前确定所有输出路径 + +#### 2.3 实现方式 + +```python +def extract_and_write_objects( + pdf_path, figure_inventory, table_inventory, + asset_root, render_root, ..., +): + import os + from concurrent.futures import ThreadPoolExecutor, as_completed + + max_workers = min(2, os.cpu_count() or 4) + + # 串行阶段：目录创建、清理旧文件 → 安全，不做线程中 + _prepare_directories(...) + + # 串行预构建 task list，确定所有稳定 ID + tasks: list[Callable] = [] + + for match in figure_inventory.get("matched_figures", []): + fig_id = _resolve_figure_id(match) # 稳定 ID + asset_path = figures_asset_dir / f"{fig_id}.jpg" + render_path = figures_render_dir / f"{fig_id}.md" + tasks.append( + lambda m=match: _crop_one_figure(m, pdf_path, asset_path, render_path, ...) + ) + + for table in table_inventory.get("tables", []): + tbl_id = _resolve_table_id(table) # 稳定 ID + asset_path = tables_asset_dir / f"{tbl_id}.jpg" + render_path = tables_render_dir / f"{tbl_id}.md" + tasks.append( + lambda t=table: _crop_one_table(t, pdf_path, asset_path, render_path, ...) + ) + + # 并行裁剪：每个 task 自己打开 fitz.Document + with ThreadPoolExecutor(max_workers=max_workers) as pool: + futures = [pool.submit(t) for t in tasks] + for f in as_completed(futures): + f.result() # 让异常透出 +``` + +#### 2.4 页面缓存 atomic write + +```python +_RENDER_LOCKS: dict[int, threading.Lock] = {} + +def _render_page_cached(doc_path: Path, page_num: int, cache_dir: Path) -> Path | None: + """线程安全页面渲染：每个 page 独立锁 + 临时文件 atomic replace。""" + lock = _RENDER_LOCKS.setdefault(page_num, threading.Lock()) + cache_path = cache_dir / f"page_{page_num:03d}.jpg" + + if cache_path.exists(): + return cache_path + + with lock: + # double-check: 另一个线程可能已经写完 + if cache_path.exists(): + return cache_path + + import fitz + doc = fitz.open(str(doc_path)) + try: + page = doc[page_num - 1] + pix = page.get_pixmap(dpi=200) + # 写临时文件 → atomic replace + tmp = cache_path.with_suffix(".tmp.jpg") + pix.save(str(tmp)) + tmp.replace(cache_path) + return cache_path + finally: + doc.close() +``` + +#### 2.5 每个 crop task 打开自己的 PDF + +```python +def _crop_one_figure(match, pdf_path, asset_path, render_path, ...): + import fitz + + doc = fitz.open(str(pdf_path)) + try: + page = doc[page_num - 1] + pix = page.get_pixmap(clip=bbox, dpi=200) + tmp = asset_path.with_suffix(".tmp.jpg") + pix.save(str(tmp)) + tmp.replace(asset_path) + finally: + doc.close() + + md = render_figure_object_markdown(...) + md_tmp = render_path.with_suffix(".tmp.md") + md_tmp.write_text(md) + md_tmp.replace(render_path) + + return True +``` + +### 3. Rebuild 阶段化（A1） + +#### 3.1 拆函数 + +`_rebuild_one_paper()` 内部拆成 5 个阶段函数： + +```python +# Phase 1: Span +meta = _phase_span(raw_blocks, artifacts, meta, vault, key, source_pdf_path) +page_pdf_lines = extract_pdf_lines_normalized(source_pdf_path) + +# Phase 2: Structure + Metadata +structured, doc_structure, resolved = _phase_structure( + raw_blocks, source_meta, artifacts, vault, key +) + +# Phase 3: Figures + Tables + Objects +figure_inventory, table_inventory = _phase_figures_tables( + structured, page_pdf_lines, source_pdf_path, artifacts +) + +# Phase 4: Render + Health +markdown = _phase_render(structured, resolved, figure_inventory, table_inventory, + doc_structure, reader_payload, meta) + +# Phase 5: Index + Meta +_phase_index_and_meta(structured, resolved, markdown, meta, meta_patches, + artifacts, paper_root, vault, key) +``` + +#### 3.2 必须保护的顺序 + +**审核意见：不能重排 figure/table 管道顺序。** 以下顺序是代码注释明确要求的： + +```text +1. build_figure_inventory() +2. write_back_figure_roles() +3. residual_author_bio_pass() + post_ref_bio_cleanup() +4. synthesize_reader_figures() +5. build_table_inventory() +6. resolve_media_asset_conflicts() +7. attach_ownership_conflicts() +8. apply_object_writebacks() ← 必须在 write_figure_inventory 之前 +9. write_figure_inventory() +10. write_back_table_roles() +11. write_table_inventory() +12. write_structured_blocks_jsonl() +13. extract_and_write_objects() +``` + +### 4. Rebuild 进度 + +- 串行模式：现有 `progress_bar`（tqdm）不变 +- 并行模式：`as_completed` + `print(f"[{i}/{n}] {key} — total {elapsed:.1f}s", file=sys.stderr)` +- `--json` 时禁用人类进度 + +### 5. Embed 进度 + +**审核意见：原来 spec 写错位置。** 当前 embed build 循环不在 `embedding/builder.py`，在 `commands/embed.py`。 + +当前已有机器进度协议（`EMBED_START / EMBED_PROGRESS / EMBED_DONE`）到 stdout，插件端依赖它们。 + +正确方案： + +```python +# commands/embed.py — 在现有机器协议之上叠加人类进度 + +from paperforge.worker._progress import progress_bar + +def run_embed(vault, papers, ...): + total = len(papers) + print(f"EMBED_START:{total}", flush=True) # 保留机器协议 + + for i, paper in enumerate(progress_bar(papers, desc="Embedding", disable=args.json)): + key = paper["key"] + chunks = chunker.chunk(paper["fulltext"]) + count = embed_paper(vault, key, chunks) + print(f"EMBED_PROGRESS:{i+1}:{total}:{key}:{count}", flush=True) # 保留机器协议 + + print("EMBED_DONE", flush=True) # 保留机器协议 +``` + +改动范围： +- `commands/embed.py`：加 tqdm 到 stderr，受 `--json` 开关控制 +- `embedding/builder.py`：不动 +- 机器输出协议：不变 + +### 6. Maintenance 显示封装（A4） + +**审核意见：不能直接删除 `_compute_display_fields()`，`compute_maintenance_manifest()` 也在用它。** + +正确方案： + +```python +@dataclass +class OCRMaintenanceRow: + ... + + def __post_init__(self): + df = self.compute_display_fields(...) + self.display_action = df["display_action"] + self.display_label = df["display_label"] + self.display_severity = df["display_severity"] + self.display_group = df["display_group"] + self.visible_in_maintenance = df["visible_in_maintenance"] + self.show_in_base = df["show_in_base"] + + @staticmethod + def compute_display_fields(...) -> dict: + """和 compute_maintenance_manifest() 共享的静态方法。""" + ... # 现有 10-clause if-elif 逻辑 +``` + +改动： +- `compute_display_fields()` 转为 `@staticmethod`，签名不变 +- `OCRMaintenanceRow.__post_init__()` 调用 `compute_display_fields()` +- `compute_maintenance_manifest()` 也调用 `OCRMaintenanceRow.compute_display_fields()` +- 外部独立函数 `_compute_display_fields()` 删除 + +说明：本次只封装 display fields。`fulltext_drift_state` 等字段保持构造后赋值，不在此次改动范围内。 + +### 7. CLI 文档补齐（B1） + +**审核意见：命令名必须从实际 `build_parser()` 逐项核对。** 以下是当前实际命令名（非猜测）： + +| 命令 | 实际 CLI 示例 | 文档状态 | +|------|-------------|---------| +| 运行 OCR | `paperforge ocr --key KEY` | 已有 | +| 重建 derived | `paperforge ocr rebuild [KEY...] [--all] [--status ...] [--parallel N]` | 需补 | +| 同步 | `paperforge sync` | 已有 | +| 状态 | `paperforge status` | 已有 | +| 检索 | `paperforge search ` | 已有 | +| 内容发现 | `paperforge content-discovery ` | 需补 | +| 范围获取 | `paperforge scoped-fetch ` | 需补 | +| 论文导航 | `paperforge paper-navigation ` | 需补 | +| 论文查找 | `paperforge paper-lookup ` | 需补 | +| 阅读日志 | `paperforge reading-log --lookup/--write/--render/--validate/--import` | 需补 | +| 项目日志 | `paperforge project-log --write/--list/--render/--project/--payload` | 需补 | +| 论文状态 | `paperforge paper-status ` | 需补 | +| Agent 上下文 | `paperforge agent-context ` | 需补 | +| 运行时健康 | `paperforge runtime-health` | 需补 | +| 嵌入构建 | `paperforge embed build [--force] [--resume]` | 需补 | +| 嵌入状态 | `paperforge embed status` | 需补 | +| 停止嵌入 | `paperforge embed stop` | 需补 | +| 说明 | 以上 CLI 示例必须从 `build_parser()` 逐项核对，不要手写猜测命令名 | — | + +文档格式：每命令一行用途 + 实际 CLI 示例（从 parser real_defaults 验证）。 + +## 非功能约束 + +| 属性 | 目标 | +|------|------| +| 向后兼容 | 串行路径行为 0 变化 | +| 并行下性能 | 40 篇 paper ≥ 2x 串行（**目标，非保证**；无 profiler 数据） | +| 内存 | 默认 4 workers 应在笔记本内存可承受范围内；后续用 40-paper benchmark 实测后考虑提高默认值 | +| 错误隔离 | 单篇失败不阻止其他论文；主进程 try/except f.result() 兜底 | +| 测试 | 每阶段函数独立测试 + 串行/并行路径 smoke test | + +## 设计决策明细 + +| # | 决策 | 理由 | +|---|------|------| +| D1 | `max_workers=4` 固定值 | I/O 争用受限，非纯 CPU 密集 | +| D2 | 并行默认开 | 用户预期 "rebuild 就是快的" | +| D3 | CLI `--parallel`/`--no-parallel` 互斥组 | 语义最清晰 | +| D4 | `_rebuild_one_paper()` 返回 dict，业务异常已内部捕获 | 隔离异常传递 | +| D5 | **主进程仍 try/except f.result()** | 子进程 crash/unpicklable/OOM 无法内部拦截 | +| D6 | 独立 checkpoint_dir + `.done.` 标记 | 零库依赖，零竞态 | +| D7 | 每个 crop task 打开自己的 fitz.Document | **审核后修正**：避免共享 PDF 的线程安全风险 | +| D8 | 页面缓存用 per-page 锁 + atomic replace | **审核后修正**：避免读到半成品缓存 | +| D9 | orphan/table ID 串行预分配后再并行 | 避免并行递增计数器 | +| D10 | rebuild 和 embed 进度各自独立策略 | rebuild 用 as_completed 打印；embed 保持机器协议 + tqdm | +| D11 | `compute_display_fields()` 转为 `@staticmethod` | 同时被 `__post_init__` 和 `compute_maintenance_manifest()` 调用 | +| D12 | CLI 文档从 `build_parser()` 逐项核对 | 避免手写错误命令名 | +| D13 | embed 进度改 `commands/embed.py`，不动 `embedding/builder.py` | **审核后修正**：真正 build loop 在 commands 层 | +| D14 | 保留 `EMBED_START/EMBED_PROGRESS/EMBED_DONE` stdout 协议 | 插件端依赖此协议 | diff --git a/paperforge/cli.py b/paperforge/cli.py index 25bef9cc..3c15bd1f 100644 --- a/paperforge/cli.py +++ b/paperforge/cli.py @@ -249,6 +249,11 @@ def build_parser() -> argparse.ArgumentParser: rebuild_parser.add_argument("--status", metavar="STATUS", help="Filter by OCR status (done, done_degraded, failed)") rebuild_parser.add_argument("--dry-run", action="store_true", help="List papers that would be rebuilt without executing") rebuild_parser.add_argument("--resume", action="store_true", help="Skip papers already in checkpoint") + parallel_group = rebuild_parser.add_mutually_exclusive_group() + parallel_group.add_argument("--parallel", type=int, nargs="?", const=4, default=4, metavar="N", + help="Number of parallel workers (default: 4)") + parallel_group.add_argument("--no-parallel", dest="parallel", action="store_const", const=0, + help="Disable parallel processing (serial)") # context (Phase 26: traceable AI context packs) p_context = sub.add_parser("context", help="Generate traceable AI context pack for paper(s)") diff --git a/paperforge/commands/embed.py b/paperforge/commands/embed.py index 1708496f..1c6fd8ac 100644 --- a/paperforge/commands/embed.py +++ b/paperforge/commands/embed.py @@ -20,6 +20,7 @@ from paperforge.embedding.preflight import _preflight_check from paperforge.memory.chunker import chunk_fulltext from paperforge.memory.state_snapshot import write_vector_runtime from paperforge.worker.asset_index import read_index +from paperforge.worker._progress import progress_bar def run(args: argparse.Namespace) -> int: @@ -183,7 +184,8 @@ def run(args: argparse.Namespace) -> int: ) i = 0 - for entry in done_papers: + papers_iter = progress_bar(done_papers, desc="Embedding", disable=args.json) + for entry in papers_iter: key = entry.get("zotero_key") fulltext_rel = entry.get("fulltext_path", "") if not fulltext_rel: diff --git a/paperforge/commands/ocr.py b/paperforge/commands/ocr.py index 4cb8e81e..c8f0e2dd 100644 --- a/paperforge/commands/ocr.py +++ b/paperforge/commands/ocr.py @@ -323,6 +323,7 @@ def _run_ocr_rebuild( status_filter: str | None = None, dry_run: bool = False, resume: bool = False, + parallel_workers: int = 4, ) -> int: """Rebuild OCR-derived artifacts from existing raw blocks.""" from paperforge.worker.ocr_maintenance import collect_maintenance_rows @@ -345,11 +346,10 @@ def _run_ocr_rebuild( print("No papers matched for rebuild.") return 0 - # Resume: skip keys already in checkpoint - cp = vault / "System" / "PaperForge" / ".ocr_rebuild_checkpoint.json" - if resume and cp.exists(): - import json - done = set(json.loads(cp.read_text(encoding="utf-8"))) + # Resume: skip keys already in checkpoint (.done.* markers) + cp_dir = vault / "System" / "PaperForge" / ".ocr_rebuild_checkpoint" + if resume and cp_dir.exists(): + done = {p.name.removeprefix(".done.") for p in cp_dir.glob(".done.*")} skipped = [k for k in keys if k in done] keys = [k for k in keys if k not in done] if skipped: @@ -366,7 +366,12 @@ def _run_ocr_rebuild( return 0 from paperforge.worker._progress import progress_bar - result = run_derived_rebuild_for_keys(vault, keys, progress_bar=progress_bar, checkpoint=cp) + result = run_derived_rebuild_for_keys( + vault, keys, + progress_bar=progress_bar, + checkpoint_dir=cp_dir if resume else None, + parallel=parallel_workers, + ) count = result.get("rebuild_count", 0) print(f"Done. Rebuilt {count} paper(s).") return 0 @@ -417,6 +422,7 @@ def run(args: argparse.Namespace) -> int: ) if ocr_action == "rebuild": + parallel_workers = 0 if getattr(args, "no_parallel", False) else max(1, int(getattr(args, "parallel", 4) or 4)) return _run_ocr_rebuild( vault, keys=getattr(args, "keys", None) or None, @@ -424,6 +430,7 @@ def run(args: argparse.Namespace) -> int: status_filter=getattr(args, "status", None), dry_run=getattr(args, "dry_run", False), resume=getattr(args, "resume", False), + parallel_workers=parallel_workers, ) if key: diff --git a/paperforge/worker/ocr_maintenance.py b/paperforge/worker/ocr_maintenance.py index eeba8e2b..e2f3e266 100644 --- a/paperforge/worker/ocr_maintenance.py +++ b/paperforge/worker/ocr_maintenance.py @@ -46,6 +46,97 @@ class OCRMaintenanceRow: fulltext_drift_state: str = "UNKNOWN" fulltext_drift_reason: str = "" show_in_base: bool = True + def __post_init__(self) -> None: + df = self.compute_display_fields( + status=self.status, health_overall=self.health, + version=self.version, can_redo=self.can_redo, + can_rebuild=self.can_rebuild, error_stage=self.error_stage, + error_summary=self.error_summary, + degraded_reasons=self.degraded_reasons, + ) + for k, v in df.items(): + setattr(self, k, v) + + @staticmethod + def compute_display_fields( + status: str, + health_overall: str, + version: str, + can_redo: bool, + can_rebuild: bool, + error_stage: str = "", + error_summary: str = "", + degraded_reasons: list[str] | None = None, + ) -> dict: + """Map raw OCR state to display fields for the maintenance tab.""" + is_degraded = health_overall in ("yellow", "red") or status == "done_degraded" + + if status in ("pending",): + return dict(display_action="none", display_label="等待处理", display_label_key="", + display_reason="", display_reason_key="", + display_group="hidden", display_severity="normal", + visible_in_maintenance=False, show_in_base=True) + if status in ("running", "queued", "processing"): + return dict(display_action="none", display_label="处理中", display_label_key="", + display_reason="", display_reason_key="", + display_group="hidden", display_severity="normal", + visible_in_maintenance=False, show_in_base=True) + if status in ("failed", "error", "fatal_error", "done_incomplete", "retryable_error") and can_redo: + return dict(display_action="retry_ocr", display_label="重试 OCR", + display_label_key="maintenance_action_retry_ocr", + display_reason="上次处理未完成，可以重新尝试", + display_reason_key="maintenance_reason_retry", + display_group="retry", display_severity="actionable", + visible_in_maintenance=True, show_in_base=True) + if version == "v1" and can_redo: + return dict(display_action="upgrade_legacy", display_label="升级旧结果", + display_label_key="maintenance_action_upgrade_legacy", + display_reason="旧版本结果仍然可用，升级后可获得更好的章节、图表和问答效果", + display_reason_key="maintenance_reason_legacy", + display_group="legacy_optional", display_severity="optional", + visible_in_maintenance=True, show_in_base=True) + if is_degraded and can_rebuild: + return dict(display_action="rebuild_result", display_label="重建结果", + display_label_key="", + display_reason="已有OCR数据，可重建获得更稳定的结果", + display_reason_key="", + display_group="rebuild", display_severity="actionable", + visible_in_maintenance=True, show_in_base=True) + if is_degraded and not can_rebuild and can_redo: + return dict(display_action="retry_ocr", display_label="重试 OCR", + display_label_key="", + display_reason="降级结果无法重建，可重新OCR", + display_reason_key="", + display_group="retry", display_severity="actionable", + visible_in_maintenance=True, show_in_base=True) + if status == "nopdf": + return dict(display_action="add_pdf", display_label="补充 PDF", + display_label_key="", + display_reason="请去 Zotero 添加 PDF 文件", + display_reason_key="", + display_group="external_action", display_severity="external", + visible_in_maintenance=False, show_in_base=True) + if status == "blocked": + return dict(display_action="configure_ocr", display_label="配置 OCR", + display_label_key="", + display_reason="请配置 PaddleOCR API Token", + display_reason_key="", + display_group="external_action", display_severity="external", + visible_in_maintenance=False, show_in_base=True) + if status == "done" and not is_degraded: + return dict(display_action="none", display_label="已完成", display_label_key="", + display_reason="", display_reason_key="", + display_group="hidden", display_severity="normal", + visible_in_maintenance=False, show_in_base=True) + if not can_redo and not can_rebuild: + return dict(display_action="none", display_label="已完成", display_label_key="", + display_reason="", display_reason_key="", + display_group="hidden", display_severity="normal", + visible_in_maintenance=False, show_in_base=False) + return dict(display_action="none", display_label="已完成", display_label_key="", + display_reason="", display_reason_key="", + display_group="hidden", display_severity="normal", + visible_in_maintenance=False, show_in_base=True) def to_dict(self) -> dict: return { @@ -183,85 +274,7 @@ def _recommended_action(meta: dict, has_raw: bool, has_source_meta: bool) -> str return "" -def _compute_display_fields( - status: str, - health_overall: str, - version: str, - can_redo: bool, - can_rebuild: bool, - error_stage: str = "", - error_summary: str = "", - degraded_reasons: list[str] | None = None, -) -> dict: - """Map raw OCR state to display fields for the maintenance tab.""" - is_degraded = health_overall in ("yellow", "red") or status == "done_degraded" - - if status in ("pending",): - return dict(display_action="none", display_label="等待处理", display_label_key="", - display_reason="", display_reason_key="", - display_group="hidden", display_severity="normal", - visible_in_maintenance=False, show_in_base=True) - if status in ("running", "queued", "processing"): - return dict(display_action="none", display_label="处理中", display_label_key="", - display_reason="", display_reason_key="", - display_group="hidden", display_severity="normal", - visible_in_maintenance=False, show_in_base=True) - if status in ("failed", "error", "fatal_error", "done_incomplete", "retryable_error") and can_redo: - return dict(display_action="retry_ocr", display_label="重试 OCR", - display_label_key="maintenance_action_retry_ocr", - display_reason="上次处理未完成，可以重新尝试", - display_reason_key="maintenance_reason_retry", - display_group="retry", display_severity="actionable", - visible_in_maintenance=True, show_in_base=True) - if version == "v1" and can_redo: - return dict(display_action="upgrade_legacy", display_label="升级旧结果", - display_label_key="maintenance_action_upgrade_legacy", - display_reason="旧版本结果仍然可用，升级后可获得更好的章节、图表和问答效果", - display_reason_key="maintenance_reason_legacy", - display_group="legacy_optional", display_severity="optional", - visible_in_maintenance=True, show_in_base=True) - if is_degraded and can_rebuild: - return dict(display_action="rebuild_result", display_label="重建结果", - display_label_key="", - display_reason="已有OCR数据，可重建获得更稳定的结果", - display_reason_key="", - display_group="rebuild", display_severity="actionable", - visible_in_maintenance=True, show_in_base=True) - if is_degraded and not can_rebuild and can_redo: - return dict(display_action="retry_ocr", display_label="重试 OCR", - display_label_key="", - display_reason="降级结果无法重建，可重新OCR", - display_reason_key="", - display_group="retry", display_severity="actionable", - visible_in_maintenance=True, show_in_base=True) - if status == "nopdf": - return dict(display_action="add_pdf", display_label="补充 PDF", - display_label_key="", - display_reason="请去 Zotero 添加 PDF 文件", - display_reason_key="", - display_group="external_action", display_severity="external", - visible_in_maintenance=False, show_in_base=True) - if status == "blocked": - return dict(display_action="configure_ocr", display_label="配置 OCR", - display_label_key="", - display_reason="请配置 PaddleOCR API Token", - display_reason_key="", - display_group="external_action", display_severity="external", - visible_in_maintenance=False, show_in_base=True) - if status == "done" and not is_degraded: - return dict(display_action="none", display_label="已完成", display_label_key="", - display_reason="", display_reason_key="", - display_group="hidden", display_severity="normal", - visible_in_maintenance=False, show_in_base=True) - if not can_redo and not can_rebuild: - return dict(display_action="none", display_label="已完成", display_label_key="", - display_reason="", display_reason_key="", - display_group="hidden", display_severity="normal", - visible_in_maintenance=False, show_in_base=False) - return dict(display_action="none", display_label="已完成", display_label_key="", - display_reason="", display_reason_key="", - display_group="hidden", display_severity="normal", - visible_in_maintenance=False, show_in_base=True) +_compute_display_fields = OCRMaintenanceRow.compute_display_fields def compute_maintenance_manifest(vault: Path) -> dict[str, str]: @@ -299,7 +312,7 @@ def compute_maintenance_manifest(vault: Path) -> dict[str, str]: rec_action = _recommended_action(meta, has_raw, has_source_meta) # Compute display fields for hash - df = _compute_display_fields( + df = OCRMaintenanceRow.compute_display_fields( status=status, health_overall=health_overall, version=version, can_redo=can_redo, can_rebuild=can_rebuild, error_stage=_error_stage(meta), diff --git a/paperforge/worker/ocr_objects.py b/paperforge/worker/ocr_objects.py index 2d36c1dc..85f7240a 100644 --- a/paperforge/worker/ocr_objects.py +++ b/paperforge/worker/ocr_objects.py @@ -1,14 +1,29 @@ from __future__ import annotations import contextlib +import os import re +import threading from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path from typing import Any from paperforge.worker.ocr_math import normalize_ocr_math_text +_RENDER_LOCKS: dict[tuple[str, int], threading.Lock] = {} +_RENDER_LOCKS_LOCK = threading.Lock() + + +def _get_render_lock(page_cache_dir: Path, page_num: int) -> threading.Lock: + key = (str(page_cache_dir), page_num) + with _RENDER_LOCKS_LOCK: + if key not in _RENDER_LOCKS: + _RENDER_LOCKS[key] = threading.Lock() + return _RENDER_LOCKS[key] + + def render_figure_object_markdown(figure: dict[str, Any]) -> str: caption = figure.get("caption", "") image_relpath = figure.get("image_relpath", "") @@ -101,8 +116,6 @@ def _crop_asset_from_pdf( rotation_deg: int = 0, ) -> bool: - - if dst.exists(): with contextlib.suppress(Exception): dst.unlink() @@ -134,26 +147,35 @@ def _crop_asset_from_pdf( try: if page_width > 0 and page_height > 0 and page_cache_dir is not None and not rotation_deg: - try: - from paperforge.worker.ocr import crop_block_asset, render_pdf_page_cached - except ImportError: - return False - - try: - page_image_path = page_cache_dir / f"page_{page_num:03d}.jpg" - rendered = render_pdf_page_cached( - doc, - page_num, - target_width=page_width, - target_height=page_height, - destination=page_image_path, - ) - if not rendered: + lock = _get_render_lock(page_cache_dir, page_num) + with lock: + cached_page_image = _find_cached_page_image(page_cache_dir, page_num) + if cached_page_image is not None: + try: + from paperforge.worker.ocr import crop_block_asset + except ImportError: + return False + ok = crop_block_asset(cached_page_image, [int(v) for v in bbox], dst) + return ok + try: + from paperforge.worker.ocr import crop_block_asset, render_pdf_page_cached + except ImportError: + return False + try: + page_image_path = page_cache_dir / f"page_{page_num:03d}.jpg" + rendered = render_pdf_page_cached( + doc, + page_num, + target_width=page_width, + target_height=page_height, + destination=page_image_path, + ) + if not rendered: + return False + ok = crop_block_asset(rendered, [int(v) for v in bbox], dst) + return ok + except Exception: return False - ok = crop_block_asset(rendered, [int(v) for v in bbox], dst) - return ok - except Exception: - return False try: import fitz @@ -190,6 +212,163 @@ def _crop_asset_from_pdf( doc.close() +def _write_figure_object_task( + data: dict[str, Any], + *, + pdf_path: Path | None, + page_cache_dir: Path | None, + asset_dir: Path, + render_dir: Path, +) -> None: + """Extract crop and write markdown for one figure/cluster object.""" + fig_id = data["fig_id"] + page = data["page"] + page_width = data["page_width"] + page_height = data["page_height"] + crop_bbox = data["crop_bbox"] + rotation_deg = data.get("rotation_deg", 0) + caption = data.get("caption", "") + asset_path_rel = data["asset_path_rel"] + asset_path_abs = asset_dir / f"{fig_id}.jpg" + + was_cropped = False + if crop_bbox and all(v > 0 for v in crop_bbox): + was_cropped = _crop_asset_from_pdf( + pdf_path, + page, + crop_bbox, + asset_path_abs, + page_width=page_width, + page_height=page_height, + page_cache_dir=page_cache_dir, + pdf_doc=None, + pdf_doc_provider=None, + rotation_deg=rotation_deg, + ) + + if not was_cropped: + for asset_info in data.get("matched_assets", []): + bbox = asset_info.get("bbox", [0, 0, 0, 0]) + if pdf_path and bbox and all(v > 0 for v in bbox): + if _crop_asset_from_pdf( + pdf_path, + page, + bbox, + asset_path_abs, + page_width=page_width, + page_height=page_height, + page_cache_dir=page_cache_dir, + pdf_doc=None, + pdf_doc_provider=None, + rotation_deg=rotation_deg, + ): + was_cropped = True + break + + md = render_figure_object_markdown( + { + "figure_id": fig_id, + "page": page, + "caption": caption, + "image_relpath": asset_path_rel, + "confidence": data.get("confidence", 0.5), + "was_cropped": was_cropped, + } + ) + _write_object_markdown(md, render_dir / f"{fig_id}.md") + + +def _write_table_object_task( + data: dict[str, Any], + *, + pdf_path: Path | None, + page_cache_dir: Path | None, + asset_dir: Path, + render_dir: Path, +) -> None: + """Extract crop and write markdown for one table object.""" + tbl_id = data["tbl_id"] + page = data["page"] + page_width = data["page_width"] + page_height = data["page_height"] + crop_bbox = data["crop_bbox"] + rotation_deg = data.get("rotation_deg", 0) + caption = data.get("caption", "") + asset_path_rel = data["asset_path_rel"] + asset_path_abs = asset_dir / f"{tbl_id}.jpg" + + was_cropped = False + if data.get("has_asset") and pdf_path and crop_bbox and all(v > 0 for v in crop_bbox): + was_cropped = _crop_asset_from_pdf( + pdf_path, + page, + crop_bbox, + asset_path_abs, + page_width=page_width, + page_height=page_height, + page_cache_dir=page_cache_dir, + pdf_doc=None, + pdf_doc_provider=None, + rotation_deg=rotation_deg, + ) + + md = render_table_object_markdown( + { + "table_id": tbl_id, + "page": page, + "caption": caption, + "image_relpath": asset_path_rel, + "confidence": 0.85 if was_cropped else 0.4, + "formal_table_number": data.get("formal_table_number"), + "note_texts": data.get("note_texts", []), + "note_match_reason": data.get("note_match_reason", ""), + } + ) + _write_object_markdown(md, render_dir / f"{tbl_id}.md") + + +def _write_orphan_object_task( + data: dict[str, Any], + *, + pdf_path: Path | None, + page_cache_dir: Path | None, + asset_dir: Path, + render_dir: Path, +) -> None: + """Extract crop and write markdown for one orphan object.""" + orphan_id = data["orphan_id"] + page = data["page"] + page_width = data["page_width"] + page_height = data["page_height"] + bbox = data["bbox"] + asset_path_rel = data["asset_path_rel"] + asset_path_abs = asset_dir / f"{orphan_id}.jpg" + + if pdf_path and bbox and all(v > 0 for v in bbox): + _crop_asset_from_pdf( + pdf_path, + page, + bbox, + asset_path_abs, + page_width=page_width, + page_height=page_height, + page_cache_dir=page_cache_dir, + pdf_doc=None, + pdf_doc_provider=None, + ) + + md = render_figure_object_markdown( + { + "figure_id": orphan_id, + "page": page, + "caption": "", + "image_relpath": asset_path_rel, + "confidence": 0.3, + } + ) + _write_object_markdown(md, render_dir / f"{orphan_id}.md") + + def extract_and_write_objects( pdf_path: Path | None, figure_inventory: dict[str, Any], @@ -238,260 +417,194 @@ def extract_and_write_objects( def _page_dims(page_num: int) -> tuple[int, int]: return page_dimensions_by_page.get(page_num, (0, 0)) - shared_pdf_doc: Any | None = None - shared_pdf_open_attempted = False + # ---- Serial phase: build task list with stable IDs ---- + tasks: list[tuple[str, dict[str, Any]]] = [] - def _get_shared_pdf_doc() -> Any | None: - nonlocal shared_pdf_doc, shared_pdf_open_attempted - if shared_pdf_doc is not None: - return shared_pdf_doc - if shared_pdf_open_attempted: - return None - if pdf_path is None or not pdf_path.exists(): - shared_pdf_open_attempted = True - return None - shared_pdf_open_attempted = True - try: - import fitz - shared_pdf_doc = fitz.open(str(pdf_path)) - except Exception: - return None - return shared_pdf_doc + # Matched figures + for i, match in enumerate(figure_inventory.get("matched_figures", [])): + fig_id = match.get("figure_id", f"figure_{i + 1:03d}") + caption_text = match.get("text", "") + page = match.get("page", 0) + page_width, page_height = _page_dims(page) + rotation_deg = int(match.get("rotation_correction_deg", 0) or 0) - try: - # Process matched figures - for i, match in enumerate(figure_inventory.get("matched_figures", [])): - fig_id = match.get("figure_id", f"figure_{i + 1:03d}") - caption_text = match.get("text", "") - page = match.get("page", 0) - page_width, page_height = _page_dims(page) - asset_path_rel = f"assets/figures/{fig_id}.jpg" - asset_path_abs = figures_asset_dir / f"{fig_id}.jpg" - - rotation_deg = int(match.get("rotation_correction_deg", 0) or 0) - - # Build crop bbox: cluster_bbox, else union of matched_assets, - # expanded to include owned figure_inner_text blocks (same page) - crop_bbox = match.get("cluster_bbox") or [0, 0, 0, 0] - if not (len(crop_bbox) == 4 and all(v > 0 for v in crop_bbox)): - asset_bboxes = [ - a.get("bbox", [0, 0, 0, 0]) - for a in match.get("matched_assets", []) - if len(a.get("bbox") or []) >= 4 and all(v > 0 for v in a.get("bbox", [0, 0, 0, 0])) + crop_bbox = match.get("cluster_bbox") or [0, 0, 0, 0] + if not (len(crop_bbox) == 4 and all(v > 0 for v in crop_bbox)): + asset_bboxes = [ + a.get("bbox", [0, 0, 0, 0]) + for a in match.get("matched_assets", []) + if len(a.get("bbox") or []) >= 4 and all(v > 0 for v in a.get("bbox", [0, 0, 0, 0])) + ] + if asset_bboxes: + crop_bbox = [ + min(b[0] for b in asset_bboxes), + min(b[1] for b in asset_bboxes), + max(b[2] for b in asset_bboxes), + max(b[3] for b in asset_bboxes), ] - if asset_bboxes: - crop_bbox = [ - min(b[0] for b in asset_bboxes), - min(b[1] for b in asset_bboxes), - max(b[2] for b in asset_bboxes), - max(b[3] for b in asset_bboxes), - ] - if structured_blocks and len(crop_bbox) == 4 and all(v > 0 for v in crop_bbox): - for blk in structured_blocks: - if blk.get("role") != "figure_inner_text": - continue - if str(blk.get("_object_owner_id", "")) != str(fig_id): - continue - if int(blk.get("page", 0) or 0) != int(page or 0): - continue - bb = blk.get("bbox") or [0, 0, 0, 0] - if len(bb) < 4 or not all(v > 0 for v in bb): - continue - crop_bbox = [ - min(crop_bbox[0], bb[0]), - min(crop_bbox[1], bb[1]), - max(crop_bbox[2], bb[2]), - max(crop_bbox[3], bb[3]), - ] + if structured_blocks and len(crop_bbox) == 4 and all(v > 0 for v in crop_bbox): + for blk in structured_blocks: + if blk.get("role") != "figure_inner_text": + continue + if str(blk.get("_object_owner_id", "")) != str(fig_id): + continue + if int(blk.get("page", 0) or 0) != int(page or 0): + continue + bb = blk.get("bbox") or [0, 0, 0, 0] + if len(bb) < 4 or not all(v > 0 for v in bb): + continue + crop_bbox = [ + min(crop_bbox[0], bb[0]), + min(crop_bbox[1], bb[1]), + max(crop_bbox[2], bb[2]), + max(crop_bbox[3], bb[3]), + ] - was_cropped = False - if len(crop_bbox) == 4 and all(v > 0 for v in crop_bbox): - was_cropped = _crop_asset_from_pdf( - pdf_path, - page, - crop_bbox, - asset_path_abs, - page_width=page_width, - page_height=page_height, + tasks.append(( + "figure", + { + "fig_id": fig_id, + "page": page, + "page_width": page_width, + "page_height": page_height, + "crop_bbox": crop_bbox, + "rotation_deg": rotation_deg, + "matched_assets": match.get("matched_assets", []), + "caption": caption_text, + "asset_path_rel": f"assets/figures/{fig_id}.jpg", + "confidence": match.get("confidence", 0.5), + }, + )) + + # Unresolved figure clusters + for i, cluster in enumerate(figure_inventory.get("unresolved_clusters", [])): + cluster_id = cluster.get("cluster_id") or f"unresolved_cluster_{i + 1:03d}" + page = cluster.get("page", 0) + page_width, page_height = _page_dims(page) + bbox = cluster.get("cluster_bbox", [0, 0, 0, 0]) + + tasks.append(( + "figure", + { + "fig_id": cluster_id, + "page": page, + "page_width": page_width, + "page_height": page_height, + "crop_bbox": bbox, + "rotation_deg": 0, + "matched_assets": [], + "caption": "", + "asset_path_rel": f"assets/figures/{cluster_id}.jpg", + "confidence": 0.45, + }, + )) + + # Pre-allocate orphan IDs (stable ordering across figure + table orphans) + num_figure_orphans = len(figure_inventory.get("unmatched_assets", [])) + num_table_orphans = len(table_inventory.get("unmatched_assets", [])) + orphan_index = 0 + + # Figure unmatched assets as orphans + for asset in figure_inventory.get("unmatched_assets", []): + orphan_index += 1 + orphan_id = f"orphan_{orphan_index:03d}" + page = asset.get("page", 0) + page_width, page_height = _page_dims(page) + bbox = asset.get("bbox", [0, 0, 0, 0]) + + tasks.append(( + "orphan", + { + "orphan_id": orphan_id, + "page": page, + "page_width": page_width, + "page_height": page_height, + "bbox": bbox, + "asset_path_rel": f"assets/orphans/{orphan_id}.jpg", + }, + )) + + # Tables + for i, table in enumerate(table_inventory.get("tables", [])): + tbl_id = f"table_{i + 1:03d}" + caption_text = table.get("caption_text", "") + page = table.get("page", 0) + page_width, page_height = _page_dims(page) + asset_bbox = table.get("asset_bbox", [0, 0, 0, 0]) + crop_bbox = table.get("render_bbox") or asset_bbox + rotation_deg = table.get("render_rotation_deg", 0) or 0 + + tasks.append(( + "table", + { + "tbl_id": tbl_id, + "page": page, + "page_width": page_width, + "page_height": page_height, + "crop_bbox": crop_bbox, + "rotation_deg": rotation_deg, + "caption": caption_text, + "has_asset": table.get("has_asset", False), + "asset_path_rel": f"assets/tables/{tbl_id}.jpg", + "formal_table_number": table.get("formal_table_number") or table.get("table_number"), + "note_texts": table.get("note_texts", []), + "note_match_reason": table.get("note_match_reason", ""), + }, + )) + + # Table unmatched assets as orphans (continue from figure orphan index) + for asset in table_inventory.get("unmatched_assets", []): + orphan_index += 1 + orphan_id = f"orphan_{orphan_index:03d}" + page = asset.get("page", 0) + page_width, page_height = _page_dims(page) + bbox = asset.get("bbox", [0, 0, 0, 0]) + + tasks.append(( + "orphan", + { + "orphan_id": orphan_id, + "page": page, + "page_width": page_width, + "page_height": page_height, + "bbox": bbox, + "asset_path_rel": f"assets/orphans/{orphan_id}.jpg", + }, + )) + + # ---- Parallel phase: dispatch crops and markdown writes ---- + max_workers = min(2, os.cpu_count() or 4) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [] + for task_type, data in tasks: + if task_type == "figure": + futures.append(executor.submit( + _write_figure_object_task, + data, + pdf_path=pdf_path, page_cache_dir=page_cache_dir, - pdf_doc_provider=_get_shared_pdf_doc, - rotation_deg=rotation_deg, - ) - if not was_cropped: - for asset_info in match.get("matched_assets", []): - bbox = asset_info.get("bbox", [0, 0, 0, 0]) - if ( - pdf_path - and bbox - and all(v > 0 for v in bbox) - and _crop_asset_from_pdf( - pdf_path, - page, - bbox, - asset_path_abs, - page_width=page_width, - page_height=page_height, - page_cache_dir=page_cache_dir, - pdf_doc_provider=_get_shared_pdf_doc, - rotation_deg=rotation_deg, - ) - ): - was_cropped = True - break - - md = render_figure_object_markdown( - { - "figure_id": fig_id, - "page": page, - "caption": caption_text, - "image_relpath": asset_path_rel, - "confidence": match.get("confidence", 0.5), - "was_cropped": was_cropped, - } - ) - _write_object_markdown(md, figures_render_dir / f"{fig_id}.md") - - # Process unresolved figure clusters (multi-panel without reliable legend) - for i, cluster in enumerate(figure_inventory.get("unresolved_clusters", [])): - cluster_id = cluster.get("cluster_id") or f"unresolved_cluster_{i + 1:03d}" - page = cluster.get("page", 0) - bbox = cluster.get("cluster_bbox", [0, 0, 0, 0]) - page_width, page_height = _page_dims(page) - asset_path_rel = f"assets/figures/{cluster_id}.jpg" - asset_path_abs = figures_asset_dir / f"{cluster_id}.jpg" - - if pdf_path and bbox and all(v > 0 for v in bbox): - _crop_asset_from_pdf( - pdf_path, - page, - bbox, - asset_path_abs, - page_width=page_width, - page_height=page_height, + asset_dir=figures_asset_dir, + render_dir=figures_render_dir, + )) + elif task_type == "table": + futures.append(executor.submit( + _write_table_object_task, + data, + pdf_path=pdf_path, page_cache_dir=page_cache_dir, - pdf_doc_provider=_get_shared_pdf_doc, - ) - - md = render_figure_object_markdown( - { - "figure_id": cluster_id, - "page": page, - "caption": "", - "image_relpath": asset_path_rel, - "confidence": 0.45, - } - ) - _write_object_markdown(md, figures_render_dir / f"{cluster_id}.md") - - # Process unmatched assets as orphans - orphan_count = 0 - for asset in figure_inventory.get("unmatched_assets", []): - orphan_count += 1 - orphan_id = f"orphan_{orphan_count:03d}" - bbox = asset.get("bbox", [0, 0, 0, 0]) - page = asset.get("page", 0) - page_width, page_height = _page_dims(page) - asset_path_rel = f"assets/orphans/{orphan_id}.jpg" - asset_path_abs = orphans_asset_dir / f"{orphan_id}.jpg" - - if pdf_path and bbox and all(v > 0 for v in bbox): - _crop_asset_from_pdf( - pdf_path, - page, - bbox, - asset_path_abs, - page_width=page_width, - page_height=page_height, + asset_dir=tables_asset_dir, + render_dir=tables_render_dir, + )) + elif task_type == "orphan": + futures.append(executor.submit( + _write_orphan_object_task, + data, + pdf_path=pdf_path, page_cache_dir=page_cache_dir, - pdf_doc_provider=_get_shared_pdf_doc, - ) + asset_dir=orphans_asset_dir, + render_dir=figures_render_dir, + )) - md = render_figure_object_markdown( - { - "figure_id": orphan_id, - "page": page, - "caption": "", - "image_relpath": asset_path_rel, - "confidence": 0.3, - } - ) - _write_object_markdown(md, figures_render_dir / f"{orphan_id}.md") - - # Process tables - for i, table in enumerate(table_inventory.get("tables", [])): - tbl_id = f"table_{i + 1:03d}" - caption_text = table.get("caption_text", "") - page = table.get("page", 0) - page_width, page_height = _page_dims(page) - asset_bbox = table.get("asset_bbox", [0, 0, 0, 0]) - asset_path_rel = f"assets/tables/{tbl_id}.jpg" - asset_path_abs = tables_asset_dir / f"{tbl_id}.jpg" - - was_cropped = False - # Use render_bbox + render_rotation_deg for rotated tables - _crop_bbox = table.get("render_bbox") or asset_bbox - _rot_deg = table.get("render_rotation_deg", 0) or 0 - if table.get("has_asset") and pdf_path and _crop_bbox and all(v > 0 for v in _crop_bbox): - was_cropped = _crop_asset_from_pdf( - pdf_path, - page, - _crop_bbox, - asset_path_abs, - page_width=page_width, - page_height=page_height, - page_cache_dir=page_cache_dir, - pdf_doc_provider=_get_shared_pdf_doc, - rotation_deg=_rot_deg, - ) - - md = render_table_object_markdown( - { - "table_id": tbl_id, - "page": page, - "caption": caption_text, - "image_relpath": asset_path_rel, - "confidence": 0.85 if was_cropped else 0.4, - "formal_table_number": table.get("formal_table_number") or table.get("table_number"), - "note_texts": table.get("note_texts", []), - "note_match_reason": table.get("note_match_reason", ""), - } - ) - _write_object_markdown(md, tables_render_dir / f"{tbl_id}.md") - - # Process unmatched table assets as orphans - for asset in table_inventory.get("unmatched_assets", []): - orphan_count += 1 - orphan_id = f"orphan_{orphan_count:03d}" - bbox = asset.get("bbox", [0, 0, 0, 0]) - page = asset.get("page", 0) - page_width, page_height = _page_dims(page) - asset_path_rel = f"assets/orphans/{orphan_id}.jpg" - asset_path_abs = orphans_asset_dir / f"{orphan_id}.jpg" - - if pdf_path and bbox and all(v > 0 for v in bbox): - _crop_asset_from_pdf( - pdf_path, - page, - bbox, - asset_path_abs, - page_width=page_width, - page_height=page_height, - page_cache_dir=page_cache_dir, - pdf_doc_provider=_get_shared_pdf_doc, - ) - - md = render_figure_object_markdown( - { - "figure_id": orphan_id, - "page": page, - "caption": "", - "image_relpath": asset_path_rel, - "confidence": 0.3, - } - ) - _write_object_markdown(md, figures_render_dir / f"{orphan_id}.md") - finally: - if shared_pdf_doc is not None: - with contextlib.suppress(Exception): - shared_pdf_doc.close() + for future in as_completed(futures): + future.result() diff --git a/paperforge/worker/ocr_rebuild.py b/paperforge/worker/ocr_rebuild.py index 1f3d84a2..b12d46bc 100644 --- a/paperforge/worker/ocr_rebuild.py +++ b/paperforge/worker/ocr_rebuild.py @@ -2,10 +2,13 @@ from __future__ import annotations import datetime import json +import logging from pathlib import Path from paperforge.core.io import read_json, write_json +logger = logging.getLogger(__name__) + CURRENT_SPAN_BACKFILL_VERSION = "2026-07-01.1" CURRENT_SPAN_VISUAL_CONTAINER_VERSION = "2026-06-26.6" MIN_SPAN_BACKFILL_COVERAGE = 0.90 @@ -116,37 +119,53 @@ def select_papers_for_derived_rebuild(papers: list[dict]) -> list[str]: return [p["zotero_key"] for p in papers if p.get("derived_stale") and not p.get("raw_upgradable")] -def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None, checkpoint: Path | None = None) -> dict: - """Run derived-layer rebuild for the given paper keys without raw OCR rerun. - Rebuilds: structured blocks, metadata, figure/table inventories, objects, - render outputs, and health — from stored raw blocks only. - If checkpoint is provided, appends each rebuilt key so interrupted runs - can skip completed work via --resume. +def _filter_completed_keys(checkpoint_dir: Path | None, keys: list[str]) -> list[str]: + """Return keys that do not have a .done. marker in checkpoint_dir.""" + if not checkpoint_dir: + return keys + cp = Path(checkpoint_dir) + if not cp.exists(): + return keys + done = {p.name.removeprefix(".done.") for p in cp.glob(".done.*")} + return [k for k in keys if k not in done] + + +def _write_done_marker(checkpoint_dir: Path | None, key: str) -> None: + """Write a completion marker for a successfully rebuilt paper.""" + if not checkpoint_dir: + return + (Path(checkpoint_dir) / f".done.{key}").touch() + + +def _rebuild_one_paper(vault: Path, key: str) -> dict: + """Rebuild derived artifacts for a single paper. Module-level for pickle. + + Returns dict with status ('ok', 'skipped') and details. """ from paperforge.worker._utils import pipeline_paths, read_jsonl from paperforge.worker.ocr import validate_ocr_meta from paperforge.worker.ocr_artifacts import artifact_paths_for_root ocr_root = pipeline_paths(vault)["ocr"] - rebuilt_count = 0 + artifacts = artifact_paths_for_root(ocr_root, key) + paper_root = artifacts.paper_root - keys_iter = progress_bar(keys, desc="OCR rebuild") if progress_bar else keys - for key in keys_iter: - artifacts = artifact_paths_for_root(ocr_root, key) - paper_root = artifacts.paper_root - if not paper_root.exists(): - continue + if not paper_root.exists(): + return {"key": key, "status": "skipped", "reason": "no_paper_dir"} + if not artifacts.blocks_raw.exists(): + return {"key": key, "status": "skipped", "reason": "no_raw_blocks"} - # Read stored raw blocks - if not artifacts.blocks_raw.exists(): - continue - all_raw_blocks = list(read_jsonl(artifacts.blocks_raw)) + all_raw_blocks = list(read_jsonl(artifacts.blocks_raw)) + ocr_meta = read_json(artifacts.meta_json) if artifacts.meta_json.exists() else {} + + # ── Phase 1: clean raw blocks and span backfill ── + def _phase1_span_backfill() -> dict: + """Reject overlapping fallback blocks, backfill span_metadata from PDF. + Returns (span_meta_patch, source_pdf_path). Modifies all_raw_blocks in place.""" + nonlocal all_raw_blocks - # Reject pdf_text_layer_fallback blocks whose text overlaps >=80% with - # a non-backfill block on the same page. These are column fragments - # the PDF text layer spilled during the initial OCR backfill. from paperforge.worker.ocr_pdf_spans import ( _BACKFILL_OVERLAP_REJECT_THRESHOLD, _backfill_coverage_in_existing, @@ -175,10 +194,7 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None block["_ocr_raw_error_type"] = "backfill_overlaps_existing_text_block" block["_text_source"] = "pdf_text_layer_fallback_rejected" - # Backfill span_metadata from source PDF - ocr_meta = read_json(artifacts.meta_json) if artifacts.meta_json.exists() else {} source_pdf_path = _resolve_source_pdf_for_rebuild(vault, key, ocr_meta) - span_meta_patch: dict[str, object] = {} covered_count, eligible_count, coverage = _compute_span_backfill_coverage(all_raw_blocks) @@ -223,17 +239,19 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None status="rerun_backfill", ) - # Extract normalized PDF rawdict lines for asset-internal figure number recovery + return {"span_meta_patch": span_meta_patch, "source_pdf_path": source_pdf_path} + + # ── Phase 2: PDF lines, enrich meta, build structured blocks ── + def _phase2_build_structured(source_pdf_path: Path | None) -> dict: + """Extract PDF lines, enrich source metadata, build structured blocks. + Returns {structured, doc_structure, resolved, source_meta, page_pdf_lines_by_page}.""" from paperforge.worker.ocr_pdf_spans import extract_pdf_lines_normalized + page_pdf_lines_by_page = extract_pdf_lines_normalized(source_pdf_path) - # Read source metadata. If legacy/old OCR papers are missing canonical - # bibliographic metadata, enrich source_metadata.json from the formal - # Literature-hub note frontmatter before rebuilding OCR-derived layers. _enrich_meta_from_paper_note(vault, key, artifacts.source_metadata) source_meta = read_json(artifacts.source_metadata) if artifacts.source_metadata.exists() else {} - # Rebuild structured blocks from paperforge.worker.ocr_blocks import build_structured_blocks, write_structured_blocks_jsonl structured, doc_structure = build_structured_blocks( @@ -241,12 +259,11 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None source_metadata=source_meta, structure_output_dir=artifacts.blocks_structured.parent, ) - # Write role-level span profiles + from paperforge.worker.ocr_profiles import write_role_span_profiles write_role_span_profiles(structured, artifacts.blocks_structured.parent) - # Rebuild resolved metadata from paperforge.worker.ocr_metadata import ( extract_frontmatter_candidates_from_blocks, resolve_metadata, @@ -264,7 +281,21 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None ) write_resolved_metadata(metadata_dir / "resolved_metadata.json", resolved) - # Rebuild figure inventory + return { + "structured": structured, + "doc_structure": doc_structure, + "resolved": resolved, + "source_meta": source_meta, + "page_pdf_lines_by_page": page_pdf_lines_by_page, + } + + # ── Phase 3: figure/table inventories, bio passes, writebacks ── + def _phase3_figure_tables( + structured: list[dict], + page_pdf_lines_by_page: dict[int, list[dict]], + source_meta: dict, + ) -> dict: + """Build figure and table inventories, run bio passes, resolve conflicts.""" from paperforge.worker.ocr_figures import ( build_figure_inventory, write_back_figure_roles, @@ -274,7 +305,6 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None figure_inventory = build_figure_inventory(structured, page_pdf_lines_by_page=page_pdf_lines_by_page) write_back_figure_roles(figure_inventory, structured) - # Author bio passes (Pass B + Pass C) from paperforge.worker.ocr_bio import ( residual_author_bio_pass, post_ref_bio_cleanup, @@ -289,7 +319,7 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None if ref_start_page is not None: post_ref_bio_cleanup(figure_inventory, structured, ref_start_page=ref_start_page) prune_figure_inventory_after_bio(figure_inventory) - # Rebuild reader figures + from paperforge.worker.ocr_figure_reader import synthesize_reader_figures reader_payload = synthesize_reader_figures(figure_inventory, structured_blocks=structured) @@ -297,7 +327,6 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None reader_figures_dir.mkdir(parents=True, exist_ok=True) write_json(reader_figures_dir / "reader_figures.json", reader_payload) - # Rebuild table inventory from paperforge.worker.ocr_tables import build_table_inventory, write_back_table_roles, write_table_inventory table_inventory = build_table_inventory(structured) @@ -306,8 +335,6 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None resolve_media_asset_conflicts(figure_inventory, table_inventory) attach_ownership_conflicts(figure_inventory, table_inventory) - # Apply object writeback seam (ownership evidence, contained/side-adjacent text, consumed-block contract) - # Must run BEFORE write_figure_inventory so claims land in the persisted inventory. from paperforge.worker.ocr_object_writeback import apply_object_writebacks apply_object_writebacks( @@ -320,15 +347,32 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None write_back_table_roles(table_inventory, structured) write_table_inventory(artifacts.blocks_structured.parent / "table_inventory.json", table_inventory) - # Re-persist structured blocks with writeback roles (table_html, figure_asset) # ponytail: writes entire list again; if throughput matters, write only changed blocks - write_structured_blocks_jsonl(artifacts.blocks_structured, structured) + from paperforge.worker.ocr_blocks import write_structured_blocks_jsonl as _write_structured_blocks_jsonl + _write_structured_blocks_jsonl(artifacts.blocks_structured, structured) - # Rebuild object artifacts + return { + "figure_inventory": figure_inventory, + "table_inventory": table_inventory, + "reader_payload": reader_payload, + } + + # ── Phase 4: objects, render, health ── + def _phase4_render_health( + structured: list[dict], + resolved: dict, + figure_inventory: dict, + table_inventory: dict, + reader_payload: dict, + doc_structure: dict, + ocr_meta: dict, + source_pdf_path: Path | None, + ) -> str: + """Extract object artifacts, render fulltext markdown, build health report. + Returns markdown string.""" from paperforge.worker.ocr_objects import extract_and_write_objects - ocr_meta = read_json(artifacts.meta_json) if artifacts.meta_json.exists() else {} - source_pdf_path = Path(ocr_meta.get("source_pdf", "")) if ocr_meta.get("source_pdf") else None + _source_pdf_path = Path(ocr_meta.get("source_pdf", "")) if ocr_meta.get("source_pdf") else None page_dimensions_by_page: dict[int, tuple[int, int]] = {} for block in structured: page = int(block.get("page", 0) or 0) @@ -338,7 +382,7 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None page_dimensions_by_page[page] = (width, height) extract_and_write_objects( - pdf_path=source_pdf_path, + pdf_path=_source_pdf_path, figure_inventory=figure_inventory, table_inventory=table_inventory, asset_root=paper_root / "assets", @@ -347,7 +391,6 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None structured_blocks=structured, ) - # Rebuild render output from paperforge.worker.ocr_render import render_fulltext_markdown, write_render_outputs rebuild_page_count = ocr_meta.get("page_count", 0) or 0 @@ -364,7 +407,6 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None reader_payload=reader_payload, ) - # Rebuild health from paperforge.worker.ocr_health import build_ocr_health, build_ocr_raw_integrity_health, write_ocr_health health_report = build_ocr_health( @@ -380,12 +422,20 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None health_report["ocr_raw_integrity"] = build_ocr_raw_integrity_health(all_raw_blocks) write_ocr_health(paper_root / "health", health_report) - # Persist decision log from paperforge.worker.ocr_decisions import collect_decisions, write_decision_log write_decision_log(paper_root / "health" / "decision_log.jsonl", collect_decisions(structured)) - # Rebuild role index + return markdown + + # ── Phase 5: indexes, version flags, write meta ── + def _phase5_finalize( + resolved: dict, + structured: list[dict], + markdown: str, + span_meta_patch: dict, + ) -> None: + """Rebuild indexes, apply version flags, write meta.json.""" from paperforge.worker.ocr_index import build_role_indexes, write_role_index role_indexes = build_role_indexes( @@ -393,20 +443,19 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None resolved_metadata=resolved, ) write_role_index(paper_root / "index", role_indexes) - # Rebuild structure tree + from paperforge.retrieval.structure_tree import build_structure_tree, write_structure_tree structure_tree = build_structure_tree(structured) write_structure_tree(paper_root / "index", structure_tree) - # Update version state in meta.json meta = ocr_meta meta.update(span_meta_patch) meta = _apply_post_rebuild_version_flags(meta) - # Rebuild regenerated the derived outputs; validate from a clean - # optimistic status instead of short-circuiting on a stale - # done_incomplete value from a previous render. meta["ocr_status"] = "done" + + from paperforge.worker.ocr_render import write_render_outputs + meta = write_render_outputs( render_root=paper_root / "render", user_fulltext=artifacts.compat_fulltext, @@ -414,21 +463,97 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None meta=meta, rebuild_increment=True, ) - # Re-validate and clear stale errors (e.g. page marker mismatch from pre-fix render) paths_dict = {"ocr": pipeline_paths(vault)["ocr"]} _status, _err = validate_ocr_meta(paths_dict, meta) meta["ocr_status"] = _status meta["error"] = _err if _err else "" write_json(artifacts.meta_json, meta) - rebuilt_count += 1 - if checkpoint: - done = [] - if checkpoint.exists(): - done = json.loads(checkpoint.read_text(encoding="utf-8")) - done.append(key) - checkpoint.write_text(json.dumps(done, ensure_ascii=False), encoding="utf-8") + # ── Execute phases ── + phase1_result = _phase1_span_backfill() + span_meta_patch = phase1_result["span_meta_patch"] + source_pdf_path = phase1_result["source_pdf_path"] + phase2_result = _phase2_build_structured(source_pdf_path) + structured = phase2_result["structured"] + doc_structure = phase2_result["doc_structure"] + resolved = phase2_result["resolved"] + source_meta = phase2_result["source_meta"] + page_pdf_lines_by_page = phase2_result["page_pdf_lines_by_page"] + + phase3_result = _phase3_figure_tables(structured, page_pdf_lines_by_page, source_meta) + figure_inventory = phase3_result["figure_inventory"] + table_inventory = phase3_result["table_inventory"] + reader_payload = phase3_result["reader_payload"] + + markdown = _phase4_render_health( + structured, resolved, figure_inventory, table_inventory, + reader_payload, doc_structure, ocr_meta, source_pdf_path, + ) + + _phase5_finalize(resolved, structured, markdown, span_meta_patch) + + return {"key": key, "status": "ok"} + + +def _run_parallel_rebuild(vault: Path, keys: list[str], workers: int, checkpoint_dir: Path | None) -> list[dict]: + """Run rebuild in parallel using a process pool.""" + from concurrent.futures import ProcessPoolExecutor, as_completed + + results: list[dict] = [] + with ProcessPoolExecutor(max_workers=workers) as executor: + futures = {executor.submit(_rebuild_one_paper, vault, k): k for k in keys} + for future in as_completed(futures): + key = futures[future] + try: + result = future.result() + if result.get("status") == "ok": + _write_done_marker(checkpoint_dir, key) + results.append(result) + except Exception as e: + results.append({"key": key, "status": "failed", "error": str(e)}) + return results + + +def run_derived_rebuild_for_keys( + vault: Path, + keys: list[str], + progress_bar=None, + checkpoint_dir: Path | None = None, + parallel: int = 4, +) -> dict: + """Run derived-layer rebuild for the given paper keys without raw OCR rerun. + + Rebuilds: structured blocks, metadata, figure/table inventories, objects, + render outputs, and health — from stored raw blocks only. + + If checkpoint_dir is provided, .done. marker files track progress so + interrupted runs can skip completed work via --resume. + + Args: + vault: Vault root path. + keys: Paper keys to rebuild. + progress_bar: Optional progress bar wrapper (tqdm-style). + checkpoint_dir: Directory for .done. completion markers. + parallel: Number of parallel workers (0 = serial). Default 4. + """ + keys = _filter_completed_keys(checkpoint_dir, keys) + if not keys: + return {"rebuild_count": 0} + + workers = int(parallel) if parallel else 0 + + if workers > 0 and len(keys) > 1: + results = _run_parallel_rebuild(vault, keys, workers, checkpoint_dir) + return {"rebuild_count": sum(1 for r in results if r.get("status") == "ok")} + + rebuilt_count = 0 + keys_iter = progress_bar(keys, desc="OCR rebuild") if progress_bar else keys + for key in keys_iter: + result = _rebuild_one_paper(vault, key) + if result.get("status") == "ok": + rebuilt_count += 1 + _write_done_marker(checkpoint_dir, key) return {"rebuild_count": rebuilt_count} diff --git a/tests/test_ocr_maintenance.py b/tests/test_ocr_maintenance.py index 6db0796e..927fa6c3 100644 --- a/tests/test_ocr_maintenance.py +++ b/tests/test_ocr_maintenance.py @@ -18,8 +18,8 @@ def _call(status: str, health: str = "-", version: str = "-", can_redo: bool = False, can_rebuild: bool = False, error_stage: str = "", error_summary: str = "", degraded_reasons: list[str] | None = None) -> dict: - from paperforge.worker.ocr_maintenance import _compute_display_fields - return _compute_display_fields( + from paperforge.worker.ocr_maintenance import OCRMaintenanceRow + return OCRMaintenanceRow.compute_display_fields( status=status, health_overall=health, version=version, can_redo=can_redo, can_rebuild=can_rebuild, error_stage=error_stage, error_summary=error_summary,