Coverage for src / quber / playground / benchmark / pages.py: 0%

1"""Backfill `financebench.documents.page_count` from the PDFs themselves. 

2 

3The FinanceBench release files carry no page counts, so this reads each 

4document's PDF from the bucket (through the quber S3 cache) and records its 

5page count. Run it after `load.py` — the loader rebuilds the schema, which 

6clears the column. A missing PDF logs a warning and leaves the count NULL; 

7it is never a silent skip. 

8""" 

9 

10from __future__ import annotations 

11 

12from typing import Optional 

13 

14import pymupdf 

15from loguru import logger 

16 

17from quber.playground import db 

18 

19BUCKET = "qubera-docs" 

20 

21 

22def backfill() -> tuple[int, int]: 

23 """Set page_count for every document whose PDF is in the bucket. 

24 

25 Returns (updated, missing) counts. 

26 """ 

27 from quber.files.cache import resolve_document 

28 

29 with db.connect() as conn: 

30 rows = conn.execute( 

31 "SELECT doc_name, s3_key FROM financebench.documents ORDER BY doc_name" 

32 ).fetchall() 

33 updated = missing = 0 

34 for doc_name, s3_key in rows: 

35 try: 

36 local = resolve_document(f"s3://{BUCKET}/{s3_key}") 

37 with pymupdf.open(local) as pdf: 

38 n = pdf.page_count 

39 except Exception as exc: 

40 logger.warning("No page count for {} ({}): {}", doc_name, s3_key, exc) 

41 missing += 1 

42 continue 

43 conn.execute( 

44 "UPDATE financebench.documents SET page_count = %s WHERE doc_name = %s", 

45 (n, doc_name), 

46 ) 

47 updated += 1 

48 

49 logger.success("Page counts: {} documents updated, {} missing PDFs", updated, missing) 

50 return updated, missing 

51 

52 

53def main(argv: Optional[list[str]] = None) -> None: 

54 backfill() 

55 

56 

57if __name__ == "__main__": 

58 main()