Coverage for src / quber / core / extractors / set_of_mark / assemble.py: 81%

1"""Turn a cell grid and its geometry into a finished `ExtractedTable`. 

2 

3Everything after the cells are in hand is the same work whichever tool produced 

4them: correct the structure against the table's image, ground every corrected 

5cell in a measured box, verify the statuses that carry no box, and record what 

6the header area printed that no cell carries. 

7 

8Two callers hand cells in. The Set-of-Mark pipeline hands in a region-constrained 

9Camelot grid. The page-scan workflow hands in the grid a scan read off a table 

10printed as an image, where there was no text layer for Camelot to read. Both get 

11back the same object, filled the same way, because the correction and grounding 

12steps read the table's image and the page's own text and never ask which tool 

13found the cells. 

14""" 

15 

16from __future__ import annotations 

17 

18import asyncio 

19from dataclasses import dataclass 

20from pathlib import Path 

21from typing import List, Literal, Optional, Tuple 

22 

23from quber.agents.llm_client import LLMClient 

24from quber.agents.status_inspector import StatusInspector 

25from quber.core.extractors.base import ( 

26 ExtractedTable, 

27 MergedCellBox, 

28 grid_fingerprint, 

29 grounded_grid, 

30 table_address, 

31) 

32from quber.core.extractors.camelot.acquire import grid_to_markdown 

33from quber.core.extractors.camelot.correspondence.correction import correct_structure, printed_title 

34from quber.core.extractors.set_of_mark.inspection import inspect_gap_cells 

35from quber.core.extractors.set_of_mark.merge_grounding import ( 

36 find_dropped_header_text, 

37 locate_markers, 

38 log_ungrounded, 

39 log_ungrounded_cells, 

40 resolve_corrected_grid, 

41 resolve_merges, 

42) 

43 

44Box = Tuple[float, float, float, float] 

45NormBox = Tuple[float, float, float, float] 

46TableKind = Literal["text_table", "chart", "image_table"] 

47 

48 

49@dataclass 

50class TableAssembly: 

51 """Run-scoped inputs the assembly needs, the same for every table in a run.""" 

52 

53 source: Path 

54 llm: LLMClient 

55 correct_sem: asyncio.Semaphore 

56 dpi: int 

57 # Verifies each unboxed cell's proposed status against the table image; None 

58 # disables inspection and the proposed statuses stand. 

59 inspector: Optional[StatusInspector] = None 

60 

61 

62@dataclass 

63class CamelotOrigin: 

64 """The Camelot extraction a table's cells came from. 

65 

66 Absent for cells no Camelot pass produced, whose `bbox`, `flavor` and 

67 `camelot_accuracy` then keep their defaults, because there is no extraction 

68 for those fields to describe. 

69 """ 

70 

71 bbox: Optional[Box] 

72 flavor: Literal["lattice", "stream", "unknown"] 

73 accuracy: float 

74 

75 

76async def assemble_table( 

77 assembly: TableAssembly, 

78 page: int, 

79 ordinal: int, 

80 page_image: Path, 

81 page_dims: Tuple[float, float], 

82 cells: List[List[str]], 

83 cell_boxes: List[List[Optional[Box]]], 

84 bbox: Optional[Box], 

85 title: str = "", 

86 som_region: Optional[NormBox] = None, 

87 kind: TableKind = "text_table", 

88 camelot: Optional[CamelotOrigin] = None, 

89 ground_values: bool = True, 

90) -> ExtractedTable: 

91 """One table, corrected and grounded, from its cells and their boxes. 

92 

93 `bbox` is the table's box in PDF points with a bottom-left page origin. It 

94 scopes the image crop the correction reads and the page text the grounding 

95 guard allows, so an adjacent table can never bleed in. Without it there is 

96 nothing to crop, so the cells are rendered as they arrived and the fields the 

97 correction fills stay empty. 

98 

99 `cell_boxes` is shaped exactly like `cells`, in the same frame as `bbox`, and 

100 may hold None wherever no box was measured. `title` is the identity read off 

101 the page image before the cells were captured; it stands only when no 

102 correction ran, and only if the page prints it. A correction's title is the 

103 printed name copied off the page, or empty when none is printed, and that is 

104 what the table carries. 

105 """ 

