CoolFace
Apppublic

breakpointsoftware/document-parser

sourceHugging Faceupdated 2mo agoView on Hugging Face
0likes
document_processing.py58 linesDownload Raw Back to root
1from __future__ import annotations2 3import base644import io5from pathlib import Path6 7 8SUPPORTED_EXTENSIONS = {".txt", ".pdf", ".docx", ".jpg", ".jpeg", ".png"}9IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"}10 11 12def extract_text(file_path: Path) -> str:13    suffix = file_path.suffix.lower()14 15    if suffix == ".txt":16        return file_path.read_text(encoding="utf-8", errors="ignore")17 18    if suffix == ".pdf":19        from pypdf import PdfReader20 21        reader = PdfReader(str(file_path))22        pages = [page.extract_text() or "" for page in reader.pages]23        return "\n".join(pages)24 25    if suffix == ".docx":26        from docx import Document27 28        document = Document(str(file_path))29        paragraphs = [paragraph.text for paragraph in document.paragraphs]30        return "\n".join(paragraphs)31 32    raise ValueError(f"Unsupported file type: {file_path.suffix}")33 34 35def load_documents(folder: Path) -> list[Path]:36    if not folder.exists():37        raise FileNotFoundError(f"Input folder does not exist: {folder}")38 39    files = [path for path in folder.rglob("*") if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS]40    return sorted(files)41 42 43def is_image_document(file_path: Path) -> bool:44    return file_path.suffix.lower() in IMAGE_EXTENSIONS45 46 47def to_data_uri(file_path: Path) -> str:48    mime_type = {49        ".jpg": "image/jpeg",50        ".jpeg": "image/jpeg",51        ".png": "image/png",52    }.get(file_path.suffix.lower())53 54    if not mime_type:55        raise ValueError(f"Unsupported image file type: {file_path.suffix}")56 57    encoded = base64.b64encode(file_path.read_bytes()).decode("ascii")58    return f"data:{mime_type};base64,{encoded}"