refactor(tree): extract _embed_l2_subtree to reduce embed_all complexity
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
+23
-9
@@ -481,15 +481,7 @@ class TreeIndex:
|
|||||||
if l1.embedding is None:
|
if l1.embedding is None:
|
||||||
l1.embedding = embed_fn(l1.summary)[0].astype(np.float32)
|
l1.embedding = embed_fn(l1.summary)[0].astype(np.float32)
|
||||||
for l2 in l1.children:
|
for l2 in l1.children:
|
||||||
if l2.embedding is None:
|
self._embed_l2_subtree(l2, embed_fn)
|
||||||
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.metadata.embed_model = model_name
|
self.metadata.embed_model = model_name
|
||||||
self.metadata.embed_dim = embed_dim
|
self.metadata.embed_dim = embed_dim
|
||||||
logger.info(
|
logger.info(
|
||||||
@@ -498,6 +490,28 @@ class TreeIndex:
|
|||||||
embed_dim=embed_dim,
|
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
|
||||||
|
|
||||||
# ------------------------------------------------------------------ #
|
# ------------------------------------------------------------------ #
|
||||||
# 嵌入矩阵提取
|
# 嵌入矩阵提取
|
||||||
# ------------------------------------------------------------------ #
|
# ------------------------------------------------------------------ #
|
||||||
|
|||||||
Reference in New Issue
Block a user