目录

前言

新增PdfNewExtractor类

替换ExtractProcessor类

最终结果


前言


dify的1.1.3版本知识库pdf解析实现使用pypdfium2提取文本,主要存在以下问题:
1. 文本提取能力有限,对表格和图片支持不足
2. 缺乏专门的中文处理优化
3. 没有文档结构分析
4. 缺少文档质量评估
建议优化方案:
1. 使用pdfplumber替代pypdfium2
2. 增加OCR支持
3. 优化中文处理逻辑
4. 添加文档结构分析
5. 实现智能表格识别
6. 增加缓存机制
7. 优化大文件处理

导入包pdfplumber和pytesseract

pip install pdfplumber
pip install pytesseract

新增PdfNewExtractor类


新增一个PdfNewExtractor处理类替代老的PdfExtractor

from collections.abc import Iterator
from typing import Optional, cast
import pdfplumber
import pytesseract
from PIL import Image
import io

from core.rag.extractor.blob.blob import Blob
from core.rag.extractor.extractor_base import BaseExtractor
from core.rag.models.document import Document
from extensions.ext_storage import storage

class PdfNewExtractor(BaseExtractor):
    """Enhanced PDF loader with improved text extraction, OCR support, and structure analysis.

    Args:
        file_path: Path to the PDF file to load.
        file_cache_key: Optional cache key for storing extracted text.
        enable_ocr: Whether to enable OCR for text extraction from images.
    """

    def __init__(self, file_path: str, file_cache_key: Optional[str] = None, enable_ocr: bool = False):
        """Initialize with file path and optional settings."""
        self._file_path = file_path
        self._file_cache_key = file_cache_key
        self._enable_ocr = enable_ocr

    def extract(self) -> list[Document]:
        """Extract text from PDF with caching support."""
        plaintext_file_exists = False
        if self._file_cache_key:
            try:
                text = cast(bytes, storage.load(self._file_cache_key)).decode("utf-8")
                plaintext_file_exists = True
                return [Document(page_content=text)]
            except FileNotFoundError:
                pass

        documents = list(self.load())
        text_list = []
        for document in documents:
            text_list.append(document.page_content)
        text = "\n\n".join(text_list)

        # Save plaintext file for caching
        if not plaintext_file_exists and self._file_cache_key:
            storage.save(self._file_cache_key, text.encode("utf-8"))

        return documents

    def load(self) -> Iterator[Document]:
        """Lazy load PDF pages with enhanced text extraction."""
        blob = Blob.from_path(self._file_path)
        yield from self.parse(blob)

    def parse(self, blob: Blob) -> Iterator[Document]:
        """Parse PDF with enhanced features including OCR and structure analysis."""
        with blob.as_bytes_io() as file_obj:
            with pdfplumber.open(file_obj) as pdf:
                for page_number, page in enumerate(pdf.pages):
                    # Extract text with layout preservation and encoding detection
                    content = page.extract_text(layout=True)
                    # Try to detect and fix encoding issues
                    try:
                        # First try to decode as UTF-8
                        content = content.encode('utf-8').decode('utf-8')
                    except UnicodeError:
                        try:
                            # If UTF-8 fails, try GB18030 (common Chinese encoding)
                            content = content.encode('utf-8').decode('gb18030', errors='ignore')
                        except UnicodeError:
                            # If all else fails, use a more lenient approach
                            content = content.encode('utf-8', errors='ignore').decode('utf-8', errors='ignore')
                    
                    # Extract tables if present
                    tables = page.extract_tables()
                    if tables:
                        table_text = "\n\nTables:\n"
                        for table in tables:
                            # Convert table to text format
                            table_text += "\n" + "\n".join(
                                ["\t".join([str(cell) if cell else "" for cell in row]) 
                                 for row in table]
                            )
                        content += table_text

                    # Perform OCR if enabled and text content is limited or contains potential encoding issues
                    if self._enable_ocr and (len(content.strip()) < 100 or any('\ufffd' in line for line in content.splitlines())):
                        image = page.to_image()
                        img_bytes = io.BytesIO()
                        image.original.save(img_bytes, format='PNG')
                        img_bytes.seek(0)
                        pil_image = Image.open(img_bytes)
                        # Use multiple language models and improve OCR accuracy
                        ocr_text = pytesseract.image_to_string(
                            pil_image,
                            lang='chi_sim+chi_tra+eng',  # Support both simplified and traditional Chinese
                            config='--psm 3 --oem 3'  # Use more accurate OCR mode
                        )
                        if ocr_text.strip():
                            # Clean and normalize OCR text
                            ocr_text = ocr_text.replace('\x0c', '').strip()
                            content = f"{content}\n\nOCR Text:\n{ocr_text}"

                    metadata = {
                        "source": blob.source,
                        "page": page_number,
                        "has_tables": bool(tables)
                    }
                    
                    yield Document(page_content=content, metadata=metadata)

替换ExtractProcessor类


在ExtractProcessor中把两处extractor = PdfExtractor(file_path),替换成extractor = PdfNewExtractor(file_path)。
分别在代码144行和148行

最终结果


经过测试,优化效果完美

Logo

中国智能体开发者社区,聚焦智能体与大模型开发,提供前沿资讯、实用工具链、开源项目及行业案例。通过技术沙龙、开发者大赛等活动,促进经验交流与协作,助力开发者快速构建创新智能应用。

更多推荐