106 page_w, page_h = page_dims 

107 correction = await correct_structure( 

108 cells, 

109 page_image, 

110 bbox, 

111 str(assembly.source), 

112 page, 

113 assembly.llm, 

114 assembly.correct_sem, 

115 assembly.dpi, 

116 ground_values, 

117 ) 

118 # The pre-correction grid paired with per-cell geometry; values are immutable 

119 # through correction, so it stays joinable to the final markdown. Boxes 

120 # normalize to the page frame here so the whole table reads in one frame. 

121 cell_grid = grounded_grid(cells, cell_boxes, page_w, page_h) 

122 

123 table_id = table_address(assembly.source, page, ordinal) 

124 fingerprint = grid_fingerprint(cells) 

125 origin_bbox = camelot.bbox if camelot else None 

126 flavor: Literal["lattice", "stream", "unknown"] = camelot.flavor if camelot else "unknown" 

127 accuracy = camelot.accuracy if camelot else 0.0 

128 

129 if correction is None: 

130 return ExtractedTable( 

131 table_id=table_id, 

132 content_fingerprint=fingerprint, 

133 title=await printed_title(title, str(assembly.source), page), 

134 markdown=grid_to_markdown(cells), 

135 page=page, 

136 source=str(assembly.source), 

137 som_region=som_region, 

138 kind=kind, 

139 cell_grid=cell_grid, 

140 corrected_grid=cell_grid, 

141 bbox=origin_bbox, 

142 flavor=flavor, 

143 camelot_accuracy=accuracy, 

144 ) 

145 

146 # Cells the agent combined resolve by the source ADDRESSES it read off the 

147 # printed coordinate frame: address -> cell box, text-validated. The 

148 # grounding stage then closes the whole table: corrected_grid gives every 

149 # corrected cell its measured box, so nothing downstream re-derives geometry. 

150 merged: List[MergedCellBox] = [] 

151 if correction.cell_merges: 

152 merged = resolve_merges(correction.cell_merges, cell_grid, correction.markdown) 

153 log_ungrounded(merged, page) 

154 corrected_grid = resolve_corrected_grid( 

155 correction.markdown, merged, cell_grid, correction.footnote_refs, correction.region_text 

156 ) 

157 # Classified statuses are text-evidence hypotheses; the inspector verifies 

158 # each against the table image, and whatever it cannot positively confirm 

159 # downgrades to `unverified` for user inspection. 

160 await inspect_gap_cells( 

161 corrected_grid, 

162 correction.markdown, 

163 str(assembly.source), 

164 page, 

165 page_image, 

166 bbox, 

167 assembly.dpi, 

168 assembly.inspector, 

169 ) 

170 log_ungrounded_cells(corrected_grid, page) 

171 dropped_text = find_dropped_header_text( 

172 corrected_grid, 

173 cell_grid, 

174 " ".join( 

175 [correction.title, correction.caption, correction.units] 

176 + [f"{f.marker} {f.text}".strip() for f in correction.footnotes] 

177 ), 

178 page, 

179 ) 

180 return ExtractedTable( 

181 table_id=table_id, 

182 content_fingerprint=fingerprint, 

183 title=correction.title, 

184 caption=correction.caption, 

185 markdown=correction.markdown, 

186 footnotes=correction.footnotes, 

187 footnote_refs=correction.footnote_refs, 

188 footnote_marks=locate_markers( 

189 correction.footnote_marks, 

190 correction.footnote_refs, 

191 correction.markdown, 

192 correction.footnotes, 

193 # Both titles: the correction sometimes strips a marker suffix the 

194 # image-read title retains. 

195 f"{correction.title} {title} {correction.caption}", 

196 ), 

197 units=correction.units, 

198 header_rows=correction.header_rows, 

199 llm_corrected=correction.llm_corrected, 

200 corrected_grid=corrected_grid, 

201 merged_cells=merged, 

202 dropped_text=dropped_text, 

203 page=page, 

204 source=str(assembly.source), 

205 som_region=som_region, 

206 kind=kind, 

207 cell_grid=cell_grid, 

208 bbox=origin_bbox, 

209 flavor=flavor, 

210 camelot_accuracy=accuracy, 

211 )