关于RAG文件格式转换

需求：pdf格式转换为md，doc，docx转换为md，用于在知识库中建立索引

遇到的问题，pdf中包含了一些图片和cad格式图，转换为md后数据丢失，解决方法是md+png+svg。

pymupdf4llm 转换pdf，生成md、png、svg。

LibreOffice 转换doc，docx为md。

分工合作，最终的效果还可以。我也尝试了docling，但pdf里面的cad图有丢失，也可能是我脚本写的不对。最终用pynupdf4llm达到了效果。

代码如下：

#!/usr/bin/env python3
"""
批量将 data/ 下的 PDF / DOC / DOCX 转为 Markdown，输出到 data_md/。

- PDF  : PyMuPDF4LLM 生成 Markdown + PNG（位图，用于多模态理解）
         同时用 PyMuPDF 导出 SVG（矢量 + 可检索文字，用于文本索引）
- DOCX : Docling（保留标题、列表、表格、公式）
- DOC  : 先用 LibreOffice 转 DOCX，再用 Docling

输出结构示例：
    data_md/
      └─ 子目录结构与 data/ 一致
          ├─ 报告.md
          ├─ images/               # PNG（多模态理解用）
          │     ├─ page-1-image-0.png
          │     └─ ...
          └─ svg/                  # SVG（矢量 + 可检索文字）
                ├─ page-1.svg
                └─ ...
"""

import os
import shutil
import subprocess
import sys
import traceback
from pathlib import Path

import pymupdf
import pymupdf4llm

# --------------------------------------------------------------------------
# 路径与配置（基于脚本自身位置，双击运行也不会错）
# --------------------------------------------------------------------------
BASE = Path(__file__).resolve().parent
INPUT_DIR = BASE / "data"
OUTPUT_DIR = BASE / "data_md"
SUPPORTED_EXTS = {".pdf", ".docx", ".doc"}

PDF_DPI = 300          # CAD 图纸建议 300；普通文档 150 足够
EXPORT_SVG = True      # 是否导出 SVG（矢量 + 可检索文字）


# --------------------------------------------------------------------------
# 自动探测 LibreOffice（处理 .doc 和 DOCX 中的 DrawingML）
# --------------------------------------------------------------------------
def setup_libreoffice():
    if os.environ.get("DOCLING_LIBREOFFICE_CMD"):
        return
    for name in ("soffice", "libreoffice"):
        p = shutil.which(name)
        if p:
            os.environ["DOCLING_LIBREOFFICE_CMD"] = p
            return
    candidates = [
        r"C:\Program Files\LibreOffice\program\soffice.exe",
        r"C:\Program Files (x86)\LibreOffice\program\soffice.exe",
        "/Applications/LibreOffice.app/Contents/MacOS/soffice",
        "/usr/bin/soffice",
    ]
    for p in candidates:
        if Path(p).exists():
            os.environ["DOCLING_LIBREOFFICE_CMD"] = p
            return
    print("提示：未找到 LibreOffice，DOCX 中的 DrawingML 元素可能无法提取。")


setup_libreoffice()

# Docling 转换器全局复用（延迟导入，避免未安装时直接报错）
_converter = None


def get_converter():
    global _converter
    if _converter is None:
        from docling.document_converter import DocumentConverter
        _converter = DocumentConverter()
    return _converter


