Compare commits

..

7 Commits

4 changed files with 811 additions and 58 deletions

View File

@@ -1,5 +1,5 @@
from elasticsearch import Elasticsearch
from src.add_chunk_cli_pdf_img import update_positon_img_id_in_elasticsearch
# 初始化 Elasticsearch 用户名elastic密码infini_rag_flow
es = Elasticsearch(
[{'host': '127.0.0.1', 'port': 1200, 'scheme': 'http'}],
@@ -58,6 +58,39 @@ def update_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id, new_img_id):
else:
return {"code": 100, "message": "Failed to update img_id"}
def get_index_mapping(tenant_id):
"""
获取指定索引的 mapping 信息
:param tenant_id: 租户 ID
:return: mapping 信息
"""
index_name = f"ragflow_{tenant_id}"
try:
mapping = es.indices.get_mapping(index=index_name)
# 将 ObjectApiResponse 转换为普通字典
mapping_dict = dict(mapping)
return {"code": 0, "message": "", "data": mapping_dict}
except Exception as e:
return {"code": 500, "message": str(e), "data": {}}
# 在主函数中调用示例
if __name__ == "__main__":
# ... 现有代码 ...
# 获取 mapping 信息
tenant_id = "9c73df5a3ebc11f08410c237296aa408"
mapping_result = get_index_mapping(tenant_id)
if mapping_result["code"] == 0:
print("索引 mapping 信息:")
import json
# 使用 default=str 处理不能直接序列化的对象
print(json.dumps(mapping_result["data"], indent=2, ensure_ascii=False, default=str))
else:
print(f"获取 mapping 失败: {mapping_result['message']}")
def list_chunk_information(tenant_id, dataset_id, doc_id=None, chunk_id=None, size=1000):
"""
@@ -121,17 +154,33 @@ if __name__ == "__main__":
tenant_id = "9c73df5a3ebc11f08410c237296aa408"
dataset_id = "0e6127da574a11f0a59c7e7439a490f8" # dataset_id = kb_id
doc_id = "cbf576385bc911f08f23fedc3996e479"
doc_id = "323113d8670c11f0b4255ea1d23c381a"
doc_id = "323113d8670c11f0b4255ea1d23c381a"
doc_id = "5cdab2fa67cb11f0a21592edb0e63cad" #
chunk_id = "f035247f7de579b0" #
chunk_id = "b2d53baddbfde97c" #
chunk_id = "e46a067c1edf939a"
new_img_id = "10345832587311f0919f3a2728512a4b-f035247f7de579b0" #"new_img_id_12345"
new_img_id = "0e6127da574a11f0a59c7e7439a490f8-b2d53baddbfde97c"
#new_img_id = "0e6127da574a11f0a59c7e7439a490f8-b2d53baddbfde97c"
#new_img_id ="c5142bce5ac611f0ae707a8b5ba029cb-thumbnail_fb3cbc165ac611f0b5897a8b5ba029cb.png"
pos= [3, 317, 397, 123, 182]
# 获取 mapping 信息
tenant_id = "9c73df5a3ebc11f08410c237296aa408"
mapping_result = get_index_mapping(tenant_id)
if mapping_result["code"] == 0:
print("索引 mapping 信息:")
import json
print(json.dumps(mapping_result["data"], indent=2, ensure_ascii=False))
else:
print(f"获取 mapping 失败: {mapping_result['message']}")
#chunk_list = list_chunk_information(tenant_id, dataset_id, doc_id=doc_id)
update_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id,new_img_id)
# update_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id,new_img_id)
update_positon_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id, pos, new_img_id)
# if chunk_list["code"] == 0:
# print(f"找到 {len(chunk_list['data'])} 个 chunks")
# for chunk in chunk_list['data']:

256
chunk_pos.py Normal file
View File

@@ -0,0 +1,256 @@
from elasticsearch import Elasticsearch
#from src.add_chunk_cli_pdf_img import update_positon_img_id_in_elasticsearch
# 初始化 Elasticsearch 用户名elastic密码infini_rag_flow
from dotenv import load_dotenv # 新增
import os
import json
# 加载 .env 文件中的环境变量
load_dotenv()
# 初始化 Elasticsearch
es = Elasticsearch(
[{
'host': os.getenv("ELASTIC_HOST"),
'port': int(os.getenv("ELASTIC_PORT")),
'scheme': 'http'
}],
basic_auth=(
os.getenv("ELASTIC_USERNAME"),
os.getenv("ELASTIC_PASSWORD")
)
)
def get_index_mapping(tenant_id):
"""
获取指定索引的 mapping 信息
:param tenant_id: 租户 ID
:return: mapping 信息
"""
index_name = f"ragflow_{tenant_id}"
try:
mapping = es.indices.get_mapping(index=index_name)
# 将 ObjectApiResponse 转换为普通字典
mapping_dict = dict(mapping)
return {"code": 0, "message": "", "data": mapping_dict}
except Exception as e:
return {"code": 500, "message": str(e), "data": {}}
def update_positon_in_elasticsearch(tenant_id, doc_id, chunk_id, positions):
"""
在 Elasticsearch 中更新指定文档块的position and img_id。
:param tenant_id: 租户 ID
:param doc_id: 文档 ID
:param chunk_id: 文档块 ID
:param new_img_id: 新的 img_id
:param position: 位置信息
:return: 更新结果
"""
if not positions:
return
position_int = []
for pos in positions:
if len(pos) != 5:
continue # Skip invalid positions
pn, left, right, top, bottom = pos
# 使用元组格式与原始RAGFlow保持一致
position_int.append((int(pn + 1), int(left), int(right), int(top), int(bottom)))
if position_int: # Only add if we have valid positions
# 仅添加精确位置信息,不修改排序字段
# 构建索引名称
index_name = f"ragflow_{tenant_id}"
# 构建查询条件
query = {
"bool": {
"must": [
{"term": {"doc_id": doc_id}},
{"term": {"_id": chunk_id}}
]
}
}
# 搜索目标文档
result = es.search(index=index_name, body={"query": query})
# 检查是否找到目标文档
if result['hits']['total']['value'] == 0:
print(f"在 Elasticsearch 中未找到文档: index={index_name}, doc_id={doc_id}, chunk_id={chunk_id}")
return {"code": 102, "message": f"Can't find this chunk {chunk_id}"}
# 获取目标文档的 ID
hit = result['hits']['hits'][0]
doc_id_in_es = hit['_id']
# 构建更新请求 - 只更新存在的字段
update_body = {"doc": {}}
update_body["doc"]["position_int"] = position_int
update_body["doc"]["page_num_int"] = [position_int[0][0]]
update_body["doc"]["top_int"] = [position_int[0][3]]
# 更新文档
update_result = es.update(
index=index_name,
id=doc_id_in_es,
body=update_body,
refresh=True # 确保更新立即可见
)
print(f"Elasticsearch 更新结果: index={index_name}, id={doc_id_in_es}, result={update_result}")
def update_positon_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id, position, new_img_id):
"""
在 Elasticsearch 中更新指定文档块的position and img_id。
:param tenant_id: 租户 ID
:param doc_id: 文档 ID
:param chunk_id: 文档块 ID
:param new_img_id: 新的 img_id
:param position: 位置信息
:return: 更新结果
"""
try:
# 构建索引名称
index_name = f"ragflow_{tenant_id}"
# 构建查询条件
query = {
"bool": {
"must": [
{"term": {"doc_id": doc_id}},
{"term": {"_id": chunk_id}}
]
}
}
# 搜索目标文档
result = es.search(index=index_name, body={"query": query})
# 检查是否找到目标文档
if result['hits']['total']['value'] == 0:
print(f"在 Elasticsearch 中未找到文档: index={index_name}, doc_id={doc_id}, chunk_id={chunk_id}")
return {"code": 102, "message": f"Can't find this chunk {chunk_id}"}
# 获取目标文档的 ID
hit = result['hits']['hits'][0]
doc_id_in_es = hit['_id']
# 构建更新请求 - 只更新存在的字段
update_body = {"doc": {}}
#只有当 new_img_id 存在时才更新 img_id
if new_img_id is not None:
update_body["doc"]["img_id"] = new_img_id
# 只有当 position 存在时才更新 positions
if position is not None:
update_body["doc"]["positions"] = position
# 如果没有需要更新的字段,直接返回成功
if not update_body["doc"]:
print("没有需要更新的字段")
return {"code": 0, "message": "No fields to update"}
# 更新文档
update_result = es.update(
index=index_name,
id=doc_id_in_es,
body=update_body,
refresh=True # 确保更新立即可见
)
print(f"Elasticsearch 更新结果: index={index_name}, id={doc_id_in_es}, result={update_result}")
# 验证更新
verify_doc = es.get(index=index_name, id=doc_id_in_es)
# 检查 img_id 是否已更新(如果提供了 new_img_id
img_id_updated = True
if new_img_id is not None:
img_id_updated = verify_doc['_source'].get('img_id') == new_img_id
if img_id_updated:
print(f"成功更新 img_id 为: {new_img_id}")
else:
print(f"更新验证失败,当前 img_id: {verify_doc['_source'].get('img_id')}")
# 检查 position 是否已更新(如果提供了 position
position_updated = True
if position is not None:
position_updated = verify_doc['_source'].get('positions') == position
if position_updated:
print(f"成功更新 position 为: {position}")
else:
print(f"更新验证失败,当前 position: {verify_doc['_source'].get('positions')}")
# 统一返回结果
if img_id_updated and position_updated:
return {"code": 0, "message": ""}
else:
return {"code": 100, "message": "Failed to verify update"}
except Exception as e:
print(f"更新 Elasticsearch 时发生错误: {str(e)}")
return {"code": 101, "message": f"Error updating img_id: {str(e)}"}
# 示例调用 - 列出特定文档的所有 chunks
if __name__ == "__main__":
try:
print(es.info())
except Exception as e:
print("连接失败:", e)
# 单位电脑
tenant_id = "d669205e57a211f0b9e7324e7f243034"
new_img_id ="10345832587311f0919f3a2728512a4b-bd04866cd05337281"
doc_id="ea8d75966df811f0925ac6e8db75f472"
chunk_id="4a4927560a7e6d80"
# 添加以下代码来检查索引映射
# mapping_result = get_index_mapping(tenant_id)
# print("Positions field mapping:", mapping_result["data"][f"ragflow_{tenant_id}"]["mappings"]["properties"]["positions"])
# 左,右 -->
#上, 下| 上面最小,下面最大
pos = [[4, 0, 100, 200, 510]]
#pos_string = json.dumps(pos) # 转换为 JSON 字符串
update_positon_in_elasticsearch(tenant_id, doc_id, chunk_id, pos)
#update_positon_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id, pos, "")

View File

@@ -9,8 +9,10 @@ import tempfile
from elasticsearch import Elasticsearch
from minio import Minio
from minio.error import S3Error
from find_text_in_pdf_enhanced import find_text_in_pdf
import time
from get_pos_pdf import smart_fuzzy_find_text_batch, find_text_positions_batch
from dotenv import load_dotenv # 新增
@@ -47,7 +49,149 @@ MINIO_CONFIG = {
"secure": False
}
def update_positon_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id, position, new_img_id):
from elasticsearch.helpers import bulk
def bulk_update_elasticsearch(tenant_id, updates):
"""
批量更新Elasticsearch中的文档
:param tenant_id: 租户ID
:param updates: 更新信息列表每个元素包含doc_id, chunk_id, positions, new_img_id
:return: 更新结果
"""
try:
index_name = f"ragflow_{tenant_id}"
# 构建批量操作列表
actions = []
for update_info in updates:
doc_id = update_info['doc_id']
chunk_id = update_info['chunk_id']
positions = update_info.get('positions', [])
new_img_id = update_info.get('new_img_id')
# 构建查询条件来找到文档
query = {
"bool": {
"must": [
{"term": {"doc_id": doc_id}},
{"term": {"_id": chunk_id}}
]
}
}
# 搜索目标文档
result = es.search(index=index_name, body={"query": query})
# 检查是否找到目标文档
if result['hits']['total']['value'] == 0:
print(f"在 Elasticsearch 中未找到文档: index={index_name}, doc_id={doc_id}, chunk_id={chunk_id}")
continue
# 获取目标文档的 ID
hit = result['hits']['hits'][0]
doc_id_in_es = hit['_id']
# 构建更新请求 - 只更新存在的字段
doc_update = {}
# 只有当 new_img_id 存在时才更新 img_id
if new_img_id is not None:
doc_update["img_id"] = new_img_id
# 只有当 positions 存在时才更新 positions
if positions:
position_int = []
for pos in positions:
if len(pos) != 5:
continue # Skip invalid positions
pn, left, right, top, bottom = pos
# 使用元组格式与原始RAGFlow保持一致
position_int.append((int(pn + 1), int(left), int(right), int(top), int(bottom)))
if position_int:
doc_update["position_int"] = position_int
doc_update["page_num_int"] = [position_int[0][0]]
doc_update["top_int"] = [position_int[0][3]]
# 如果没有需要更新的字段,跳过
if not doc_update:
print(f"没有需要更新的字段 for chunk {chunk_id}")
continue
# 添加到批量操作列表
action = {
"_op_type": "update",
"_index": index_name,
"_id": doc_id_in_es,
"doc": doc_update
}
actions.append(action)
# 执行批量更新
if actions:
results = bulk(es, actions, refresh=True)
print(f"批量更新完成,成功处理 {results[0]} 个操作")
return {"code": 0, "message": f"Successfully updated {results[0]} documents"}
else:
print("没有需要执行的更新操作")
return {"code": 0, "message": "No updates to perform"}
except Exception as e:
print(f"批量更新 Elasticsearch 时发生错误: {str(e)}")
return {"code": 101, "message": f"Error in bulk update: {str(e)}"}
# 修改 process_pdf_txt_pairs 函数以使用批量更新
def process_pdf_txt_pairs_bulk(pdf_dict, txt_dict, dataset):
"""处理PDF-TXT文件对使用批量更新提高效率"""
# 收集所有需要更新的信息
all_updates = []
for name, pdf_path in pdf_dict.items():
display_name = os.path.basename(pdf_path)
document = upload_or_get_document(dataset, pdf_path, display_name)
print(f"选择的文档: {document.name}ID: {document.id}")
if not document:
continue
txt_path = txt_dict.get(name)
if txt_path:
chunks_info = process_txt_chunks(dataset.id, document, txt_path)
time.sleep(1) # 等待chunk处理完成
if chunks_info:
chunks_info = get_positions_from_chunk(pdf_path, chunks_info)
# 收集更新信息而不是立即更新
for chunk_info in chunks_info:
print(f"Chunk ID: {chunk_info['id']}, Text: {chunk_info['text'][:30]}..., Has Image: {chunk_info['has_image']}, Positions: {chunk_info['positions']}")
update_info = {
'doc_id': document.id,
'chunk_id': chunk_info['id'],
'positions': chunk_info['positions']
}
if chunk_info['has_image']:
# 如果有图片准备更新img_id
update_info['new_img_id'] = f"{dataset.id}-{chunk_info['id']}"
# 如果没有图片new_img_id为None不会更新img_id字段
all_updates.append(update_info)
# 执行批量更新
if all_updates:
result = bulk_update_elasticsearch(elastic_tenant_id, all_updates)
print(f"批量更新结果: {result}")
def update_positon_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id, positions, new_img_id):
"""
在 Elasticsearch 中更新指定文档块的position and img_id。
@@ -59,6 +203,7 @@ def update_positon_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id, position
:return: 更新结果
"""
try:
# 构建索引名称
index_name = f"ragflow_{tenant_id}"
@@ -87,13 +232,27 @@ def update_positon_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id, position
# 构建更新请求 - 只更新存在的字段
update_body = {"doc": {}}
# 只有当 new_img_id 存在时才更新 img_id
#只有当 new_img_id 存在时才更新 img_id
if new_img_id is not None:
update_body["doc"]["img_id"] = new_img_id
# 只有当 position 存在时才更新 positions
if position is not None:
update_body["doc"]["positions"] = position
if positions :
position_int = []
for pos in positions:
if len(pos) != 5:
continue # Skip invalid positions
pn, left, right, top, bottom = pos
# 使用元组格式与原始RAGFlow保持一致
position_int.append((int(pn + 1), int(left), int(right), int(top), int(bottom)))
if position_int:
update_body["doc"]["position_int"] = position_int
update_body["doc"]["page_num_int"] = [position_int[0][0]]
update_body["doc"]["top_int"] = [position_int[0][3]]
# 如果没有需要更新的字段,直接返回成功
if not update_body["doc"]:
@@ -110,32 +269,7 @@ def update_positon_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id, position
print(f"Elasticsearch 更新结果: index={index_name}, id={doc_id_in_es}, result={update_result}")
# 验证更新
verify_doc = es.get(index=index_name, id=doc_id_in_es)
# 检查 img_id 是否已更新(如果提供了 new_img_id
img_id_updated = True
if new_img_id is not None:
img_id_updated = verify_doc['_source'].get('img_id') == new_img_id
if img_id_updated:
print(f"成功更新 img_id 为: {new_img_id}")
else:
print(f"更新验证失败,当前 img_id: {verify_doc['_source'].get('img_id')}")
# 检查 position 是否已更新(如果提供了 position
position_updated = True
if position is not None:
position_updated = verify_doc['_source'].get('positions') == position
if position_updated:
print(f"成功更新 position 为: {position}")
else:
print(f"更新验证失败,当前 position: {verify_doc['_source'].get('positions')}")
# 统一返回结果
if img_id_updated and position_updated:
return {"code": 0, "message": ""}
else:
return {"code": 100, "message": "Failed to verify update"}
except Exception as e:
@@ -143,6 +277,7 @@ def update_positon_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id, position
return {"code": 101, "message": f"Error updating img_id: {str(e)}"}
def get_minio_client():
"""创建MinIO客户端"""
return Minio(
@@ -427,43 +562,57 @@ def get_positions_from_chunk(pdf_path, chunks_info):
try:
# 提取所有chunk的文本内容用于批量查找
chunk_texts = [chunk_info['text'] for chunk_info in chunks_info]
print(f"批量查找文本块: {chunk_texts}")
# 使用智能模糊查找获取位置信息
batch_positions = smart_fuzzy_find_text_batch(pdf_path, chunk_texts, similarity_threshold=0.7)
matches = find_text_in_pdf(
pdf_path,
chunk_texts,
threshold=60
)
print(f"匹配结果: {matches}")
# 将位置信息与chunks_info关联并确保数据类型正确
for i, chunk_info in enumerate(chunks_info):
positions = batch_positions[i] if i < len(batch_positions) else []
# 处理位置信息
processed_positions = []
for pos in positions:
if isinstance(pos, dict):
# 创建新的位置字典,确保所有坐标都是整数
processed_pos = {
'x0': int(round(float(pos['x0']))) if pos.get('x0') is not None else 0,
'y0': int(round(float(pos['y0']))) if pos.get('y0') is not None else 0,
'x1': int(round(float(pos['x1']))) if pos.get('x1') is not None else 0,
'y1': int(round(float(pos['y1']))) if pos.get('y1') is not None else 0,
'page': int(pos['page']) if pos.get('page') is not None else 0
}
processed_positions.append(processed_pos)
# 更新chunk_info中的positions
chunk_info['positions'] = processed_positions
# 确保 chunk_info 包含 'positions' 键
if 'positions' not in chunk_info:
chunk_info['positions'] = []
print(f"处理第 {i+1} 个chunk: {chunk_info['text']}")
print(f"更新前位置: {chunk_info['positions']}")
if isinstance(matches, list) and i < len(matches):
chunk_info['positions']=[mat['position_int'] for mat in matches[i] if 'position_int' in mat]
# # 如果matches是列表且索引有效
# if isinstance(matches[i], dict) and 'position_int' in matches[i]:
# chunk_info['positions'] = matches[i]['position_int']
# print(f"更新后位置: {chunk_info['positions']}")
# else:
# chunk_info['positions'] = []
# print(f"未找到有效位置信息,设置为空列表")
else:
chunk_info['positions'] = []
print(f"匹配结果无效或索引越界,设置为空列表")
# 验证更新结果
print("最终chunks_info状态:")
for i, chunk_info in enumerate(chunks_info):
print(f" Chunk {i+1}: ID={chunk_info['id']}, Positions={chunk_info['positions']}")
return chunks_info
except Exception as e:
print(f"获取PDF文本位置信息时出错: {str(e)}")
# 出错时为每个chunk添加空的位置信息
for chunk_info in chunks_info:
chunk_info['positions'] = []
# 确保 chunk_info 包含 'positions' 键
if 'positions' not in chunk_info:
chunk_info['positions'] = []
return chunks_info
def process_pdf_txt_pairs(pdf_dict, txt_dict, dataset):
"""处理PDF-TXT文件对"""
for name, pdf_path in pdf_dict.items():
@@ -476,6 +625,8 @@ def process_pdf_txt_pairs(pdf_dict, txt_dict, dataset):
txt_path = txt_dict.get(name)
if txt_path:
chunks_info=process_txt_chunks(dataset.id,document, txt_path)
time.sleep(1)
if chunks_info:
chunks_info=get_positions_from_chunk(pdf_path, chunks_info)
for chunk_info in chunks_info:
@@ -490,7 +641,6 @@ def process_pdf_txt_pairs(pdf_dict, txt_dict, dataset):
def main():
"""主函数处理PDF和TXT文件对
dataset.id = bucket_name
@@ -511,8 +661,8 @@ def main():
print("未选择数据集。")
return
process_pdf_txt_pairs(pdf_dict, txt_dict, dataset)
# 使用批量处理函数替代原来的处理函数
process_pdf_txt_pairs_bulk(pdf_dict, txt_dict, dataset)

