Studio: splice VLM figure captions next to their 'Figure N:' line

Captions were appended at the bottom of the page text, so the chunk
containing 'Figure 1: Asymmetries ...' got chunked separately from
'**Figure**: Flowchart with ...' on the same page. Retrieval surfaced
the caption-text chunk but the VLM description landed in a different
chunk, leaving the LLM without the visual content right next to the
figure label.

Splice each VLM caption right after the matching 'Figure N:' (or
'Table N:') line as '**Figure N description**: ...', so:

  - The figure-boundary chunker now keeps both the original in-PDF
    caption AND the VLM description in the same chunk (which starts
    with 'Figure N:').
  - Multi-figure pages get per-figure attribution — the prefix
    'Figure N description' lets the LLM tell two figures on the same
    page apart, even though the bbox renderer still emits one image
    per page today (multi-figure clustering is a follow-up).
  - When the page text has no figure lines (DOCX/HTML/TXT or rare
    PDF layouts) the old end-of-page appendix is kept as a fallback.
This commit is contained in:
Roland Tannous 2026-05-27 17:02:36 +04:00
commit 8145f1d527

View file

@ -3,9 +3,20 @@
from __future__ import annotations
import re
from dataclasses import dataclass, field
from pathlib import Path
# Same shape as the chunker's figure-boundary regex but captures the
# figure label (Figure / Fig. / Table / Tab.) AND the number so we can
# attribute a VLM caption back to a specific figure on a multi-figure
# page.
_FIGURE_LINE_RE = re.compile(
r"^(?P<lead>\**)(?P<label>Figure|Fig\.|Table|Tab\.)\s+"
r"(?P<num>[A-Z]?\.?\d+(?:\.\d+)?)(?P<trail>\**[\.:])",
re.MULTILINE | re.IGNORECASE,
)
@dataclass(frozen = True)
class ParsedPage:
@ -44,6 +55,63 @@ class UnsupportedFormatError(ValueError):
pass
def _normalize_label(raw: str) -> str:
head = raw.lower()
if head.startswith("fig"):
return "Figure"
if head.startswith("tab"):
return "Table"
return raw.capitalize()
def _splice_inline_at_figure_lines(text: str, captions: list[str]) -> str:
"""Splice each caption right after the matching 'Figure N:' line.
Captions are consumed in order against the figure caption lines
appearing in the page text. Each spliced block is prefixed with
"**Figure N description**:" so a retrieved chunk lets the LLM
distinguish between multiple figures on the same page.
Fallback when no figure caption lines exist (or fewer than we have
captions): leftover captions are appended at the end of the page
text as generic "**Figure**: ..." entries.
"""
if not captions:
return text
matches = list(_FIGURE_LINE_RE.finditer(text))
if not matches:
appendix = "\n\n".join(f"**Figure**: {c}" for c in captions)
return f"{text}\n\n{appendix}"
parts: list[str] = []
cursor = 0
caps = iter(captions)
used = 0
for m in matches:
# Insertion point = end of the line containing the figure label.
line_end = text.find("\n", m.end())
if line_end == -1:
line_end = len(text)
parts.append(text[cursor:line_end])
try:
cap = next(caps)
except StopIteration:
cursor = line_end
continue
used += 1
label = f"{_normalize_label(m.group('label'))} {m.group('num')}"
parts.append(f"\n\n**{label} description**: {cap}")
cursor = line_end
parts.append(text[cursor:])
leftover = list(caps)
body = "".join(parts)
if leftover:
appendix = "\n\n".join(f"**Figure**: {c}" for c in leftover)
body = f"{body}\n\n{appendix}"
return body
def inline_image_captions(
pages: list[ParsedPage],
images: list[ParsedImage],
@ -51,10 +119,13 @@ def inline_image_captions(
) -> list[ParsedPage]:
"""Splice per-image captions into the markdown of the pages they came from.
Mirrors PR #5351's chat-composer pattern: figure captions become
inline text in the page markdown so the chunker indexes them like
any other content. Captions appear at the end of the page's text
block as ``**Figure**: `` lines.
Each caption lands right after the page's matching ``Figure N:`` or
``Table N:`` line as ``**Figure N description**: `` so the chunker
keeps the VLM description adjacent to the figure's existing in-PDF
caption text. When a page has more figure caption lines than we have
captioned images, the extra figure lines are left alone; when there
are more captions than figure lines (or no figure lines at all),
leftovers fall back to an end-of-page ``**Figure**: `` appendix.
``captions`` is parallel to ``images`` (same length, same order).
Empty or whitespace-only captions are skipped. Images without a
@ -64,11 +135,8 @@ def inline_image_captions(
if not images or not captions:
return list(pages)
if len(captions) != len(images):
# Defensive: caller mismatch shouldn't happen but we don't want
# to lose pages over it.
return list(pages)
# Bucket captions per page_number (None bucket → single-page docs).
per_page: dict[int | None, list[str]] = {}
for img, cap in zip(images, captions):
cleaned = (cap or "").strip()
@ -83,17 +151,15 @@ def inline_image_captions(
null_bucket = per_page.get(None, [])
for page in pages:
captions_for_this = per_page.get(page.page_number, [])
# If this is the single-page case (no page_number) also flush
# the null-bucket so DOCX/HTML/TXT pick up captions correctly.
if page.page_number is None and null_bucket:
captions_for_this = captions_for_this + null_bucket
if not captions_for_this:
out.append(page)
continue
appendix = "\n\n".join(f"**Figure**: {cap}" for cap in captions_for_this)
new_text = _splice_inline_at_figure_lines(page.text, captions_for_this)
out.append(
ParsedPage(
text = f"{page.text}\n\n{appendix}",
text = new_text,
page_number = page.page_number,
)
)