1 line
53 KiB
JSON
1 line
53 KiB
JSON
{"repo": "bit-datalab/bench-u", "files": {"systems/quest/db/indexer/zendb_indexer.py": {"f2da981b5d37c058345f27f23c7979eb44697f0c": {"t": 1785809913.3198915, "n": 46007, "h": "87fd8425c81030b7", "c": "#!/usr/bin/env python3\r\n\"\"\"\r\nZenDB 完整实现 - 基于论文设计的学术文档分析系统\r\n\"\"\"\r\n\r\nMAX_SENTENCES = 10\r\n\r\nimport pdfplumber\r\nimport sqlite3\r\nimport openai\r\nimport re\r\nimport json\r\nimport pandas as pd\r\nimport numpy as np\r\nfrom pathlib import Path\r\nfrom sklearn.cluster import KMeans\r\nfrom sklearn.preprocessing import StandardScaler\r\nfrom dataclasses import dataclass\r\nfrom typing import List, Dict, Any, Optional\r\nfrom datetime import datetime\r\nimport time\r\nfrom openai import OpenAI\r\nimport random\r\nfrom quest.core.datapack.doc import ZenDBDoc\r\nfrom quest.core.llm.llm_query import LLMInfo\r\nfrom quest.core.nlp.doc_summary import tfidf_summary\r\nfrom quest.db.indexer.single_indexer import SingleIndexer\r\nfrom quest.db.querier.querier import OpenGaussQuerier\r\nimport wordninja\r\nimport pyparsing as pp\r\nimport os\r\nimport pickle\r\nfrom sklearn.metrics.pairwise import cosine_similarity\r\nimport faiss\r\n\r\n# 设置DeepSeek API密钥和配置\r\nDEEPSEEK_API_KEY = 'sk-DrStrjrlXohoCI2iBc14Eb41A38c428b88A0973aA29fCc3b'\r\nDEEPSEEK_BASE_URL = 'https://aihubmix.com/v1'\r\n\r\n# 初始化DeepSeek客户端\r\nclient = OpenAI(\r\n api_key=DEEPSEEK_API_KEY,\r\n base_url=DEEPSEEK_BASE_URL\r\n)\r\n\r\n@dataclass\r\nclass SHTNode:\r\n \"\"\"语义层次树节点\"\"\"\r\n node_id: int\r\n doc_id: int\r\n name: str\r\n granularity: int\r\n context: str\r\n summary: str\r\n start_pos: int\r\n end_pos: int\r\n parent: Optional['SHTNode'] = None\r\n children: List['SHTNode'] = None\r\n visual_pattern: Dict[str, Any] = None\r\n full_context: str = None # 新增字段，包含本节点及所有子节点的全文\r\n \r\n def __post_init__(self):\r\n if self.children is None:\r\n self.children = []\r\n if self.full_context is None:\r\n self.full_context = self.context\r\n \r\n def to_dict(self):\r\n return {\r\n \"node_id\": self.node_id,\r\n \"doc_id\": self.doc_id,\r\n \"name\": self.name,\r\n \"granularity\": self.granularity,\r\n \"context\": self.context,\r\n \"summary\": self.summary,\r\n \"start_pos\": self.start_pos,\r\n \"end_pos\": self.end_pos,\r\n \"visual_pattern\": self.visual_pattern,\r\n # \"full_context\": self.full_context, # 不需要存full_text，太大\r\n \"children\": [child.to_dict() for child in self.children],\r\n }\r\n\r\n @staticmethod\r\n def from_dict(data, parent=None):\r\n node = SHTNode(\r\n node_id= int(data[\"node_id\"]),\r\n doc_id=int(data[\"doc_id\"]),\r\n name=data[\"name\"],\r\n granularity=data[\"granularity\"],\r\n context=data[\"context\"],\r\n summary=data.get(\"summary\", \"\"),\r\n start_pos= int(data.get(\"start_pos\", 0)),\r\n end_pos= int(data.get(\"end_pos\", 0)) ,\r\n parent=parent,\r\n children=[],\r\n visual_pattern=data.get(\"visual_pattern\"),\r\n # full_context=data.get(\"full_context\")\r\n )\r\n node.children = [SHTNode.from_dict(child, parent=node) for child in data.get(\"children\", [])]\r\n return node\r\n\r\n\r\n def get_ancestors(self):\r\n \"\"\"获取祖先节点路径\"\"\"\r\n ancestors = []\r\n current = self.parent\r\n while current:\r\n ancestors.append(current.name)\r\n current = current.parent\r\n return ancestors[::-1]\r\n \r\n def is_leaf(self):\r\n \"\"\"判断是否为叶子节点\"\"\"\r\n return len(self.children) == 0\r\n\r\nfrom quest.conf.settings import ABS_PROJECT_ROOT_PATH, count_tokens\r\n\r\nclass ZenDBDocIndexer(SingleIndexer):\r\n \"\"\"ZenDBIndexer主类\"\"\"\r\n \r\n def __init__(self, table_name: str, type: str = \"ZenDBDoc\", root_save_path = os.path.join(ABS_PROJECT_ROOT_PATH, \"data/single_index\"), need_clean_chunk = False, embedding_model = None, **kwargs):\r\n print(\"🚀 建立ZenDB-singleIndex索引...\")\r\n self.embedding_size = embedding_model.emb_size \r\n self.embedding_model = embedding_model\r\n self.need_clean_chunk = need_clean_chunk\r\n self.table_name = table_name\r\n self.type = type\r\n\r\n self.table_save_path = os.path.join(root_save_path, f\"{self.table_name}.json\")\r\n self.embedding_save_path = os.path.join(root_save_path, f\"{self.table_name}_embeddings.pkl\")\r\n self.sht_tables = {} # doc_id: int -> SHT nodes\r\n self.user_tables = {} # table_name -> SQLite table\r\n self.templates = [] # 跨文档模板缓存\r\n # 新增embedding相关成员变量\r\n self.node_embeddings: Dict[int, np.ndarray] = {} # node_id -> embedding向量\r\n # self.embedding_cache: Dict[str, np.ndarray] = {} # 缓存相同文本的embedding\r\n self.query_embedding_cache: Dict[str, np.ndarray] = {} # 缓存相同文本的embedding\r\n # self.querier = OpenGaussQuerier() \r\n\r\n def build_indexer(self, docs: List[ZenDBDoc]) -> None:\r\n # 提取并存储docs_meta\r\n self.docs_meta = {doc[\"doc_id\"]: doc.metadata for doc in docs}\r\n # self.querier.build_cache_table(self.table_name, self.docs_meta)\r\n\r\n self.build_sht_for_all_docs(docs)\r\n # 为所有节点生成embedding\r\n self._generate_embeddings_for_all_docs()\r\n self.save_indexer()\r\n return\r\n\r\n def save_indexer(self):\r\n # 保存SHT树结构\r\n path = self.table_save_path\r\n sht_dict = {doc_id: root.to_dict() for doc_id, root in self.sht_tables.items()}\r\n data = {\r\n \"docs_meta\": self.docs_meta,\r\n \"sht_trees\": sht_dict\r\n }\r\n with open(path, \"w\", encoding=\"utf-8\") as f:\r\n json.dump(data, f, ensure_ascii=False, indent=2)\r\n \r\n # 保存embedding数据\r\n with open(self.embedding_save_path, \"wb\") as f:\r\n pickle.dump(self.node_embeddings, f)\r\n\r\n def load_indexer(self):\r\n # 加载SHT树结构\r\n path = self.table_save_path\r\n with open(path, \"r\", encoding=\"utf-8\") as f:\r\n data = json.load(f)\r\n self.docs_meta = data[\"docs_meta\"]\r\n # 把字典中所有key转换成int类型\r\n self.docs_meta = {int(k): v for k, v in self.docs_meta.items()}\r\n self.sht_tables = {\r\n int(doc_id) : SHTNode.from_dict(root_dict)\r\n for doc_id, root_dict in data[\"sht_trees\"].items()\r\n }\r\n \r\n # 加载embedding数据\r\n if os.path.exists(self.embedding_save_path):\r\n with open(self.embedding_save_path, \"rb\") as f:\r\n self.node_embeddings = pickle.load(f)\r\n return\r\n\r\n def _generate_embeddings_for_all_docs(self):\r\n \"\"\"为所有文档的所有节点生成embedding\"\"\"\r\n print(\"\\n🧠 生成节点embedding...\")\r\n total_nodes = 0\r\n for doc_id, root in self.sht_tables.items():\r\n self._generate_embeddings(root)\r\n total_nodes += self._count_nodes(root)\r\n print(f\" ✓ 完成 {total_nodes} 个节点的embedding生成\")\r\n\r\n def _generate_embeddings(self, node: SHTNode) -> None:\r\n \"\"\"为节点生成embedding - 使用批量嵌入提高效率\"\"\"\r\n # 获取所有节点的context和node_id列表\r\n context_list, node_id_list = self.level_traverse(node.doc_id)\r\n \r\n if not context_list:\r\n return\r\n \r\n # 批量生成embedding\r\n embeddings = self.embedding_model.embed_documents(context_list)\r\n \r\n # 构建node_id到embedding的映射\r\n for node_id, embedding in zip(node_id_list, embeddings):\r\n self.node_embeddings[node_id] = embedding\r\n\r\n def _get_embedding(self, node_id: int) -> np.ndarray:\r\n \"\"\"获取节点的context的embedding向量\"\"\"\r\n return self.node_embeddings.get(node_id, None)\r\n\r\n def _get_cached_query_embedding(self, query: str):\r\n \"\"\"\r\n 获取缓存的query embedding，如果不存在则计算并缓存\r\n \r\n Args:\r\n query: 查询字符串\r\n \r\n Returns:\r\n query的embedding向量\r\n \"\"\"\r\n cache_key = query[:256] # 使用前256个字符作为缓存键\r\n \r\n if cache_key not in self.query_embedding_cache:\r\n # 计算embedding并缓存\r\n self.query_embedding_cache[cache_key] = self.embedding_model.embed_query(query)\r\n \r\n return self.query_embedding_cache[cache_key]\r\n\r\n\r\n def _semantic_similarity_search(self, query: str, node_ids: List[int], topk: int) -> List[int]:\r\n \"\"\"在指定节点集合中进行语义相似度搜索（使用faiss加速）\"\"\"\r\n if not node_ids:\r\n return []\r\n \r\n query_embedding = self._get_cached_query_embedding(query)\r\n # query_embedding = self.embedding_model.embed_query(query)\r\n # 收集节点embedding\r\n embeddings = []\r\n valid_node_ids = []\r\n for node_id in node_ids:\r\n emb = self._get_embedding(node_id)\r\n if emb is not None:\r\n embeddings.append(emb)\r\n valid_node_ids.append(node_id)\r\n else:\r\n raise ValueError(f\"节点 {node_id} 没有embedding\")\r\n if not embeddings:\r\n return []\r\n # 转为np数组\r\n xb = np.array(embeddings).astype('float32')\r\n xq = np.array([query_embedding]).astype('float32')\r\n # 归一化（余弦相似度）\r\n faiss.normalize_L2(xb)\r\n faiss.normalize_L2(xq)\r\n # 建立faiss索引\r\n index = faiss.IndexFlatIP(xb.shape[1])\r\n index.add(xb)\r\n D, I = index.search(xq, min(topk, xb.shape[0]))\r\n # 返回对应的node_id\r\n result_ids = [valid_node_ids[i] for i in I[0]]\r\n return result_ids\r\n\r\n def build_indexer(self, docs: List[ZenDBDoc]) -> None:\r\n # 提取并存储docs_meta\r\n self.docs_meta = {doc[\"doc_id\"]: doc.metadata for doc in docs}\r\n # self.querier.build_cache_table(self.table_name, self.docs_meta)\r\n\r\n self.build_sht_for_all_docs(docs)\r\n # 为所有节点生成embedding\r\n self._generate_embeddings_for_all_docs()\r\n self.save_indexer()\r\n return\r\n\r\n def save_indexer(self):\r\n # 保存SHT树结构\r\n path = self.table_save_path\r\n sht_dict = {doc_id: root.to_dict() for doc_id, root in self.sht_tables.items()}\r\n data = {\r\n \"docs_meta\": self.docs_meta,\r\n \"sht_trees\": sht_dict\r\n }\r\n with open(path, \"w\", encoding=\"utf-8\") as f:\r\n json.dump(data, f, ensure_ascii=False, indent=2)\r\n \r\n # 保存embedding数据\r\n with open(self.embedding_save_path, \"wb\") as f:\r\n pickle.dump(self.node_embeddings, f)\r\n\r\n def load_indexer(self):\r\n # 加载SHT树结构\r\n path = self.table_save_path\r\n with open(path, \"r\", encoding=\"utf-8\") as f:\r\n data = json.load(f)\r\n self.docs_meta = data[\"docs_meta\"]\r\n # 把字典中所有key转换成int类型\r\n self.docs_meta = {int(k): v for k, v in self.docs_meta.items()}\r\n self.sht_tables = {\r\n int(doc_id) : SHTNode.from_dict(root_dict)\r\n for doc_id, root_dict in data[\"sht_trees\"].items()\r\n }\r\n \r\n # 加载embedding数据\r\n if os.path.exists(self.embedding_save_path):\r\n with open(self.embedding_save_path, \"rb\") as f:\r\n self.node_embeddings = pickle.load(f)\r\n return\r\n\r\n def get_file_name_by_id(self, doc_id: int) -> str:\r\n \"\"\"\r\n 根据文档ID获取对应的文件名\r\n\r\n Args:\r\n doc_id: 文档编号\r\n\r\n Returns:\r\n str: 对应的文件名\r\n \"\"\"\r\n if doc_id in self.docs_meta:\r\n return self.docs_meta[doc_id].get(\"file_name\", \"\")\r\n return \"\"\r\n\r\n def get_chunks_by_docid(self, doc_id):\r\n chunk_list, id_list = self.level_traverse(doc_id)\r\n return chunk_list\r\n\r\n def level_traverse(self, doc_id):\r\n \"\"\"\r\n 层序遍历指定文档的分块树，返回 (context_list, node_id_list)。\r\n \"\"\"\r\n root = self.sht_tables.get(doc_id)\r\n if root is None:\r\n return [], []\r\n context_list = []\r\n node_id_list = []\r\n queue = [] \r\n queue.extend(root.children) # 不要把root节点的context返回到结果中，因为它包含全文。\r\n while queue:\r\n node = queue.pop(0)\r\n context_list.append(node.context)\r\n node_id_list.append(node.node_id)\r\n queue.extend(node.children)\r\n return context_list, node_id_list\r\n\r\n\r\n def print_sht_tree(self, doc_id):\r\n \"\"\"\r\n 以前序遍历的形式递归地打印整个SHT分块树的结构，按节点所在的深度进行缩进。\r\n \r\n Args:\r\n doc_id: 文档编号\r\n \"\"\"\r\n root = self.sht_tables.get(doc_id)\r\n if root is None:\r\n print(f\"文档 {doc_id} 的SHT树不存在\")\r\n return\r\n \r\n def _print_node(node, depth=0):\r\n \"\"\"递归打印节点及其子节点\"\"\"\r\n indent = \" \" * depth\r\n print(f\"{indent}node_id: {node.node_id}\")\r\n for child in node.children:\r\n _print_node(child, depth + 1)\r\n \r\n print(f\"文档 {doc_id} 的SHT树结构:\")\r\n _print_node(root)\r\n\r\n\r\n def llm_judge_contain_attr(self, node, query):\r\n \"\"\"\r\n 调用LLM，根据当前节点的summary判断当前节点是否可能包含query中需要提取的属性，\r\n llm只给出true或false的结果以及结果的confidence(0到100的整数)，\r\n 返回 (bool, int)，如果格式不对则默认(False, 100)，并统计token消耗。\r\n \"\"\"\r\n try:\r\n # 构造提示\r\n prompt = (\r\n f\"Determine if the following section contains the attribute defined in the query based on the summary of the section.\\n\"\r\n f\"Query: {query}\\n\"\r\n f\"Section Title: {node.name}\\n\"\r\n f\"Section Summary: {node.summary}\\n\"\r\n f\"Answer format: true or false, confidence from 0 to 100. e.g. (true, 100)\"\r\n )\r\n # 统计输入输出token（简易估算）\r\n input_tokens = count_tokens(prompt)\r\n response = self._llm_query(prompt)\r\n output_tokens = count_tokens(response)\r\n # 更新token计数和调用次数计数\r\n LLMInfo.add_input_tokens(input_tokens)\r\n LLMInfo.add_output_tokens(output_tokens)\r\n LLMInfo.add_query_times(1)\r\n\r\n # 解析返回结果\r\n m = re.search(r\"(true|false)\\s*[,，]\\s*(\\d{1,3})\", response.lower(), re.IGNORECASE)\r\n if m:\r\n result = m.group(1).lower() == 'true'\r\n confidence = min(max(int(m.group(2)), 0), 100)\r\n return result, confidence\r\n except Exception:\r\n pass\r\n # 默认值\r\n return False, 0\r\n\r\n def pre_order_dfs_search_related_node(self, node, query):\r\n \"\"\"\r\n 基于DFS的前序遍历，判断并收集可能包含查询属性的节点：\r\n - 若 node 为空或 llm_judge_contain_attr 返回 False，返回 False\r\n - 若 llm_judge_contain_attr 返回 True，则对子节点递归调用，\r\n 当且仅当所有子节点递归均返回 False 时，将当前节点加入候选\r\n 返回是否该 subtree 可能包含属性的 bool\r\n 高级剪枝策略：利用topK，如果当前判断包含属性的块数已经达到topK，则直接返回。\r\n \"\"\"\r\n if node is None:\r\n return False\r\n contain, confidence = self.llm_judge_contain_attr(node, query)\r\n if not contain:\r\n return False\r\n # 标记子节点中是否有命中\r\n child_hit = False\r\n for child in node.children:\r\n if self.pre_order_dfs_search_related_node(child, query):\r\n child_hit = True\r\n # 若子节点都未命中，则将当前节点作为候选\r\n if not child_hit:\r\n if not hasattr(self, 'candidate_nodes_with_confidence'):\r\n self.candidate_nodes = []\r\n self.candidate_nodes.append((node, confidence))\r\n return True\r\n\r\n\r\n def topk_semantic_beam_search_related_node(self, node: SHTNode, query: str, beam_topk: int) -> bool:\r\n \"\"\"基于层次化束搜索的节点搜索方法\"\"\"\r\n self.candidate_nodes = []\r\n only_one_1_level_title_flag = False\r\n if beam_topk <= 0:\r\n return False\r\n \r\n \r\n # 第一阶段：收集所有granularity=1的节点\r\n # granularity_1_nodes = []\r\n # self._collect_nodes_by_granularity(node, 1, granularity_1_nodes)\r\n # 直接取node(即root节点)的所有子节点作为granularity=1的节点\r\n # jztodo 如果1级标题只有1个，则把2级标题作为granularity_1_nodes。\r\n granularity_1_nodes = node.children\r\n if len(granularity_1_nodes) == 1:\r\n granularity_1_nodes = granularity_1_nodes[0].children\r\n only_one_1_level_title_flag = True\r\n \r\n if len(granularity_1_nodes) == 0 or not granularity_1_nodes:\r\n # 如果没有granularity=1的节点，直接认为整篇文章不包含attr，返回False\r\n return False\r\n \r\n # 按语义相似度排序granularity=1的节点，取topK个进行LLM判断\r\n node_ids = [n.node_id for n in granularity_1_nodes]\r\n sorted_node_ids = self._semantic_similarity_search(query, node_ids, beam_topk)\r\n \r\n # 第二阶段：LLM判断语义相似度最高的topK个granularity=1节点\r\n V1 = [] # 存储被LLM判定为true的granularity=1节点\r\n for node_id in sorted_node_ids:\r\n # 找到对应的节点\r\n target_node = None\r\n for n in granularity_1_nodes:\r\n if n.node_id == node_id:\r\n target_node = n\r\n break\r\n \r\n if target_node:\r\n contain, confidence = self.llm_judge_contain_attr(target_node, query)\r\n if contain:\r\n V1.append(target_node)\r\n \r\n # 第三、四阶段：层次化收集子节点\r\n all_candidate_nodes = V1.copy() # 从V1开始\r\n current_level_nodes = V1\r\n \r\n while current_level_nodes:\r\n # 收集当前层级所有节点的子节点\r\n next_level_nodes = []\r\n for parent_node in current_level_nodes:\r\n next_level_nodes.extend(parent_node.children)\r\n \r\n if not next_level_nodes:\r\n break\r\n \r\n # 按语义相似度排序，取topK个\r\n child_node_ids = [n.node_id for n in next_level_nodes]\r\n sorted_child_ids = self._semantic_similarity_search(query, child_node_ids, beam_topk)\r\n \r\n # 获取topK个子节点\r\n topk_children = []\r\n for node_id in sorted_child_ids:\r\n for child in next_level_nodes:\r\n if child.node_id == node_id:\r\n topk_children.append(child)\r\n break\r\n \r\n # 添加到候选列表\r\n all_candidate_nodes.extend(topk_children)\r\n current_level_nodes = topk_children\r\n \r\n # 第五阶段：合并所有候选节点，按语义相似度排序，取最终topK\r\n if all_candidate_nodes:\r\n # 去重（避免重复节点）\r\n unique_nodes = []\r\n seen_ids = set()\r\n for node in all_candidate_nodes:\r\n if node.node_id not in seen_ids:\r\n unique_nodes.append(node)\r\n seen_ids.add(node.node_id)\r\n \r\n # 按语义相似度排序 jztodo 如果1级标题就1个，把它也加入到列表中\r\n final_node_ids = [n.node_id for n in unique_nodes]\r\n if only_one_1_level_title_flag:\r\n final_node_ids.append(granularity_1_nodes[0].node_id)\r\n final_sorted_ids = self._semantic_similarity_search(query, final_node_ids, beam_topk)\r\n \r\n # 取topK个节点，进行最终的LLM判断并添加到候选列表\r\n for i, node_id in enumerate(final_sorted_ids[:beam_topk]):\r\n # 找到对应的节点\r\n target_node = None\r\n for n in unique_nodes:\r\n if n.node_id == node_id:\r\n target_node = n\r\n unique_nodes.remove(n)\r\n break\r\n \r\n if target_node:\r\n self.candidate_nodes.append(target_node)\r\n \r\n return len(self.candidate_nodes) > 0\r\n\r\n def _collect_nodes_by_granularity(self, node: SHTNode, target_granularity: int, result_list: List[SHTNode]):\r\n \"\"\"递归收集指定granularity的节点\"\"\"\r\n if node.granularity == target_granularity:\r\n result_list.append(node)\r\n \r\n # 继续递归搜索子节点\r\n for child in node.children:\r\n self._collect_nodes_by_granularity(child, target_granularity, result_list)\r\n\r\n\r\n def topk_query_all_nodes(self, doc_id: int, query: str, topk: int) -> List[tuple[str, int]]:\r\n \"\"\"\r\n 对指定文档的所有节点进行语义相似度搜索，返回最相关的topK个节点\r\n \r\n Args:\r\n doc_id: 文档ID\r\n query: 查询文本\r\n topk: 返回的节点数量\r\n \r\n Returns:\r\n List[tuple[str, int]]: [(node_context, node_id), ...] 按相似度降序排列\r\n \"\"\"\r\n # 获取根节点\r\n root = self.sht_tables.get(doc_id)\r\n if root is None:\r\n return []\r\n \r\n # 使用level_traverse获取所有节点\r\n context_list, node_id_list = self.level_traverse(doc_id)\r\n \r\n if not node_id_list:\r\n return []\r\n \r\n # 对所有节点进行语义相似度搜索\r\n sorted_node_ids = self._semantic_similarity_search(query, node_id_list, topk)\r\n \r\n # 构造结果列表\r\n results = []\r\n node_id_to_context = dict(zip(node_id_list, context_list))\r\n \r\n for node_id in sorted_node_ids:\r\n context = node_id_to_context.get(node_id, \"\")\r\n results.append((context, node_id))\r\n \r\n return results\r\n\r\n\r\n\r\n\r\n def get_relative_chunks_text_with_id(self, doc_id: int, query: str, topk: int) -> List[tuple[str, int]]:\r\n \"\"\"\r\n 查询指定文档中与query相关的topk个文本块\r\n 返回 [(chunk_text, chunk_id), ...]，chunk_id为块起始位置\r\n \"\"\"\r\n # 获取根节点\r\n root = self.sht_tables.get(doc_id)\r\n if root is None:\r\n return []\r\n \r\n # 重置候选列表\r\n self.candidate_nodes = []\r\n \r\n # 只有根节点的一级子节点会参与判断，其他节点都是直接用束搜索取topk。\r\n has_candidate_flag = self.topk_semantic_beam_search_related_node(root, query, topk) \r\n # 如果has_candidate_flag为False，退化成对所有树节点进行topK语义相似度搜索。先使用level_traversal获得所有节点，然后进行topK语义相似度搜索。 \r\n if not has_candidate_flag:\r\n return self.topk_query_all_nodes(doc_id, query, topk) \r\n\r\n topk_nodes = self.candidate_nodes \r\n \r\n # 构造结果列表\r\n results = []\r\n for node in topk_nodes:\r\n results.append((node.context, node.node_id))\r\n \r\n \r\n return results\r\n\r\n def depreacated_topk_semantic_dfs_search_related_node(self, node: SHTNode, query: str, remaining_topk: int) -> bool:\r\n \"\"\"新的基于语义相似度的DFS搜索方法\"\"\"\r\n if remaining_topk <= 0:\r\n return False\r\n \r\n # 对当前节点进行LLM判断\r\n contain, confidence = self.llm_judge_contain_attr(node, query)\r\n \r\n if contain and confidence >= 70:\r\n # 如果当前节点包含属性，加入候选集合\r\n if not hasattr(self, 'candidate_nodes_with_confidence'):\r\n self.candidate_nodes = []\r\n \r\n if len(self.candidate_nodes) < remaining_topk:\r\n self.candidate_nodes.append((node, confidence))\r\n remaining_topk -= 1\r\n \r\n # 如果还需要更多候选节点且有子节点，继续搜索\r\n if remaining_topk > 0 and node.children:\r\n # 使用语义相似度对子节点排序\r\n child_ids = [child.node_id for child in node.children]\r\n sorted_child_ids = self._semantic_similarity_search(query, child_ids, len(child_ids))\r\n \r\n # 按相似度顺序处理子节点\r\n for child_id in sorted_child_ids:\r\n if remaining_topk <= 0:\r\n break\r\n \r\n # 找到对应的子节点\r\n child_node = None\r\n for child in node.children:\r\n if child.node_id == child_id:\r\n child_node = child\r\n break\r\n \r\n if child_node:\r\n if self.depreacated_topk_semantic_dfs_search_related_node(child_node, query, remaining_topk):\r\n remaining_topk = max(0, remaining_topk - len(self.candidate_nodes))\r\n \r\n return contain\r\n\r\n def _clean_text(self, text):\r\n \"\"\"深度清洗文本内容，去除PDF转换乱码和无效字符\"\"\"\r\n if not text:\r\n return \"\"\r\n \r\n # 1. 移除PDF字符编码乱码 (cid:XX) 格式\r\n text = re.sub(r'\\(cid:\\d+\\)', '', text)\r\n \r\n # 2. 移除其他常见PDF乱码模式\r\n text = re.sub(r'[\\x00-\\x08\\x0b\\x0c\\x0e-\\x1f\\x7f]', '', text) # 控制字符\r\n text = re.sub(r'�', '', text) # Unicode替换字符\r\n \r\n # 3. 移除重复的空白字符\r\n text = re.sub(r'\\s+', ' ', text) # 多个空格合并为一个\r\n text = re.sub(r'\\n\\s*\\n\\s*\\n', '\\n\\n', text) # 多个连续换行最多保留两个\r\n \r\n # 4. 清理行首行尾空格\r\n lines = text.split('\\n')\r\n cleaned_lines = []\r\n for line in lines:\r\n line = line.strip()\r\n if line: # 跳过空行\r\n cleaned_lines.append(line)\r\n \r\n text = '\\n'.join(cleaned_lines)\r\n \r\n # 5. 移除明显的页眉页脚模式（单独的数字可能是页码）\r\n lines = text.split('\\n')\r\n filtered_lines = []\r\n for line in lines:\r\n # 跳过可能的页码（单独的数字或短横线）\r\n if re.match(r'^\\s*[-\\d\\s]+\\s*$', line) and len(line.strip()) < 5:\r\n continue\r\n # 跳过只有特殊字符的行\r\n if re.match(r'^[\\s\\-_=~`!@#$%^&*(){}[\\]|\\\\:\";\\'<>?,.]+$', line):\r\n continue\r\n filtered_lines.append(line)\r\n \r\n text = '\\n'.join(filtered_lines)\r\n \r\n # 6. 最终格式化\r\n text = text.strip()\r\n \r\n # 确保文本有意义（至少有一些字母数字字符）\r\n if len(re.findall(r'[a-zA-Z0-9]', text)) < 10:\r\n return \"\"\r\n \r\n return text\r\n\r\n\r\n def _extract_phrases(self, doc):\r\n def smart_split_phrase(phrase):\r\n # 匹配编号+正文\r\n m = re.match(r'^((?:\\d+\\.)*\\d+)\\s*([A-Za-z]+.*)', phrase)\r\n if m:\r\n prefix = m.group(1)\r\n rest = m.group(2)\r\n # 对正文部分分词\r\n rest_split = ' '.join(wordninja.split(rest))\r\n return f\"{prefix} {rest_split}\"\r\n # 只对全英文长串分词\r\n if re.match(r'^[A-Za-z]{10,}$', phrase):\r\n return ' '.join(wordninja.split(phrase))\r\n return phrase\r\n \r\n \"\"\"根据视觉特征连续合并短语，符合论文定义\"\"\"\r\n words = doc[\"words\"]\r\n if (not words) or (doc[\"file_type\"] == \"txt\") or (doc[\"file_type\"] == \"md\"):\r\n # 针对txt和md，按md标题来分phrases\r\n line = doc[\"text\"]\r\n phrases = []\r\n lines = [line.strip() for line in doc[\"text\"].splitlines() if line.strip()]\r\n\r\n normal_features = {\r\n \"size\": 12, # 默认\r\n \"fontname\": \"normal_text\",\r\n \"type\": {\"bold\": False, \"underline\": False},\r\n \"all_cap\": line.isupper(),\r\n \"num_st\": line[:1].isdigit(),\r\n \"alpha_st\": line[:1].isalpha(),\r\n \"center\": False\r\n }\r\n\r\n title_features = {\r\n \"size\": 12, # 默认\r\n \"fontname\": \"title\",\r\n \"type\": {\"bold\": True, \"underline\": False},\r\n \"all_cap\": line.isupper(),\r\n \"num_st\": line[:1].isdigit(),\r\n \"alpha_st\": line[:1].isalpha(),\r\n \"center\": False\r\n }\r\n\r\n #####################################\r\n\r\n\r\n i = 0\r\n while i < len(lines):\r\n line = lines[i]\r\n # 检查是否为markdown标题\r\n header_match = re.match(r'^(#{1,6})\\s+(.+)$', line)\r\n \r\n if header_match:\r\n # 如果发现是markdown标题\r\n header_level = len(header_match.group(1)) # 获取#的数量\r\n header_text = line # 获取完整标题，包括前面的#\r\n \r\n # 设置title特征，fontname为title{level}\r\n features = title_features.copy()\r\n features[\"fontname\"] = f\"title{header_level}\"\r\n features[\"type\"][\"bold\"] = True\r\n \r\n # 添加到phrases\r\n phrases.append((header_text, features))\r\n i += 1\r\n else:\r\n # 不是标题，收集连续的非标题文本\r\n normal_text = line\r\n normal_features_copy = normal_features.copy()\r\n \r\n # 向前收集连续的非标题行\r\n j = i + 1\r\n while j < len(lines):\r\n next_line = lines[j]\r\n # 检查下一行是否为标题\r\n if re.match(r'^(#{1,6})\\s+(.+)$', next_line):\r\n break # 遇到标题行，停止合并\r\n normal_text += '\\n' + next_line # 合并非标题行\r\n j += 1\r\n \r\n # 添加合并后的正文\r\n phrases.append((normal_text, normal_features_copy))\r\n \r\n # 跳过已处理的行\r\n i = j\r\n\r\n return phrases\r\n # for line in lines:\r\n # features = normal_features\r\n\r\n # # jztodo 如果发现是markdown标题，先把之前的normal_lines(如果有的话)合并成一个phrase，然后再把title_features的fontname设置成title{n}，n是markdown标题的级别, type -> bold设置成True，然后append到phrases里面\r\n # # 如果不是标题，则和其他的line合并。 \r\n # phrases.append((line, features))\r\n\r\n # return phrases\r\n\r\n\r\n #################################\r\n\r\n phrases = []\r\n current_phrase = \"\"\r\n current_features = None\r\n current_visual = None\r\n\r\n def get_word_visual(word):\r\n # 只考虑字体名、大小、类型（加粗/下划线）\r\n fontname = word.get(\"fontname\", \"\")\r\n size = word.get(\"size\", 12)\r\n is_bold = int(\"bold\" in fontname.lower() or \"Bold\" in fontname)\r\n is_underline = int(\"underline\" in fontname.lower())\r\n return (size, fontname, is_bold, is_underline)\r\n\r\n for i, word in enumerate(words):\r\n word_text = word.get(\"text\", \"\")\r\n visual = get_word_visual(word)\r\n\r\n if current_phrase and visual != current_visual:\r\n # 生成短语级视觉特征\r\n phrase_text = current_phrase.strip()\r\n # 自动在编号和正文之间加空格\r\n phrase_text = smart_split_phrase(phrase_text)\r\n \r\n features = {\r\n \"size\": current_visual[0],\r\n \"fontname\": current_visual[1],\r\n \"type\": {\r\n \"bold\": bool(current_visual[2]),\r\n \"underline\": bool(current_visual[3])\r\n },\r\n \"all_cap\": phrase_text.isupper(),\r\n \"num_st\": phrase_text[:1].isdigit(),\r\n \"alpha_st\": phrase_text[:1].isalpha(),\r\n \"center\": False # 可根据x0/x1判断是否居中\r\n }\r\n phrases.append((phrase_text, features))\r\n current_phrase = word_text\r\n current_visual = visual\r\n else:\r\n current_phrase += (\"\" if current_phrase else \"\") + word_text\r\n current_visual = visual\r\n\r\n # 最后一个短语\r\n if current_phrase:\r\n phrase_text = current_phrase.strip()\r\n features = {\r\n \"size\": current_visual[0],\r\n \"fontname\": current_visual[1],\r\n \"type\": {\r\n \"bold\": bool(current_visual[2]),\r\n \"underline\": bool(current_visual[3])\r\n },\r\n \"all_cap\": phrase_text.isupper(),\r\n \"num_st\": phrase_text[:1].isdigit(),\r\n \"alpha_st\": phrase_text[:1].isalpha(),\r\n \"center\": False\r\n }\r\n phrases.append((phrase_text, features))\r\n\r\n\r\n return phrases\r\n \r\n def _build_sht(self, doc):\r\n return self._build_markdown_sht(doc)\r\n\r\n\r\n\r\n def _build_markdown_sht(self, doc):\r\n # jzmod\r\n \"\"\"构建语义层级树\"\"\"\r\n print(f\"\\n🏗️ 为文档 '{doc['name']}' 构建SHT...\")\r\n \r\n # 1. 提取短语\r\n phrases = self._extract_phrases(doc)\r\n print(f\" ✓ 提取 {len(phrases)} 个短语\")\r\n \r\n # 2. 识别标题短语\r\n headers = self._identify_headers(phrases)\r\n print(f\" ✓ 识别 {len(headers)} 个标题\")\r\n \r\n # 3. 构建树结构\r\n root = self._build_tree_structure(headers, doc)\r\n print(f\" ✓ 构建完成，根节点包含 {len(root.children)} 个子节点\")\r\n \r\n # 4. 生成摘要\r\n self._generate_summaries(root)\r\n print(f\" ✓ 生成节点摘要\")\r\n \r\n return root\r\n\r\n def _identify_headers(self, phrases):\r\n # 直接根据features判断哪些是标题。\r\n headers = []\r\n for phrase, features in phrases:\r\n if self._is_markdown_header(features):\r\n headers.append((phrase, features))\r\n\r\n return headers\r\n\r\n def _build_pdf_sht(self, doc):\r\n \"\"\"构建语义层次树\"\"\"\r\n print(f\"\\n🏗️ 为文档 '{doc['name']}' 构建SHT...\")\r\n \r\n # 1. 提取短语\r\n phrases = self._extract_phrases(doc)\r\n print(f\" ✓ 提取 {len(phrases)} 个短语\")\r\n \r\n # 2. 识别标题短语\r\n headers = self._identify_pdf_headers(phrases)\r\n print(f\" ✓ 识别 {len(headers)} 个标题\")\r\n \r\n # 3. 构建树结构\r\n root = self._build_pdf_tree_structure(headers, doc)\r\n print(f\" ✓ 构建完成，根节点包含 {len(root.children)} 个子节点\")\r\n \r\n # 4. 生成摘要\r\n self._generate_summaries(root)\r\n print(f\" ✓ 生成节点摘要\")\r\n \r\n return root\r\n \r\n def _identify_pdf_headers(self, phrases, k=10, n_clusters=8):\r\n \"\"\"根据视觉特征聚类+LLM识别标题短语，符合论文定义\"\"\"\r\n if not phrases:\r\n return []\r\n\r\n # 先用启发式规则过滤明显不是标题的短语\r\n # phrases = [(phrase, info) for phrase, info in phrases if self._is_likely_header(phrase)]\r\n if not phrases:\r\n return []\r\n\r\n # 1. 提取视觉特征向量\r\n features = []\r\n for phrase, info in phrases:\r\n size = info['size']\r\n fontname = hash(info['fontname']) % 10000\r\n bold = int(info['type']['bold'])\r\n underline = int(info['type']['underline'])\r\n features.append([size, fontname, bold, underline])\r\n\r\n # 2. 聚类\r\n kmeans = KMeans(n_clusters=min(n_clusters, len(phrases)), random_state=42)\r\n cluster_ids = kmeans.fit_predict(features)\r\n\r\n # 3. 给每个phrase加cluster id\r\n for idx, (phrase, info) in enumerate(phrases):\r\n info['cluster'] = int(cluster_ids[idx])\r\n\r\n # 4. 按cluster分组\r\n cluster_groups = {}\r\n for phrase, info in phrases:\r\n cluster = info['cluster']\r\n if cluster not in cluster_groups:\r\n cluster_groups[cluster] = []\r\n cluster_groups[cluster].append((phrase, info))\r\n\r\n headers = []\r\n headers_set = set()\r\n for cluster, cluster_phrases in cluster_groups.items():\r\n # 采样前再用启发式过滤\r\n # filtered_cluster_phrases = [p for p in cluster_phrases if self._is_likely_header(p[0])]\r\n filtered_cluster_phrases = cluster_phrases\r\n if not filtered_cluster_phrases:\r\n continue\r\n sample = random.sample(filtered_cluster_phrases, min(k, len(filtered_cluster_phrases)))\r\n header_votes = 0\r\n for phrase, info in sample:\r\n prompt = (\r\n \"You are analyzing a document. \"\r\n \"Section headers in a document typically indicate the start of a new section or subsection, such as chapter titles, main headings, or important topic divisions. \"\r\n \"Section headers are often short, descriptive, and may be formatted differently from the main text. \"\r\n \"Given the following phrase from the document:\\n\\n\"\r\n f'\"{phrase}\"\\n\\n'\r\n \"Is this phrase a section or subsection header in the document? \"\r\n \"Answer only 'yes' or 'no'.\"\r\n)\r\n\r\n try:\r\n response = self._llm_query(prompt)\r\n #print(f\"LLM判断: '{phrase}' -> {response}\")\r\n if response.strip().lower().startswith('y'):\r\n header_votes += 1\r\n except Exception as e:\r\n print(f\"LLM异常: {e}\")\r\n continue\r\n if header_votes > len(sample) / 2:\r\n for phrase, info in filtered_cluster_phrases:\r\n if phrase not in headers_set:\r\n headers.append((phrase, cluster))\r\n headers_set.add(phrase)\r\n return headers\r\n \r\n def _is_markdown_header(self, features):\r\n if \"title\" in features[\"fontname\"]:\r\n return True\r\n else:\r\n return False\r\n\r\n def _is_likely_header(self, phrase):\r\n \"\"\"启发式规则判断是否为标题\"\"\"\r\n phrase = phrase.strip()\r\n \r\n # 标题特征\r\n\r\n if len(phrase) < 5 or len(phrase) > 100:\r\n return False\r\n \r\n # 包含常见标题模式\r\n header_patterns = [\r\n r'^\\d+\\.?\\s+', # 数字开头\r\n r'^[A-Z][A-Z\\s]+$', # 全大写\r\n r'(abstract|introduction|related work|methodology|experiments|conclusion)', # 常见标题词\r\n r'^(figure|table|algorithm)\\s+\\d+', # 图表标题\r\n ]\r\n \r\n for pattern in header_patterns:\r\n if re.search(pattern, phrase, re.IGNORECASE):\r\n return True\r\n \r\n return False\r\n \r\n def _deprecated2_build_pdf_tree_structure(self, headers, doc):\r\n \"\"\"构建树形结构\"\"\"\r\n # 创建根节点\r\n root = SHTNode(\r\n node_id= 0 , # f\"{doc['id']}_root\",\r\n doc_id=doc['doc_id'],\r\n name=f\"Document: {doc['name']}\",\r\n granularity=0,\r\n context=doc['text'][:500], # 前500字符作为上下文\r\n summary=\"\",\r\n start_pos=0,\r\n end_pos=len(doc['text'])\r\n )\r\n \r\n if not headers:\r\n return root\r\n \r\n # 按文档中的出现顺序排序标题\r\n sorted_headers = sorted(headers, key=lambda x: doc['text'].find(x[0]))\r\n \r\n # 构建层次结构\r\n node_stack = [root]\r\n \r\n for i, (header_text, cluster) in enumerate(sorted_headers):\r\n # 确定层级\r\n granularity = self._determine_pdf_granularity(header_text)\r\n \r\n # 找到合适的父节点\r\n while len(node_stack) > 1 and node_stack[-1].granularity >= granularity:\r\n node_stack.pop()\r\n \r\n parent = node_stack[-1]\r\n \r\n # 计算文本位置\r\n start_pos = doc['text'].find(header_text)\r\n if start_pos == -1:\r\n start_pos = 0\r\n \r\n # 下一个标题的位置作为结束位置\r\n if i + 1 < len(sorted_headers):\r\n next_header = sorted_headers[i + 1][0]\r\n end_pos = doc['text'].find(next_header)\r\n if end_pos == -1:\r\n end_pos = len(doc['text'])\r\n else:\r\n end_pos = len(doc['text'])\r\n \r\n # 提取该节点的上下文\r\n context = doc['text'][start_pos:end_pos]\r\n if self.need_clean_chunk:\r\n context = self._clean_text(context)\r\n \r\n # 创建节点\r\n node = SHTNode(\r\n node_id= i+1 , # f\"{doc['id']}_node_{i}\",\r\n doc_id=doc['doc_id'],\r\n name=header_text.strip(),\r\n granularity=granularity,\r\n context=context,\r\n summary=\"\",\r\n start_pos=start_pos,\r\n end_pos=end_pos,\r\n parent=parent\r\n )\r\n parent.children.append(node)\r\n node_stack.append(node)\r\n # 合并全文context\r\n self._merge_full_context(root)\r\n return root\r\n\r\n\r\n def _build_tree_structure(self, headers, doc):\r\n \"\"\"构建树形结构\"\"\"\r\n # 创建根节点\r\n root = SHTNode(\r\n node_id= 0 , # f\"{doc['id']}_root\",\r\n doc_id=doc['doc_id'],\r\n name=f\"Document: {doc['name']}\",\r\n granularity=0,\r\n context=doc['text'], # 根节点-全文作为上下文。\r\n summary=\"\",\r\n start_pos=0,\r\n end_pos=len(doc['text'])\r\n )\r\n \r\n if not headers:\r\n return root\r\n \r\n # 构建标题位置映射，处理重复标题问题\r\n header_positions = self._find_all_header_positions(headers, doc['text'])\r\n \r\n # 按实际位置排序标题\r\n sorted_headers = sorted(header_positions, key=lambda x: x[2]) # 按position排序\r\n \r\n # 构建层次结构\r\n node_stack = [root]\r\n \r\n for i, (header_text, cluster, start_pos) in enumerate(sorted_headers):\r\n # 确定层级\r\n granularity = self._determine_granularity(header_text)\r\n \r\n # 找到合适的父节点\r\n while len(node_stack) > 1 and node_stack[-1].granularity >= granularity:\r\n node_stack.pop()\r\n \r\n parent = node_stack[-1]\r\n \r\n # 下一个标题的位置作为结束位置\r\n if i + 1 < len(sorted_headers):\r\n end_pos = sorted_headers[i + 1][2] # 使用下一个标题的实际位置\r\n else:\r\n end_pos = len(doc['text'])\r\n \r\n # 提取该节点的上下文\r\n context = doc['text'][start_pos:end_pos]\r\n if self.need_clean_chunk:\r\n context = self._clean_text(context)\r\n \r\n # 创建节点\r\n node = SHTNode(\r\n node_id= i+1 , # f\"{doc['id']}_node_{i}\", 根节点是0\r\n doc_id=doc['doc_id'],\r\n name=header_text.strip(),\r\n granularity=granularity,\r\n context=context,\r\n summary=\"\",\r\n start_pos=start_pos,\r\n end_pos=end_pos,\r\n parent=parent\r\n )\r\n parent.children.append(node)\r\n node_stack.append(node)\r\n # 合并全文context\r\n self._merge_full_context(root)\r\n return root\r\n\r\n def _find_all_header_positions(self, headers, text):\r\n \"\"\"\r\n 找到所有标题在文本中的实际位置，处理重复标题问题\r\n \r\n Args:\r\n headers: [(header_text, cluster), ...]\r\n text: 文档全文\r\n \r\n Returns:\r\n [(header_text, cluster, position), ...] 按出现顺序排序\r\n \"\"\"\r\n header_positions = []\r\n used_positions = set() # 记录已使用的位置，避免重复\r\n \r\n for header_text, cluster in headers:\r\n # 找到该标题的所有出现位置\r\n start = 0\r\n while True:\r\n pos = text.find(header_text, start)\r\n if pos == -1:\r\n break\r\n \r\n # 如果这个位置还没被使用过，则添加\r\n if pos not in used_positions:\r\n header_positions.append((header_text, cluster, pos))\r\n used_positions.add(pos)\r\n break # 找到第一个未使用的位置就跳出\r\n \r\n start = pos + 1 # 继续查找下一个位置\r\n \r\n return header_positions\r\n\r\n\r\n def _merge_full_context(self, node):\r\n \"\"\"递归合并本节点及所有子节点的context到full_context\"\"\"\r\n if not node.children:\r\n node.full_context = node.context\r\n else:\r\n merged = node.context + '\\n'\r\n for child in node.children:\r\n self._merge_full_context(child)\r\n merged += child.full_context + '\\n'\r\n node.full_context = merged\r\n\r\n def _determine_granularity(self, header_text):\r\n header_match = re.match(r'^(#{1,6})\\s+(.+)$', header_text )\r\n header_level = len(header_match.group(1)) # 获取#的数量\r\n return header_level\r\n\r\n def _determine_pdf_granularity(self, header_text):\r\n header = header_text.strip()\r\n # 匹配编号（如 1、1.1、1.1.1、1.1.1.1）\r\n m = re.match(r'^((\\d+)(\\.\\d+)*)', header)\r\n if m:\r\n # 层级数 = 点的数量 + 1\r\n level = m.group(1).count('.') + 1\r\n return level\r\n elif header.isupper() and len(header) < 50: # ABSTRACT 等全大写短标题\r\n return 1\r\n else:\r\n return 1 # 默认层级\r\n\r\n\r\n\r\n def _generate_summaries(self, node):\r\n \"\"\"为节点生成摘要\"\"\"\r\n if not node:\r\n return\r\n \r\n # 生成节点摘要\r\n if node.context:\r\n # 提取关键句子作为摘要\r\n sentences = self._extract_key_sentences(node.context, max_sentences=MAX_SENTENCES)\r\n ancestors_path = \" → \".join(node.get_ancestors())\r\n \r\n node.summary = f\"\"\"Title: {node.name}\r\nAncestors: {ancestors_path}\r\nKey Content: {sentences}\"\"\"\r\n \r\n # 递归处理子节点\r\n for child in node.children:\r\n self._generate_summaries(child)\r\n \r\n def _extract_key_sentences(self, text, max_sentences=MAX_SENTENCES):\r\n \"\"\"提取关键句子\"\"\"\r\n if not text:\r\n return \"\" \r\n summary = tfidf_summary(text, max_sentences)\r\n return summary\r\n \r\n\r\n\r\n def _find_node_by_id(self, root, node_id):\r\n \"\"\"根据ID找到节点\"\"\"\r\n if root.node_id == node_id:\r\n return root\r\n \r\n for child in root.children:\r\n result = self._find_node_by_id(child, node_id)\r\n if result:\r\n return result\r\n \r\n return None\r\n \r\n # 辅助模块\r\n def _llm_query(self, prompt, model=\"gpt-4.1\", max_retries=3):\r\n \"\"\"LLM调用封装\"\"\"\r\n for attempt in range(max_retries):\r\n try:\r\n response = client.chat.completions.create(\r\n model=model,\r\n messages=[{\"role\": \"user\", \"content\": prompt}],\r\n temperature=0,\r\n max_tokens=500\r\n )\r\n return response.choices[0].message.content.strip()\r\n except Exception as e:\r\n if attempt < max_retries - 1:\r\n print(f\" ⚠️ LLM调用失败，重试 {attempt + 1}/{max_retries}: {e}\")\r\n time.sleep(1)\r\n else:\r\n raise e\r\n \r\n # 主要工作流程\r\n def build_sht_for_all_docs(self, docs):\r\n \"\"\"为所有文档构建SHT\"\"\"\r\n print(\"\\n🏗️ 为所有文档构建语义层次树...\")\r\n \r\n for doc in docs:\r\n sht_root = self._build_sht(doc)\r\n self.sht_tables[doc['doc_id']] = sht_root\r\n \r\n \r\n print(f\"✓ 完成 {len(self.sht_tables)} 个文档的SHT构建\")\r\n\r\n def get_docs_id(self) -> list[int]:\r\n \"\"\"\r\n 获取表中所有文档的doc_id\r\n\r\n Returns:\r\n list[int]: 文档ID列表\r\n \"\"\"\r\n return list(self.docs_meta.keys())\r\n\r\n\r\n def _count_nodes(self, node):\r\n \"\"\"计算节点总数\"\"\"\r\n count = 1\r\n for child in node.children:\r\n count += self._count_nodes(child)\r\n return count\r\n \r\n def _get_max_depth(self, node):\r\n \"\"\"获取最大深度\"\"\"\r\n if not node.children:\r\n return node.granularity\r\n \r\n return max(self._get_max_depth(child) for child in node.children)\r\n\r\n"}}}} |