Spaces:
Sleeping
Sleeping
| import hashlib | |
| import json | |
| import tempfile | |
| import os | |
| import io | |
| from pathlib import Path | |
| from typing import List, Dict | |
| from llama_index.core import Document | |
| import pandas as pd | |
| # Advanced Docling Imports for Table Extraction | |
| from llama_index.readers.docling import DoclingReader | |
| from docling.datamodel.base_models import InputFormat | |
| from docling.datamodel.pipeline_options import PdfPipelineOptions, TableStructureOptions, TableFormerMode | |
| from docling.document_converter import DocumentConverter | |
| from docling.document_converter import PdfFormatOption | |
| class PDFProcessor: | |
| def __init__(self, chunk_size: int = 1000, chunk_overlap: int = 200): | |
| self.chunk_size = chunk_size | |
| self.chunk_overlap = chunk_overlap | |
| # Initialize advanced pipeline options for accurate table discovery | |
| pipeline_options = PdfPipelineOptions() | |
| pipeline_options.do_table_structure = True | |
| # NOTE: OCR is disabled by default for 10x speed boost on text-based PDFs. | |
| pipeline_options.do_ocr = False | |
| pipeline_options.table_structure_options.mode = TableFormerMode.ACCURATE | |
| self.doc_converter = DocumentConverter( | |
| format_options={ | |
| InputFormat.PDF: PdfFormatOption( | |
| pipeline_options=pipeline_options | |
| ) | |
| } | |
| ) | |
| def get_pdf_hash(self, pdf_file) -> str: | |
| """ | |
| Generates an MD5 hash for the PDF file object to serve as a cache key. | |
| """ | |
| pos = pdf_file.tell() | |
| pdf_file.seek(0) | |
| file_hash = hashlib.md5(pdf_file.read()).hexdigest() | |
| pdf_file.seek(pos) | |
| return file_hash | |
| def load_docling_documents(self, pdf_file, cache_path: Path = None) -> Dict: | |
| """ | |
| Uses Docling for unified RAG and Table Extraction in a single pass. | |
| Returns a dict with 'documents' (LlamaIndex) and 'tables' (List of DataFrames). | |
| """ | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf", dir=os.getcwd()) as tmp: | |
| pdf_file.seek(0) | |
| tmp.write(pdf_file.read()) | |
| tmp_path = Path(tmp.name) | |
| try: | |
| # 1. Single-Pass Conversion (Core Optimization - Truly One Pass) | |
| result = self.doc_converter.convert(tmp_path) | |
| doc = result.document # This is a DoclingDocument (v2) | |
| # 2. Extract LlamaIndex Documents from the result manually | |
| # This replaces DoclingReader and avoids double-conversion | |
| json_content = result.document.model_dump_json() | |
| # We create a single LlamaIndex Document with the full JSON content. | |
| # Use our existing hash generator for consistency. | |
| pdf_file.seek(0) | |
| file_hash = self.get_pdf_hash(pdf_file) | |
| documents = [Document( | |
| text=json_content, | |
| metadata={ | |
| "filename": pdf_file.name, | |
| "dl_doc_hash": file_hash | |
| } | |
| )] | |
| # 3. Extract structured tables (Uses Docling v2 high-speed export) | |
| tables = [] | |
| for i, table in enumerate(doc.tables): | |
| try: | |
| df = table.export_to_dataframe(doc=doc) | |
| if not df.empty: | |
| # Find page number for labeling | |
| page_no = "?" | |
| if table.prov and len(table.prov) > 0: | |
| page_no = table.prov[0].page_no | |
| tables.append({ | |
| "id": i + 1, | |
| "label": f"Table {i+1} (Page {page_no})", | |
| "df": df | |
| }) | |
| except Exception as e: | |
| print(f"Table Extraction Error [Table {i+1}]: {e}") | |
| # PERSIST: Save Markdown and Tables | |
| if cache_path: | |
| try: | |
| cache_path.mkdir(parents=True, exist_ok=True) | |
| # Markdown Export (Instant from existing result) | |
| md_content = result.document.export_to_markdown() | |
| with open(cache_path / "content.md", "w", encoding="utf-8") as f: | |
| f.write(md_content) | |
| # Store tables as JSON for persistence | |
| if tables: | |
| serialized_tables = [{ | |
| "id": t["id"], | |
| "label": t["label"], | |
| "data": t["df"].to_dict(orient="records") | |
| } for t in tables] | |
| with open(cache_path / "tables.json", "w", encoding="utf-8") as f: | |
| json.dump(serialized_tables, f, indent=2) | |
| except Exception as e: | |
| print(f"Persistence Error: {e}") | |
| return { | |
| "documents": documents, | |
| "tables": tables | |
| } | |
| finally: | |
| try: | |
| if tmp_path.exists(): | |
| tmp_path.unlink() | |
| except Exception: | |
| pass | |
| if __name__ == "__main__": | |
| # Test | |
| pass | |