import fitz # PyMuPDF - keep for fallback import tiktoken import os from typing import List, Dict, Any from pathlib import Path import logging # SmolDocling imports import torch from docling_core.types.doc import DoclingDocument from docling_core.types.doc.document import DocTagsDocument from transformers import AutoProcessor, AutoModelForVision2Seq from pdf2image import convert_from_path from PIL import Image import tempfile logger = logging.getLogger(__name__) class PDFProcessor: def __init__(self): self.encoding = tiktoken.encoding_for_model("gpt-3.5-turbo") self.max_tokens = int(os.getenv("MAX_TOKENS", "180000")) self.chunk_size = int(os.getenv("CHUNK_SIZE", "8000")) # Initialize SmolDocling model self.device = "cuda" if torch.cuda.is_available() else "cpu" self.model_path = "ds4sd/SmolDocling-256M-preview" # Load SmolDocling model and processor try: logger.info(f"Loading SmolDocling model on {self.device}") self.processor = AutoProcessor.from_pretrained(self.model_path) self.model = AutoModelForVision2Seq.from_pretrained( self.model_path, torch_dtype=torch.bfloat16 if self.device == "cuda" else torch.float32, _attn_implementation="flash_attention_2" if self.device == "cuda" else "eager", ).to(self.device) logger.info("SmolDocling model loaded successfully") except Exception as e: logger.error(f"Failed to load SmolDocling model: {str(e)}") logger.info("Falling back to PyMuPDF for extraction") self.model = None self.processor = None def extract_text(self, pdf_path: Path) -> str: """Extract text from PDF using SmolDocling (with PyMuPDF fallback)""" try: if self.model is None or self.processor is None: logger.info("Using PyMuPDF fallback for text extraction") return self._extract_text_pymupdf(pdf_path) logger.info(f"Extracting text from PDF using SmolDocling: {pdf_path}") # Convert PDF to images images = convert_from_path(str(pdf_path)) all_text = "" for page_num, image in enumerate(images): logger.info(f"Processing page {page_num + 1}/{len(images)}") # Create input messages for SmolDocling messages = [ { "role": "user", "content": [ {"type": "image"}, {"type": "text", "text": "Convert this page to docling."} ] }, ] # Prepare inputs prompt = self.processor.apply_chat_template(messages, add_generation_prompt=True) inputs = self.processor(text=prompt, images=[image], return_tensors="pt") inputs = inputs.to(self.device) # Generate outputs with torch.no_grad(): generated_ids = self.model.generate(**inputs, max_new_tokens=8192) prompt_length = inputs.input_ids.shape[1] trimmed_generated_ids = generated_ids[:, prompt_length:] doctags = self.processor.batch_decode( trimmed_generated_ids, skip_special_tokens=False, )[0].lstrip() # Convert DocTags to text try: doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctags], [image]) doc = DoclingDocument(name=f"Page_{page_num + 1}") doc.load_from_doctags(doctags_doc) # Export as markdown and extract text page_text = doc.export_to_markdown() # Add page separator and content all_text += f"\n--- Page {page_num + 1} ---\n{page_text}\n" except Exception as e: logger.warning(f"Failed to convert DocTags for page {page_num + 1}: {str(e)}") # Fallback: extract text directly from DocTags page_text = self._extract_text_from_doctags(doctags) all_text += f"\n--- Page {page_num + 1} ---\n{page_text}\n" logger.info(f"Successfully extracted text from {len(images)} pages using SmolDocling") return all_text.strip() except Exception as e: logger.error(f"Error extracting text with SmolDocling: {str(e)}") logger.info("Falling back to PyMuPDF") return self._extract_text_pymupdf(pdf_path) def _extract_text_pymupdf(self, pdf_path: Path) -> str: """Fallback method using PyMuPDF""" try: doc = fitz.open(pdf_path) text = "" for page_num in range(len(doc)): page = doc.load_page(page_num) page_text = page.get_text() text += f"\n--- Page {page_num + 1} ---\n{page_text}\n" doc.close() return text.strip() except Exception as e: logger.error(f"Error extracting text from PDF with PyMuPDF: {str(e)}") raise Exception(f"Failed to extract text from PDF: {str(e)}") def _extract_text_from_doctags(self, doctags: str) -> str: """Extract plain text from DocTags as fallback""" try: # Simple text extraction from DocTags import re # Remove XML-like tags and extract text content text = re.sub(r'<[^>]+>', '', doctags) text = re.sub(r'\s+', ' ', text) # Normalize whitespace return text.strip() except Exception as e: logger.warning(f"Failed to extract text from DocTags: {str(e)}") return "Failed to extract text from this page" def count_tokens(self, text: str) -> int: """Count tokens in text""" return len(self.encoding.encode(text)) def chunk_text(self, text: str) -> List[Dict[str, Any]]: """Split text into chunks based on token limits with proper management""" chunks = [] # Calculate total tokens for the entire document total_tokens = self.count_tokens(text) logger.info(f"Total document tokens: {total_tokens}") # If text is within token limit, return as single chunk if total_tokens <= self.chunk_size: return [{ "chunk_id": 0, "text": text, "tokens": total_tokens, "page_range": "all", "original_length": len(text) }] # Split by pages first pages = text.split("--- Page ") current_chunk = "" current_tokens = 0 chunk_id = 0 start_page = 1 total_processed_tokens = 0 logger.info(f"Processing {len(pages)-1} pages into chunks") for i, page in enumerate(pages): if i == 0: # Skip empty first split continue page_text = f"--- Page {page}" page_tokens = self.count_tokens(page_text) # If single page exceeds chunk size, split it further if page_tokens > self.chunk_size: logger.info(f"Page {i} has {page_tokens} tokens, splitting further") # Save current chunk if it has content if current_chunk: chunks.append({ "chunk_id": chunk_id, "text": current_chunk, "tokens": current_tokens, "page_range": f"{start_page}-{i-1}", "original_length": len(current_chunk) }) total_processed_tokens += current_tokens chunk_id += 1 # Split large page into smaller chunks page_chunks = self._split_large_page(page_text, page_tokens, chunk_id, i) chunks.extend(page_chunks) chunk_id += len(page_chunks) total_processed_tokens += page_tokens # Reset for next chunk current_chunk = "" current_tokens = 0 start_page = i + 1 # If adding this page would exceed chunk size, save current chunk elif current_tokens + page_tokens > self.chunk_size: if current_chunk: chunks.append({ "chunk_id": chunk_id, "text": current_chunk, "tokens": current_tokens, "page_range": f"{start_page}-{i-1}", "original_length": len(current_chunk) }) total_processed_tokens += current_tokens chunk_id += 1 # Start new chunk with current page current_chunk = page_text current_tokens = page_tokens start_page = i else: # Add page to current chunk if current_chunk: current_chunk += "\n" + page_text else: current_chunk = page_text current_tokens += page_tokens # Add final chunk if it has content if current_chunk: chunks.append({ "chunk_id": chunk_id, "text": current_chunk, "tokens": current_tokens, "page_range": f"{start_page}-{len(pages)-1}", "original_length": len(current_chunk) }) total_processed_tokens += current_tokens logger.info(f"Created {len(chunks)} chunks, total processed tokens: {total_processed_tokens}") # Verify we didn't lose content if abs(total_processed_tokens - total_tokens) > 100: # Allow small variance logger.warning(f"Token count mismatch: original={total_tokens}, processed={total_processed_tokens}") return chunks def _split_large_page(self, page_text: str, page_tokens: int, start_chunk_id: int, page_num: int) -> List[Dict[str, Any]]: """Split a large page into smaller chunks""" chunks = [] lines = page_text.split('\n') current_chunk = "" current_tokens = 0 chunk_id = start_chunk_id logger.info(f"Splitting page {page_num} with {page_tokens} tokens into smaller chunks") for line in lines: line_tokens = self.count_tokens(line) if current_tokens + line_tokens > self.chunk_size: if current_chunk: chunks.append({ "chunk_id": chunk_id, "text": current_chunk, "tokens": current_tokens, "page_range": f"page-{page_num}-part-{chunk_id-start_chunk_id+1}", "original_length": len(current_chunk) }) chunk_id += 1 current_chunk = line current_tokens = line_tokens else: current_chunk += "\n" + line if current_chunk else line current_tokens += line_tokens # Add final chunk if current_chunk: chunks.append({ "chunk_id": chunk_id, "text": current_chunk, "tokens": current_tokens, "page_range": f"page-{page_num}-part-{chunk_id-start_chunk_id+1}", "original_length": len(current_chunk) }) logger.info(f"Split page {page_num} into {len(chunks)} chunks") return chunks def get_text_preview(self, text: str, max_chars: int = 500) -> str: """Get a preview of the text""" if len(text) <= max_chars: return text return text[:max_chars] + "..."