pdf_processor.py

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] + "..."