90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""DOCX → Markdown via mammoth (preserves headings, lists, tables)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import re
|
|
from pathlib import Path
|
|
|
|
from . import ParsedImage, ParsedPage, ParseResult
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Force common DOCX heading style names to Markdown headings.
|
|
_STYLE_MAP = """
|
|
p[style-name='Title'] => h1.title:fresh
|
|
p[style-name='Subtitle'] => h2.subtitle:fresh
|
|
p[style-name='Heading 1'] => h1:fresh
|
|
p[style-name='Heading 2'] => h2:fresh
|
|
p[style-name='Heading 3'] => h3:fresh
|
|
p[style-name='Heading 4'] => h4:fresh
|
|
p[style-name='Heading 5'] => h5:fresh
|
|
p[style-name='Heading 6'] => h6:fresh
|
|
"""
|
|
|
|
|
|
def _html_to_markdown(html: str) -> str:
|
|
from markdownify import markdownify
|
|
|
|
md = markdownify(html, heading_style = "ATX", strip = ["script", "style"])
|
|
md = re.sub(r"\n{3,}", "\n\n", md)
|
|
return md.strip()
|
|
|
|
|
|
def extract(path: Path, *, want_images: bool = False) -> ParseResult:
|
|
import mammoth
|
|
|
|
images: list[ParsedImage] = []
|
|
|
|
if want_images:
|
|
# Capture bytes; suppress src so base64 doesn't land in Markdown.
|
|
def _convert(image):
|
|
with image.open() as image_bytes:
|
|
blob = image_bytes.read()
|
|
mime = (image.content_type or "application/octet-stream").lower()
|
|
images.append(
|
|
ParsedImage(
|
|
image_bytes = blob,
|
|
mime_type = mime,
|
|
page_number = None,
|
|
nearest_caption = "",
|
|
)
|
|
)
|
|
return {"src": ""}
|
|
|
|
convert_image = mammoth.images.img_element(_convert)
|
|
else:
|
|
convert_image = mammoth.images.img_element(lambda _image: {"src": ""})
|
|
|
|
with open(path, "rb") as fp:
|
|
result = mammoth.convert_to_html(
|
|
fp,
|
|
convert_image = convert_image,
|
|
style_map = _STYLE_MAP,
|
|
)
|
|
for message in result.messages:
|
|
logger.debug("mammoth %s: %s", getattr(message, "type", "msg"), message.message)
|
|
|
|
markdown = _html_to_markdown(result.value)
|
|
if want_images and images:
|
|
# Approximate caption: first 1500 chars of doc.
|
|
caption_pool = markdown[:1500]
|
|
images = [
|
|
ParsedImage(
|
|
image_bytes = img.image_bytes,
|
|
mime_type = img.mime_type,
|
|
page_number = img.page_number,
|
|
nearest_caption = caption_pool,
|
|
)
|
|
for img in images
|
|
]
|
|
if not markdown:
|
|
return ParseResult(pages = [], images = images)
|
|
return ParseResult(
|
|
pages = [ParsedPage(text = markdown, page_number = None)],
|
|
images = images,
|
|
)
|