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:
parent
0be7ca39a4
commit
8145f1d527
1 changed files with 77 additions and 11 deletions
|
|
@ -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,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue