718 lines
28 KiB
Python
718 lines
28 KiB
Python
import json
|
||
import re
|
||
from typing import Dict, List
|
||
from langchain_text_splitters import CharacterTextSplitter, MarkdownHeaderTextSplitter,MarkdownTextSplitter,TokenTextSplitter,RecursiveCharacterTextSplitter
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
from markdown_process.save_markdown import Save_Markdown_to_Milvus
|
||
from markdown_process.convert_header import convert_numbered_headings, custom_convert_numbered_headings,parse_markdown,header_postprocess,postprocess_chunks
|
||
from markdown_it import MarkdownIt
|
||
from markdown_it.tree import SyntaxTreeNode
|
||
def reset_chunk(chunks):
|
||
id_list = []
|
||
for res in chunks:
|
||
if res['Header_1'] == res['text']:
|
||
id = res['id']
|
||
id_list.append(id)
|
||
# print(id_list)
|
||
for iditem in id_list:
|
||
result = [item for item in chunks if item['id'] == iditem][0]
|
||
result_ = [item for item in chunks if item['id'] == iditem + 1][0]
|
||
result_['Header_2'] = result_['Header_1']
|
||
result_['Header_1'] = result['Header_1']
|
||
result_['text'] = result['text'] + '\n' + result_['text']
|
||
result_devide = [record for record in chunks if not record['Header_1'] == record['text']]
|
||
for index, item in enumerate(result_devide):
|
||
item['id'] = index + 1
|
||
return result_devide
|
||
def remove_markdown_images(text):
|
||
"""
|
||
移除Markdown格式的图像链接
|
||
|
||
参数:
|
||
text: 包含Markdown图像链接的原始字符串
|
||
|
||
返回:
|
||
移除图像链接后的文本
|
||
"""
|
||
# 匹配  或 [](path) 格式
|
||
pattern = r'!?\[\]\(/root/output/[^/]+/images/[a-f0-9]+\.(?:jpg|png|gif|bmp)\)'
|
||
return re.sub(pattern, '', text)
|
||
def remove_markdown_images(text):
|
||
"""
|
||
移除Markdown格式的图像链接
|
||
|
||
参数:
|
||
text: 包含Markdown图像链接的原始字符串
|
||
|
||
返回:
|
||
移除图像链接后的文本
|
||
"""
|
||
# 匹配  或 [](path) 格式
|
||
pattern = r'http://\d{0,9}\.\d{0,9}\.\d{0,9}\.\d{0,9}:\d+/[^\s\'"<>]+'
|
||
return re.sub(pattern, '', text)
|
||
|
||
def extract_urls(text):
|
||
# 定义匹配 URL 的正则表达式模式
|
||
pattern = r'http://\d{0,9}\.\d{0,9}\.\d{0,9}\.\d{0,9}:\d{0,9}/output/images/[a-f0-9]+\.(?:jpg|jpeg|png|gif)'
|
||
|
||
# 查找所有匹配的 URL
|
||
urls = re.findall(pattern, text)
|
||
|
||
# # 可选：过滤只保留特定扩展名的文件（如.jpg, .png等）
|
||
# image_urls = [url for url in urls if url.lower().endswith(('.jpg', '.jpeg', '.png', '.gif'))]
|
||
|
||
return urls
|
||
|
||
def extract_and_convert_image_paths(text):
|
||
"""
|
||
从字符串中提取图像路径并转换为新格式
|
||
|
||
参数:
|
||
text: 包含图像路径的原始字符串
|
||
|
||
返回:
|
||
包含原始路径和转换后路径的字典列表
|
||
"""
|
||
# 正则表达式匹配图像路径
|
||
pattern = r'(/root/output/[^/]+/images/[a-f0-9]+\.(?:jpg|png|gif|bmp))'
|
||
matches = re.findall(pattern, text)
|
||
|
||
converted_paths = []
|
||
|
||
for original_path in matches:
|
||
# 分解原始路径
|
||
path_parts = original_path.split('/')
|
||
|
||
# 提取关键部分
|
||
# 格式: /root/output/<文件夹>/images/<图片名>
|
||
folder_name = path_parts[3] # 获取"生成式人工智能传播风险：理论内涵、形成机理与治理策略_景奎"这部分
|
||
image_name = path_parts[-1] # 获取图片文件名
|
||
|
||
# 构建新路径
|
||
new_path = f"mineru/{folder_name}/images/{image_name}"
|
||
|
||
converted_paths.append({
|
||
'original_path': original_path,
|
||
'converted_path': new_path
|
||
})
|
||
|
||
return converted_paths
|
||
|
||
class MarkdownAwareTextSplitter(RecursiveCharacterTextSplitter):
|
||
def __init__(self, **kwargs):
|
||
# 优先用Markdown的分隔符
|
||
separators = [
|
||
# 保持表格完整 (匹配Markdown表格)
|
||
r'\n\s*\|', # 表格行
|
||
r'\n#+\s+', # 标题
|
||
r'\n\*\*\s+', # 分隔线
|
||
r'\n\n', # 双换行
|
||
r'\n', # 单换行
|
||
r'\.\s+', # 句号+空格
|
||
r'\s+', # 任意空白
|
||
]
|
||
super().__init__(separators=separators, **kwargs)
|
||
|
||
def split_text(self, text: str) -> List[str]:
|
||
# 预处理：保护表格结构
|
||
protected_tables = []
|
||
def _protect_tables(match):
|
||
protected_tables.append(match.group(0))
|
||
return f"__TABLE_{len(protected_tables)-1}__"
|
||
|
||
# 用临时标记替换表格
|
||
table_pattern = r'(\n\|.+\|.+\n(\|.+\|.+\n)+)'
|
||
text_with_protected_tables = re.sub(
|
||
table_pattern,
|
||
_protect_tables,
|
||
text,
|
||
flags=re.MULTILINE
|
||
)
|
||
|
||
# 常规切分
|
||
chunks = super().split_text(text_with_protected_tables)
|
||
|
||
# 恢复表格
|
||
final_chunks = []
|
||
for chunk in chunks:
|
||
for i, table in enumerate(protected_tables):
|
||
chunk = chunk.replace(f"__TABLE_{i}__", table)
|
||
final_chunks.append(chunk)
|
||
|
||
return final_chunks
|
||
|
||
def split_regular(resource,mdDocs,logger,database_name,collection_name,chunk_type,chunk_size,chunk_overlap=50,max_level=3):
|
||
|
||
if chunk_type =="03":
|
||
text_splitter =TokenTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||
texts = text_splitter.split_text(mdDocs)
|
||
milvus = Save_Markdown_to_Milvus(database_name=database_name,collection_name=collection_name)
|
||
for text in texts:
|
||
convert_path_list=extract_and_convert_image_paths(text)
|
||
text = remove_markdown_images(text)
|
||
if convert_path_list:
|
||
img_path = ",".join(convert_path_list)
|
||
vector = OpenaiAPI.get_embeddings(text)
|
||
data = [[text],[vector],[resource],[img_path]]
|
||
milvus.save_hybrid_md(data=data)
|
||
results.append({"id":j+i+1,"text":text,"resource":resource,"img_path":img_path})
|
||
milvus.close()
|
||
else:
|
||
#自定义分解
|
||
md_header_json = parse_markdown(mdDocs,max_level=max_level)
|
||
# print(md_header_json)
|
||
#合并处理
|
||
md_header_list = header_postprocess(md_header_json)
|
||
logger.info("markdown splitting...")
|
||
|
||
milvus = Save_Markdown_to_Milvus(database_name=database_name,collection_name=collection_name)
|
||
results = []
|
||
for i ,md in enumerate(md_header_list):
|
||
if md["text"]:
|
||
header_1 = md["Header_1"]
|
||
header_2 = md["Header_2"]
|
||
header_3 = md["Header_3"]
|
||
text = md["text"]
|
||
|
||
# print("Header_1:",header_1)
|
||
# print("Header_2:",header_2)
|
||
# print("Header_3:",header_3)
|
||
# print("text:",text_)
|
||
# print("======================================")
|
||
if chunk_type =="01":
|
||
text_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=chunk_size,
|
||
chunk_overlap=chunk_overlap,
|
||
length_function=len,
|
||
separators=[" ","\n\n"]
|
||
)
|
||
texts = text_splitter.split_text(text)
|
||
for j,txt in enumerate(texts):
|
||
img_list = extract_urls(txt)
|
||
txt = remove_markdown_images(txt)
|
||
if img_list:
|
||
img_path = ",".join(img_list)
|
||
# convert_path_list=extract_and_convert_image_paths(text)
|
||
# text = remove_markdown_images(text)
|
||
# if convert_path_list:
|
||
# img_path = ",".join(convert_path_list)
|
||
# print(j)
|
||
# print(texts)
|
||
# print("======================================")
|
||
vector = OpenaiAPI.get_embeddings(txt)
|
||
data = [[header_1],[header_2],[header_3],[txt],[vector],[resource],[img_path]]
|
||
# milvus.save_md(data=data)
|
||
milvus.save_hybrid_md(data=data)
|
||
results.append({"id":j+i+1,"Header_1":header_1,"Header_2":header_2,"Header_3":header_3,"text":txt,"resource":resource,"img_path":img_path})
|
||
|
||
else:
|
||
img_list = extract_urls(txt)
|
||
txt = remove_markdown_images(txt)
|
||
if img_list:
|
||
img_path = ",".join(img_list)
|
||
# convert_path_list=extract_and_convert_image_paths(text)
|
||
# text = remove_markdown_images(text)
|
||
# if convert_path_list:
|
||
# img_path = ",".join(convert_path_list)
|
||
vector = OpenaiAPI.get_embeddings(text)
|
||
data = [[header_1],[header_2],[header_3],[text],[vector],[resource],[img_path]]
|
||
milvus.save_hybrid_md(data=data)
|
||
results.append({"id":i+1,"Header_1":header_1,"Header_2":header_2,"Header_3":header_3,"text":txt,"resource":resource,"img_path":img_path})
|
||
milvus.close()
|
||
|
||
return results
|
||
|
||
def split_text(resource, mdDocs, logger, chunk_type, chunk_size, chunk_overlap=50,max_level=3):
|
||
"""
|
||
拆分文本的函数
|
||
|
||
Args:
|
||
resource: 资源标识
|
||
mdDocs: Markdown文档内容
|
||
logger: 日志记录器
|
||
chunk_type: 拆分类型("01"或"03")
|
||
"01": 智能分割(基于递归字符分割)
|
||
"03": 按字符分割(基于token分割)
|
||
chunk_size: 块长度
|
||
chunk_overlap: 块重叠大小(默认为50)
|
||
|
||
Returns:
|
||
list: 拆分后的文本块列表，每个元素是包含文本和元数据的字典
|
||
"""
|
||
#print(2222222222222222222222222222222222)
|
||
if chunk_type == "03":
|
||
logger.info(f"chunk size: {chunk_size},chunk_overlap: {chunk_overlap},chunk_type: {chunk_type}")
|
||
return _split_by_token(resource, mdDocs, logger, chunk_size, chunk_overlap)
|
||
else:
|
||
return _split_custom(resource, mdDocs, logger, chunk_type, chunk_size, chunk_overlap,max_level)
|
||
|
||
def split_page_chunk(resource,content_chunk_list,chunk_size,chunk_overlap,logger,max_level,chunk_type):
|
||
def process_content_list(parsed_data: List[Dict],max_level) -> List[Dict]:
|
||
"""
|
||
处理 parsed_data：
|
||
1. 将 text_level=1/2 的标题添加到后续无 text_level 的文本前
|
||
2. 合并 type=image 的项到前一个 chunk
|
||
"""
|
||
processed_data = []
|
||
current_titles = {level: "" for level in range(1, max_level + 1)} # 各级标题缓存
|
||
|
||
if isinstance(max_level, str):
|
||
try:
|
||
max_level = int(max_level)
|
||
except ValueError:
|
||
max_level = 3 # 默认值
|
||
|
||
for item in parsed_data:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
|
||
raw_text_level = item.get("text_level")
|
||
text_level = None
|
||
if raw_text_level is not None:
|
||
if isinstance(raw_text_level, str):
|
||
try:
|
||
text_level = int(raw_text_level.strip()) if raw_text_level.strip() else None
|
||
except ValueError:
|
||
text_level = None
|
||
elif isinstance(raw_text_level, int):
|
||
text_level = raw_text_level
|
||
|
||
item_type = item.get("type", "text")
|
||
text = item.get("text", "")
|
||
page_idx = item.get("page_idx", -1)
|
||
id = item.get("id")
|
||
|
||
if text_level is not None and 1 <= text_level <= max_level:
|
||
# 更新当前层级及更低层级的标题（例如：遇到2级标题会清除3级的缓存）
|
||
for level in range(text_level, max_level + 1):
|
||
current_titles[level] = text + "\n" if level == text_level else ""
|
||
continue # 标题本身不单独生成 chunk
|
||
|
||
# 合并图片逻辑
|
||
if item_type == "image":
|
||
image_path = item.get("image_path","")
|
||
image_caption = item.get("image_caption","")
|
||
if image_path:
|
||
text = f""
|
||
if processed_data:
|
||
processed_data[-1]["text"] += f""
|
||
continue
|
||
elif item_type =="table":
|
||
table_caption = "".join(item.get("table_caption", ""))
|
||
table_body = item.get("table_body", "")
|
||
text = table_caption + "\n" + table_body
|
||
# 普通文本：添加当前所有层级的标题（从高到低）
|
||
combined_title = "".join(
|
||
current_titles[level]
|
||
for level in range(1, max_level + 1)
|
||
if current_titles[level]
|
||
)
|
||
full_text = combined_title + text
|
||
|
||
# 添加到结果（仅当有内容时）
|
||
if full_text.strip():
|
||
processed_data.append({
|
||
"id":id,
|
||
"text": full_text,
|
||
"page_idx": page_idx+1,
|
||
"type": item_type,
|
||
"image_path":image_path if item_type == "image" else ""
|
||
})
|
||
|
||
return processed_data
|
||
results = []
|
||
processed_data = process_content_list(content_chunk_list,max_level=max_level)
|
||
for i,data in enumerate(processed_data):
|
||
if chunk_type == "01" or "02":
|
||
chunks = _split_recursively(data["text"], chunk_size, chunk_overlap)
|
||
else:
|
||
chunks = _split_by_token(resource=resource,md_text=data.get["text"],logger=logger,chunk_size=chunk_size,chunk_overlap=chunk_overlap)
|
||
for j, chunk in enumerate(chunks):
|
||
result = _process_chunk(chunk, resource, {}, str(data.get("id")) +f"-{j}", logger,page_idx=data.get("page_idx"))
|
||
results.append(result)
|
||
print(111111111111111111)
|
||
return results
|
||
|
||
def _split_by_token(resource, md_text,logger, chunk_size, chunk_overlap):
|
||
"""
|
||
使用markdown解析器保持结构完整的切分方法，支持chunk重叠
|
||
|
||
参数:
|
||
resource: 关联的资源标识
|
||
md_text: 要切分的Markdown文本
|
||
chunk_size: 每个chunk的目标大小
|
||
chunk_overlap: chunk之间的重叠大小
|
||
|
||
返回:
|
||
包含text和resource字段的字典列表
|
||
"""
|
||
logger.info("开始按字符分片")
|
||
if chunk_overlap >= chunk_size:
|
||
raise ValueError("chunk_overlap必须小于chunk_size")
|
||
|
||
md = MarkdownIt()
|
||
tokens = md.parse(md_text)
|
||
root = SyntaxTreeNode(tokens)
|
||
if chunk_size >= 8192:
|
||
raise ValueError("chunk_size不能大于4096")
|
||
chunks = []
|
||
current_chunk = ""
|
||
overlap_buffer = ""
|
||
|
||
def get_node_text(node):
|
||
"""递归获取节点的完整文本内容"""
|
||
if node.type == "text":
|
||
return node.content
|
||
elif node.type == "inline" and node.children:
|
||
return "".join(get_node_text(child) for child in node.children)
|
||
elif node.children:
|
||
# 对于块级元素，添加适当的换行
|
||
contents = []
|
||
for child in node.children:
|
||
content = get_node_text(child)
|
||
if content:
|
||
contents.append(content)
|
||
return "\n".join(contents)
|
||
return ""
|
||
|
||
def process_node(node):
|
||
nonlocal current_chunk, overlap_buffer
|
||
content = get_node_text(node)
|
||
|
||
# 处理当前内容
|
||
remaining_content = content
|
||
while remaining_content:
|
||
# 计算还能添加到当前chunk的内容量
|
||
remaining_space = chunk_size - len(current_chunk)
|
||
|
||
if remaining_space <= 0:
|
||
# 当前chunk已满，保存并准备重叠部分
|
||
chunks.append(current_chunk)
|
||
overlap_buffer = current_chunk[-chunk_overlap:] if chunk_overlap > 0 else ""
|
||
current_chunk = overlap_buffer + remaining_content[:remaining_space]
|
||
remaining_content = remaining_content[remaining_space:]
|
||
else:
|
||
# 添加到当前chunk
|
||
add_content = remaining_content[:remaining_space]
|
||
current_chunk += add_content
|
||
remaining_content = remaining_content[remaining_space:]
|
||
|
||
# 处理所有节点
|
||
for child in root.children:
|
||
process_node(child)
|
||
|
||
# 添加最后一个chunk
|
||
if current_chunk:
|
||
chunks.append(current_chunk)
|
||
|
||
# 过滤空chunk并添加resource信息
|
||
return [{"text": text.strip(), "resource": resource}
|
||
for text in chunks if text.strip()]
|
||
|
||
# def _split_by_token(resource, text, logger, chunk_size, chunk_overlap):
|
||
# """按token分割文本"""
|
||
# logger.info("开始按字符分割")
|
||
# text_splitter = TokenTextSplitter(chunk_size=1024, chunk_overlap=100)
|
||
# texts = text_splitter.split_text(text)
|
||
# return [{"text": text, "resource": resource} for text in texts]
|
||
|
||
|
||
def _split_custom(resource, mdDocs, logger, chunk_type, chunk_size, chunk_overlap,max_level):
|
||
"""自定义分割处理"""
|
||
logger.info("开始按自定义分割")
|
||
results = []
|
||
# if chunk_type=="02":
|
||
if chunk_type=="02" or chunk_type=="01":
|
||
# md_text = custom_convert_numbered_headings(mdDocs, max_level)
|
||
# md_text = '\n'.join(md_text)
|
||
chunks = _split_markdown(mdDocs,max_level)
|
||
# chunks = _split_markdown(md_text,max_level)
|
||
for idx, doc in enumerate(chunks, start=1):
|
||
# 确保所有Header字段都存在
|
||
headers = {f"Header_{i}": doc.metadata.get(f"Header_{i}", "")
|
||
for i in range(1, max_level + 1)}
|
||
# if headers.get("Header_1") =="目 录" or headers.get("Header_1") =="目录" or headers.get("Header_1") =="目 录":
|
||
# continue
|
||
result = _process_chunk(doc.page_content, resource, headers, idx, logger)
|
||
results.append(result)
|
||
results = reset_chunk(results)
|
||
|
||
return results
|
||
|
||
|
||
def _split_markdown(text,max_level,chunk_size=2048,chunk_overlap=50):
|
||
"""
|
||
根据 max_level 动态生成 headers_to_split_on 配置
|
||
规则：
|
||
- max_level=1 → 只按 # 分割
|
||
- max_level=2 → 按 # 和 ## 分割
|
||
- 以此类推
|
||
"""
|
||
# 动态生成 headers_to_split_on 配置
|
||
headers_to_split_on = [
|
||
(f"{'#' * i}", f"Header_{i}") for i in range(1, max_level + 1)
|
||
]
|
||
|
||
# 初始化分割器
|
||
splitter = MarkdownHeaderTextSplitter(
|
||
headers_to_split_on=headers_to_split_on,
|
||
strip_headers=False # 废弃标题文本
|
||
)
|
||
chunks = splitter.split_text(text)
|
||
char_splitter = CharacterTextSplitter(
|
||
chunk_size=chunk_size,
|
||
chunk_overlap=chunk_overlap,
|
||
length_function=len,
|
||
separator="" # 优先按段落分割
|
||
)
|
||
# 检查并处理过大的分块
|
||
|
||
final_chunks = []
|
||
header_chunks = []
|
||
for chunk in chunks:
|
||
# content = chunk.page_content # 只获取文本内容，忽略metadata
|
||
if len(chunk.page_content) <=8000:
|
||
final_chunks.append(chunk)
|
||
header_chunks.append(chunk.metadata) #dict
|
||
else:
|
||
# 直接分割文本内容
|
||
split_texts = char_splitter.split_documents([chunk])
|
||
final_chunks.extend(split_texts)
|
||
|
||
return final_chunks
|
||
|
||
# 执行分割
|
||
# return splitter.split_text(text)
|
||
|
||
def _split_recursively(text, chunk_size, chunk_overlap):
|
||
"""递归字符分割"""
|
||
text_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=1024,
|
||
chunk_overlap=100,
|
||
length_function=len,
|
||
# separators=["\n\n\n","(?<!#)\n\n"," "],
|
||
separators=[""],
|
||
keep_separator=True,
|
||
is_separator_regex=True
|
||
)
|
||
semantic_chunks = text_splitter.split_text(text)
|
||
|
||
# 2. 再用 CharacterTextSplitter 按字符数精确分割
|
||
char_splitter = CharacterTextSplitter(
|
||
chunk_size=1024,#2048
|
||
chunk_overlap=50,
|
||
length_function=len,
|
||
separator="" # 按纯字符切分
|
||
)
|
||
|
||
final_chunks = []
|
||
for chunk in semantic_chunks:
|
||
final_chunks.extend(char_splitter.split_text(chunk))
|
||
|
||
return final_chunks
|
||
# return text_splitter.split_text(text)
|
||
|
||
|
||
def _split_with_images(resource, text, logger, chunk_size, chunk_overlap):
|
||
"""处理没有章节信息的文本分割"""
|
||
results = []
|
||
chunks = _split_recursively(text, chunk_size, chunk_overlap)
|
||
|
||
for i, chunk in enumerate(chunks):
|
||
img_list = extract_urls(chunk)
|
||
logger.info(f"img_list: {img_list}")
|
||
# clean_text = remove_markdown_images(chunk)
|
||
img_path = ",".join(img_list) if img_list else ""
|
||
|
||
results.append({
|
||
"id": i+1,
|
||
"text": text,#clean_text,
|
||
"resource": resource,
|
||
"img_path": img_path
|
||
})
|
||
logger.info(f"results: {results}")
|
||
return results
|
||
|
||
|
||
def _process_chunk(text, resource, headers, chunk_id, logger,page_idx=None,origin_text=None):
|
||
"""处理单个文本块，提取图片和清理文本"""
|
||
img_list = extract_urls(text)
|
||
# logger.info(f"img_list: {img_list}")
|
||
# clean_text = remove_markdown_images(text)
|
||
img_path = ",".join(img_list) if img_list else ""
|
||
if img_path:
|
||
logger.info(f"Get Image Path: {img_path}")
|
||
|
||
return {
|
||
"id": chunk_id,
|
||
"Header_1": headers.get("Header_1", ""),
|
||
"Header_2": headers.get("Header_2", ""),
|
||
"Header_3": headers.get("Header_3", ""),
|
||
"text":text.replace("# ",""),
|
||
"resource": resource,
|
||
"img_path": img_path,
|
||
"page":page_idx,
|
||
"origin_text":text
|
||
}
|
||
|
||
def save_md(results, collection_name,database_name):
|
||
"""
|
||
保存Markdown到Milvus数据库的函数
|
||
:param results: 拆分后的文本块列表
|
||
:param database_name: 数据库名称
|
||
:param collection_name: 集合名称
|
||
:return: 保存后的结果列表
|
||
"""
|
||
texts = [item.get("text") for item in results]
|
||
# vectors = OpenaiAPI.get_embeddings(texts)
|
||
vectors = OpenaiAPI.batch_embeddings(texts)
|
||
|
||
milvus = Save_Markdown_to_Milvus(database_name=database_name, collection_name=collection_name)
|
||
all_data = []
|
||
for idx, item in enumerate(results):
|
||
text = item.get("text")
|
||
resource = item.get("resource")
|
||
img_path = item.get("img_path", "null")
|
||
vector = vectors[idx].get("embedding") # 获取对应的嵌入向量
|
||
|
||
data_item = {
|
||
"text": text,
|
||
"vector": vector,
|
||
"resource": resource,
|
||
"img_path": img_path,
|
||
"page_idx": item.get("page_idx"),
|
||
"origin_id": str(item.get("id")),
|
||
"origin_text":item.get("origin_text")
|
||
}
|
||
|
||
# 添加标题信息（如果有）
|
||
if "Header_1" in item:
|
||
data_item.update({
|
||
"Header_1": item.get("Header_1"),
|
||
"Header_2": item.get("Header_2"),
|
||
"Header_3": item.get("Header_3")
|
||
})
|
||
all_data.append(data_item)
|
||
|
||
# milvus.save_hybrid_md(data=[data_item])
|
||
milvus.save_hybrid_md(data=all_data)
|
||
|
||
milvus.close()
|
||
return results
|
||
# milvus = Save_Markdown_to_Milvus(database_name=database_name, collection_name=collection_name)
|
||
|
||
# for item in results:
|
||
# text = item.get("text")
|
||
# resource = item.get("resource")
|
||
# img_path = item.get("img_path", "null")
|
||
# vector = OpenaiAPI.get_embeddings(text)
|
||
|
||
# if "Header_1" in item: # 有标题的情况
|
||
# data = [{"text":text,
|
||
# "vector":vector,
|
||
# "resource":resource,
|
||
# "img_path":img_path,
|
||
# "Header_1":item.get("Header_1"),
|
||
# "Header_2":item.get("Header_2"),
|
||
# "Header_3":item.get("Header_3"),
|
||
# "page_idx":item.get("page_idx"),
|
||
# "origin_id":str(item.get("id"))}]
|
||
|
||
# else: # 没有标题的情况
|
||
# data = [{"text":text,
|
||
# "vector":vector,
|
||
# "resource":resource,
|
||
# "img_path":img_path,
|
||
# "page_idx":item.get("page_idx"),
|
||
# "origin_id":str(item.get("id"))
|
||
# }]
|
||
|
||
# milvus.save_hybrid_md(data=data)
|
||
|
||
# milvus.close()
|
||
# return results
|
||
|
||
def hierarchical_save(results,collection_name,database_name):
|
||
texts = [item.get("text") for item in results]
|
||
# vectors = OpenaiAPI.get_embeddings(texts)
|
||
vectors = OpenaiAPI.batch_embeddings(texts)
|
||
|
||
milvus = Save_Markdown_to_Milvus(database_name=database_name, collection_name=collection_name)
|
||
all_data = []
|
||
for idx, item in enumerate(results):
|
||
text = item.get("text")
|
||
resource = item.get("resource")
|
||
img_path = item.get("img_path", "null")
|
||
vector = vectors[idx].get("embedding") # 获取对应的嵌入向量
|
||
|
||
data_item = {
|
||
"text": text,
|
||
"vector": vector,
|
||
"resource": resource,
|
||
"img_path": img_path,
|
||
"page_idx": item.get("page_idx"),
|
||
"origin_id": str(item.get("id")),
|
||
"origin_text": item.get("origin_text"),
|
||
}
|
||
|
||
# 添加标题信息（如果有）
|
||
if "Header_1" in item:
|
||
data_item.update({
|
||
"Header_1": item.get("Header_1"),
|
||
"Header_2": item.get("Header_2"),
|
||
"Header_3": item.get("Header_3")
|
||
})
|
||
all_data.append(data_item)
|
||
|
||
# milvus.save_hybrid_md(data=[data_item])
|
||
milvus.hierarchical_insert(data=all_data,collection_name=collection_name,database_name=database_name)
|
||
|
||
milvus.close()
|
||
return results
|
||
|
||
def split_insert(resource, mdDocs, logger, database_name, collection_name, chunk_type, chunk_length, chunk_overlap=50,max_level=3):
|
||
"""
|
||
原始函数的包装函数，保持接口不变
|
||
"""
|
||
# 先拆分文本
|
||
results = split_text(resource, mdDocs, logger, chunk_type, chunk_length, chunk_overlap,max_level)
|
||
# 然后保存到数据库
|
||
return save_md(results, database_name, collection_name)
|
||
|
||
|
||
def split_by_md_header(mdDocs,logger,database_name,collection_name):
|
||
"""
|
||
Split with metadata like "Metadata (("Header_1", "xxx"), ("Header_2", "xxxx")) page_content="xxx..."
|
||
:param mdDocs: markdown text
|
||
:return:
|
||
"""
|
||
mdDocs = convert_numbered_headings(mdDocs)
|
||
headers_to_split_on = [("#", "Header_1"), ("##", "Header_2"), ("###", "Header_3")]
|
||
md_header_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on, return_each_line=True)
|
||
md_header_splits = md_header_splitter.split_text(mdDocs)
|
||
# md_text_splitter = MarkdownTextSplitter(chunk_size=512,)
|
||
#存入向量数据库
|
||
milvus = Save_Markdown_to_Milvus(database_name=database_name,collection_name=collection_name)
|
||
final_splits = []
|
||
for _, i in enumerate(md_header_splits):
|
||
logger.info(f"Index: {_}, Metadata: {i.metadata} Content: {i.page_content}")
|
||
|
||
if i.metadata:
|
||
header_1 = i.metadata["Header_1"] if "Header_1" in i.metadata.keys() else None
|
||
header_2 = i.metadata["Header_2"] if "Header_2" in i.metadata.keys() else None
|
||
header_3 = i.metadata["Header_3"] if "Header_3" in i.metadata.keys() else None
|
||
text = i.page_content
|
||
# metadata = {"title": header_1, "sub_title": header_2, "sub_sub_title": header_3,"text":text}
|
||
vector = OpenaiAPI.get_embeddings(text)
|
||
if vector is None:
|
||
logger.error("向量生成失败")
|
||
raise ValueError("向量生成失败")
|
||
|
||
data = [[header_1],[header_2],[header_3],[text],[vector]]
|
||
logger.info(f"Index:{_},header_1:{header_1},header_2:{header_2},header_3:{header_3},text:{text}")
|
||
# data={"text":text,"header_1":header_1,"header_2":header_2,"header_3":header_3,"vector":vector}
|
||
result = milvus.save_md(data=data)
|
||
# print(f"Index:{_},header_1:{header_1},header_2:{header_2},header_3:{header_3},text:{text}")
|
||
# print("---------------------")
|
||
return md_header_splits
|