iomgaa
6bdb802f01
- 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
267 lines
10 KiB
Python
267 lines
10 KiB
Python
"""
|
||
端到端推理管线
|
||
==============
|
||
串联 预处理 → 检索 → 生成 的完整推理流程。
|
||
提供 ``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")
|
||
)
|