mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-10 00:08:58 +00:00
Studio RAG: fix RTL/Indic PDF corruption and dropped DOCX tables (#6780)
* Studio RAG: fix RTL/Indic PDF corruption and dropped DOCX tables The RAG parser prefers pymupdf4llm.to_markdown for PDFs, but that rebuilds text from positioned glyphs and mangles complex-shaping scripts (RTL Arabic/Hebrew come back as shaped Presentation Forms, Indic matras drop to U+FFFD) and can silently drop most of a heavy-RTL page. _pdf now compares the Markdown against PyMuPDF's logical-order get_text() per page and falls back to it when the Markdown looks corrupted (shaped Presentation Forms or U+FFFD above a small floor/ratio) or holds far fewer letters than the raw layer. Latin PDFs are unaffected and keep their Markdown tables/headings. _docx walked document.paragraphs, which excludes table cells, so DOCX tables were dropped entirely. It now walks body content in document order via iter_inner_content, emitting each table row as pipe-joined cells (deduped across merged cells); the preview locator already anchors on pipes. Adds parser tests for the corruption and incompleteness fallbacks and for DOCX table extraction. These mirror the chat document-extractor guard raised in the unslothai/ unsloth#5351 review; the RAG parser is a separate module and needed its own fix. * RAG DOCX: keep empty table cells and collapse in-cell newlines Skipping empty cells shifted later cells left and broke column alignment across rows; a cell with internal paragraphs (newlines) also broke the pipe-joined row. Keep every cell (dropping the row only when all are empty) and normalize each cell with " ".join(split()) so multi-paragraph cells stay on one row. Adds a test for both. * RAG DOCX: dedup merged table cells on the <w:tc> element directly Store the shared <w:tc> lxml element in the seen set instead of its id(); it is hashable and compares by the underlying node, so it dedups spanned/merged cells the same way without relying on id(). Adds a merged-cell test. * RAG DOCX: align merged cells, pad skipped grid columns, flatten nested tables * RAG DOCX: walk cells in document order so nested tables keep in-cell position * RAG DOCX: dedup vertically merged cells so a spanning label is indexed once --------- Co-authored-by: danielhanchen <michaelhan2050@gmail.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
ac6ba96f9e
commit
62e9644266
2 changed files with 307 additions and 5 deletions
|
|
@ -12,6 +12,7 @@ from __future__ import annotations
|
|||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from html.parser import HTMLParser
|
||||
|
||||
|
|
@ -69,6 +70,39 @@ def _html(raw: str) -> list[Page]:
|
|||
return [_page("\n".join(parser.out), 1)]
|
||||
|
||||
|
||||
# pymupdf4llm rebuilds text from positioned glyphs, which mangles complex-shaping
|
||||
# scripts (RTL Arabic/Hebrew emerge as shaped Presentation Forms, Indic matras drop to
|
||||
# U+FFFD) and can silently drop most of a heavy-RTL page. When Markdown trips these
|
||||
# signals we fall back to PyMuPDF's logical-order get_text(). Thresholds mirror the chat
|
||||
# extractor guard (unslothai/unsloth#5351 review).
|
||||
_SHAPED_PRESENTATION_FORMS = re.compile("[\ufb1d-\ufdff\ufe70-\ufefc]")
|
||||
_PDF_FALLBACK_MIN_BAD_GLYPHS = 5
|
||||
_PDF_FALLBACK_BAD_GLYPH_RATIO = 0.0005
|
||||
_PDF_INCOMPLETE_RATIO = 0.75
|
||||
_PDF_INCOMPLETE_MIN_LETTERS = 200
|
||||
|
||||
|
||||
def _markdown_corrupted(text: str) -> bool:
|
||||
"""True when pymupdf4llm's glyph reconstruction mangled the text: shaped RTL
|
||||
Presentation Forms or U+FFFD replacements above a small floor/ratio (so a lone
|
||||
legitimate shaped glyph does not force the fallback)."""
|
||||
if not text:
|
||||
return False
|
||||
threshold = max(_PDF_FALLBACK_MIN_BAD_GLYPHS, _PDF_FALLBACK_BAD_GLYPH_RATIO * len(text))
|
||||
shaped = len(_SHAPED_PRESENTATION_FORMS.findall(text))
|
||||
return shaped > threshold or text.count("\ufffd") > threshold
|
||||
|
||||
|
||||
def _markdown_incomplete(markdown: str, plain: str) -> bool:
|
||||
"""True when ``markdown`` holds far fewer letters than the raw ``get_text`` layer -- a
|
||||
coarse guard for heavy-RTL pages pymupdf4llm silently drops without shaped glyphs."""
|
||||
plain_letters = sum(1 for c in plain if c.isalnum())
|
||||
if plain_letters < _PDF_INCOMPLETE_MIN_LETTERS:
|
||||
return False
|
||||
markdown_letters = sum(1 for c in markdown if c.isalnum())
|
||||
return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters
|
||||
|
||||
|
||||
def _pdf_markdown(doc) -> list[str] | None:
|
||||
"""Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index
|
||||
i maps to page i+1. Returns None when the lib is missing, extraction fails, or the
|
||||
|
|
@ -100,9 +134,19 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
|
|||
try:
|
||||
md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None
|
||||
for i, page in enumerate(doc):
|
||||
# Prefer layout-aware Markdown (keeps tables/headings legible for retrieval);
|
||||
# fall back to plain text when Markdown is off, unavailable, or empty here.
|
||||
text = (md[i] if md else "") or page.get_text("text") or ""
|
||||
plain = page.get_text("text") or ""
|
||||
candidate = md[i] if md else ""
|
||||
# Prefer layout-aware Markdown (keeps tables/headings legible for retrieval),
|
||||
# but drop to PyMuPDF's logical-order text when Markdown is off/empty or when
|
||||
# pymupdf4llm mangled it (RTL/Indic) or dropped most of the page.
|
||||
if (
|
||||
candidate
|
||||
and not _markdown_corrupted(candidate)
|
||||
and not _markdown_incomplete(candidate, plain)
|
||||
):
|
||||
text = candidate
|
||||
else:
|
||||
text = plain
|
||||
pages.append(_page(text, i + 1))
|
||||
if want_images:
|
||||
for img in page.get_images(full = True):
|
||||
|
|
@ -308,12 +352,61 @@ def render_pdf_pages(
|
|||
doc.close()
|
||||
|
||||
|
||||
def _docx_table_rows(table) -> list[str]:
|
||||
"""Each row as pipe-joined cell text (the locator splits anchors on pipes).
|
||||
Columns stay aligned to the layout grid (merged cells fill their spanned slots,
|
||||
skipped leading/trailing grid columns become empty fields). Cells are walked in
|
||||
document order so a nested table, and any text after it, flattens in place."""
|
||||
from docx.table import Table
|
||||
from docx.text.paragraph import Paragraph
|
||||
|
||||
rows: list[str] = []
|
||||
seen: set = set() # <w:tc> already emitted; dedups merges spanning columns or rows
|
||||
for row in table.rows:
|
||||
cells: list[str] = [""] * getattr(row, "grid_cols_before", 0)
|
||||
trailing: list[str] = [] # nested rows + any post-nested text, kept in order
|
||||
for cell in row.cells:
|
||||
# A merged cell shares one <w:tc> across the columns and rows it spans:
|
||||
# emit its text once, then placeholders, so columns and rows stay aligned.
|
||||
if cell._tc in seen:
|
||||
cells.append("")
|
||||
continue
|
||||
seen.add(cell._tc)
|
||||
# Paragraph text before the first nested table is the aligned field; the
|
||||
# nested table and anything after it flatten below the row, in order.
|
||||
field: list[str] = []
|
||||
after_table = False
|
||||
for item in cell.iter_inner_content():
|
||||
if isinstance(item, Table):
|
||||
after_table = True
|
||||
trailing.extend(_docx_table_rows(item))
|
||||
elif isinstance(item, Paragraph):
|
||||
text = " ".join(item.text.split()) # collapse in-cell newlines
|
||||
if text:
|
||||
(trailing if after_table else field).append(text)
|
||||
cells.append(" ".join(field)) # empty cells kept so columns line up
|
||||
cells.extend([""] * getattr(row, "grid_cols_after", 0))
|
||||
if any(c.strip() for c in cells):
|
||||
rows.append(" | ".join(cells))
|
||||
rows.extend(trailing)
|
||||
return rows
|
||||
|
||||
|
||||
def _docx(path: str) -> list[Page]:
|
||||
import docx
|
||||
from docx.table import Table
|
||||
from docx.text.paragraph import Paragraph
|
||||
|
||||
document = docx.Document(path)
|
||||
text = "\n".join(p.text for p in document.paragraphs)
|
||||
return [_page(text, None)]
|
||||
lines: list[str] = []
|
||||
# Walk body content in document order: paragraphs alone drop tables entirely.
|
||||
for block in document.iter_inner_content():
|
||||
if isinstance(block, Paragraph):
|
||||
if block.text.strip():
|
||||
lines.append(block.text)
|
||||
elif isinstance(block, Table):
|
||||
lines.extend(_docx_table_rows(block))
|
||||
return [_page("\n".join(lines), None)]
|
||||
|
||||
|
||||
def parse(path: str, *, want_images: bool = False):
|
||||
|
|
|
|||
|
|
@ -86,3 +86,212 @@ def test_pdf_markdown_falls_back_when_lib_missing(tmp_path, monkeypatch):
|
|||
_table_pdf(pdf)
|
||||
pages = parsers.parse(str(pdf))
|
||||
assert pages and "Quarter" in pages[0].text
|
||||
|
||||
|
||||
def _long_text_pdf(path):
|
||||
import pymupdf
|
||||
|
||||
doc = pymupdf.open()
|
||||
page = doc.new_page()
|
||||
body = "The quick brown fox jumps over the lazy dog. " * 12 # >200 letters
|
||||
page.insert_textbox(pymupdf.Rect(40, 40, 550, 750), body, fontsize = 11)
|
||||
doc.save(str(path))
|
||||
doc.close()
|
||||
|
||||
|
||||
def test_pdf_markdown_corruption_falls_back_to_plain(tmp_path, monkeypatch):
|
||||
# pymupdf4llm can emit shaped RTL Presentation Forms for Arabic/Hebrew; the parser
|
||||
# detects that and uses PyMuPDF's logical-order text instead of the mangled Markdown.
|
||||
from core.rag import config, parsers
|
||||
|
||||
monkeypatch.setattr(config, "PDF_MARKDOWN", True)
|
||||
shaped = "".join(chr(c) for c in range(0xFE8D, 0xFEA0)) * 20 # heavy shaped forms
|
||||
monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: [shaped] * doc.page_count)
|
||||
pdf = tmp_path / "table.pdf"
|
||||
_table_pdf(pdf)
|
||||
text = "\n".join(p.text for p in parsers.parse(str(pdf)))
|
||||
assert "Quarter" in text # real logical-order text recovered
|
||||
assert not parsers._markdown_corrupted(text) # shaped garbage not carried through
|
||||
|
||||
|
||||
def test_pdf_markdown_incomplete_falls_back_to_plain(tmp_path, monkeypatch):
|
||||
# If pymupdf4llm silently drops most of a page, the parser prefers the fuller raw layer.
|
||||
from core.rag import config, parsers
|
||||
|
||||
monkeypatch.setattr(config, "PDF_MARKDOWN", True)
|
||||
monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: ["x"] * doc.page_count)
|
||||
pdf = tmp_path / "long.pdf"
|
||||
_long_text_pdf(pdf)
|
||||
text = "\n".join(p.text for p in parsers.parse(str(pdf)))
|
||||
assert "quick brown fox" in text # fuller raw layer used, not the near-empty Markdown
|
||||
|
||||
|
||||
def _docx_with_table(path):
|
||||
import docx
|
||||
|
||||
document = docx.Document()
|
||||
document.add_paragraph("Intro before table.")
|
||||
table = document.add_table(rows = 2, cols = 2)
|
||||
table.cell(0, 0).text = "NAME"
|
||||
table.cell(0, 1).text = "SCORE"
|
||||
table.cell(1, 0).text = "Alice"
|
||||
table.cell(1, 1).text = "97pts"
|
||||
document.add_paragraph("Outro after table.")
|
||||
document.save(str(path))
|
||||
|
||||
|
||||
def test_docx_extracts_table_cells(tmp_path):
|
||||
# document.paragraphs alone drops tables; the parser walks body content in order so
|
||||
# table cells survive (pipe-joined, which the preview locator anchors on).
|
||||
pytest.importorskip("docx")
|
||||
from core.rag import parsers
|
||||
|
||||
docx_path = tmp_path / "t.docx"
|
||||
_docx_with_table(docx_path)
|
||||
text = "\n".join(p.text for p in parsers.parse(str(docx_path)))
|
||||
assert all(v in text for v in ("NAME", "SCORE", "Alice", "97pts")) # cells kept
|
||||
assert "Alice | 97pts" in text # row cells joined
|
||||
assert text.index("Intro") < text.index("NAME") < text.index("Outro") # order kept
|
||||
|
||||
|
||||
def test_docx_table_keeps_columns_and_collapses_cell_newlines(tmp_path):
|
||||
# Empty cells are kept (so columns stay aligned across rows) and a cell's internal
|
||||
# newlines are collapsed to spaces (so a multi-paragraph cell can't break the row).
|
||||
pytest.importorskip("docx")
|
||||
import docx
|
||||
|
||||
from core.rag import parsers
|
||||
|
||||
document = docx.Document()
|
||||
table = document.add_table(rows = 2, cols = 3)
|
||||
table.cell(0, 0).text = "A"
|
||||
table.cell(0, 1).text = "" # empty middle cell
|
||||
table.cell(0, 2).text = "C"
|
||||
multiline = table.cell(1, 0)
|
||||
multiline.text = "line1"
|
||||
multiline.add_paragraph("line2") # cell now holds an internal newline
|
||||
table.cell(1, 1).text = "mid"
|
||||
table.cell(1, 2).text = "end"
|
||||
path = tmp_path / "aligned.docx"
|
||||
document.save(str(path))
|
||||
|
||||
text = "\n".join(p.text for p in parsers.parse(str(path)))
|
||||
assert "A | | C" in text # empty cell preserved -> columns line up
|
||||
assert "line1 line2 | mid | end" in text # internal newline collapsed to a space
|
||||
|
||||
|
||||
def test_docx_table_merged_cell_keeps_grid_alignment(tmp_path):
|
||||
# A horizontally merged cell repeats across the spanned columns: emit its text once
|
||||
# then a placeholder, so the row keeps as many fields as its siblings (columns stay
|
||||
# aligned) without duplicating the merged text.
|
||||
pytest.importorskip("docx")
|
||||
import docx
|
||||
|
||||
from core.rag import parsers
|
||||
|
||||
document = docx.Document()
|
||||
table = document.add_table(rows = 2, cols = 3)
|
||||
table.cell(0, 0).text = "WIDE"
|
||||
table.cell(0, 2).text = "END"
|
||||
table.cell(0, 0).merge(table.cell(0, 1)) # span the first two columns
|
||||
table.cell(1, 0).text = "a"
|
||||
table.cell(1, 1).text = "b"
|
||||
table.cell(1, 2).text = "c"
|
||||
path = tmp_path / "merged.docx"
|
||||
document.save(str(path))
|
||||
|
||||
text = "\n".join(p.text for p in parsers.parse(str(path)))
|
||||
assert text.count("WIDE") == 1 # merged cell not duplicated across spanned columns
|
||||
assert "WIDE | | END" in text # placeholder keeps 3 fields, aligned with "a | b | c"
|
||||
assert "a | b | c" in text
|
||||
|
||||
|
||||
def test_docx_table_pads_omitted_grid_columns(tmp_path):
|
||||
# A row that skips leading grid columns exposes the gap via grid_cols_before; pad it
|
||||
# with empty fields so the value stays under the right header instead of shifting left.
|
||||
pytest.importorskip("docx")
|
||||
import docx
|
||||
from docx.oxml.ns import qn
|
||||
|
||||
from core.rag import parsers
|
||||
|
||||
document = docx.Document()
|
||||
table = document.add_table(rows = 2, cols = 3)
|
||||
table.cell(0, 0).text = "H1"
|
||||
table.cell(0, 1).text = "H2"
|
||||
table.cell(0, 2).text = "H3"
|
||||
tr = table.rows[1]._tr # drop the first cell and mark it skipped via <w:gridBefore>
|
||||
tr.remove(tr.tc_lst[0])
|
||||
trPr = tr.get_or_add_trPr()
|
||||
trPr.insert(0, trPr.makeelement(qn("w:gridBefore"), {qn("w:val"): "1"}))
|
||||
table.rows[1].cells[0].text = "X" # sits in column 2
|
||||
path = tmp_path / "gap.docx"
|
||||
document.save(str(path))
|
||||
|
||||
text = "\n".join(p.text for p in parsers.parse(str(path)))
|
||||
assert " | X | " in text # leading gap padded so X lines up under H2, not H1
|
||||
|
||||
|
||||
def test_docx_flattens_nested_table(tmp_path):
|
||||
# cell.text ignores tables nested inside a cell; walk cell.tables so nested rows are
|
||||
# not silently dropped from the indexed text.
|
||||
pytest.importorskip("docx")
|
||||
import docx
|
||||
|
||||
from core.rag import parsers
|
||||
|
||||
document = docx.Document()
|
||||
outer = document.add_table(rows = 1, cols = 1).cell(0, 0)
|
||||
outer.text = "outer"
|
||||
nested = outer.add_table(rows = 1, cols = 2)
|
||||
nested.cell(0, 0).text = "NESTED-A"
|
||||
nested.cell(0, 1).text = "NESTED-B"
|
||||
path = tmp_path / "nested.docx"
|
||||
document.save(str(path))
|
||||
|
||||
text = "\n".join(p.text for p in parsers.parse(str(path)))
|
||||
assert "NESTED-A | NESTED-B" in text # nested table flattened, not dropped
|
||||
|
||||
|
||||
def test_docx_nested_table_keeps_in_cell_order(tmp_path):
|
||||
# A cell holding paragraph, nested table, paragraph must serialize in that order
|
||||
# (cell.text alone would emit both paragraphs before the nested rows).
|
||||
pytest.importorskip("docx")
|
||||
import docx
|
||||
|
||||
from core.rag import parsers
|
||||
|
||||
document = docx.Document()
|
||||
cell = document.add_table(rows = 1, cols = 1).cell(0, 0)
|
||||
cell.text = "before"
|
||||
nested = cell.add_table(rows = 1, cols = 2)
|
||||
nested.cell(0, 0).text = "NESTED-A"
|
||||
nested.cell(0, 1).text = "NESTED-B"
|
||||
cell.add_paragraph("after")
|
||||
path = tmp_path / "nested_order.docx"
|
||||
document.save(str(path))
|
||||
|
||||
text = "\n".join(p.text for p in parsers.parse(str(path)))
|
||||
assert text.index("before") < text.index("NESTED-A") < text.index("after")
|
||||
|
||||
|
||||
def test_docx_table_vertical_merge_emitted_once(tmp_path):
|
||||
# A vertically merged cell maps every continuation row back to the origin <w:tc>;
|
||||
# emit it once and leave placeholders below so a row-spanning label isn't repeated.
|
||||
pytest.importorskip("docx")
|
||||
import docx
|
||||
|
||||
from core.rag import parsers
|
||||
|
||||
document = docx.Document()
|
||||
table = document.add_table(rows = 3, cols = 2)
|
||||
table.cell(0, 0).merge(table.cell(1, 0)).merge(table.cell(2, 0)).text = "SECTION"
|
||||
table.cell(0, 1).text = "r0"
|
||||
table.cell(1, 1).text = "r1"
|
||||
table.cell(2, 1).text = "r2"
|
||||
path = tmp_path / "vmerge.docx"
|
||||
document.save(str(path))
|
||||
|
||||
text = "\n".join(p.text for p in parsers.parse(str(path)))
|
||||
assert text.count("SECTION") == 1 # not repeated on each spanned row
|
||||
assert "SECTION | r0" in text and " | r1" in text and " | r2" in text
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue