Spaces:
Sleeping
Sleeping
File size: 3,521 Bytes
4f25e4a 54a9b55 4f25e4a 54a9b55 4f25e4a 54a9b55 4f25e4a 54a9b55 4f25e4a 54a9b55 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | import logging
import pdfplumber
from pathlib import Path
from dataclasses import dataclass
logger = logging.getLogger(__name__)
@dataclass
class PageContent:
page_number: int
text: str
char_count: int
@dataclass
class DocumentContent:
file_path: str
file_name: str
total_pages: int
pages: list[PageContent]
@property
def full_text(self) -> str:
return "\n\n".join(p.text for p in self.pages if p.text)
@property
def total_chars(self) -> int:
return sum(p.char_count for p in self.pages)
def extract_pdf(
file_path: str | Path,
parser: str = "pymupdf4llm",
) -> DocumentContent:
"""Extract text from a PDF, returning a DocumentContent with per-page text.
Args:
file_path: Path to the PDF file.
parser: "pymupdf4llm" (default) uses structured markdown extraction
that preserves headers, tables, and lists. Falls back to
"pdfplumber" automatically if pymupdf4llm fails.
Pass "pdfplumber" explicitly to always use the flat extractor.
"""
path = Path(file_path)
if not path.exists():
raise FileNotFoundError(f"PDF not found: {path}")
if path.suffix.lower() != ".pdf":
raise ValueError(f"Expected a .pdf file, got: {path.suffix}")
if parser == "pymupdf4llm":
try:
return extract_pdf_structured(path)
except Exception as exc:
logger.warning(
"pymupdf4llm extraction failed for %s (%s) — falling back to pdfplumber.",
path.name, exc,
)
return _extract_pdf_pdfplumber(path)
return _extract_pdf_pdfplumber(path)
def extract_pdf_structured(file_path: str | Path) -> DocumentContent:
"""Extract a PDF to per-page markdown using pymupdf4llm.
pymupdf4llm preserves document structure as markdown:
- Section headers become ## / ### headings
- Tables become markdown tables
- Bullet lists become markdown lists
Args:
file_path: Path to an existing .pdf file.
Returns:
DocumentContent whose page texts are markdown strings.
"""
import pymupdf4llm # lazy import — optional dependency
path = Path(file_path)
page_dicts: list[dict] = pymupdf4llm.to_markdown(str(path), page_chunks=True)
pages: list[PageContent] = []
for i, page_dict in enumerate(page_dicts, start=1):
text = page_dict.get("text", "").strip()
pages.append(PageContent(page_number=i, text=text, char_count=len(text)))
return DocumentContent(
file_path=str(path.resolve()),
file_name=path.name,
total_pages=len(page_dicts),
pages=pages,
)
def _extract_pdf_pdfplumber(path: Path) -> DocumentContent:
"""Extract flat text from a PDF using pdfplumber (original implementation)."""
pages: list[PageContent] = []
with pdfplumber.open(path) as pdf:
total_pages = len(pdf.pages)
for i, page in enumerate(pdf.pages, start=1):
raw = page.extract_text() or ""
text = _clean(raw)
pages.append(PageContent(page_number=i, text=text, char_count=len(text)))
return DocumentContent(
file_path=str(path.resolve()),
file_name=path.name,
total_pages=total_pages,
pages=pages,
)
def _clean(text: str) -> str:
lines = (line.strip() for line in text.splitlines())
non_empty = (line for line in lines if line)
return "\n".join(non_empty)
|