diff --git a/app/tree/index.py b/app/tree/index.py index 7e0e8e2..6db5e09 100644 --- a/app/tree/index.py +++ b/app/tree/index.py @@ -481,15 +481,7 @@ class TreeIndex: if l1.embedding is None: l1.embedding = embed_fn(l1.summary)[0].astype(np.float32) for l2 in l1.children: - if l2.embedding is None: - l2.embedding = embed_fn(l2.description)[0].astype(np.float32) - # L3 批量 embed - need_embed = [l3 for l3 in l2.children if l3.embedding is None] - if need_embed: - texts = [l3.description for l3 in need_embed] - embs = embed_fn(texts).astype(np.float32) # [N, D] - for l3, emb in zip(need_embed, embs, strict=True): - l3.embedding = emb + self._embed_l2_subtree(l2, embed_fn) self.metadata.embed_model = model_name self.metadata.embed_dim = embed_dim logger.info( @@ -498,6 +490,28 @@ class TreeIndex: embed_dim=embed_dim, ) + def _embed_l2_subtree( + self, + l2: L2Node, + embed_fn: Callable[[str | list[str]], np.ndarray], + ) -> None: + """对单个 L2 节点及其 L3 子节点执行 embedding(仅处理 embedding 为 None 的节点)。 + + 参数: + l2: 待 embed 的 L2 节点。 + embed_fn: EmbeddingModel.embed 方法,接受 str 或 List[str], + 返回 [N, D] ndarray。 + """ + if l2.embedding is None: + l2.embedding = embed_fn(l2.description)[0].astype(np.float32) + # L3 批量 embed + need_embed = [l3 for l3 in l2.children if l3.embedding is None] + if need_embed: + texts = [l3.description for l3 in need_embed] + embs = embed_fn(texts).astype(np.float32) # [N, D] + for l3, emb in zip(need_embed, embs, strict=True): + l3.embedding = emb + # ------------------------------------------------------------------ # # 嵌入矩阵提取 # ------------------------------------------------------------------ #