Create-knowledge-find-decode/scripts/embed_trees.py

"""把 live 分类树节点名向量化（火山 multimodal）→ scope_trees/trees_embeddings.npy。

行序与 trees_index.json 一一对应。须先跑 dump_trees.py，且能连火山。
用法：python scripts/embed_trees.py [env_file]
"""
from __future__ import annotations

import json
import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np

from core.embedding import ArkEmbedConfig, embed_text

OUT = Path("scope_trees")


def main(env_file: str = ".env") -> None:
    cfg = ArkEmbedConfig.from_env(env_file)
    index = json.loads((OUT / "trees_index.json").read_text(encoding="utf-8"))
    names = [it["name"] for it in index]
    n = len(names)
    embs: list[list[float] | None] = [None] * n

    def work(i: int) -> int:
        embs[i] = embed_text(names[i], cfg)
        return i

    done = 0
    with ThreadPoolExecutor(max_workers=8) as ex:
        for _ in ex.map(work, range(n)):
            done += 1
            if done % 100 == 0 or done == n:
                print(f"  embedded {done}/{n}")

    arr = np.asarray(embs, dtype=np.float32)
    if arr.shape != (n, cfg.dim):
        raise AssertionError(f"shape {arr.shape} != {(n, cfg.dim)}")
    np.save(OUT / "trees_embeddings.npy", arr)
    print("saved", arr.shape, "->", OUT / "trees_embeddings.npy")


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else ".env")