chore: track claude skills, tools, templates, reference code and research-wiki

- Add all claude skills (brainstorming, commit, debugging, TDD, etc.) - Add claude hooks (pre-commit-guard, post-edit-quality) - Add research templates (experiment plan, research brief, etc.) - Add claude tools (arxiv/semantic_scholar/openalex fetch, wiki, exa) - Add TRM4 reference implementation as algorithm fidelity baseline - Add research-wiki content (plans, index, graph, query_pack) - Update .gitignore to exclude .graphify_version runtime state
2026-07-06 20:59:03 -04:00
parent 0616d16956
commit 6bdb802f01
98 changed files with 19321 additions and 0 deletions
@@ -0,0 +1,266 @@
+"""
+端到端推理管线
+==============
+串联 预处理 → 检索 → 生成 的完整推理流程。
+提供 ``build_index()`` 和 ``query()`` 两个高层接口。
+
+使用方式::
+
+    from video_tree_trm.config import Config
+    from video_tree_trm.pipeline import Pipeline
+
+    cfg = Config.load("config/default.yaml")
+    pipeline = Pipeline(cfg)
+
+    # 构建（或从缓存加载）树索引
+    tree = pipeline.build_index("data/my_doc.txt", modality="text")
+
+    # 问答
+    answer = pipeline.query("文档的主要结论是什么？", tree)
+    print(answer)
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from utils.logger_system import ensure, log_msg
+from video_tree_trm.answer_generator import AnswerGenerator
+from video_tree_trm.config import Config
+from video_tree_trm.embeddings import EmbeddingModel
+from video_tree_trm.llm_client import LLMClient
+from video_tree_trm.recursive_retriever import RecursiveRetriever
+from video_tree_trm.text_tree_builder import TextTreeBuilder
+from video_tree_trm.tree_index import TreeIndex
+from video_tree_trm.video_tree_builder import VideoTreeBuilder
+
+
+class Pipeline:
+    """端到端推理管线（预处理 → 检索 → 生成）。
+
+    将所有子模块按配置串联，对外暴露两个接口：
+    - ``build_index()``: 从原始文件构建 TreeIndex，支持磁盘缓存。
+    - ``query()``: 对已有 TreeIndex 执行问答，返回生成答案字符串。
+
+    属性:
+        config: 全局配置对象。
+        embed_model: 文本嵌入模型（冻结）。
+        llm: 文本大语言模型客户端。
+        vlm: 视觉语言模型客户端。
+        retriever: TRM 递归检索器（eval 模式）。
+        generator: 答案生成器。
+    """
+
+    def __init__(self, config: Config) -> None:
+        """初始化端到端推理管线。
+
+        参数:
+            config: 通过 ``Config.load()`` 加载的全局配置对象。
+
+        实现细节:
+            - 若 ``config.retriever.checkpoint`` 非 None，加载预训练权重。
+            - 检索器始终切换到 eval 模式（关闭 Dropout 等训练行为）。
+        """
+        self.config = config
+
+        # Phase 1: 初始化各子模块（embed_model 懒加载，仅 query/embed 时触发）
+        self._embed_model: Optional[EmbeddingModel] = None
+        self.llm = LLMClient(config.llm)
+        self.vlm = LLMClient(config.vlm)
+        self.retriever = RecursiveRetriever(config.retriever)
+
+        # Phase 2: 可选加载检查点
+        if config.retriever.checkpoint:
+            ensure(
+                os.path.isfile(config.retriever.checkpoint),
+                f"检查点文件不存在: {config.retriever.checkpoint}",
+            )
+            state_dict = torch.load(config.retriever.checkpoint, map_location="cpu")
+            self.retriever.load_state_dict(state_dict)
+            log_msg(
+                "INFO",
+                "检索器权重已加载",
+                checkpoint=config.retriever.checkpoint,
+            )
+
+        self.retriever.eval()
+        self.generator = AnswerGenerator(self.llm, self.vlm)
+
+        log_msg(
+            "INFO",
+            "Pipeline 初始化完成",
+            modality_embed=config.embed.model_name,
+            has_checkpoint=bool(config.retriever.checkpoint),
+        )
+
+    @property
+    def embed_model(self) -> EmbeddingModel:
+        """懒加载 EmbeddingModel，仅在首次访问时初始化（index 阶段不触发）。"""
+        if self._embed_model is None:
+            log_msg("INFO", "懒加载 EmbeddingModel", model=self.config.embed.model_name)
+            self._embed_model = EmbeddingModel(self.config.embed)
+        return self._embed_model
+
+    def build_index(self, source_path: str, modality: str) -> TreeIndex:
+        """构建并缓存 TreeIndex（JSON 格式，含 embedding）。
+
+        参数:
+            source_path: 原始文件路径（文本文件或视频文件）。
+            modality:    模态类型，"text" 或 "video"。
+
+        返回:
+            构建完成的 TreeIndex 对象（已 embed）。
+
+        实现细节:
+            - 缓存路径: ``{cache_dir}/{stem}_{modality}.json``。
+            - 缓存命中时直接反序列化返回（自动恢复 embedding 若有）。
+            - 缓存未命中时调用 VLM 生成描述文本，执行 embedding，保存为 JSON。
+        """
+        ensure(
+            modality in ("text", "video"),
+            f"modality 须为 'text' 或 'video'，实际={modality}",
+        )
+
+        # Phase 1: 缓存路径计算
+        stem = Path(source_path).stem
+        cache_dir = Path(self.config.tree.cache_dir)
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        cache_path = str(cache_dir / f"{stem}_{modality}.json")
+
+        if os.path.isfile(cache_path):
+            log_msg("INFO", "缓存命中，直接加载 TreeIndex", cache_path=cache_path)
+            tree = TreeIndex.load_json(cache_path)
+            # 若缓存中已有 embedding，直接返回；否则按需 embed
+            if tree.is_embedded:
+                return tree
+            log_msg("INFO", "缓存中无 embedding，开始执行 embed_all")
+            self._embed_tree(tree, cache_path=cache_path)
+            return tree
+
+        # Phase 2: 构建树索引（纯 VLM 文字描述）
+        log_msg(
+            "INFO",
+            "缓存未命中，开始构建 TreeIndex",
+            source_path=source_path,
+            modality=modality,
+        )
+
+        if modality == "text":
+            with open(source_path, encoding="utf-8") as f:
+                text = f.read()
+            builder = TextTreeBuilder(self.llm, self.config.tree)
+            tree = builder.build(text, source_path)
+        else:
+            builder = VideoTreeBuilder(self.vlm, self.config.tree)
+            tree = builder.build(source_path)
+
+        # Phase 3: 执行 embedding 并保存（含 embedding）
+        self._embed_tree(tree, cache_path=cache_path)
+
+        return tree
+
+    def _embed_tree(self, tree: TreeIndex, cache_path: Optional[str] = None) -> None:
+        """对树的所有节点执行 embedding，可选回写缓存。
+
+        参数:
+            tree: 待 embed 的 TreeIndex（embedding=None 的节点）。
+            cache_path: 若非 None，embed 完成后回写到此路径（JSON 格式，含 embedding）。
+
+        实现细节:
+            调用 TreeIndex.embed_all，传入 EmbeddingModel.embed 作为 embed_fn。
+            embed_all 内部按 L2 分组批量处理 L3，减少 API 调用次数。
+            若 cache_path 非 None，保存时 include_embedding=True。
+        """
+        log_msg("INFO", "开始对树执行 embedding")
+        tree.embed_all(
+            embed_fn=self.embed_model.embed,
+            model_name=self.config.embed.model_name,
+            embed_dim=self.embed_model.dim,
+        )
+        if cache_path is not None:
+            tree.save_json(cache_path, include_embedding=True)
+            log_msg("INFO", "embed_all 完成，缓存已更新（含 embedding）", cache_path=cache_path)
+        else:
+            log_msg("INFO", "embed_all 完成（仅内存，未写磁盘）")
+
+    def _load_or_build_video_tree(self, video_path: str) -> TreeIndex:
+        """根据视频路径优先从缓存加载 TreeIndex，若无缓存则在线构建。
+
+        参数:
+            video_path: 视频文件路径或 youtube_id。
+
+        返回:
+            加载或构建完成的 TreeIndex 对象。
+        """
+        # 如果传入的是 youtube_id，尝试拼凑路径
+        if not os.path.isfile(video_path):
+            video_path_full = os.path.join("data/videomme/videos", f"{video_path}.mp4")
+            if os.path.isfile(video_path_full):
+                video_path = video_path_full
+
+        return self.build_index(video_path, modality="video")
+
+    def query(
+        self,
+        question: str,
+        tree: TreeIndex | str,
+        modality: Optional[str] = None,
+        cache_path: Optional[str] = None,
+    ) -> str:
+        """执行端到端问答。
+
+        参数:
+            question: 用户查询字符串。
+            tree:     TreeIndex 对象，或树 JSON 路径，或视频路径。
+            modality: 当 tree 为字符串且无法自动推断时，指定模态 ("text" 或 "video")。
+            cache_path: 若非 None，embed 完成后回写到此路径。
+
+        返回:
+            生成的答案字符串。
+        """
+        # Phase 0: 处理输入，确保得到 TreeIndex 对象
+        if isinstance(tree, str):
+            if tree.endswith(".json"):
+                log_msg("INFO", "直接从 JSON 路径加载 TreeIndex", path=tree)
+                tree_obj = TreeIndex.load_json(tree)
+                # 若 cache_path 未指定，使用 tree 的 JSON 路径
+                if cache_path is None:
+                    cache_path = tree
+            elif modality == "video" or tree.endswith(".mp4"):
+                log_msg("INFO", "根据视频路径获取 TreeIndex", path=tree)
+                tree_obj = self._load_or_build_video_tree(tree)
+            else:
+                # 默认为文本
+                log_msg("INFO", "根据文本路径获取 TreeIndex", path=tree)
+                tree_obj = self.build_index(tree, modality="text")
+        else:
+            tree_obj = tree
+
+        # Phase 1: 确保树已 embed
+        if not tree_obj.is_embedded:
+            log_msg("INFO", "树尚未 embed，触发 embed_all 并回写缓存", cache_path=cache_path)
+            self._embed_tree(tree_obj, cache_path=cache_path)
+
+        # Phase 2: 嵌入查询
+        q: torch.Tensor = self.embed_model.embed_tensor(question)  # [1, D]
+
+        # Phase 3: 递归检索
+        with torch.no_grad():
+            result = self.retriever(q, tree_obj)
+
+        log_msg(
+            "INFO",
+            "检索完成",
+            num_rounds=result["num_rounds"],
+            num_paths=len(result["paths"]),
+            question=question[:50],
+        )
+
+        # Phase 4: 生成答案
+        return self.generator.generate(
+            question, result["paths"], tree_obj, frame_hits=result.get("frame_hits")
+        )