feat: 添加Markdown文件处理器以支持分块处理
新增markdown_processor.py文件,提供Markdown文件的分块处理功能。该处理器能够根据一级标题将内容分块,并使用LLM处理每个块,最终将处理后的内容保存为新的Markdown文件。此功能适用于处理大文件,避免超出LLM的输入限制。
This commit is contained in:
parent
18774df43d
commit
490a6d6eec
137
markdown_processor.py
Normal file
137
markdown_processor.py
Normal file
@ -0,0 +1,137 @@
|
|||||||
|
import os
|
||||||
|
import re
|
||||||
|
import argparse
|
||||||
|
import openai
|
||||||
|
from typing import List, Dict, Any
|
||||||
|
|
||||||
|
# 设置OpenAI API
|
||||||
|
openai.api_key = "sk-no-key-required"
|
||||||
|
openai.api_base = "http://localhost:1234/v1" # LM-Studio默认地址
|
||||||
|
|
||||||
|
def read_markdown_file(file_path: str) -> str:
|
||||||
|
"""读取Markdown文件内容"""
|
||||||
|
try:
|
||||||
|
with open(file_path, 'r', encoding='utf-8') as file:
|
||||||
|
return file.read()
|
||||||
|
except Exception as e:
|
||||||
|
print(f"读取文件时出错: {e}")
|
||||||
|
return ""
|
||||||
|
|
||||||
|
def split_by_headers(content: str, max_length: int) -> List[str]:
|
||||||
|
"""根据一级标题(#)将内容分块,确保每个块小于max_length"""
|
||||||
|
# 查找所有一级标题的位置
|
||||||
|
pattern = r'^# .+$'
|
||||||
|
headers = [(m.start(), m.group()) for m in re.finditer(pattern, content, re.MULTILINE)]
|
||||||
|
|
||||||
|
if not headers:
|
||||||
|
# 如果没有一级标题,将整个内容作为一个块
|
||||||
|
return [content] if len(content) <= max_length else chunk_content(content, max_length)
|
||||||
|
|
||||||
|
chunks = []
|
||||||
|
current_chunk = ""
|
||||||
|
last_pos = 0
|
||||||
|
|
||||||
|
# 遍历所有标题
|
||||||
|
for i, (pos, header) in enumerate(headers):
|
||||||
|
# 获取当前标题到下一个标题之间的内容
|
||||||
|
if i == 0 and pos > 0:
|
||||||
|
# 处理文件开头到第一个标题之间的内容
|
||||||
|
current_chunk = content[:pos]
|
||||||
|
|
||||||
|
# 当前标题的内容
|
||||||
|
section_end = headers[i+1][0] if i+1 < len(headers) else len(content)
|
||||||
|
section = content[pos:section_end]
|
||||||
|
|
||||||
|
# 检查添加当前部分是否会超过最大长度
|
||||||
|
if len(current_chunk) + len(section) <= max_length:
|
||||||
|
current_chunk += section
|
||||||
|
else:
|
||||||
|
# 如果当前块不为空,添加到chunks
|
||||||
|
if current_chunk:
|
||||||
|
chunks.append(current_chunk)
|
||||||
|
|
||||||
|
# 如果单个部分超过最大长度,需要进一步分割
|
||||||
|
if len(section) > max_length:
|
||||||
|
sub_chunks = chunk_content(section, max_length)
|
||||||
|
chunks.extend(sub_chunks)
|
||||||
|
current_chunk = ""
|
||||||
|
else:
|
||||||
|
current_chunk = section
|
||||||
|
|
||||||
|
last_pos = section_end
|
||||||
|
|
||||||
|
# 添加最后一个块
|
||||||
|
if current_chunk:
|
||||||
|
chunks.append(current_chunk)
|
||||||
|
|
||||||
|
return chunks
|
||||||
|
|
||||||
|
def chunk_content(content: str, max_length: int) -> List[str]:
|
||||||
|
"""将内容分割成固定大小的块"""
|
||||||
|
chunks = []
|
||||||
|
for i in range(0, len(content), max_length):
|
||||||
|
chunks.append(content[i:i + max_length])
|
||||||
|
return chunks
|
||||||
|
|
||||||
|
def process_chunk_with_llm(chunk: str, model: str = "gpt-3.5-turbo") -> str:
|
||||||
|
"""使用LLM处理每个块"""
|
||||||
|
try:
|
||||||
|
response = openai.ChatCompletion.create(
|
||||||
|
model=model,
|
||||||
|
messages=[
|
||||||
|
{"role": "system", "content": "你是一个有用的助手,请处理以下Markdown内容。"},
|
||||||
|
{"role": "user", "content": chunk}
|
||||||
|
],
|
||||||
|
temperature=0.7,
|
||||||
|
)
|
||||||
|
return response.choices[0].message.content
|
||||||
|
except Exception as e:
|
||||||
|
print(f"处理块时出错: {e}")
|
||||||
|
return chunk # 出错时返回原始内容
|
||||||
|
|
||||||
|
def save_markdown_file(content: str, output_path: str) -> None:
|
||||||
|
"""保存处理后的Markdown内容到文件"""
|
||||||
|
try:
|
||||||
|
with open(output_path, 'w', encoding='utf-8') as file:
|
||||||
|
file.write(content)
|
||||||
|
print(f"已保存处理后的文件到: {output_path}")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"保存文件时出错: {e}")
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description='处理Markdown文件')
|
||||||
|
parser.add_argument('input_file', help='输入的Markdown文件路径')
|
||||||
|
parser.add_argument('output_file', help='输出的Markdown文件路径')
|
||||||
|
parser.add_argument('--max_length', type=int, default=4000, help='每个块的最大长度')
|
||||||
|
parser.add_argument('--model', default='gpt-3.5-turbo', help='使用的LLM模型名称')
|
||||||
|
parser.add_argument('--api_base', default='http://localhost:1234/v1', help='API基础URL')
|
||||||
|
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
# 设置API基础URL
|
||||||
|
openai.api_base = args.api_base
|
||||||
|
|
||||||
|
# 读取文件
|
||||||
|
content = read_markdown_file(args.input_file)
|
||||||
|
if not content:
|
||||||
|
return
|
||||||
|
|
||||||
|
# 分块
|
||||||
|
chunks = split_by_headers(content, args.max_length)
|
||||||
|
print(f"文件已分成 {len(chunks)} 个块")
|
||||||
|
|
||||||
|
# 处理每个块
|
||||||
|
processed_chunks = []
|
||||||
|
for i, chunk in enumerate(chunks):
|
||||||
|
print(f"处理块 {i+1}/{len(chunks)}...")
|
||||||
|
processed_chunk = process_chunk_with_llm(chunk, args.model)
|
||||||
|
processed_chunks.append(processed_chunk)
|
||||||
|
|
||||||
|
# 合并处理后的内容
|
||||||
|
final_content = '\n'.join(processed_chunks)
|
||||||
|
|
||||||
|
# 保存结果
|
||||||
|
save_markdown_file(final_content, args.output_file)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
Loading…
x
Reference in New Issue
Block a user