211 lines
9.1 KiB
Python
211 lines
9.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
tables.py — selective table-cell extraction for the exam extractor (PLAN.md §B).
|
|
|
|
Two sources, unified into one cell-grid schema:
|
|
* STANDARD — the Tesseract+TableFormer backbone already emits `tables[].data.table_cells`
|
|
(text + row/col offsets + spans + bbox). Free, cached, every run. Good on ruled tables;
|
|
but it MISSES some data tables and OCRs them as loose tokens (REPORT.md p5).
|
|
* GRANITE — Granite-Docling-258M VLM emits `<otsl>` grids in DocTags (clean rows/cols even
|
|
where the backbone scrambles them). GPU cost, so used SELECTIVELY: only on pages the router
|
|
flags (a standard table present, or dense picture/checkbox), routed through dsync's GPU lock
|
|
+ Redis cache. Recipe (REPORT.md): {"to_formats":["doctags","json"], "pipeline":"vlm",
|
|
"vlm_pipeline_model":"granite_docling"}.
|
|
|
|
Unified table = {page, n_rows, n_cols, grid (2D text), cells, caption, source, is_furniture}.
|
|
"""
|
|
import re, json, os, glob, base64, urllib.request
|
|
|
|
# ----------------------------------------------------------------- OTSL (Granite DocTags)
|
|
OTSL_BLOCK = re.compile(r"<otsl>(.*?)</otsl>", re.S)
|
|
CAPTION = re.compile(r"<caption>(?:<loc_\d+>)*(.*?)</caption>", re.S)
|
|
CELL_TOK = re.compile(r"<(fcel|ecel|ched|rhed|lcel|ucel|xcel|nl)>([^<]*)")
|
|
HEADER_TAGS = {"ched", "rhed"}
|
|
|
|
|
|
def parse_otsl(doctags):
|
|
"""Parse every <otsl> block in a DocTags string into unified tables."""
|
|
out = []
|
|
for block in OTSL_BLOCK.findall(doctags):
|
|
cap = None
|
|
mc = CAPTION.search(block)
|
|
if mc:
|
|
cap = re.sub(r"\s+", " ", mc.group(1)).strip()
|
|
body = CAPTION.sub("", block)
|
|
body = re.sub(r"<loc_\d+>", "", body)
|
|
rows, cur = [], []
|
|
for tag, txt in CELL_TOK.findall(body):
|
|
if tag == "nl":
|
|
rows.append(cur); cur = []
|
|
else:
|
|
cur.append({"text": txt.strip(), "header": tag in HEADER_TAGS,
|
|
"empty": tag == "ecel"})
|
|
if cur:
|
|
rows.append(cur)
|
|
rows = [r for r in rows if r]
|
|
if not rows:
|
|
continue
|
|
n_cols = max(len(r) for r in rows)
|
|
grid = [[c["text"] for c in r] + [""] * (n_cols - len(r)) for r in rows]
|
|
out.append({"page": None, "n_rows": len(rows), "n_cols": n_cols, "grid": grid,
|
|
"caption": cap, "source": "granite-otsl",
|
|
"is_furniture": is_furniture(grid, cap)})
|
|
return out
|
|
|
|
|
|
# ----------------------------------------------------------------- standard TableFormer
|
|
def tables_from_standard(doc):
|
|
out = []
|
|
for t in doc.get("tables", []):
|
|
data = t.get("data", {}) or {}
|
|
cells = data.get("table_cells", []) or []
|
|
nr, nc = data.get("num_rows") or 0, data.get("num_cols") or 0
|
|
grid = [["" for _ in range(nc)] for _ in range(nr)]
|
|
for c in cells:
|
|
r0, c0 = c.get("start_row_offset_idx"), c.get("start_col_offset_idx")
|
|
if r0 is not None and c0 is not None and r0 < nr and c0 < nc and c.get("text"):
|
|
grid[r0][c0] = c["text"]
|
|
prov = t.get("prov") or []
|
|
page = prov[0].get("page_no") if prov else None
|
|
cap = " ".join(x.get("text", "") for x in (t.get("captions") or []) if isinstance(x, dict)) or None
|
|
out.append({"page": page, "n_rows": nr, "n_cols": nc, "grid": grid,
|
|
"caption": cap, "source": "docling-standard",
|
|
"is_furniture": is_furniture(grid, cap)})
|
|
return out
|
|
|
|
|
|
# ----------------------------------------------------------------- furniture filter
|
|
FURNITURE_RE = re.compile(r"examiner|do not write|leave\s+blank|question\s*mark|"
|
|
r"for marker|total marks?$", re.I)
|
|
|
|
|
|
def is_furniture(grid, caption=None):
|
|
"""A table that is exam scaffolding (mark grid / 'For Examiner's Use'), not question data."""
|
|
blob = " ".join(cell for row in grid for cell in row) + " " + (caption or "")
|
|
if FURNITURE_RE.search(blob):
|
|
return True
|
|
# a single-column strip of question numbers / blanks = a mark grid
|
|
flat = [c for row in grid for c in row if c.strip()]
|
|
if flat and all(re.fullmatch(r"\d{1,2}", c.strip()) for c in flat):
|
|
return True
|
|
return False
|
|
|
|
|
|
# ----------------------------------------------------------------- Granite via dsync
|
|
VLM_OPTS = {"to_formats": ["doctags", "json"], "pipeline": "vlm",
|
|
"vlm_pipeline_model": "granite_docling", "image_export_mode": "placeholder"}
|
|
|
|
|
|
def _serve_vlm(pdf_b64, fname, page):
|
|
import dsync
|
|
opts = {**VLM_OPTS, "page_range": [page, page]}
|
|
body = {"options": opts,
|
|
"sources": [{"kind": "file", "base64_string": pdf_b64, "filename": fname}],
|
|
"target": {"kind": "inbody"}}
|
|
req = urllib.request.Request(dsync.SERVE + "/v1/convert/source",
|
|
data=json.dumps(body).encode(),
|
|
headers={"Content-Type": "application/json"})
|
|
for _ in range(4): # tolerate the single-use 404 race
|
|
try:
|
|
return json.loads(urllib.request.urlopen(req, timeout=1200).read())
|
|
except urllib.error.HTTPError as e:
|
|
if e.code == 404:
|
|
import time; time.sleep(3); continue
|
|
raise
|
|
raise RuntimeError("serve vlm: repeated 404")
|
|
|
|
|
|
def _doctags_of(resp):
|
|
doc = resp.get("document") or {}
|
|
return doc.get("doctags_content") or doc.get("doc_tags") or doc.get("doctags") or ""
|
|
|
|
|
|
def granite_tables(pdf, pages, *, cached_glob=None, retries=4):
|
|
"""Run Granite-Docling on the given pages via dsync (GPU lock + OOM retry + Redis cache),
|
|
parse <otsl>, tag each table with its page. Falls back to cached *.doctags if serve fails."""
|
|
import dsync, time
|
|
cache = _load_cached_doctags(cached_glob) if cached_glob else {}
|
|
r = dsync._redis()
|
|
b64 = base64.b64encode(open(pdf, "rb").read()).decode()
|
|
fname = os.path.basename(pdf)
|
|
sha = dsync._sha(pdf)
|
|
out = []
|
|
for pg in pages:
|
|
key = f"docling:vlm:{sha}:p{pg}"
|
|
doctags = None
|
|
if r and (hit := r.get(key)):
|
|
doctags = hit if isinstance(hit, str) else hit.decode()
|
|
if doctags is None:
|
|
delay = 5
|
|
for attempt in range(retries):
|
|
with dsync._GpuLock(r):
|
|
resp = _serve_vlm(b64, fname, pg)
|
|
if dsync._is_oom(resp):
|
|
print(f"[granite] p{pg} OOM, backoff {delay}s ({attempt+1}/{retries})")
|
|
time.sleep(delay); delay = min(delay * 2, 120); continue
|
|
doctags = _doctags_of(resp)
|
|
if r and doctags:
|
|
r.set(key, doctags, ex=dsync.CACHE_TTL)
|
|
break
|
|
if not doctags and pg in cache:
|
|
print(f"[granite] p{pg} serve empty -> cached doctags")
|
|
doctags = cache[pg]
|
|
for tbl in parse_otsl(doctags or ""):
|
|
tbl["page"] = pg
|
|
out.append(tbl)
|
|
return out
|
|
|
|
|
|
def _load_cached_doctags(glob_pat):
|
|
"""Map page_no -> doctags text from files named *p<N>.doctags."""
|
|
cache = {}
|
|
for fn in glob.glob(glob_pat):
|
|
m = re.search(r"p(\d+)\.doctags$", fn)
|
|
if m:
|
|
cache[int(m.group(1))] = open(fn, encoding="utf-8", errors="replace").read()
|
|
return cache
|
|
|
|
|
|
# ----------------------------------------------------------------- routing + attach
|
|
def candidate_pages(doc):
|
|
"""Pages the router sends to Granite: a standard table, or a dense picture/checkbox page."""
|
|
pages = set()
|
|
for t in doc.get("tables", []):
|
|
prov = t.get("prov") or []
|
|
if prov and prov[0].get("page_no"):
|
|
pages.add(prov[0]["page_no"])
|
|
chk = {}
|
|
for it in doc.get("texts", []):
|
|
if it.get("label", "").startswith("checkbox"):
|
|
prov = it.get("prov") or []
|
|
if prov and prov[0].get("page_no"):
|
|
chk[prov[0]["page_no"]] = chk.get(prov[0]["page_no"], 0) + 1
|
|
pages |= {p for p, n in chk.items() if n >= 2}
|
|
return sorted(pages)
|
|
|
|
|
|
def attach_to_questions(tables, parts):
|
|
"""Assign each non-furniture table to the nearest preceding part on its page (by y); if no
|
|
geometry, attach to the first part on that page. Records table refs on the part."""
|
|
data_tables = [t for t in tables if not t["is_furniture"]]
|
|
by_page = {}
|
|
for lab, v in parts.items():
|
|
by_page.setdefault(v.get("page"), []).append((lab, v))
|
|
for i, t in enumerate(data_tables):
|
|
t["id"] = i
|
|
cands = by_page.get(t["page"], [])
|
|
if not cands:
|
|
t["for_part"] = None; continue
|
|
# best-effort: the part highest on the page (largest bbox top = the page's question stem),
|
|
# else the earliest part label. (Tables sit under the stem; we don't carry table y here.)
|
|
with_geo = [(lab, v) for lab, v in cands if v.get("bbox")]
|
|
if with_geo:
|
|
lab = max(with_geo, key=lambda kv: (kv[1]["bbox"] or {}).get("t", 0))[0]
|
|
else:
|
|
lab = sorted(cands, key=lambda kv: kv[0])[0][0]
|
|
t["for_part"] = lab
|
|
parts[lab].setdefault("tables", []).append(
|
|
{"id": i, "n_rows": t["n_rows"], "n_cols": t["n_cols"],
|
|
"caption": t["caption"], "source": t["source"]})
|
|
return data_tables
|