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
1415 lines
52 KiB
Markdown
1415 lines
52 KiB
Markdown
# 技术方案（TD）— Video-Tree-TRM
|
||
|
||
## 技术决策
|
||
|
||
- 单机本地执行，无服务端/数据库，所有数据 pickle/JSON 序列化到本地文件。
|
||
- 节点选择使用 **Cross-Attention**（学习 W_q/W_k/W_v/W_o 投影），替代简单 cosine 路由，更强表达力。
|
||
- L_level 推理模块使用 **MLP-based**（RMSNorm + SwiGLU），因操作对象为单向量 `[B, D]`，非序列，无需 self-attention。
|
||
- 三个可学习组件（CrossAttentionSelector, ReasoningModule, q_head）**跨层级共享权重**，与 TRM 原设计一致。
|
||
- 文本嵌入器（text_embed）**冻结不训练**，TreeIndex 中所有 embedding 为预计算静态值。
|
||
- 训练分两阶段：Phase 1 纯导航监督（单轮），Phase 2 加入 ACT halt（多轮）。
|
||
- MVP 优先文本模态（LongBench），视频模态（VideoMME）后续扩展。
|
||
- 配置管理：dataclass（无默认值，纯类型定义）+ YAML（全量配置）+ .env（敏感信息），优先级 CLI args > .env > YAML，三者统一归口到 dataclass。
|
||
|
||
---
|
||
|
||
## 目录
|
||
|
||
- [技术决策](#技术决策)
|
||
- [模块设计](#1-模块设计)
|
||
- [tree_index.py](#11-tree_indexpy--统一数据结构)
|
||
- [embeddings.py](#12-embeddingspy--嵌入服务)
|
||
- [llm_client.py](#13-llm_clientpy--llmvlm-客户端)
|
||
- [text_tree_builder.py](#14-text_tree_builderpy--文本树构建)
|
||
- [video_tree_builder.py](#15-video_tree_builderpy--视频树构建)
|
||
- [recursive_retriever.py](#16-recursive_retrieverpy--trm-递归检索器)
|
||
- [losses.py](#17-lossespy--损失函数)
|
||
- [answer_generator.py](#18-answer_generatorpy--答案生成)
|
||
- [pipeline.py](#19-pipelinepy--端到端管线)
|
||
- [config.py](#110-configpy--配置管理)
|
||
- [训练管线](#2-训练管线)
|
||
- [实验计划](#3-实验计划)
|
||
- [文件结构与依赖](#4-文件结构与依赖)
|
||
|
||
---
|
||
|
||
## 1. 模块设计
|
||
|
||
### 1.1 tree_index.py — 统一数据结构
|
||
|
||
**文件**: `video_tree_trm/tree_index.py`
|
||
**职责**: 定义三层树索引的节点类型、序列化/反序列化、嵌入矩阵提取。
|
||
|
||
> **延迟 Embedding 设计**：build 阶段所有节点 `embedding=None`，`IndexMeta.embed_model/embed_dim` 也为 None。首次检索前由 `Pipeline._embed_tree()` 调用 `embed_all()` 统一填充。
|
||
|
||
```python
|
||
@dataclass
|
||
class IndexMeta:
|
||
source_path: str # 原始文件路径
|
||
modality: str # "text" | "video"
|
||
embed_model: Optional[str] = None # build 时为 None，embed_all 后填充
|
||
embed_dim: Optional[int] = None # build 时为 None，embed_all 后填充
|
||
created_at: str # ISO 时间戳（自动生成）
|
||
|
||
@dataclass
|
||
class L3Node:
|
||
id: str
|
||
description: str # 视频=VLM帧描述, 文本=原始段落
|
||
embedding: Optional[ndarray] # [D]，build 时为 None，embed_all 后填充
|
||
raw_content: Optional[str] # 原始文本（文本模式）
|
||
frame_path: Optional[str] # 帧图像路径（视频模式）
|
||
timestamp: Optional[float] # 帧时间戳（视频模式）
|
||
|
||
@dataclass
|
||
class L2Node:
|
||
id: str
|
||
description: str # 1-2句片段描述
|
||
embedding: Optional[ndarray] # [D]，build 时为 None
|
||
time_range: Optional[Tuple[float, float]]
|
||
children: List[L3Node]
|
||
|
||
@dataclass
|
||
class L1Node:
|
||
id: str
|
||
summary: str # 2-3句聚合摘要
|
||
embedding: Optional[ndarray] # [D]，build 时为 None
|
||
time_range: Optional[Tuple[float, float]]
|
||
children: List[L2Node]
|
||
|
||
@dataclass
|
||
class TreeIndex:
|
||
metadata: IndexMeta
|
||
roots: List[L1Node]
|
||
```
|
||
|
||
**关键方法**:
|
||
|
||
```python
|
||
class TreeIndex:
|
||
@property
|
||
def is_embedded(self) -> bool:
|
||
"""所有 L1/L2/L3 节点的 embedding 均非 None 时返回 True"""
|
||
|
||
def embed_all(
|
||
self,
|
||
embed_fn: Callable[[Union[str, List[str]]], ndarray],
|
||
model_name: str,
|
||
embed_dim: int,
|
||
) -> None:
|
||
"""批量 embed 所有节点，更新 metadata。
|
||
- L3 按 L2 分组批量调用（减少 API 调用次数）
|
||
- L1/L2 各单独 embed
|
||
- 仅对 embedding=None 的节点执行（支持增量更新）"""
|
||
|
||
def l1_embeddings(self) -> ndarray:
|
||
"""返回所有 L1 嵌入矩阵 [N1, D]（需先 embed_all）"""
|
||
|
||
def l2_embeddings_of(self, l1_idx: int) -> ndarray:
|
||
"""返回指定 L1 下所有 L2 子节点嵌入 [N2, D]"""
|
||
|
||
def l3_embeddings_of(self, l1_idx: int, l2_idx: int) -> ndarray:
|
||
"""返回指定 L2 下所有 L3 子节点嵌入 [N3, D]"""
|
||
|
||
def get_node(self, l1: int, l2: int, l3: int) -> L3Node:
|
||
"""按路径索引获取 L3 节点"""
|
||
|
||
# JSON 序列化（主格式，无 embedding，适合缓存和人工查看）
|
||
def save_json(self, path: str) -> None:
|
||
"""序列化为 JSON 文件（不含 embedding 向量）"""
|
||
|
||
@classmethod
|
||
def load_json(cls, path: str) -> "TreeIndex":
|
||
"""从 JSON 文件加载（embedding=None，需后续 embed_all）"""
|
||
|
||
# pickle 序列化（向后兼容，含 embedding）
|
||
def save(self, path: str) -> None:
|
||
"""pickle 序列化（含 embedding 向量）"""
|
||
|
||
@classmethod
|
||
def load(cls, path: str) -> "TreeIndex":
|
||
"""从 pickle 文件加载"""
|
||
```
|
||
|
||
**依赖**: numpy, pickle, json（标准库）
|
||
|
||
---
|
||
|
||
### 1.2 embeddings.py — 嵌入服务
|
||
|
||
**文件**: `video_tree_trm/embeddings.py`
|
||
**职责**: 封装文本嵌入器，支持本地 sentence-transformers 和远程 OpenAI 兼容 API 双后端，冻结不训练。
|
||
|
||
```python
|
||
class EmbeddingModel:
|
||
"""文本嵌入器封装（冻结），支持本地/远程双后端。"""
|
||
|
||
def __init__(self, config: EmbedConfig):
|
||
"""
|
||
根据 config.backend 初始化:
|
||
- "local": 加载 sentence-transformers 模型，冻结参数
|
||
- "remote": 初始化 OpenAI 兼容 API 客户端
|
||
"""
|
||
|
||
@property
|
||
def dim(self) -> int:
|
||
"""嵌入维度 D"""
|
||
|
||
def embed(self, texts: Union[str, List[str]]) -> ndarray:
|
||
"""
|
||
文本 → 嵌入向量 (L2 归一化)
|
||
Args:
|
||
texts: 单条或批量文本
|
||
Returns:
|
||
[N, D] ndarray（单条时 N=1，每行 L2 范数为 1.0）
|
||
"""
|
||
|
||
def embed_tensor(self, texts: Union[str, List[str]]) -> Tensor:
|
||
"""同 embed()，返回 torch.Tensor [N, D]（float32）"""
|
||
|
||
# 内部方法
|
||
def _embed_local(self, texts: List[str]) -> ndarray:
|
||
"""sentence-transformers 本地推理，torch.no_grad() + normalize_embeddings=True"""
|
||
|
||
def _embed_remote(self, texts: List[str]) -> ndarray:
|
||
"""OpenAI 兼容 API: client.embeddings.create() → 提取向量 → L2 归一化"""
|
||
```
|
||
|
||
**远程模式示例** (GPUStack qwen3-embedding):
|
||
```python
|
||
# .env
|
||
EMBED_API_KEY=sk-xxx
|
||
EMBED_API_URL=http://gpu-host:8080/v1
|
||
|
||
# config/default.yaml
|
||
embed:
|
||
backend: "remote"
|
||
model_name: "qwen3-embedding-4b"
|
||
embed_dim: 2048
|
||
device: "cpu" # 远程模式不使用
|
||
api_key: "" # 从 .env 覆盖
|
||
api_url: "" # 从 .env 覆盖
|
||
```
|
||
|
||
**依赖**: sentence-transformers（本地模式）, openai SDK（远程模式）, torch, numpy
|
||
|
||
---
|
||
|
||
### 1.3 llm_client.py — LLM/VLM 客户端
|
||
|
||
**文件**: `video_tree_trm/llm_client.py`
|
||
**职责**: 统一封装 LLM（纯文本）和 VLM（多模态）API 调用，仅支持 OpenAI-compatible 单一接口，通过配置 `api_url` + `model` 切换服务商（Qwen DashScope、OpenAI、本地推理服务等）。
|
||
|
||
```python
|
||
class LLMClient:
|
||
"""OpenAI-compatible LLM/VLM 统一客户端。"""
|
||
|
||
def __init__(self, config: Union[LLMConfig, VLMConfig]) -> None:
|
||
"""
|
||
初始化客户端。
|
||
Args:
|
||
config: LLMConfig 或 VLMConfig，含 api_key、api_url、model 等参数。
|
||
Raises:
|
||
ValueError: api_key 或 api_url 为空时抛出。
|
||
实现:
|
||
openai.OpenAI(api_key=config.api_key, base_url=config.api_url)
|
||
"""
|
||
|
||
def chat(self, prompt: str, max_tokens: Optional[int] = None) -> str:
|
||
"""纯文本单轮对话，返回生成文本。max_tokens=None 时取 config.max_tokens。"""
|
||
|
||
def chat_with_images(
|
||
self, prompt: str, images: List[str], max_tokens: Optional[int] = None
|
||
) -> str:
|
||
"""
|
||
多模态对话（VLM）。
|
||
Args:
|
||
prompt: 文本指令。
|
||
images: 图像列表，每项为本地文件路径或已编码的 data URI 字符串。
|
||
Returns:
|
||
生成文本。
|
||
实现:
|
||
本地路径 → _encode_image() 转 base64 → _build_messages() 拼 content → API 调用。
|
||
"""
|
||
|
||
def batch_chat(self, prompts: List[str], max_tokens: Optional[int] = None) -> List[str]:
|
||
"""ThreadPoolExecutor(max_workers=8) 并发调用 chat()，保序返回。"""
|
||
|
||
# ── 私有辅助 ──
|
||
|
||
def _encode_image(self, path_or_b64: str) -> str:
|
||
"""
|
||
本地路径 → "data:image/{jpeg|png};base64,<data>"。
|
||
已含 "base64," 标记则直接返回（不重复编码）。
|
||
"""
|
||
|
||
def _build_messages(
|
||
self, prompt: str, images: Optional[List[str]] = None
|
||
) -> List[Dict]:
|
||
"""
|
||
无图像: [{"role": "user", "content": prompt}]
|
||
有图像: content 为列表，image_url 项在前，text 项在后。
|
||
"""
|
||
```
|
||
|
||
**消息结构（OpenAI-compatible）**:
|
||
|
||
```python
|
||
# 纯文本
|
||
[{"role": "user", "content": prompt}]
|
||
|
||
# 多模态（图在 text 之前）
|
||
[{"role": "user", "content": [
|
||
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}},
|
||
{"type": "text", "text": prompt}
|
||
]}]
|
||
```
|
||
|
||
**502/503 自动重试机制**:
|
||
|
||
```python
|
||
# 模块级辅助函数
|
||
def _call_with_retry(fn, label: str):
|
||
"""对 API 调用执行指数退避重试（仅重试 502/503）。
|
||
- 最多重试 20 次（约等待 20+ 分钟）
|
||
- 首次等待 60s，每次翻倍，上限 300s
|
||
"""
|
||
wait = 60
|
||
for attempt in range(1, 21):
|
||
try:
|
||
return fn()
|
||
except openai.InternalServerError as exc:
|
||
if exc.status_code not in {502, 503}:
|
||
raise
|
||
time.sleep(wait)
|
||
wait = min(wait * 2, 300)
|
||
```
|
||
|
||
**代理绕过**:
|
||
|
||
```python
|
||
self._client = openai.OpenAI(
|
||
api_key=config.api_key,
|
||
base_url=config.api_url,
|
||
http_client=httpx.Client(proxy=None), # 显式绕过系统代理，直连内网地址
|
||
)
|
||
```
|
||
|
||
**与 TD 原设计的差异**:
|
||
|
||
| 项目 | 原设计 | 实际实现 |
|
||
|------|--------|---------|
|
||
| 构造器签名 | `(backend, api_key, model, **kwargs)` | `(config: Union[LLMConfig, VLMConfig])` |
|
||
| 后端区分 | `"qwen" \| "openai" \| "ollama"` 分支 | 统一走 OpenAI-compatible，无后端分支 |
|
||
| `max_tokens` 默认值 | 函数参数硬编码 `= 256` | `= None`，None 时取 `config.max_tokens` |
|
||
| `batch_chat` 并发 | "并发或顺序" | `ThreadPoolExecutor(max_workers=8)` |
|
||
| **重试** | 无 | `_call_with_retry()` 502/503 指数退避重试（新增） |
|
||
| **代理** | 无 | `httpx.Client(proxy=None)` 绕过系统代理（新增） |
|
||
|
||
**依赖**: openai SDK（≥1.0）, httpx, python-dotenv（间接，via config）
|
||
|
||
---
|
||
|
||
### 1.4 text_tree_builder.py — 文本树构建
|
||
|
||
> **状态**: ✅ 已实现 | 测试: `tests/unit/test_text_tree_builder.py`（43 个用例全部通过）
|
||
|
||
**文件**: `video_tree_trm/text_tree_builder.py`
|
||
**职责**: 长文本 → TreeIndex，实现 L2 轴心构建策略。
|
||
|
||
> **注意**: 构造器不接受 `EmbeddingModel`（延迟 embedding 设计）。所有节点 `embedding=None`，由 `Pipeline.embed_all()` 在检索前统一填充。
|
||
|
||
#### 公共接口
|
||
|
||
```python
|
||
class TextTreeBuilder:
|
||
"""文本模态树构建器"""
|
||
|
||
def __init__(self, llm: LLMClient, config: TreeConfig):
|
||
self.llm = llm # LLM 客户端
|
||
self.config = config # TreeConfig（关键字段: max_paragraphs_per_l2）
|
||
# ⚠ 无 embed_model：embedding 延迟到 Pipeline.embed_all()
|
||
|
||
def build(self, text: str, source_path: str) -> TreeIndex:
|
||
"""
|
||
完整构建流程:
|
||
Phase 1: _segment_text() → sections: List[List[str]]
|
||
Phase 2: llm.batch_chat() → 所有 L2 摘要并发生成（一次调用）
|
||
Phase 3: 逐层组装 L3 → L2 → L1 节点
|
||
Phase 5: 组装 TreeIndex + 写日志
|
||
"""
|
||
```
|
||
|
||
#### 内部方法
|
||
|
||
```python
|
||
def _segment_text(self, text: str) -> List[List[str]]:
|
||
"""调度: _detect_toc() → True → _segment_with_regex()
|
||
→ False → _segment_with_llm()
|
||
返回: sections[i] = [para_1, para_2, ...]
|
||
外层 = L1 章节，内层 = 该章节下所有段落（扁平）
|
||
注: build() 负责按 max_paragraphs_per_l2 等长分块为 L2 组"""
|
||
|
||
def _detect_toc(self, text: str) -> bool:
|
||
"""检测文本是否含 # 或 ## 开头的 Markdown 标题行（正则: ^#{1,2}\s+\S）"""
|
||
|
||
def _segment_with_regex(self, text: str) -> List[List[str]]:
|
||
"""正则解析 #/## 标题边界:
|
||
# → L1 切换（flush 当前 section）
|
||
## → 段落分隔（flush 当前段落，标题文本作为一段）
|
||
空行 → 段落分隔
|
||
### 及以下 → 视为普通段落内容"""
|
||
|
||
def _segment_with_llm(self, text: str) -> List[List[str]]:
|
||
"""LLM 单次调用语义分段，返回只有一个外层元素的 list（整篇视为单 L1）
|
||
Prompt: '将以下文本分成若干语义段落...只返回 JSON 数组...'
|
||
解析: json.loads()，失败时通过 ensure() 抛 ValueError
|
||
支持 LLM 返回值被 markdown 代码块包裹（正则提取）"""
|
||
|
||
def _collect_paragraphs(self, text: str) -> List[str]:
|
||
"""保底策略: 按双换行符切分段落（_segment_with_regex 无结果时兜底）"""
|
||
|
||
def _build_l2(self, paragraphs: List[str], l2_id: str) -> L2Node:
|
||
"""段落组 → L2Node（不含 children，由 build() 填充）
|
||
LLM prompt: _L2_PROMPT.format(text="\n\n".join(paragraphs))
|
||
注: 实际由 build() 统一调 batch_chat() 批量处理，此方法仅供单独调用"""
|
||
|
||
def _build_l3_from_paragraphs(
|
||
self, paragraphs: List[str], l1_i: int, l2_j: int
|
||
) -> List[L3Node]:
|
||
"""段落列表批量嵌入 → L3Node 列表（不调用 LLM）
|
||
description = raw_content = 原始段落文本
|
||
embed.embed(paragraphs) 一次调用获取全部向量 [N, D]
|
||
节点 ID: f"l1_{l1_i}_l2_{l2_j}_l3_{k}" """
|
||
|
||
def _build_l1(self, l2_children: List[L2Node], l1_id: str) -> L1Node:
|
||
"""聚合所有 L2 描述 → L1Node（含 children）
|
||
LLM prompt: _L1_PROMPT.format(l2_descriptions="1. ...\n2. ...")
|
||
节点 ID: f"l1_{l1_i}" """
|
||
```
|
||
|
||
#### 关键实现决策
|
||
|
||
| 决策 | 说明 |
|
||
|------|------|
|
||
| **批量 LLM** | `build()` 收集所有 L2 段落组后调用 `llm.batch_chat()` 一次并发生成所有 L2 摘要，避免串行延迟 |
|
||
| **L2 等长分块** | 当段落数超过 `max_paragraphs_per_l2` 时，`_chunk(lst, size)` 等长切块（固定步长无重叠），同一 `#` 章节下可产生多个 L2 |
|
||
| **L3 无 LLM** | L3 直接复用原始段落文本（`description == raw_content`），`embed.embed()` 批量调用 |
|
||
| **节点 ID** | `l1_{i}` / `l1_{i}_l2_{j}` / `l1_{i}_l2_{j}_l3_{k}`，全局唯一 |
|
||
| **Prompt 常量** | `_L2_PROMPT`, `_L1_PROMPT`, `_SEG_PROMPT` 定义在模块顶层 |
|
||
|
||
#### Prompt 设计
|
||
|
||
```python
|
||
_L2_PROMPT = "用1-2句话描述以下段落的核心内容，与同级小节形成区分:\n\n{text}"
|
||
_L1_PROMPT = "用2-3句话总结以下小节的核心内容:\n\n{l2_descriptions}"
|
||
_SEG_PROMPT = "将以下文本分成若干语义段落，每段为完整语义单元。\n只返回 JSON 数组，格式: [\"段落1\", ...]，不要其他内容。\n文本:\n\n{text}"
|
||
```
|
||
|
||
**依赖**: `tree_index`, `embeddings`, `llm_client`, `utils.logger_system`
|
||
|
||
---
|
||
|
||
### 1.5 video_tree_builder.py — 视频树构建
|
||
|
||
**文件**: `video_tree_trm/video_tree_builder.py`
|
||
**职责**: 长视频 → TreeIndex，实现 L2 轴心构建策略 + VLM 帧描述。
|
||
**状态**: ✅ 已实现
|
||
|
||
> **注意**: 构造器不接受 `EmbeddingModel`（延迟 embedding 设计）。所有节点 `embedding=None`，由 `Pipeline.embed_all()` 在检索前统一填充。支持本地文件路径和 YouTube URL 两种输入。
|
||
|
||
```python
|
||
class VideoTreeBuilder:
|
||
"""视频模态树构建器"""
|
||
|
||
def __init__(self, vlm: LLMClient, config: TreeConfig):
|
||
self.vlm = vlm
|
||
self.config = config
|
||
# ⚠ 无 embed_model：embedding 延迟到 Pipeline.embed_all()
|
||
|
||
def build(self, video_path: str) -> TreeIndex:
|
||
"""
|
||
支持本地文件路径或 YouTube URL。
|
||
URL 模式: _resolve_stream() 获取 CDN 直链，_get_video_duration() 获取时长
|
||
本地模式: OpenCV 直接读取
|
||
|
||
完整构建流程（ThreadPoolExecutor 异步事件循环）:
|
||
Step 0: URL 处理（若 video_path 为 URL）
|
||
Step 1: _segment_video → L1 时间区间列表
|
||
Step 2: 收集全局 L2 任务列表，预计算每个 L1 的 L2 数量
|
||
Step 3: ThreadPoolExecutor(max_workers=concurrency) 一次性提交所有 L2 任务（非阻塞）
|
||
Step 4: 事件循环（cfwait FIRST_COMPLETED）：
|
||
L2 完成 → 立即提交 L3 任务（_build_l3_task）
|
||
L3 完成 → 检查 L1 就绪 → 立即提交 L1 任务
|
||
L1 完成 → 收集结果
|
||
Step 5: 有序重建 l1_nodes，组装 TreeIndex（全部 embedding=None）
|
||
|
||
主线程单线程操作 l1_l2_buckets，无竞争，无需 Lock。
|
||
"""
|
||
|
||
# ── URL 流式辅助方法（静态方法）──
|
||
|
||
@staticmethod
|
||
def _is_url(path_or_url: str) -> bool:
|
||
"""判断输入是否为 http/https URL"""
|
||
|
||
@staticmethod
|
||
def _source_stem(video_path: str) -> str:
|
||
"""提取短标识符用于帧缓存目录：
|
||
YouTube URL → 视频 ID（v= 参数）；本地文件 → stem（限 64 字符）"""
|
||
|
||
@staticmethod
|
||
def _resolve_stream(url: str) -> str:
|
||
"""yt-dlp -g 获取 YouTube CDN 直链（不下载，仅元数据）"""
|
||
|
||
@staticmethod
|
||
def _get_video_duration(url: str) -> float:
|
||
"""yt-dlp --dump-json 获取视频时长（cv2 在 HTTP 流上 CAP_PROP_FRAME_COUNT 不可靠）"""
|
||
|
||
# ── 内部方法 ──
|
||
|
||
def _segment_video(
|
||
self, video_path: str, duration_hint: Optional[float] = None
|
||
) -> List[Tuple[float, float]]:
|
||
"""固定步长切分 L1 区间。
|
||
本地文件: cv2 读取总时长；HTTP 流: 使用 duration_hint"""
|
||
|
||
def _get_l2_clips(self, l1_range: Tuple[float, float]) -> List[Tuple[float, float]]:
|
||
"""将 L1 区间按 l2_clip_duration 等分为 L2 clips"""
|
||
|
||
def _extract_frames(
|
||
self, video_path: str, time_range: Tuple[float, float], fps: float,
|
||
source_id: Optional[str] = None
|
||
) -> List[Tuple[str, float]]:
|
||
"""L3 专用：按 fps 密集提取帧到 {cache_dir}/frames/{source_id}/
|
||
已存在的帧文件自动跳过（缓存复用）"""
|
||
|
||
def _build_l2_video(
|
||
self, video_path: str, clip_range: Tuple[float, float], l2_id: str,
|
||
source_id: Optional[str] = None
|
||
) -> L2Node:
|
||
"""稀疏均匀 seek l2_representative_frames 帧 → VLM 描述（1-2句）
|
||
embedding=None"""
|
||
|
||
def _build_l3_video(
|
||
self, frames: List[Tuple[str, float]], l2_description: str, l1_i: int, l2_j: int
|
||
) -> List[L3Node]:
|
||
"""注入 L2 上下文的 VLM 批量帧描述
|
||
- 主路径: 一次 VLM 调用，要求返回 JSON 数组
|
||
- 降级路径: JSON 解析失败时逐帧调用
|
||
所有节点 embedding=None"""
|
||
|
||
def _build_l3_task(
|
||
self,
|
||
video_path: str,
|
||
l2_node: L2Node,
|
||
clip_range: Tuple[float, float],
|
||
source_id: str,
|
||
l1_i: int,
|
||
l2_j: int,
|
||
) -> L2Node:
|
||
"""L3 线程任务单元：提取帧 + _build_l3_video，返回已填充 children 的 L2Node。
|
||
由事件循环在 L2 完成后自动提交（非阻塞），线程安全（内部独立持有 VideoCapture）。"""
|
||
|
||
def _call_vlm_batch(self, prompt, frame_paths, n, l1_i, l2_j) -> List[str]:
|
||
"""批量 VLM 调用 + JSON 解析失败时降级逐帧"""
|
||
|
||
def _parse_json_descriptions(self, raw: str, expected_n: int) -> Optional[List[str]]:
|
||
"""从 VLM 输出解析 JSON 数组，长度不匹配返回 None"""
|
||
|
||
def _build_l1_video(
|
||
self, l2_children: List[L2Node], l1_id: str, l1_range: Tuple[float, float]
|
||
) -> L1Node:
|
||
"""拼接 L2 描述 → vlm.chat()（纯文本）生成 2-3 句摘要；embedding=None"""
|
||
```
|
||
|
||
**关键配置参数**（`config.tree`）:
|
||
|
||
| 参数 | 默认值 | 说明 |
|
||
|------|--------|------|
|
||
| `l1_segment_duration` | 600.0s | L1 切分步长 |
|
||
| `l2_clip_duration` | 60.0s | L2 clip 时长 |
|
||
| `l3_fps` | 1.0 | L3 帧提取速率（帧/秒） |
|
||
| `l2_representative_frames` | 10 | L2 稀疏代表帧数 |
|
||
| `cache_dir` | `cache/trees` | 帧图像持久化目录 |
|
||
| `concurrency` | 16 | 视频内 L2/L3 任务并发数（ThreadPoolExecutor max_workers） |
|
||
|
||
**依赖**: tree_index, embeddings, llm_client, opencv-python（帧提取）
|
||
|
||
---
|
||
|
||
### 1.6 recursive_retriever.py — TRM 递归检索器
|
||
|
||
**文件**: `video_tree_trm/recursive_retriever.py`
|
||
**职责**: 核心可训练模型。Cross-Attention 节点选择 + MLP 推理 + ACT halt。
|
||
**状态**: ✅ 已实现 | 测试: `tests/unit/test_recursive_retriever.py`（17 个用例全部通过）
|
||
|
||
#### 1.6.1 CrossAttentionSelector
|
||
|
||
```python
|
||
class CrossAttentionSelector(nn.Module):
|
||
"""跨层节点选择器（共享，用于 L1/L2/L3 三个阶段）"""
|
||
|
||
def __init__(self, embed_dim: int, num_heads: int):
|
||
self.W_q = Linear(embed_dim, embed_dim)
|
||
self.W_k = Linear(embed_dim, embed_dim)
|
||
self.W_v = Linear(embed_dim, embed_dim)
|
||
self.W_o = Linear(embed_dim, embed_dim)
|
||
self.num_heads = num_heads
|
||
self.head_dim = embed_dim // num_heads
|
||
self.scale = self.head_dim ** -0.5
|
||
|
||
def forward(
|
||
self, state: Tensor, candidates: Tensor
|
||
) -> Tuple[Tensor, Tensor, Tensor]:
|
||
"""
|
||
Args:
|
||
state: [B, D] — 当前 q+z 融合状态
|
||
candidates: [B, N, D] — 该层候选节点嵌入
|
||
Returns:
|
||
selected_info: [B, D] — attention 加权节点信息（可微）
|
||
attn_weights: [B, N] — 归一化注意力权重（用于 nav loss）
|
||
selected_idx: [B] — argmax 节点索引（用于路径记录）
|
||
"""
|
||
B, N, D = candidates.shape
|
||
|
||
Q = self.W_q(state).unsqueeze(1) # [B, 1, D]
|
||
K = self.W_k(candidates) # [B, N, D]
|
||
V = self.W_v(candidates) # [B, N, D]
|
||
|
||
# reshape → multi-head
|
||
Q = Q.view(B, 1, self.num_heads, self.head_dim).transpose(1, 2) # [B, H, 1, d]
|
||
K = K.view(B, N, self.num_heads, self.head_dim).transpose(1, 2) # [B, H, N, d]
|
||
V = V.view(B, N, self.num_heads, self.head_dim).transpose(1, 2) # [B, H, N, d]
|
||
|
||
# scaled dot-product attention
|
||
attn_out = F.scaled_dot_product_attention(Q, K, V) # [B, H, 1, d]
|
||
attn_out = attn_out.transpose(1, 2).reshape(B, 1, D)
|
||
selected_info = self.W_o(attn_out).squeeze(1) # [B, D]
|
||
|
||
# 注意力权重（对 head 维度平均，用于 loss 和可解释性）
|
||
raw_scores = (Q @ K.transpose(-2, -1)) * self.scale # [B, H, 1, N]
|
||
attn_weights = raw_scores.mean(dim=1).squeeze(1).softmax(dim=-1) # [B, N]
|
||
selected_idx = attn_weights.argmax(dim=-1) # [B]
|
||
|
||
return selected_info, attn_weights, selected_idx
|
||
```
|
||
|
||
#### 1.6.2 ReasoningModule（L-level）
|
||
|
||
```python
|
||
class ReasoningBlock(nn.Module):
|
||
"""单层 MLP 推理块"""
|
||
def __init__(self, dim: int, expansion: float):
|
||
self.norm = RMSNorm(dim)
|
||
self.ffn = SwiGLU(dim, int(dim * expansion))
|
||
|
||
def forward(self, x: Tensor) -> Tensor:
|
||
return self.norm(x + self.ffn(x)) # [B, D] → [B, D]
|
||
|
||
|
||
class ReasoningModule(nn.Module):
|
||
"""L-level 推理模块（多层 MLP，共享权重跨层级）"""
|
||
def __init__(self, dim: int, L_layers: int, expansion: float):
|
||
self.blocks = ModuleList([ReasoningBlock(dim, expansion) for _ in range(L_layers)])
|
||
|
||
def forward(self, z: Tensor, injection: Tensor) -> Tensor:
|
||
"""
|
||
Args:
|
||
z: [B, D] — 当前潜在状态
|
||
injection: [B, D] — 注入信息 (selected_info + q)
|
||
Returns:
|
||
z_new: [B, D]
|
||
"""
|
||
h = z + injection
|
||
for block in self.blocks:
|
||
h = block(h)
|
||
return h
|
||
```
|
||
|
||
#### 1.6.3 RecursiveRetriever
|
||
|
||
```python
|
||
class RecursiveRetriever(nn.Module):
|
||
"""TRM 递归检索器主模型"""
|
||
|
||
def __init__(self, config: RetrieverConfig):
|
||
self.selector = CrossAttentionSelector(config.embed_dim, config.num_heads)
|
||
self.L_level = ReasoningModule(config.embed_dim, config.L_layers, config.ffn_expansion)
|
||
self.q_head = Linear(config.embed_dim, 1) # ACT halt head
|
||
self.L_cycles = config.L_cycles
|
||
self.max_rounds = config.max_rounds
|
||
|
||
# q_head 初始化为倾向"继续"（bias = -5 → sigmoid ≈ 0）
|
||
with torch.no_grad():
|
||
self.q_head.bias.fill_(-5.0)
|
||
|
||
def forward(
|
||
self, q: Tensor, tree: TreeIndex, return_internals: bool = False
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
训练/推理统一入口。
|
||
Args:
|
||
q: [B, D] — 查询嵌入（来自冻结 text_embed）
|
||
tree: TreeIndex — 预构建树索引
|
||
return_internals: 是否返回中间状态（用于 loss 计算）
|
||
Returns:
|
||
{
|
||
"paths": List[Tuple[int, int, int]],
|
||
"num_rounds": int,
|
||
"z_final": Tensor [B, D],
|
||
# return_internals=True 时额外返回:
|
||
"attn_weights_per_step": List[Tensor], # 每步 [B, N]
|
||
"halt_logits": List[Tensor], # 每轮 [B, 1]
|
||
}
|
||
"""
|
||
z = q.clone() # [B, D]
|
||
paths = []
|
||
attn_weights_all = []
|
||
halt_logits_all = []
|
||
|
||
for round_idx in range(self.max_rounds):
|
||
path, z, step_attns = self._traverse_one_path(q, z, tree)
|
||
paths.append(path)
|
||
attn_weights_all.extend(step_attns)
|
||
|
||
halt_logit = self.q_head(z) # [B, 1]
|
||
halt_logits_all.append(halt_logit)
|
||
|
||
if not self.training and halt_logit.item() > 0 and round_idx > 0:
|
||
break
|
||
|
||
result = {
|
||
"paths": paths,
|
||
"num_rounds": len(paths),
|
||
"z_final": z,
|
||
}
|
||
if return_internals:
|
||
result["attn_weights_per_step"] = attn_weights_all
|
||
result["halt_logits"] = halt_logits_all
|
||
return result
|
||
|
||
def _traverse_one_path(
|
||
self, q: Tensor, z: Tensor, tree: TreeIndex
|
||
) -> Tuple[Tuple[int, int, int], Tensor, List[Tensor]]:
|
||
"""单次 L1 → L2 → L3 遍历"""
|
||
step_attns = []
|
||
|
||
# Phase 1: L1
|
||
M_L1 = torch.tensor(tree.l1_embeddings(), device=q.device) # [N1, D]
|
||
k1, z, attn_w = self._select_and_reason(q, z, M_L1.unsqueeze(0))
|
||
step_attns.append(attn_w)
|
||
|
||
# Phase 2: L2 (k1 的子节点)
|
||
M_L2 = torch.tensor(tree.l2_embeddings_of(k1), device=q.device) # [N2, D]
|
||
k2, z, attn_w = self._select_and_reason(q, z, M_L2.unsqueeze(0))
|
||
step_attns.append(attn_w)
|
||
|
||
# Phase 3: L3 (k2 的子节点)
|
||
M_L3 = torch.tensor(tree.l3_embeddings_of(k1, k2), device=q.device) # [N3, D]
|
||
k3, z, attn_w = self._select_and_reason(q, z, M_L3.unsqueeze(0))
|
||
step_attns.append(attn_w)
|
||
|
||
return (k1, k2, k3), z, step_attns
|
||
|
||
def _select_and_reason(
|
||
self, q: Tensor, z: Tensor, M: Tensor
|
||
) -> Tuple[int, Tensor, Tensor]:
|
||
"""
|
||
单层: Cross-Attention 选择 + L_cycles 内循环推理
|
||
Args:
|
||
q: [B, D], z: [B, D], M: [B, N, D]
|
||
Returns:
|
||
k_star: int, z_new: [B, D], attn_weights: [B, N]
|
||
"""
|
||
state = q + z
|
||
selected_info, attn_weights, selected_idx = self.selector(state, M)
|
||
|
||
z = z + selected_info
|
||
|
||
for _ in range(self.L_cycles):
|
||
z = self.L_level(z, selected_info + q)
|
||
|
||
return selected_idx.item(), z, attn_weights
|
||
```
|
||
|
||
**训练 vs 推理行为差异**:
|
||
|
||
| 行为 | 训练 | 推理 |
|
||
|------|------|------|
|
||
| 多轮循环 | 固定跑 max_rounds 轮 | halt_logit > 0 提前停止 |
|
||
| 梯度 | 全部可微 | no_grad |
|
||
| 返回值 | 含 attn_weights + halt_logits | 仅 paths + z_final |
|
||
|
||
**依赖**: torch, tree_index
|
||
|
||
---
|
||
|
||
### 1.7 losses.py — 损失函数
|
||
|
||
**文件**: `video_tree_trm/losses.py`
|
||
**职责**: 导航损失（cross-entropy）+ ACT halt 损失（Q-learning）。
|
||
**状态**: ✅ 已实现 | 测试: `tests/unit/test_losses.py`（13 个用例全部通过）
|
||
|
||
```python
|
||
class NavigationLoss(nn.Module):
|
||
"""导航监督损失：推动 attn_weights 指向正确节点"""
|
||
|
||
def forward(
|
||
self, attn_weights_list: List[Tensor], gt_path: Tuple[int, int, int]
|
||
) -> Tensor:
|
||
"""
|
||
Args:
|
||
attn_weights_list: [attn_l1, attn_l2, attn_l3]，每个 [B, N]
|
||
gt_path: (gt_l1_idx, gt_l2_idx, gt_l3_idx)
|
||
Returns:
|
||
loss: scalar
|
||
"""
|
||
loss = 0
|
||
for attn_w, gt_idx in zip(attn_weights_list, gt_path):
|
||
target = torch.tensor([gt_idx], device=attn_w.device)
|
||
log_probs = attn_w.log() # [B, N]
|
||
loss += F.nll_loss(log_probs, target) # cross-entropy
|
||
return loss / 3 # 三层平均
|
||
|
||
|
||
class ACTLoss(nn.Module):
|
||
"""ACT halt Q-learning 损失"""
|
||
|
||
def __init__(self, lambda_step: float = 0.1, gamma: float = 0.9):
|
||
self.lambda_step = lambda_step
|
||
self.gamma = gamma
|
||
|
||
def forward(
|
||
self,
|
||
halt_logits: List[Tensor], # 每轮 [B, 1]
|
||
answer_qualities: List[float], # 每轮累积的答案质量 (0~1)
|
||
) -> Tensor:
|
||
"""
|
||
Q-learning target:
|
||
若在第 t 轮停止 → Q_halt = quality_t
|
||
若继续 → Q_continue = γ * max(Q_{t+1}) - λ
|
||
"""
|
||
loss = 0
|
||
n = len(halt_logits)
|
||
for t in range(n):
|
||
halt_q = answer_qualities[t]
|
||
if t < n - 1:
|
||
continue_q = self.gamma * answer_qualities[t + 1] - self.lambda_step
|
||
else:
|
||
continue_q = halt_q - self.lambda_step # 最后一轮，继续无意义
|
||
|
||
# 目标: halt_logit > 0 当 halt_q > continue_q
|
||
target = 1.0 if halt_q >= continue_q else 0.0
|
||
pred = torch.sigmoid(halt_logits[t])
|
||
loss += F.binary_cross_entropy(pred, torch.tensor([[target]], device=pred.device))
|
||
|
||
return loss / n
|
||
```
|
||
|
||
**依赖**: torch
|
||
|
||
---
|
||
|
||
### 1.8 answer_generator.py — 答案生成
|
||
|
||
**文件**: `video_tree_trm/answer_generator.py`
|
||
**职责**: 根据检索结果组装 context，调用 LLM/VLM 生成最终答案。
|
||
**状态**: ✅ 已实现 | 测试: `tests/unit/test_answer_generator.py`（10 个用例全部通过）
|
||
|
||
```python
|
||
@dataclass
|
||
class RetrievalResult:
|
||
"""检索器输出的结构化结果"""
|
||
query: str
|
||
paths: List[Tuple[int, int, int]]
|
||
num_rounds: int
|
||
|
||
class AnswerGenerator:
|
||
def __init__(self, llm: LLMClient, vlm: LLMClient):
|
||
self.llm = llm
|
||
self.vlm = vlm
|
||
|
||
def generate(self, query: str, result: RetrievalResult, tree: TreeIndex) -> str:
|
||
"""
|
||
根据模态分发:
|
||
文本 → LLM(query, raw_text_chunks)
|
||
视频 → VLM(query, frame_images + captions)
|
||
"""
|
||
nodes = [tree.get_node(*path) for path in result.paths]
|
||
|
||
if tree.metadata.modality == "text":
|
||
context = "\n---\n".join(n.raw_content for n in nodes if n.raw_content)
|
||
return self.llm.chat(
|
||
f"根据以下上下文回答问题。\n\n上下文:\n{context}\n\n问题: {query}"
|
||
)
|
||
else:
|
||
frames = [n.frame_path for n in nodes if n.frame_path]
|
||
captions = [n.description for n in nodes]
|
||
caption_text = "\n".join(f"- {c}" for c in captions)
|
||
return self.vlm.chat_with_images(
|
||
f"根据以下关键帧回答问题。\n帧描述:\n{caption_text}\n\n问题: {query}",
|
||
images=frames,
|
||
)
|
||
```
|
||
|
||
**依赖**: tree_index, llm_client
|
||
|
||
---
|
||
|
||
### 1.9 pipeline.py — 端到端管线
|
||
|
||
**文件**: `video_tree_trm/pipeline.py`
|
||
**职责**: 串联 预处理 → 检索 → 生成 的完整推理流程。
|
||
**状态**: ✅ 已实现 | 测试: `tests/unit/test_pipeline.py`（10 个用例全部通过）
|
||
|
||
```python
|
||
class Pipeline:
|
||
"""端到端推理管线"""
|
||
|
||
def __init__(self, config: Config):
|
||
self.embed_model = EmbeddingModel(config.embed)
|
||
self.llm = LLMClient(config.llm)
|
||
self.vlm = LLMClient(config.vlm)
|
||
self.retriever = RecursiveRetriever(config.retriever)
|
||
# 可选加载检查点（checkpoint=null 时跳过）
|
||
if config.retriever.checkpoint:
|
||
state_dict = torch.load(config.retriever.checkpoint, map_location="cpu")
|
||
self.retriever.load_state_dict(state_dict)
|
||
self.retriever.eval()
|
||
self.generator = AnswerGenerator(self.llm, self.vlm)
|
||
|
||
def build_index(self, source_path: str, modality: str) -> TreeIndex:
|
||
"""构建并缓存 TreeIndex（JSON 格式，无 embedding）。
|
||
|
||
缓存路径: {cache_dir}/{stem}_{modality}.json
|
||
- 缓存命中: 直接 load_json 返回（embedding=None）
|
||
- 缓存未命中: 调用 Builder 生成文字描述，save_json 持久化
|
||
|
||
⚠ 返回的 TreeIndex embedding 全为 None，
|
||
query() 时会自动调用 _embed_tree() 填充。
|
||
"""
|
||
if modality == "text":
|
||
builder = TextTreeBuilder(self.llm, self.config.tree)
|
||
with open(source_path, encoding="utf-8") as f:
|
||
tree = builder.build(f.read(), source_path)
|
||
else:
|
||
builder = VideoTreeBuilder(self.vlm, self.config.tree)
|
||
tree = builder.build(source_path)
|
||
tree.save_json(cache_path) # JSON 持久化，无 embedding
|
||
return tree
|
||
|
||
def _embed_tree(self, tree: TreeIndex, cache_path: Optional[str] = None) -> None:
|
||
"""对树所有节点执行 embedding（内存中），可选回写缓存。
|
||
L3 按 L2 分组批量处理，L1/L2 各单独处理。"""
|
||
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) # 回写（含 embedding 的 JSON，实际不存储向量）
|
||
|
||
def query(self, question: str, tree: TreeIndex) -> str:
|
||
"""问答: question → answer。
|
||
|
||
若 tree.is_embedded 为 False（JSON 加载后），先触发 _embed_tree()。
|
||
"""
|
||
# Phase 0: 按需触发 embed_all（JSON 缓存加载后 embedding=None）
|
||
if not tree.is_embedded:
|
||
self._embed_tree(tree, cache_path=None)
|
||
|
||
# Phase 1: 嵌入查询
|
||
q = self.embed_model.embed_tensor(question) # [1, D]
|
||
|
||
# Phase 2: 递归检索
|
||
with torch.no_grad():
|
||
result = self.retriever(q, tree)
|
||
|
||
# Phase 3: 生成答案
|
||
return self.generator.generate(question, result["paths"], tree)
|
||
```
|
||
|
||
**与原设计的关键差异**:
|
||
|
||
| 项目 | 原设计 | 实际实现 |
|
||
|------|--------|---------|
|
||
| Builder 构造器 | `TextTreeBuilder(embed_model, llm, config)` | `TextTreeBuilder(llm, config)` |
|
||
| 缓存格式 | pickle（含 embedding） | JSON（无 embedding），首次 query 时内存 embed |
|
||
| `build_index` 返回 | 含 embedding 的 TreeIndex | `embedding=None` 的 TreeIndex |
|
||
| `query` 额外逻辑 | 直接检索 | 先检查 `is_embedded`，按需调用 `_embed_tree()` |
|
||
| 新增方法 | — | `_embed_tree()` |
|
||
|
||
**依赖**: 所有其他模块
|
||
|
||
---
|
||
|
||
### 1.10 config.py — 配置管理
|
||
|
||
**文件**: `video_tree_trm/config.py`
|
||
**职责**: 所有超参数的 dataclass 类型定义（无默认值）+ 多源加载。
|
||
|
||
#### 设计原则
|
||
|
||
- **Dataclass 无默认值**: 纯类型定义 + 结构化访问，YAML 必须写全，漏写即报错。
|
||
- **三层优先级**: `CLI args > .env > YAML`，高优先级覆盖低优先级。
|
||
- **统一归口**: 无论来源，最终构造唯一 `Config` dataclass 对象，代码只与 dataclass 交互。
|
||
- **敏感信息隔离**: `api_key` 等敏感字段只写在 `.env` 中，不进 YAML 和代码。
|
||
|
||
#### 加载流程
|
||
|
||
```
|
||
Step 1: 读取 YAML → base dict（全量非敏感配置）
|
||
Step 2: 读取 .env → 覆盖 dict 中对应字段（api_key 等敏感信息）
|
||
Step 3: 解析 CLI args → 最终覆盖 dict 中对应字段
|
||
Step 4: dict → Config dataclass（校验完整性，缺字段直接报错）
|
||
```
|
||
|
||
#### Dataclass 定义
|
||
|
||
```python
|
||
@dataclass
|
||
class TreeConfig:
|
||
# 文本模式
|
||
max_paragraphs_per_l2: int # 每个 L2 节点包含的最大段落数
|
||
# 视频模式
|
||
l1_segment_duration: float # L1 段时长（秒）
|
||
l2_clip_duration: float # L2 clip 时长（秒）
|
||
l3_fps: float # L3 帧提取频率
|
||
l2_representative_frames: int # L2 VLM 描述用的代表帧数
|
||
# 通用
|
||
cache_dir: str # TreeIndex 缓存目录
|
||
|
||
@dataclass
|
||
class EmbedConfig:
|
||
model_name: str # 嵌入模型名称
|
||
embed_dim: int # 嵌入维度 D
|
||
device: str # "cuda" | "cpu"
|
||
|
||
@dataclass
|
||
class LLMConfig:
|
||
backend: str # "qwen" | "openai" | "ollama"
|
||
api_key: str # 从 .env 加载，不写入 YAML
|
||
model: str # 模型名称
|
||
api_url: str # API 端点 URL
|
||
max_tokens: int # 最大生成 token 数
|
||
temperature: float # 采样温度
|
||
|
||
@dataclass
|
||
class VLMConfig:
|
||
backend: str # "qwen" | "openai" | "ollama"
|
||
api_key: str # 从 .env 加载，不写入 YAML
|
||
model: str # 模型名称
|
||
api_url: str # API 端点 URL
|
||
max_tokens: int # 最大生成 token 数
|
||
temperature: float # 采样温度
|
||
|
||
@dataclass
|
||
class RetrieverConfig:
|
||
embed_dim: int # 嵌入维度（须与 EmbedConfig.embed_dim 一致）
|
||
num_heads: int # Cross-Attention 头数
|
||
L_layers: int # ReasoningModule 层数
|
||
L_cycles: int # 每级推理迭代次数
|
||
max_rounds: int # ACT 最大遍历轮次
|
||
ffn_expansion: float # SwiGLU 扩展比
|
||
checkpoint: Optional[str] # 训练好的模型权重路径（推理时必填）
|
||
|
||
@dataclass
|
||
class TrainConfig:
|
||
lr: float # 学习率
|
||
weight_decay: float # 权重衰减
|
||
batch_size: int # 批大小
|
||
max_epochs_phase1: int # Phase 1 导航训练轮数
|
||
max_epochs_phase2: int # Phase 2 ACT 训练轮数
|
||
nav_loss_weight: float # 导航损失权重
|
||
act_loss_weight: float # ACT 损失权重
|
||
act_lambda_step: float # ACT 步数惩罚系数
|
||
act_gamma: float # ACT 折扣因子
|
||
eval_interval: int # 每 N epoch 评估一次
|
||
save_dir: str # 模型权重保存目录
|
||
dataset: str # "longbench" | "narrativeqa" | "videomme"
|
||
dataset_path: str # 数据集路径
|
||
|
||
@dataclass
|
||
class Config:
|
||
tree: TreeConfig
|
||
embed: EmbedConfig
|
||
llm: LLMConfig
|
||
vlm: VLMConfig
|
||
retriever: RetrieverConfig
|
||
train: TrainConfig
|
||
|
||
@classmethod
|
||
def load(cls, yaml_path: str, cli_args: Optional[dict] = None) -> "Config":
|
||
"""
|
||
三层合并加载:
|
||
1. 读取 YAML → base dict
|
||
2. 读取 .env → 覆盖 api_key 等敏感字段
|
||
3. cli_args → 最终覆盖
|
||
4. dict → Config（缺字段报 TypeError）
|
||
"""
|
||
...
|
||
```
|
||
|
||
#### 文件分工
|
||
|
||
| 文件 | 内容 | 提交到 Git |
|
||
|------|------|-----------|
|
||
| `config/default.yaml` | 全量非敏感配置（必须写全所有字段） | 是 |
|
||
| `.env` | 敏感信息（api_key 等） | 否 |
|
||
| `.env.example` | `.env` 模板（值留空） | 是 |
|
||
|
||
#### YAML 示例 (`config/default.yaml`)
|
||
|
||
```yaml
|
||
tree:
|
||
max_paragraphs_per_l2: 5
|
||
l1_segment_duration: 600.0
|
||
l2_clip_duration: 20.0
|
||
l3_fps: 1.0
|
||
l2_representative_frames: 3
|
||
cache_dir: "cache/trees"
|
||
|
||
embed:
|
||
model_name: "BAAI/bge-base-zh-v1.5"
|
||
embed_dim: 768
|
||
device: "cuda"
|
||
|
||
llm:
|
||
backend: "qwen"
|
||
model: "qwen-plus"
|
||
api_url: "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"
|
||
max_tokens: 256
|
||
temperature: 0.1
|
||
# api_key: 从 .env 加载，此处不写
|
||
|
||
vlm:
|
||
backend: "qwen"
|
||
model: "qwen-vl-plus"
|
||
api_url: "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"
|
||
max_tokens: 256
|
||
temperature: 0.1
|
||
# api_key: 从 .env 加载，此处不写
|
||
|
||
retriever:
|
||
embed_dim: 768
|
||
num_heads: 4
|
||
L_layers: 2
|
||
L_cycles: 4
|
||
max_rounds: 5
|
||
ffn_expansion: 2.0
|
||
checkpoint: null
|
||
|
||
train:
|
||
lr: 1.0e-4
|
||
weight_decay: 1.0e-5
|
||
batch_size: 1
|
||
max_epochs_phase1: 30
|
||
max_epochs_phase2: 20
|
||
nav_loss_weight: 1.0
|
||
act_loss_weight: 0.1
|
||
act_lambda_step: 0.1
|
||
act_gamma: 0.9
|
||
eval_interval: 5
|
||
save_dir: "checkpoints"
|
||
dataset: "longbench"
|
||
dataset_path: "data/longbench"
|
||
```
|
||
|
||
#### .env 示例
|
||
|
||
```bash
|
||
# .env — 敏感信息，不提交到 Git
|
||
LLM_API_KEY=sk-xxx
|
||
VLM_API_KEY=sk-xxx
|
||
```
|
||
|
||
**依赖**: dataclasses, yaml, python-dotenv
|
||
|
||
---
|
||
|
||
## 2. 训练管线
|
||
|
||
**文件**: `train.py`（项目根目录）
|
||
**状态**: ✅ 已实现 | 测试: `tests/unit/test_train.py`（13 个用例全部通过）
|
||
|
||
### 2.1 数据准备
|
||
|
||
```python
|
||
def prepare_training_data(config: Config) -> List[Dict]:
|
||
"""
|
||
离线预处理:
|
||
1. 加载 QA 数据集（LongBench / NarrativeQA）
|
||
2. 为每个文档构建 TreeIndex（缓存到 cache_dir）
|
||
3. 推导每个 QA 对的 ground truth 路径
|
||
Returns:
|
||
[{"query": str, "tree": TreeIndex, "gt_path": (l1, l2, l3), "answer": str}, ...]
|
||
"""
|
||
```
|
||
|
||
### 2.2 Ground Truth 路径推导
|
||
|
||
```python
|
||
def find_gt_path_text(tree: TreeIndex, answer: str) -> Optional[Tuple[int, int, int]]:
|
||
"""
|
||
文本模式: 找到与答案文本重叠度最高的 L3 节点
|
||
评分: F1(L3.raw_content, answer) — token 级别
|
||
返回: (l1_idx, l2_idx, l3_idx) 或 None
|
||
"""
|
||
best_score, best_path = 0, None
|
||
for i, l1 in enumerate(tree.roots):
|
||
for j, l2 in enumerate(l1.children):
|
||
for k, l3 in enumerate(l2.children):
|
||
score = token_f1(l3.raw_content, answer)
|
||
if score > best_score:
|
||
best_score = score
|
||
best_path = (i, j, k)
|
||
return best_path
|
||
|
||
|
||
def find_gt_path_video(tree: TreeIndex, timestamp: float) -> Optional[Tuple[int, int, int]]:
|
||
"""
|
||
视频模式: 找到最接近目标时间戳的 L3 帧
|
||
"""
|
||
for i, l1 in enumerate(tree.roots):
|
||
if l1.time_range[0] <= timestamp <= l1.time_range[1]:
|
||
for j, l2 in enumerate(l1.children):
|
||
if l2.time_range[0] <= timestamp <= l2.time_range[1]:
|
||
k = min(range(len(l2.children)),
|
||
key=lambda k: abs(l2.children[k].timestamp - timestamp))
|
||
return (i, j, k)
|
||
return None
|
||
```
|
||
|
||
### 2.3 两阶段训练策略
|
||
|
||
```
|
||
Phase 1: 导航训练（单轮, max_rounds=1）
|
||
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
|
||
目标: 训练 Selector + L_level 正确导航到目标节点
|
||
损失: NavigationLoss (cross-entropy on attn_weights)
|
||
可训练: CrossAttentionSelector, ReasoningModule
|
||
冻结: text_embed, q_head, TreeIndex embeddings
|
||
|
||
Phase 2: ACT 训练（多轮, max_rounds=5）
|
||
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
|
||
目标: 训练 q_head 判断何时停止检索
|
||
损失: NavigationLoss + λ * ACTLoss
|
||
可训练: 全部（Selector + L_level + q_head）
|
||
冻结: text_embed, TreeIndex embeddings
|
||
ACT reward: answer_quality (F1/EM) - λ_step * rounds
|
||
```
|
||
|
||
### 2.4 训练循环伪代码
|
||
|
||
```python
|
||
def train(config: Config):
|
||
# ── 初始化 ──
|
||
embed_model = EmbeddingModel(config.embed.model_name, config.embed.device)
|
||
retriever = RecursiveRetriever(config.retriever).to(config.embed.device)
|
||
nav_loss_fn = NavigationLoss()
|
||
act_loss_fn = ACTLoss(config.train.act_lambda_step, config.train.act_gamma)
|
||
|
||
dataset = prepare_training_data(config)
|
||
optimizer = AdamW(retriever.parameters(), lr=config.train.lr)
|
||
|
||
# ── Phase 1: 导航训练 ──
|
||
retriever.max_rounds = 1
|
||
for epoch in range(config.train.max_epochs_phase1):
|
||
for sample in dataset:
|
||
q = embed_model.embed_tensor(sample["query"]).to(device) # [1, D]
|
||
tree = sample["tree"]
|
||
gt_path = sample["gt_path"]
|
||
|
||
result = retriever(q, tree, return_internals=True)
|
||
# result["attn_weights_per_step"] = [attn_l1, attn_l2, attn_l3]
|
||
loss = nav_loss_fn(result["attn_weights_per_step"][:3], gt_path)
|
||
|
||
optimizer.zero_grad()
|
||
loss.backward()
|
||
optimizer.step()
|
||
|
||
# ── Phase 2: ACT 训练 ──
|
||
retriever.max_rounds = config.retriever.max_rounds
|
||
llm = LLMClient(config.llm.backend, config.llm.api_key, config.llm.model)
|
||
generator = AnswerGenerator(llm, None)
|
||
|
||
for epoch in range(config.train.max_epochs_phase2):
|
||
for sample in dataset:
|
||
q = embed_model.embed_tensor(sample["query"]).to(device)
|
||
result = retriever(q, sample["tree"], return_internals=True)
|
||
|
||
# 每轮计算答案质量
|
||
qualities = []
|
||
for round_idx in range(result["num_rounds"]):
|
||
paths_so_far = result["paths"][:round_idx + 1]
|
||
nodes = [sample["tree"].get_node(*p) for p in paths_so_far]
|
||
context = "\n".join(n.raw_content for n in nodes if n.raw_content)
|
||
answer = llm.chat(f"上下文: {context}\n问题: {sample['query']}")
|
||
quality = token_f1(answer, sample["answer"])
|
||
qualities.append(quality)
|
||
|
||
# 导航 loss（仅第一轮）
|
||
loss_nav = nav_loss_fn(result["attn_weights_per_step"][:3], sample["gt_path"])
|
||
# ACT loss
|
||
loss_act = act_loss_fn(result["halt_logits"], qualities)
|
||
# 总损失
|
||
loss = loss_nav + config.train.act_loss_weight * loss_act
|
||
|
||
optimizer.zero_grad()
|
||
loss.backward()
|
||
optimizer.step()
|
||
```
|
||
|
||
---
|
||
|
||
## 3. 实验计划
|
||
|
||
### 3.1 数据集
|
||
|
||
| 数据集 | 模态 | 样本量 | 任务类型 | 优先级 |
|
||
|--------|------|--------|----------|--------|
|
||
| LongBench | 文本 | ~5K | 长文本 QA | P0 (首发) |
|
||
| NarrativeQA | 文本 | ~30K | 叙事理解 QA | P1 |
|
||
| VideoMME | 视频 | ~2K | 视频 QA (多选) | P2 |
|
||
|
||
### 3.2 评估指标
|
||
|
||
| 指标 | 适用 | 计算方式 |
|
||
|------|------|----------|
|
||
| EM (Exact Match) | 文本 QA | 标准化后精确匹配 |
|
||
| F1 | 文本 QA | token 级 precision/recall |
|
||
| Accuracy | 视频 QA | 选项匹配正确率 |
|
||
| Avg Rounds | 全部 | 平均检索轮次（衡量效率） |
|
||
| Nav Accuracy | 全部 | 第一轮 L1/L2/L3 各层命中率 |
|
||
|
||
### 3.3 Baselines
|
||
|
||
| 方法 | 描述 |
|
||
|------|------|
|
||
| BM25 + LLM | 传统稀疏检索 baseline |
|
||
| Dense Retrieval + LLM | BGE 向量检索 + rerank |
|
||
| PageIndex (原论文) | 无 TRM 的树状导航 (cosine routing, 无推理模块) |
|
||
| Tree-TRM (原论文) | 原始 tree_trm.py 实现 |
|
||
|
||
### 3.4 消融实验
|
||
|
||
| 实验 | 变量 | 目的 |
|
||
|------|------|------|
|
||
| A1 | Cross-Attention vs Cosine 路由 | 验证 CA 选择器的增益 |
|
||
| A2 | L_cycles = {1, 2, 4, 8} | 推理深度对准确率的影响 |
|
||
| A3 | L_layers = {1, 2, 4} | 推理模块复杂度 |
|
||
| A4 | max_rounds = {1, 3, 5} | 多轮检索的边际收益 |
|
||
| A5 | 有/无 ACT halt | ACT 机制对效率的贡献 |
|
||
| A6 | num_heads = {1, 4, 8} | 注意力头数的影响 |
|
||
|
||
---
|
||
|
||
## 4. 文件结构与依赖
|
||
|
||
### 4.1 目录树
|
||
|
||
```
|
||
Video-Tree-TRM/
|
||
├── video_tree_trm/ # 主包
|
||
│ ├── __init__.py
|
||
│ ├── config.py # §1.10 配置管理
|
||
│ ├── tree_index.py # §1.1 统一数据结构
|
||
│ ├── embeddings.py # §1.2 嵌入服务
|
||
│ ├── llm_client.py # §1.3 LLM/VLM 客户端
|
||
│ ├── text_tree_builder.py # §1.4 文本树构建
|
||
│ ├── video_tree_builder.py # §1.5 视频树构建
|
||
│ ├── recursive_retriever.py # §1.6 TRM 递归检索器
|
||
│ ├── losses.py # §1.7 损失函数
|
||
│ ├── answer_generator.py # §1.8 答案生成
|
||
│ └── pipeline.py # §1.9 端到端管线
|
||
├── utils/
|
||
│ ├── __init__.py
|
||
│ └── logger_system.py # 日志系统 (log_msg, ensure, log_exception)
|
||
├── config/
|
||
│ ├── default.yaml # 默认配置（通用）
|
||
│ └── videomme.yaml # VideoMME 实验专属配置（GPUStack Qwen3-VL）
|
||
├── tests/
|
||
│ ├── conftest.py # 全局 fixture（real_config, 代理修复）
|
||
│ ├── unit/
|
||
│ │ ├── test_config.py # ✅ 已实现
|
||
│ │ ├── test_embeddings.py # ✅ 已实现
|
||
│ │ ├── test_llm_client.py # ✅ 已实现
|
||
│ │ ├── test_tree_index.py # ✅ 已实现
|
||
│ │ ├── test_text_tree_builder.py # ✅ 已实现（43 用例）
|
||
│ │ ├── test_recursive_retriever.py # ✅ 已实现（17 用例）
|
||
│ │ ├── test_losses.py # ✅ 已实现（13 用例）
|
||
│ │ ├── test_answer_generator.py # ✅ 已实现（10 用例）
|
||
│ │ ├── test_pipeline.py # ✅ 已实现（10 用例）
|
||
│ │ └── test_train.py # ✅ 已实现（13 用例）
|
||
│ ├── integration/
|
||
│ └── outputs/ # Agent 测试 MD 输出
|
||
│ └── text_tree_builder/ # ✅ build_toc_*.md
|
||
├── data/ # 数据集（不提交）
|
||
├── cache/ # TreeIndex 缓存（不提交）
|
||
├── checkpoints/ # 模型权重（不提交）
|
||
├── logs/ # 运行日志（不提交）
|
||
├── train.py # §2 训练入口
|
||
├── main.py # 推理/演示入口
|
||
├── docs/
|
||
│ ├── architecture.md # 架构设计（理念层）
|
||
│ └── TD.md # 本文档（实现层）
|
||
├── .env # API 密钥（不提交）
|
||
├── .env.example # 环境变量模板
|
||
└── requirements.txt
|
||
```
|
||
|
||
### 4.2 模块依赖关系
|
||
|
||
```
|
||
config.py ← (所有模块都依赖)
|
||
|
||
embeddings.py ← text_tree_builder.py
|
||
← video_tree_builder.py
|
||
← pipeline.py
|
||
|
||
llm_client.py ← text_tree_builder.py
|
||
← video_tree_builder.py
|
||
← answer_generator.py
|
||
← pipeline.py
|
||
|
||
tree_index.py ← text_tree_builder.py
|
||
← video_tree_builder.py
|
||
← recursive_retriever.py
|
||
← answer_generator.py
|
||
← pipeline.py
|
||
|
||
recursive_retriever.py ← pipeline.py
|
||
← train.py
|
||
|
||
losses.py ← train.py
|
||
|
||
answer_generator.py ← pipeline.py
|
||
← train.py (Phase 2, 计算 answer quality)
|
||
```
|
||
|
||
### 4.3 Python 依赖
|
||
|
||
```
|
||
# 核心
|
||
torch>=2.0
|
||
sentence-transformers>=2.2
|
||
numpy
|
||
|
||
# LLM/VLM
|
||
openai>=1.0 # 兼容 Qwen/OpenAI/Ollama 接口
|
||
python-dotenv
|
||
|
||
# 视频处理
|
||
opencv-python
|
||
|
||
# 配置
|
||
pyyaml
|
||
|
||
# 测试
|
||
pytest
|
||
pytest-cov
|
||
|
||
# 代码质量
|
||
ruff
|
||
```
|