Coverage for src / quber / playground / benchmark / verify_gold.py: 67%

1"""Check every authored gold value against the cell it claims to come from. 

2 

3The questions were authored by agents reading printed tables, so each one 

4names a table cell and quotes what that cell prints. This reads the cell out 

5of the database and compares. A question whose gold value does not match its 

6cell is dropped rather than corrected — a benchmark answer that had to be 

7repaired is not a benchmark answer. 

8 

9Input is the authoring workflow's journal; output is the gold set the 

10comparison runs against. 

11 

12Usage: 

13 python -m quber.playground.benchmark.verify_gold \\ 

14 --journal <path>/journal.jsonl --out output/benchmark/gold.json 

15""" 

16 

17from __future__ import annotations 

18 

19import argparse 

20import json 

21import re 

22from pathlib import Path 

23from typing import Any, Callable, Dict, List, Optional, Tuple 

24 

25from loguru import logger 

26 

27from quber.playground import db 

28 

29 

30def norm(s: str) -> str: 

31 return re.sub(r"\s+", " ", (s or "").strip()) 

32 

33 

34def load_documents(journal: Path) -> List[Dict[str, Any]]: 

35 docs = [] 

36 for line in journal.read_text().splitlines(): 

37 row = json.loads(line) 

38 if row.get("type") == "result" and isinstance(row.get("result"), dict): 

39 docs.append(row["result"]) 

40 return docs 

41 

42 

43def tagged_cells(content: str) -> List[Tuple[str, str]]: 

44 """The id-tagged cells of one chunk's markup, as (cell id, printed text).""" 

45 return [ 

46 (cid, norm(re.sub(r"<[^>]+>", "", inner))) 

47 for cid, inner in re.findall(r'<td id="([^"]+)"[^>]*>(.*?)</td>', content or "", re.S) 

48 ] 

49 

50 

51def cell_index(doc_key: str) -> Dict[Tuple[str, str], str]: 

52 """Every tagged cell of every chunk in one document, keyed by chunk and cell id.""" 

53 with db.connect() as conn: 

54 rows = conn.execute( 

55 """SELECT c.chunk_id, c.content FROM ade_playground.chunks c 

56 JOIN ade_playground.documents d ON d.id = c.document_id 

57 WHERE d.doc_key = %s""", 

58 (doc_key,), 

59 ).fetchall() 

60 index: Dict[Tuple[str, str], str] = {} 

61 for chunk_id, content in rows: 

62 for cid, text in tagged_cells(content): 

63 index.setdefault((chunk_id, cid), text) 

64 index.setdefault(("", cid), text) # same cell id, whichever chunk carries it 

65 return index 

66 

67 

68def occurrences(index: Dict[Tuple[str, str], str], printed: str) -> int: 

69 """How many distinct cells in the document print this exact figure. 

70 

71 Counted over cell ids, not over chunks: a table's cells are repeated in 

72 each of its line-item chunks under the same ids, so counting chunk 

73 entries would report every figure as duplicated. Reported so a reviewer 

74 can see which questions rest on a figure that appears more than once, 

75 where a right-looking answer read off the wrong row would still score. 

76 """ 

77 target = norm(printed) 

78 return len({cid for (chunk, cid), text in index.items() if not chunk and text == target}) 

79 

80 

81def verify( 

82 docs: List[Dict[str, Any]], 

83 index_for: Callable[[str], Dict[Tuple[str, str], str]] = cell_index, 

84) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: 

85 kept: List[Dict[str, Any]] = [] 

86 dropped: List[Dict[str, Any]] = [] 

87 for doc in docs: 

88 slug = doc["slug"] 

89 index = index_for(slug) 

90 for q in doc.get("value_questions", []): 

91 cid = q["gold_cell_id"] 

92 chunk = q.get("chunk_id", "") 

93 printed: Optional[str] = index.get((chunk, cid)) or index.get(("", cid)) 

94 gold = norm(q["gold_value"]) 

95 if printed is None: 

96 dropped.append({**q, "slug": slug, "why": f"cell {cid} not found in {slug}"}) 

97 continue 

98 if printed != gold: 

99 dropped.append( 

100 {**q, "slug": slug, "why": f"cell {cid} prints {printed!r}, gold says {gold!r}"} 

101 ) 

102 continue 

103 kept.append( 

104 { 

105 "slug": slug, 

106 "kind": "value", 

107 "question": q["question"], 

108 "gold_value": printed, 

109 "gold_cell_id": cid, 

110 "chunk_id": chunk, 

111 "page": q.get("page"), 

112 "unit": q.get("unit"), 

113 "period": q.get("period"), 

114 "style": q.get("style"), 

115 "uniqueness_note": q.get("uniqueness_note"), 

116 "occurrences_in_document": occurrences(index, printed), 

117 } 

118 ) 

119 for q in doc.get("prose_questions", []): 

120 kept.append( 

121 { 

122 "slug": slug, 

123 "kind": "prose", 

124 "question": q["question"], 

125 "gold_value": None, 

126 "gold_cell_id": None, 

127 "why_prose": q.get("why_prose"), 

128 } 

129 ) 

130 return kept, dropped 

131 

132 

133def main() -> None: 

134 p = argparse.ArgumentParser(description=__doc__) 

135 p.add_argument("--journal", type=Path, required=True) 

136 p.add_argument("--out", type=Path, default=Path("output/benchmark/gold.json")) 

137 args = p.parse_args() 

138 

139 docs = load_documents(args.journal) 

140 kept, dropped = verify(docs) 

141 args.out.parent.mkdir(parents=True, exist_ok=True) 

142 args.out.write_text(json.dumps(kept, indent=2)) 

143 (args.out.parent / "gold_dropped.json").write_text(json.dumps(dropped, indent=2)) 

144 

145 values = [k for k in kept if k["kind"] == "value"] 

146 repeated = [k for k in values if k["occurrences_in_document"] > 1] 

147 logger.success( 

148 "kept {v} value + {p} prose from {d} documents; dropped {x}", 

149 v=len(values), 

150 p=len(kept) - len(values), 

151 d=len(docs), 

152 x=len(dropped), 

153 ) 

154 if repeated: 

155 logger.warning( 

156 "{n} kept value questions have a gold figure printed in more than one cell", 

157 n=len(repeated), 

158 ) 

159 for d in dropped: 

160 logger.warning("dropped [{s}] {q!r}: {w}", s=d["slug"], q=d["question"][:60], w=d["why"]) 

161 

162 

163if __name__ == "__main__": 

164 main()