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

1"""Load the FinanceBench benchmark JSONL files into the `financebench` schema. 

2 

3Input: the two files from the FinanceBench open-source release — 

4`financebench_document_information.jsonl` (one row per filing) and 

5`financebench_open_source.jsonl` (one row per question, with the gold answer 

6and its evidence passages). The loader applies `schema.sql` (drop + recreate, 

7so a reload is a clean replace) and inserts everything in one pass. 

8 

9The playground corpus takes its document identity from the benchmark: 

10`doc_key_for` maps a benchmark doc_name to the playground storage key, and every 

11benchmark PDF lives at s3://qubera-docs/financebench/<doc_name>.pdf. 

12""" 

13 

14from __future__ import annotations 

15 

16import json 

17from pathlib import Path 

18from typing import Any, LiteralString, Optional, cast 

19 

20from loguru import logger 

21 

22from quber.playground import db 

23 

24SCHEMA_FILE = Path(__file__).with_name("schema.sql") 

25DEFAULT_DATA_DIR = Path(__file__).parents[5] / "financebench" / "data" 

26 

27S3_PREFIX = "financebench" 

28 

29 

30def doc_key_for(doc_name: str) -> str: 

31 """The playground storage key for a benchmark filing: lowercased, dashes 

32 for underscores. '3M_2022_10K' -> '3m-2022-10k'.""" 

33 return doc_name.lower().replace("_", "-") 

34 

35 

36def _pg_text(s: Optional[str]) -> Optional[str]: 

37 """Strip NUL characters: Postgres text fields reject 0x00, and evidence 

38 full-page text comes from the same PDF text layers that produce them.""" 

39 return s.replace("\x00", "") if s is not None else None 

40 

41 

42def _rows(path: Path) -> list[dict[str, Any]]: 

43 return [json.loads(line) for line in path.read_text().splitlines() if line.strip()] 

44 

45 

46def load(data_dir: Path = DEFAULT_DATA_DIR) -> tuple[int, int, int]: 

47 """Rebuild the financebench schema from the JSONL files. 

48 

49 Returns (documents, questions, evidence) row counts. 

50 """ 

51 docs = _rows(data_dir / "financebench_document_information.jsonl") 

52 questions = _rows(data_dir / "financebench_open_source.jsonl") 

53 

54 # The source metadata file carries one duplicated doc_name 

55 # (FOOTLOCKER_2023_annualreport, identical rows except doc_period). Keep 

56 # the first row per doc_name and report every drop — never fail silently. 

57 seen: set[str] = set() 

58 deduped = [] 

59 for d in docs: 

60 if d["doc_name"] in seen: 

61 logger.warning("Duplicate doc_name in document information, keeping first: {}", d) 

62 continue 

63 seen.add(d["doc_name"]) 

64 deduped.append(d) 

65 docs = deduped 

66 

67 with db.connect() as conn: 

68 # Same trust level as a query literal: a schema file we author. 

69 conn.execute(cast(LiteralString, SCHEMA_FILE.read_text())) 

70 

71 for d in docs: 

72 conn.execute( 

73 """INSERT INTO financebench.documents 

74 (doc_name, company, gics_sector, doc_type, doc_period, doc_link, s3_key) 

75 VALUES (%s, %s, %s, %s, %s, %s, %s)""", 

76 ( 

77 d["doc_name"], 

78 d.get("company"), 

79 d.get("gics_sector"), 

80 d.get("doc_type"), 

81 str(d.get("doc_period") or ""), 

82 d.get("doc_link"), 

83 f"{S3_PREFIX}/{d['doc_name']}.pdf", 

84 ), 

85 ) 

86 

87 n_evidence = 0 

88 for q in questions: 

89 conn.execute( 

90 """INSERT INTO financebench.questions 

91 (financebench_id, doc_name, question_type, question_reasoning, 

92 domain_question_num, question, answer, justification, dataset_subset_label) 

93 VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)""", 

94 ( 

95 q["financebench_id"], 

96 q["doc_name"], 

97 q.get("question_type"), 

98 q.get("question_reasoning"), 

99 str(q.get("domain_question_num") or ""), 

100 _pg_text(q["question"]), 

101 _pg_text(q["answer"]), 

102 _pg_text(q.get("justification")), 

103 q.get("dataset_subset_label"), 

104 ), 

105 ) 

106 for ev in q.get("evidence") or []: 

107 conn.execute( 

108 """INSERT INTO financebench.evidence 

109 (financebench_id, doc_name, evidence_page_num, evidence_text, evidence_text_full_page) 

110 VALUES (%s, %s, %s, %s, %s)""", 

111 ( 

112 q["financebench_id"], 

113 ev.get("doc_name") or q["doc_name"], 

114 ev.get("evidence_page_num"), 

115 _pg_text(ev.get("evidence_text")), 

116 _pg_text(ev.get("evidence_text_full_page")), 

117 ), 

118 ) 

119 n_evidence += 1 

120 

121 logger.success( 

122 "Loaded financebench: {d} documents, {q} questions, {e} evidence passages", 

123 d=len(docs), 

124 q=len(questions), 

125 e=n_evidence, 

126 ) 

127 return len(docs), len(questions), n_evidence 

128 

129 

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

131 import argparse 

132 

133 p = argparse.ArgumentParser(description="Load the FinanceBench JSONL files into Postgres.") 

134 p.add_argument("--data-dir", default=str(DEFAULT_DATA_DIR), help="directory holding the two JSONL files") 

135 args = p.parse_args(argv) 

136 load(Path(args.data_dir)) 

137 

138 

139if __name__ == "__main__": 

140 main()