View File

@@ -0,0 +1,298 @@
import fitz # pymupdf
import regex # 支持多行正则
from rapidfuzz import fuzz
import re
def normalize_text(text):
"""标准化文本,移除多余空白字符"""
# 将换行符、制表符等替换为空格,然后合并多个空格为一个
import re
normalized = re.sub(r'\s+', ' ', text.strip())
return normalized
def clean_text_for_fuzzy_match(text):
"""清理文本用于模糊匹配,移除特殊字符,只保留字母数字和空格"""
# 移除标点符号和特殊字符,只保留字母、数字、中文字符和空格
cleaned = re.sub(r'[^\w\s\u4e00-\u9fff]', '', text)
# 标准化空白字符
cleaned = re.sub(r'\s+', ' ', cleaned.strip())
return cleaned
def _merge_lines(lines):
"""
把多行文本合并成一段,同时记录每行 bbox 的并集。
lines: list of (text, bbox)
return: (merged_text, merged_bbox)
"""
if not lines:
return "", None
texts, bboxes = zip(*lines)
merged_text = "\n".join(texts)
# 合并 bbox取所有 bbox 的最小 x0,y0 和最大 x1,y1
x0 = min(b[0] for b in bboxes)
y0 = min(b[1] for b in bboxes)
x1 = max(b[2] for b in bboxes)
y1 = max(b[3] for b in bboxes)
# 修改:将坐标转换为整数
return merged_text, (int(x0), int(y0), int(x1), int(y1))
def _collect_lines(page):
"""
把一页的所有行按阅读顺序收集起来。
return: list of (text, bbox)
"""
lines = []
blocks = page.get_text("dict")["blocks"]
for blk in blocks:
if "lines" not in blk:
continue
for line in blk["lines"]:
line_text = "".join(span["text"] for span in line["spans"])
# 行级 bbox
x0 = min(span["bbox"][0] for span in line["spans"])
y0 = min(span["bbox"][1] for span in line["spans"])
x1 = max(span["bbox"][2] for span in line["spans"])
y1 = max(span["bbox"][3] for span in line["spans"])
# 修改:将坐标转换为整数
lines.append((line_text, (int(x0), int(y0), int(x1), int(y1))))
return lines
def find_text_in_pdf(pdf_path,
query, # 修改为支持list类型
use_regex=False,
threshold=80, # rapidfuzz 默认 0~100
page_range=None,
preprocess=True): # 添加预处理选项
"""
高级查找函数
query: 正则表达式字符串 或 普通字符串,或它们的列表
preprocess: 是否对文本进行预处理以提高模糊匹配准确性
返回: list[dict] 每个 dict 含 page, bbox, matched_text
"""
# 处理单个查询字符串的情况
if isinstance(query, str):
queries = [query]
else:
queries = query # 假设已经是列表
# 初始化结果列表
batch_results = [[] for _ in queries]
doc = fitz.open(pdf_path)
pages = range(len(doc)) if page_range is None else range(page_range[0]-1, page_range[1])
for p in pages:
page = doc.load_page(p)
lines = _collect_lines(page) # [(text, bbox), ...]
if not lines:
continue
full_text, _ = _merge_lines(lines) # 整页纯文本
# 如果启用预处理,则对整页文本进行预处理
processed_full_text = full_text
if preprocess and not use_regex:
processed_full_text = clean_text_for_fuzzy_match(full_text)
# 一次性计算所有查询的匹配结果
for idx ,q in enumerate(queries):
positions = [] # 记录匹配区间在 full_text 中的起止字符索引
results = []
if use_regex:
# regex 支持 (?s) 使 . 匹配换行
pattern = regex.compile(q)
for match in pattern.finditer(full_text):
positions.append((match.start(), match.end(), match.group()))
else:
# 模糊匹配:滑动窗口(整页 vs 查询)
# 修改:支持多个匹配结果并计算相似度分数
potential_matches = []
query_text = q
# 如果启用预处理,则对查询文本也进行预处理
if preprocess:
query_text = clean_text_for_fuzzy_match(q)
score = fuzz.partial_ratio(processed_full_text, query_text, score_cutoff=threshold)
if score >= threshold:
# 这里简单返回整页;如需精确定位,可再做二次对齐
positions.append((0, len(full_text), full_text))
# query_len = len(query_text)
# text_len = len(processed_full_text)
# # 优化:只在合理范围内进行滑动窗口匹配
# # 添加早期终止机制,一旦找到足够高的匹配就停止搜索
# best_score = 0
# for i in range(text_len - query_len + 1):
# window_text = processed_full_text[i:i + query_len]
# # 优化:只处理非空文本
# if window_text.strip():
# # 优化:使用更快速的相似度计算方法
# score = fuzz.partial_ratio(query_text, window_text)
# if score >= threshold:
# # 优化:记录当前最佳分数
# if score > best_score:
# best_score = score
# potential_matches.append((i, i + query_len, window_text, score))
# # 优化:如果找到非常高分的匹配,可以提前终止
# if score >= 95: # 如果匹配度已经很高,可以提前结束
# break
# 如果找到了潜在匹配,按分数排序并只取最高分的匹配
# if potential_matches:
# # 按分数降序排序
# potential_matches.sort(key=lambda x: x[3], reverse=True)
# # 只取分数最高的匹配
# best_match = potential_matches[0]
# positions.append((best_match[0], best_match[1], best_match[2]))
# 将字符区间映射回行
for start, end, matched_text in positions:
# 计算每一行在 full_text 中的起止字符偏移
offset = 0
matched_lines = []
for text, bbox in lines:
line_start = offset
line_end = offset + len(text)
# 检查该行是否与匹配区间有重叠 - 更严格的条件
if line_start < end and line_end > start:
matched_lines.append((text, bbox))
# 修正:正确计算偏移量,包括换行符
offset += len(text) + 1 # 加上换行符的长度
# 修正:只有当确实匹配到文本时才添加结果
if matched_lines:
_, merged_bbox = _merge_lines(matched_lines)
results.append({
"page": p,
"bbox": merged_bbox,
"matched_text": matched_text,
"position_int":[p, merged_bbox[0], merged_bbox[2], merged_bbox[1], merged_bbox[3]]
})
if results:
batch_results[idx].extend(results)
doc.close()
return batch_results
def highlight_matches(pdf_path, matches, output_path="highlighted.pdf"):
"""
把 matches 里的 bbox 用黄色高亮写入新 PDF
matches: find_text_in_pdf(...) 的返回值
"""
doc = fitz.open(pdf_path)
for m in matches:
page = doc.load_page(m["page"] - 1) # 0-based
# 修改:确保坐标为整数(虽然已经是整数了,但为了保险起见)
bbox = m["bbox"]
rect = fitz.Rect(int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]))
page.add_highlight_annot(rect) # 黄色高亮
doc.save(output_path)
doc.close()
print(f"已保存高亮 PDF{output_path}")
# ----------------- DEMO -----------------
# if __name__ == "__main__":
# pdf_path = "example.pdf"
# # 例1正则跨行匹配
# query_regex = r"条款\s*\d+\.?\s*[\s\S]*?责任限制"
# res = find_text_in_pdf(pdf_path, query_regex, use_regex=True)
# for r in res:
# print(r)
# # 例2模糊匹配一句话
# res2 = find_text_in_pdf(pdf_path, "这是一段可能不完全一样的文本", threshold=75)
# for r in res2:
# print(r)
if __name__ == "__main__":
pdf_path = 'e:\\2\\2024深化智慧城市发展推进城市全域数字化转型的指导意见.pdf'
pdf_path = 'G:\\SynologyDrive\\大模型\\RAG\\20250805党建\\中国共产党领导干部廉洁从业若干准则.pdf'
pdf_path ="F:\\Synology_nas\\SynologyDrive\\大模型\\RAG\\20250805党建\\中国共产党领导干部廉洁从业若干准则.pdf"
query = [
'''一、总体要求
以习近平新时代中国特色社会主义思想为指导完整、准确、全面贯彻新发展理念统筹发展和安全充分发挥数据的基础资源和创新引擎作用整体性重塑智慧城市技术架构、系统性变革城市管理流程、一体化推动产城深度融合全面提升城市全域数字化转型的整体性、系统性、协同性不断满足人民日益增长的美好生活需要为全面建设社会主义现代化国家提供强大动力。到2027年全国城市全域数字化转型取得明显成效形成一批横向打通、纵向贯通、各具特色的宜居、韧性、智慧城市有力支撑数字中国建设。到2030年全国城市全域数字化转型全面突破人民群众的获得感、幸福感、安全感全面提升涌现一批数字文明时代具有全球竞争力的中国式现代化城市。''',
'''二、全领域推进城市数字化转型
(一)建立城市数字化共性基础。构建统一规划、统一架构、统一标准、统一运维的城市运行和治理智能中枢,打造线上线下联动、服务管理协同的城市共性支撑平台,构建开放兼容、共性赋能、安全可靠的综合性基础环境,推进算法、模型等数字资源一体集成部署,探索建立共性组件、模块等共享协作机制。鼓励发展基于人工智能等技术的智能分析、智能调度、智能监管、辅助决策,全面支撑赋能城市数字化转型场景建设与发展。鼓励有条件的地方推进城市信息模型、时空大数据、国土空间基础信息、实景三维中国等基础平台功能整合、协同发展、应用赋能,为城市数字化转型提供统一的时空框架,因地制宜有序探索推进数字孪生城市建设,推动虚实共生、仿真推演、迭代优化的数字孪生场景落地。
(二)培育壮大城市数字经济。深入推进数字技术与一二三产业深度融合,鼓励平台企业构建多层次产业互联网服务平台。因地制宜发展智慧农业,加快工业互联网规模化应用,推动金融、物流等生产性服务业和商贸、文旅、康养等生活性服务业数字化转型,提升“上云用数赋智”水平。深化数字化转型促进中心建设,促进城市数字化转型和大中小企业融合创新协同发展。因地制宜发展新兴数字产业,加强大数据、人工智能、区块链、先进计算、未来网络、卫星遥感、三维建模等关键数字技术在城市场景中集成应用,加快技术创新成果转化,打造具有国际竞争力的数字产业集群。培育壮大数据产业,发展一批数据商和第三方专业服务机构,提高数据要素应用支撑与服务能力。''',
"""(三)促进新型产城融合发展。创新生产空间和生活空间融合的数字化场景,加强城市空间开发利用大数据分析,推进数字化赋能郊区新城,实现城市多中心、网络化、组团式发展。推动城市“数字更新”,加快街区、商圈等城市微单元基础设施智能化升级,探索利用数字技术创新应用场景,激发产城融合服务能级与数字活力。深化城市场景开放促进以城带产,提升产业聚合力。加速创新资源共享助力以产促城,发展虚拟园区和跨区域协同创新平台,增强城市数字经济就业吸附力。"""
]
query = ["""执政党的党风关系党的生死存亡。坚决惩治和有效预防腐败,是党必须始终抓好
的重大政治任务。党员领导干部廉洁从政是坚持以邓小平理论和“三个代表"重要思想为
指导,深入贯彻落实科学发展观,全面贯彻党的路线方针政策的重要保障;是新时期
从严治党,不断加强党的执政能力建设和先进性建设的重要内容;是推进改革开放和
社会主义现代化建设的基本要求;是正确行使权力、履行职责的重要基础。
党员领导干部必须具有共产主义远大理想和中国特色社会主义坚定信念,践行社
会主义核心价值体系;必须坚持全心全意为人民服务的宗旨,立党为公、执政为民;
必须在党员和人民群众中发挥表率作用,自重、自省、自警、自励;必须模范遵守党
纪国法,清正廉洁,忠于职守,正确行使权力,始终保持职务行为的廉洁性;必须弘
扬党的优良作风,求真务实,艰苦奋斗,密切联系群众。
促进党员领导干部廉洁从政,必须坚持标本兼治、综合治理、惩防并举、注重预
防的方针,按照建立健全惩治和预防腐败体系的要求,加强教育,健全制度,强化监
督,深化改革,严肃纪律,坚持自律和他律相结合。
""",
"""第三条 禁止违反公共财物管理和使用的规定,假公济私、化公为私。不准有下
列行为:
(一)用公款报销或者支付应由个人负担的费用;
(二)违反规定借用公款、公物或者将公款、公物借给他人;
(三)私存私放公款;
(四)用公款旅游或者变相用公款旅游;
(五)用公款参与高消费娱乐、健身活动和获取各种形式的俱乐部会员资格;
(六)违反规定用公款购买商业保险,缴纳住房公积金,滥发津贴、补贴、奖金
等;
(七)非法占有公共财物,或者以象征性地支付钱款等方式非法占有公共财物;
(八)挪用或者拆借社会保障基金、住房公积金等公共资金或者其他财政资金。
""",
"""
第六条禁止讲排场、比阔气、挥霍公款、铺张浪费。不准有下列行为:
(一)在公务活动中提供或者接受超过规定标准的接待,或者超过规定标准报销
招待费、差旅费等相关费用;
(二)违反规定决定或者批准兴建、装修办公楼、培训中心等楼堂馆所,超标准
配备、使用办公用房和办公用品;
(三)擅自用公款包租、占用客房供个人使用;
(四)违反规定配备、购买、更换、装饰或者使用小汽车;
(五)违反规定决定或者批准用公款或者通过摊派方式举办各类庆典活动。
第七条禁止违反规定干预和插手市场经济活动,谋取私利。不准有下列行为:
(一)干预和插手建设工程项目承发包、土地使用权出让、政府采购、房地产开
发与经营、矿产资源开发利用、中介机构服务等市场经济活动;
""",
"""第七条禁止违反规定干预和插手市场经济活动,谋取私利。不准有下列行为:
(一)干预和插手建设工程项目承发包、土地使用权出让、政府采购、房地产开
发与经营、矿产资源开发利用、中介机构服务等市场经济活动;
(二)干预和插手国有企业重组改制、兼并、破产、产权交易、清产核资、资产
评估、资产转让、重大项目投资以及其他重大经营活动等事项;
(三)干预和插手批办各类行政许可和资金借贷等事项;
四)干预和插手经济纠纷;
(五)干预和插手农村集体资金、资产和资源的使用、分配、承包、租赁等事
项。
"""
]
# 1. 找跨行正则匹配
matches = find_text_in_pdf(
pdf_path,
query,
threshold=60
)
# 修改:正确处理二维列表结果
# print(matches)
# print("------------------")
for idx,query_matches in enumerate(matches):
print(f"{idx} 个查询结果:")
print(query_matches)
#highlight_matches(pdf_path, query_matches, "example_highlighted.pdf")
for m in query_matches:
print(f"{m['page']} 页 匹配: {m['matched_text'][:50]}... 位置: {m['bbox']}, 位置_int: {m['position_int']}")
print("------------------")
# 2. 高亮并保存
# 修改:展平二维列表用于高亮
# flattened_matches = [match for query_matches in matches for match in query_matches]
# highlight_matches(pdf_path, flattened_matches, "example_highlighted.pdf")