138 lines
4.4 KiB
Python
138 lines
4.4 KiB
Python
from langchain_text_splitters import MarkdownHeaderTextSplitter, CharacterTextSplitter
|
||
import requests
|
||
import json
|
||
|
||
from openai import OpenAI
|
||
|
||
def _split_markdown(text, chunk_size=2048, chunk_overlap=50):
|
||
"""
|
||
根据 max_level 动态生成 headers_to_split_on 配置
|
||
规则：
|
||
- max_level=1 → 只按 # 分割
|
||
- max_level=2 → 按 # 和 ## 分割
|
||
- 以此类推
|
||
"""
|
||
# 动态生成 headers_to_split_on 配置（目前固定为 Header_1）
|
||
headers_to_split_on = [("#", "Header_1")]
|
||
|
||
# 初始化分割器
|
||
splitter = MarkdownHeaderTextSplitter(
|
||
headers_to_split_on=headers_to_split_on,
|
||
strip_headers=True
|
||
)
|
||
chunks = splitter.split_text(text)
|
||
|
||
char_splitter = CharacterTextSplitter(
|
||
chunk_size=chunk_size,
|
||
chunk_overlap=chunk_overlap,
|
||
length_function=len,
|
||
separator="\n\n" # 更合理的分隔符
|
||
)
|
||
|
||
final_chunks = []
|
||
start_chars = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "0", "第", "十",
|
||
"一", "二", "三", "四", "五", "六", "七", "八", "九"]
|
||
new_header = []
|
||
|
||
# 提取所有 Header_1 并筛选出“有效标题”
|
||
for chunk in chunks:
|
||
header_1 = chunk.metadata.get('Header_1', '')
|
||
if header_1 and header_1[0] in start_chars:
|
||
new_header.append(header_1)
|
||
|
||
new_header_set = set(new_header)
|
||
|
||
# 构建带标题前缀的 final_chunks
|
||
for chunk in chunks:
|
||
header_1 = chunk.metadata.get('Header_1', '')
|
||
content = chunk.page_content
|
||
if len(content) <= 20000:
|
||
chunk_with_header = f"{header_1}\n{content}" if header_1 else content
|
||
final_chunks.append(chunk_with_header)
|
||
else:
|
||
split_texts = char_splitter.split_documents([chunk])
|
||
for split_chunk in split_texts:
|
||
chunk_with_header = f"{header_1}\n{split_chunk.page_content}" if header_1 else split_chunk.page_content
|
||
final_chunks.append(chunk_with_header)
|
||
|
||
# ===== 重新分组，并记录每个组的 header =====
|
||
rechunked = []
|
||
rechunked_headers = [] # 新增：每个 rechunked 块对应的 header
|
||
current = ""
|
||
current_header = "" # 当前正在构建的块所属的标题
|
||
|
||
for chunk in final_chunks:
|
||
is_new_section = False
|
||
matched_header = ""
|
||
|
||
# 检查是否以某个 new_header 开头
|
||
for h in new_header_set:
|
||
if chunk.startswith(h + "\n"):
|
||
is_new_section = True
|
||
matched_header = h
|
||
break
|
||
|
||
if is_new_section:
|
||
# 保存上一段
|
||
if current:
|
||
rechunked.append(current)
|
||
rechunked_headers.append(current_header)
|
||
# 开启新段
|
||
current = chunk
|
||
current_header = matched_header
|
||
else:
|
||
# 合并到当前段
|
||
if current:
|
||
current += "\n" + chunk
|
||
else:
|
||
current = chunk
|
||
current_header = "" # 如果开头无标题，则 header 为空
|
||
|
||
# 添加最后一段
|
||
if current:
|
||
rechunked.append(current)
|
||
rechunked_headers.append(current_header)
|
||
|
||
return rechunked, rechunked_headers
|
||
|
||
def read_markdown_file(file_path):
|
||
"""
|
||
读取 Markdown 文件并返回其全部内容作为字符串。
|
||
|
||
参数:
|
||
file_path (str): Markdown 文件的路径（支持相对或绝对路径）
|
||
|
||
返回:
|
||
str 或 None: 成功时返回文件内容；出错时返回 None
|
||
"""
|
||
try:
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
content = f.read()
|
||
return content
|
||
except FileNotFoundError:
|
||
print(f"错误：文件未找到 - {file_path}")
|
||
except PermissionError:
|
||
print(f"错误：没有权限读取文件 - {file_path}")
|
||
except UnicodeDecodeError:
|
||
print(f"错误：文件编码不是 UTF-8，无法读取 - {file_path}")
|
||
except Exception as e:
|
||
print(f"读取文件时发生未知错误: {e}")
|
||
return None
|
||
|
||
|
||
|
||
|
||
|
||
|
||
# if __name__ == "__main__":
|
||
|
||
|
||
# # 测试代码
|
||
# file_path = "F:\zklnlp\HJ\code\output.md" # 替换为你的 Markdown 文件路径
|
||
# markdown_content = read_markdown_file(file_path)
|
||
# content, headers = _split_markdown(markdown_content, chunk_size=3000, chunk_overlap=100)
|
||
# for ins in content:
|
||
# print(ins)
|
||
# print(1111111111111111111111111)
|
||
|