from __future__ import annotations
import json
import os
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
sys.path.insert(0, "/mnt/c/Users/RyanJ/spectre-rs/scripts")
from sweep_heading_f1 import CORPUS
CACHE_PATH = Path("/tmp/spectre_extra_corpus/pmu4llm_md_cache.json")
def render_one(path: str) -> tuple[str, str, float]:
import pymupdf4llm t0 = time.time()
md = pymupdf4llm.to_markdown(path)
return path, md, time.time() - t0
def main() -> None:
cache: dict[str, str] = {}
if CACHE_PATH.exists():
cache = json.loads(CACHE_PATH.read_text())
to_render: list[str] = []
for _, path in CORPUS:
p = Path(path)
if not p.exists():
continue
key = f"{path}|{p.stat().st_mtime_ns}"
if key in cache:
continue
to_render.append(path)
print(f"to render: {len(to_render)} / {len(CORPUS)} PDFs")
if not to_render:
print("cache complete")
return
workers = min(int(os.environ.get("WORKERS", "6")), len(to_render))
print(f"workers: {workers}")
t_total = time.time()
completed = 0
with ProcessPoolExecutor(max_workers=workers) as ex:
futures = {ex.submit(render_one, p): p for p in to_render}
for fut in as_completed(futures):
try:
path, md, took = fut.result()
except Exception as e:
print(f"FAIL {futures[fut]}: {e}", flush=True)
continue
p = Path(path)
key = f"{path}|{p.stat().st_mtime_ns}"
cache[key] = md
CACHE_PATH.write_text(json.dumps(cache))
completed += 1
print(f" [{completed}/{len(to_render)}] {p.name} ({took:.1f}s)",
flush=True)
print(f"done in {time.time() - t_total:.1f}s")
if __name__ == "__main__":
main()