74 lines
2.7 KiB
Python
74 lines
2.7 KiB
Python
"""EPUB 电子书分块器 — 使用 ebooklib 提取文字后委托 TextSplitter."""
|
|
import logging
|
|
from src.core.splitters.text import TextSplitter
|
|
|
|
logger = logging.getLogger("md-vector-db")
|
|
|
|
|
|
class EPUBSplitter:
|
|
"""EPUB 分块器：ebooklib 提取各章节文字 → TextSplitter 分块.
|
|
|
|
实现 Splitter Protocol，内部组合 TextSplitter 实例。
|
|
split() 的 text 参数实际接收 EPUB 文件路径（非文本内容）。
|
|
"""
|
|
|
|
def __init__(self, max_size: int = 1000, overlap: int = 100):
|
|
self._text_splitter = TextSplitter(max_size=max_size, overlap=overlap)
|
|
|
|
def split(self, text: str, source_file: str = "") -> list[dict]:
|
|
"""从 EPUB 文件提取各章节文字并分块.
|
|
|
|
Args:
|
|
text: EPUB 文件路径（非文本内容，由 ingest_file 传入）
|
|
source_file: 来源文件名
|
|
"""
|
|
try:
|
|
import ebooklib
|
|
from ebooklib import epub
|
|
except ImportError:
|
|
raise ImportError(
|
|
"EPUB 支持需要 ebooklib 库. 请执行: uv sync --extra epub"
|
|
)
|
|
|
|
epub_path = text
|
|
try:
|
|
book = epub.read_epub(epub_path)
|
|
except Exception as e:
|
|
logger.error("EPUB 解析失败: %s — %s", epub_path, e)
|
|
raise ValueError(f"EPUB 解析失败: {e}") from e
|
|
|
|
extracted_chapters = []
|
|
for item in book.get_items_of_type(ebooklib.ITEM_DOCUMENT):
|
|
try:
|
|
content = item.get_content().decode("utf-8")
|
|
except UnicodeDecodeError:
|
|
try:
|
|
content = item.get_content().decode("utf-8-sig")
|
|
except UnicodeDecodeError:
|
|
try:
|
|
content = item.get_content().decode("latin-1")
|
|
except UnicodeDecodeError:
|
|
logger.warning("EPUB 跳过一个无法解码的章节: %s", item.get_name())
|
|
continue
|
|
|
|
# 用 BeautifulSoup 去标签（如果可用），否则保留原样
|
|
try:
|
|
from bs4 import BeautifulSoup
|
|
soup = BeautifulSoup(content, "html.parser")
|
|
for tag in soup(["script", "style"]):
|
|
tag.decompose()
|
|
text_content = soup.get_text(separator="\n")
|
|
except ImportError:
|
|
# 无 bs4 时手动去除简单标签
|
|
import re
|
|
text_content = re.sub(r"<[^>]+>", "", content)
|
|
|
|
if text_content.strip():
|
|
extracted_chapters.append(text_content)
|
|
|
|
if not extracted_chapters:
|
|
return []
|
|
|
|
full_text = "\n\n".join(extracted_chapters)
|
|
return self._text_splitter.split(full_text, source_file=source_file)
|