# --------------------------------------------------------------------------
# PDF
# --------------------------------------------------------------------------
def convert_pdf(src: Path, out_dir: Path) -> str:
    """PDF -> Markdown，PNG 写入 out_dir/images/，SVG 写入 out_dir/svg/。"""
    import re

    src = src.resolve()
    image_dir = (out_dir / "images").resolve()
    svg_dir = (out_dir / "svg").resolve()
    image_dir.mkdir(parents=True, exist_ok=True)
    svg_dir.mkdir(parents=True, exist_ok=True)

    doc = pymupdf.open(str(src))

    # ---- 第 1 步：先导出干净的 SVG（在画透明边框之前）----
    svg_files = {}
    if EXPORT_SVG:
        for page in doc:
            page_no = page.number + 1
            try:
                svg = page.get_svg_image(text_as_path=False)
                fname = f"page-{page_no}.svg"
                (svg_dir / fname).write_text(svg, encoding="utf-8")
                svg_files[page_no] = fname
            except Exception as e:
                print(f"  [WARN] page {page_no} SVG 导出失败: {e}")

    # ---- 第 2 步：给"无边框白色图形"补透明边框 ----
    for page in doc:
        try:
            for bbox in page.cluster_drawings():
                page.draw_rect(bbox, stroke_opacity=0)
        except Exception as e:
            print(f"  [WARN] page {page.number + 1} 透明边框添加失败: {e}")

    # ---- 第 3 步：PyMuPDF4LLM 提取文本 + PNG ----
    try:
        chunks = pymupdf4llm.to_markdown(
            doc,
            page_chunks=True,
            write_images=True,
            image_path=str(image_dir),
            dpi=PDF_DPI,
            image_format="png",
            image_size_limit=0.01,
            force_text=False, 
            margins=(0, 0, 0, 0),
        )
    finally:
        doc.close()

    # ---- 第 4 步：拼装 Markdown，每页末尾插入 SVG 引用 ----
    parts = []
    for i, chunk in enumerate(chunks):
        page_no = chunk.get("metadata", {}).get("page", i) + 1
        text = chunk["text"]
        svg_ref = ""
        if page_no in svg_files:
            svg_ref = f"\n\n[矢量图](svg/{svg_files[page_no]})"
        parts.append(f"<!-- page {page_no} -->\n\n{text}{svg_ref}")
    md = "\n\n".join(parts)

    # ---- 第 5 步：用正则统一重写图片引用 ----
    def _fix_img_paths(text: str) -> str:
        def repl(m):
            alt, path, close = m.group(1), m.group(2), m.group(3)
            fname = Path(path.replace("\\", "/")).name
            return f"{alt}images/{fname}{close}"
        return re.sub(r'(!\[[^\]]*\]\()([^)]+)(\))', repl, text)

    md = _fix_img_paths(md)
    return md


# --------------------------------------------------------------------------
# DOCX
# --------------------------------------------------------------------------
def convert_docx(src: Path) -> str:
    """DOCX -> Markdown (Docling)。"""
    result = get_converter().convert(str(src))
    return result.document.export_to_markdown()


# --------------------------------------------------------------------------
# DOC（旧版 Word）
# --------------------------------------------------------------------------
def convert_doc_to_docx(src: Path) -> Path:
    soffice = os.environ.get("DOCLING_LIBREOFFICE_CMD")
    if not soffice:
        raise RuntimeError(
            "未找到 LibreOffice，无法处理 .doc。"
            "请安装 LibreOffice，或手动将 .doc 另存为 .docx。"
        )
    tmp_dir = BASE / "_tmp_doc"
    tmp_dir.mkdir(exist_ok=True)
    subprocess.run(
        [soffice, "--headless", "--convert-to", "docx",
         "--outdir", str(tmp_dir), str(src)],
        check=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    return tmp_dir / (src.stem + ".docx")


# --------------------------------------------------------------------------
# 单个文件处理
# --------------------------------------------------------------------------
def process(src: Path) -> Path:
    rel = src.relative_to(INPUT_DIR)
    out_dir = OUTPUT_DIR / rel.parent
    out_dir.mkdir(parents=True, exist_ok=True)
    out_md = out_dir / (src.stem + ".md")

    ext = src.suffix.lower()

    if ext == ".pdf":
        md = convert_pdf(src, out_dir)
    elif ext == ".docx":
        md = convert_docx(src)
    elif ext == ".doc":
        docx_path = convert_doc_to_docx(src)
        try:
            md = convert_docx(docx_path)
        finally:
            try:
                docx_path.unlink()
            except OSError:
                pass
    else:
        raise ValueError(f"不支持的后缀: {ext}")

    out_md.write_text(md, encoding="utf-8")
    return out_md


# --------------------------------------------------------------------------
# 主流程
# --------------------------------------------------------------------------
def main() -> None:
    if not INPUT_DIR.exists():
        sys.exit(f"输入目录不存在: {INPUT_DIR}")

    files = sorted(
        p for p in INPUT_DIR.rglob("*")
        if p.is_file() and p.suffix.lower() in SUPPORTED_EXTS
    )
    if not files:
        print(f"{INPUT_DIR} 中没有可处理的文件（仅支持 {SUPPORTED_EXTS}）")
        return

    print(f"发现 {len(files)} 个文件，开始转换......\n")

    ok, fail = 0, 0
    for src in files:
        try:
            out = process(src)
            print(f"  [OK]   {src.relative_to(INPUT_DIR)}  ->  {out.relative_to(OUTPUT_DIR)}")
            ok += 1
        except Exception as e:
            print(f"  [FAIL] {src}: {e}")
            traceback.print_exc()
            fail += 1

    shutil.rmtree(BASE / "_tmp_doc", ignore_errors=True)
    print(f"\n完成：成功 {ok}，失败 {fail}")


if __name__ == "__main__":
    try:
        main()
    except SystemExit as e:
        print(f"\n程序退出: {e}")
    except Exception:
        print("\n===== 发生错误 =====")
        traceback.print_exc()
    finally:
        input("\n处理结束，按回车键关闭窗口...")