iomgaa/Video-Tree-TRM5
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
- 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
This commit is contained in:
iomgaa
2026-07-06 20:59:03 -04:00
parent 0616d16956
commit 6bdb802f01
98 changed files with 19321 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
reference/.env.example
+11
View File
@@ -0,0 +1,11 @@
# .env — 敏感信息，不提交到 Git
LLM_API_KEY=
LLM_MODEL=
LLM_API_URL=
VLM_API_KEY=
VLM_MODEL=
VLM_API_URL=
EMBED_BACKEND=
EMBED_MODEL=
EMBED_API_KEY=
EMBED_API_URL=
reference/BUILD_TREE_PACKAGE_MANIFEST.txt
+41
View File
@@ -0,0 +1,41 @@
Video-Tree-TRM build-tree related code package
Source: /home/undergraduate/hyt/Video-Tree-TRM
Created: 2026-07-06T22:26:58+08:00
Scope: code/config/scripts/docs/tests needed for building TreeIndex; excludes built trees, videos, frames, logs, experiments, cache, and .env secrets.
Included files:
main.py
requirements.txt
.env.example
CLAUDE.md
config/default.yaml
config/videomme.yaml
video_tree_trm/__init__.py
video_tree_trm/config.py
video_tree_trm/tree_index.py
video_tree_trm/text_tree_builder.py
video_tree_trm/video_tree_builder.py
video_tree_trm/pipeline.py
video_tree_trm/llm_client.py
video_tree_trm/embeddings.py
utils/__init__.py
utils/logger_system.py
scripts/build_tree_single.py
scripts/build_trees_batch.py
scripts/build_trees_from_mp4.sh
scripts/build_trees_from_urls.sh
scripts/build_videomme_trees.sh
scripts/download_videos.sh
scripts/_download_meta.py
docs/architecture.md
docs/TD.md
docs/EXPERIMENT.md
tests/conftest.py
tests/unit/test_config.py
tests/unit/test_tree_index.py
tests/unit/test_text_tree_builder.py
tests/unit/test_video_tree_builder.py
tests/unit/test_pipeline.py
tests/unit/test_embeddings.py
tests/unit/test_llm_client.py
tests/unit/test_embedding_persistence.py
reference/CLAUDE.md
+230
View File
@@ -0,0 +1,230 @@
# CLAUDE.md
> [!URGENT]
> **研究性项目 (Research Project)**
> 1. 本项目为 MVP（最小可行性产品），严禁过度工程化。
> 2. 你的所有思考过程和回复必须使用 **简体中文**。
## 1. 项目元数据 (Metadata)
- **核心目标**: 构建结合TRM多层训话探索能力与PageIndex树状检索能力的新型Video RAG
- **项目类型**: MVP / 研究性项目
- **后端架构**: Python 3.11
- **版本管理**: Git
- **Conda 环境**: Video-Tree-TRM (Python 3.11)
## 2. 常用命令 (Commands)
### 2.1 Conda 环境管理
> [!CRITICAL]
> **所有 Python 相关命令必须在 Video-Tree-TRM 环境中执行**
> - 使用 `conda run -n Video-Tree-TRM <command>` 确保命令在正确环境中运行
> - 或在命令前显式添加 `source activate Video-Tree-TRM &&`
> - 如果需要使用llm可以依据 '.env'环境变量文件使用
```bash
# 激活项目环境（交互式 shell）
conda activate Video-Tree-TRM
# 推荐：使用 conda run 执行命令（自动使用正确环境）
conda run -n Video-Tree-TRM pip install xxx
conda run -n Video-Tree-TRM python -m pytest xxx #注意不能conda run -n Video-Tree-TRM pytest xxx，因为这样子pytest不会调用Video-Tree-TRM
conda run -n Video-Tree-TRM python xxx
# 或者：在命令前激活环境
source activate Video-Tree-TRM && xxx
```
### 2.2 代码质量检查
```bash
# 代码格式化
conda run -n Video-Tree-TRM ruff format xxx
# 代码检查并自动修复
conda run -n Video-Tree-TRM ruff check xxx --fix
```
## 3. 标准作业程序 (Standard Operating Procedure)
> **Agent 必须严格遵守以下生命周期执行任务：**
### Phase 1: 规划与设计 (Planning)
1. **查阅规格 (Read Specs)&讨论**: 在撰写计划前，**必须**仔细阅读 `docs/` 下对应的文档与`.report`下的项目整理架构，并使用GitNexus MCP来了解整个项目的最新情况。对于不理解的地方请与人类进行多轮讨论，确保理解人类的设计意图。
2. **计划 (Plan)**: 正式编码前，**必须**使用plan模式输出开发计划，内容必须严格包含：
- **1.1 摘要 (Summary)**: 1-2句话的简单总结。
- **1.2 审查点 (User Review Required)**: 明确列出整个计划中不清楚、需要用户审查和确认的部分。若无，请注明"无"。
- **1.3 拟议变更 (Proposed Changes)**:
- 以 **文件名 + 修改内容** 的形式列出。
- 修改内容必须精确到 **函数/方法级别 (Function-level)**。
- 明确标识 `[NEW]`, `[MODIFY]`, `[DELETE]`。
- **1.4 验证计划 (Verification Plan)**: 具体描述如何验证修改是否成功（如具体的测试命令、预期日志输出等）。
4. **等待 (Wait)**: **必须** 暂停并等待用户审核开发计划。用户批准后方可进入下一阶段。
### Phase 2: 执行与验证 (Execution & Verification)
1. **编码 (Coding)**: 审核通过后，开始编写代码。
2. **验证 (Verify)**:
- **环境检查**: 确保所有命令在 Video-Tree-TRM 环境中执行（使用 `conda run -n Video-Tree-TRM`）
- **运行验证命令**:
- *失败*: 回到编码阶段修复，直到通过。
- *成功*: 进入下一步。
## 4. 核心规则 (Rules)
### 4.1 代码开发规范 (Code Style)
- **类型系统**: 强制所有函数签名包含完整类型注解 (`Union`, `Dict`, `Optional` 等)。
- **文档**: 所有模块、类、方法必须包含 **中文 Docstring** (功能、参数、返回值、关键实现细节)。
- **MVP原则**:
- **必须** 必须在`tests/`目录下编写测试代码。
- **严禁** 使用默认参数掩盖仅需逻辑（必须显式传递关键参数）。
- **必需** 运行时检查：关键维度、设备一致性必须通过 assertion 或 if 验证。
- **代码组织**:
- 使用阶段化注释 (`# Phase 1`, `# Phase 2`) 组织复杂逻辑。
- 接口返回值需包含完整诊断信息（输出、损失、统计），使用条件标志控制。
- **命名与依赖**:
- 类名 `PascalCase`，变量描述性命名，私有变量前缀 `_`。
- 导入顺序：标准库 → 第三方库 → 项目内部。
- **日志与错误处理**: 使用 `utils/logger_system.py` 的 `log_msg()`, `log_json()`, `ensure()`, `log_exception()`
- 禁用 `print()`，`log_msg("ERROR")` 不自动抛出异常，输出到 `logs/system.log` + `logs/metrics.json`
- **功能修改**:
- **必须** 不考虑向后兼容，直接修改原文件。代码简洁性优先。
### 4.2 配置管理规范
- **优先级**: CLI args > `.env` > YAML，三者统一归口到 dataclass
- **文件**: `config/default.yaml`（全量非敏感配置，必须写全）, `.env`（敏感信息，不提交）, `.env.example`（模板）
### 4.3 测试组织规范
- **目录**: `tests/{unit,integration,e2e}/test_*.py`，最低覆盖率 80%
- **运行**: `conda run -n Video-Tree-TRM pytest tests/unit/ --cov=utils --cov-report=term-missing`
#### Agent 测试输出规范
> **除了使用pytest进行单元测试，还必须构建一个完善的main.py文件将所有过程串联起来，并必须将完整执行过程保存为 Markdown 文件，供 Claude Code 智能分析，以消除格式输出正确但是逻辑错误的误差或者实际输出质量低低问题。**
| 要素 | 规范 |
|------|------|
| **输出位置** | `tests/outputs/<test_module>/<test_name>_<timestamp>.md` |
| **触发时机** | 所有涉及 Agent 执行的测试 |
| **内容要求** | 任务描述、每步 Agent 输入/输出/推理过程、工具调用、最终结果 |
| **格式要求** | 结构化 Markdown（标题、代码块、列表），人类可读 |
| **分析方式** | Claude Code 读取 MD 文件，评估推理质量、任务完成度、代码正确性 |
**示例结构**:
```markdown
# Agent 测试: <test_name>
## 任务: <task>
## Step 1: <AgentName>
- 输入: ...
- 输出: ...
- 推理: ...
## Step 2: ...
## 最终结果: ...
```
**pytest 集成**: 使用 fixture 或工具类自动保存，测试结束后输出文件路径。
## 5. 上下文获取与迷途指南 (Context & Navigation)
| 需求 | 文档路径 | 说明 |
|------|----------|------|
| 项目目标与背景 | `README.md` | 核心业务逻辑与项目定性 |
| 架构与模块设计 | `.report/CODEMAPS/{architecture,backend,data}.md` | 整体架构、分层设计、模块依赖 |
| 该项目最主要的参考项目 | `Reference/Tree-TRM`| 特定模块的详细设计 |
| 其他参考项目 | `Reference/PageIndex_ Next-Generation Vectorless, Reasoning-based RAG.md`和 `Reference/Tree-TRM`| |
## 6. 输出规范
### 6.1 语言要求
- 所有输出语言: **中文**
### 6.2 信息密度原则
- **优先使用**:
- 简洁文本描述
- 伪代码（而非完整代码）
- 表格（对比、配置、参数说明）
- 流程图（Mermaid）
- 项目符号列表
- **避免使用**:
- 大段完整代码（信息密度低，可读性差）
- 冗长的自然语言解释
- **核心原则**: 用最少的字符传递最多的信息
`
<!-- gitnexus:start -->
# GitNexus — Code Intelligence
This project is indexed by GitNexus as **Video-Tree-TRM** (860 symbols, 2247 relationships, 70 execution flows). Use the GitNexus MCP tools to understand code, assess impact, and navigate safely.
> If any GitNexus tool warns the index is stale, run `npx gitnexus analyze` in terminal first.
## Always Do
- **MUST run impact analysis before editing any symbol.** Before modifying a function, class, or method, run `gitnexus_impact({target: "symbolName", direction: "upstream"})` and report the blast radius (direct callers, affected processes, risk level) to the user.
- **MUST run `gitnexus_detect_changes()` before committing** to verify your changes only affect expected symbols and execution flows.
- **MUST warn the user** if impact analysis returns HIGH or CRITICAL risk before proceeding with edits.
- When exploring unfamiliar code, use `gitnexus_query({query: "concept"})` to find execution flows instead of grepping. It returns process-grouped results ranked by relevance.
- When you need full context on a specific symbol — callers, callees, which execution flows it participates in — use `gitnexus_context({name: "symbolName"})`.
## When Debugging
1. `gitnexus_query({query: "<error or symptom>"})` — find execution flows related to the issue
2. `gitnexus_context({name: "<suspect function>"})` — see all callers, callees, and process participation
3. `READ gitnexus://repo/Video-Tree-TRM/process/{processName}` — trace the full execution flow step by step
4. For regressions: `gitnexus_detect_changes({scope: "compare", base_ref: "main"})` — see what your branch changed
## When Refactoring
- **Renaming**: MUST use `gitnexus_rename({symbol_name: "old", new_name: "new", dry_run: true})` first. Review the preview — graph edits are safe, text_search edits need manual review. Then run with `dry_run: false`.
- **Extracting/Splitting**: MUST run `gitnexus_context({name: "target"})` to see all incoming/outgoing refs, then `gitnexus_impact({target: "target", direction: "upstream"})` to find all external callers before moving code.
- After any refactor: run `gitnexus_detect_changes({scope: "all"})` to verify only expected files changed.
## Never Do
- NEVER edit a function, class, or method without first running `gitnexus_impact` on it.
- NEVER ignore HIGH or CRITICAL risk warnings from impact analysis.
- NEVER rename symbols with find-and-replace — use `gitnexus_rename` which understands the call graph.
- NEVER commit changes without running `gitnexus_detect_changes()` to check affected scope.
## Tools Quick Reference
| Tool | When to use | Command |
|------|-------------|---------|
| `query` | Find code by concept | `gitnexus_query({query: "auth validation"})` |
| `context` | 360-degree view of one symbol | `gitnexus_context({name: "validateUser"})` |
| `impact` | Blast radius before editing | `gitnexus_impact({target: "X", direction: "upstream"})` |
| `detect_changes` | Pre-commit scope check | `gitnexus_detect_changes({scope: "staged"})` |
| `rename` | Safe multi-file rename | `gitnexus_rename({symbol_name: "old", new_name: "new", dry_run: true})` |
| `cypher` | Custom graph queries | `gitnexus_cypher({query: "MATCH ..."})` |
## Impact Risk Levels
| Depth | Meaning | Action |
|-------|---------|--------|
| d=1 | WILL BREAK — direct callers/importers | MUST update these |
| d=2 | LIKELY AFFECTED — indirect deps | Should test |
| d=3 | MAY NEED TESTING — transitive | Test if critical path |
## Resources
| Resource | Use for |
|----------|---------|
| `gitnexus://repo/Video-Tree-TRM/context` | Codebase overview, check index freshness |
| `gitnexus://repo/Video-Tree-TRM/clusters` | All functional areas |
| `gitnexus://repo/Video-Tree-TRM/processes` | All execution flows |
| `gitnexus://repo/Video-Tree-TRM/process/{name}` | Step-by-step execution trace |
## Self-Check Before Finishing
Before completing any code modification task, verify:
1. `gitnexus_impact` was run for all modified symbols
2. No HIGH/CRITICAL risk warnings were ignored
3. `gitnexus_detect_changes()` confirms changes match expected scope
4. All d=1 (WILL BREAK) dependents were updated
## CLI
- Re-index: `npx gitnexus analyze`
- Check freshness: `npx gitnexus status`
- Generate docs: `npx gitnexus wiki`
<!-- gitnexus:end -->
reference/config/default.yaml
+57
View File
@@ -0,0 +1,57 @@
tree:
max_paragraphs_per_l2: 5
l1_segment_duration: 600.0
l2_clip_duration: 60.0
l3_fps: 1.0
l2_representative_frames: 10
cache_dir: "cache/trees"
concurrency: 16
embed:
backend: "local" # "local" (sentence-transformers) | "remote" (OpenAI 兼容 API)
model_name: "BAAI/bge-base-zh-v1.5"
embed_dim: 2560
device: "cuda"
api_key: "" # 远程模式从 .env 覆盖
api_url: "" # 远程模式从 .env 覆盖
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: 2560
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
margin_loss_weight: 0.5
act_lambda_step: 0.1
act_gamma: 0.9
eval_interval: 5
save_dir: "checkpoints"
dataset: "longbench"
dataset_path: "data/longbench"
reference/config/videomme.yaml
+61
View File
@@ -0,0 +1,61 @@
tree:
max_paragraphs_per_l2: 5
l1_segment_duration: 600.0 # L1: 每段 10 分钟（长视频适配）
l2_clip_duration: 60.0 # L2: 每 clip 60 秒
l3_fps: 0.5 # L3: 0.5 帧/秒（每 2 秒一帧）
l2_representative_frames: 6 # L2 VLM 描述用的代表帧数（从10降到6以提速）
cache_dir: "data/videomme/trees" # 树索引缓存目录（相对项目根目录）
concurrency: 16 # asyncio Semaphore 上限：每视频最多 16 路同时在途 VLM 请求
embed:
backend: "local" # CPU 本地运行，无需远程嵌入服务
model_name: "BAAI/bge-base-zh-v1.5"
embed_dim: 768 # bge-base-zh-v1.5 输出维度
device: "cpu" # 本地 CPU 推理
api_key: ""
api_url: ""
llm:
backend: "openai" # GPUStack 兼容 OpenAI API
model: "gemma-4-31B"
api_url: "http://100.83.164.94:11904/v1"
max_tokens: 256
temperature: 0.1
# api_key: 从 .env 加载
vlm:
backend: "openai" # GPUStack 兼容 OpenAI API
model: "gemma-4-31B"
api_url: "http://100.83.164.94:11904/v1"
max_tokens: 512 # 5帧描述 ~300 tokens，256 会截断 JSON 触发 fallback
temperature: 0.1
# api_key: 从 .env 加载
retriever:
embed_dim: 768 # 与 bge-base-zh-v1.5 维度一致
num_heads: 4
L_layers: 2
L_cycles: 4
max_rounds: 5
ffn_expansion: 2.0
checkpoint: null
k_l1: 1
k_l2: 1
k_l3: 1
max_paths: 5
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
margin_loss_weight: 0.5
act_lambda_step: 0.1
act_gamma: 0.9
eval_interval: 5
save_dir: "data/videomme/checkpoints"
dataset: "videomme"
dataset_path: "data/videomme/splits/train.jsonl"
reference/docs/EXPERIMENT.md
+330
View File
@@ -0,0 +1,330 @@
# Video-Tree-TRM 实验计划
> **目标**: 验证结合 TRM 多层推理探索能力（Cross-Attention Selector + ACT Halt）
> 与 PageIndex 树状检索能力的 Video-Tree-TRM 在长文本和长视频问答任务上的效果。
>
> **两条并行实验线路**：
> - **线路 A** — 长文本检索（LongBench / NarrativeQA）
> - **线路 B** — 长视频检索（VideoMME）
---
## 目录
- [评测体系](#1-评测体系)
- [实验环境准备](#2-实验环境准备)
- [线路 A：长文本检索](#3-线路-a长文本检索)
- [线路 B：长视频检索](#4-线路-b长视频检索)
- [基线方法](#5-基线方法)
- [消融实验](#6-消融实验)
- [里程碑检查表](#7-里程碑检查表)
- [结果记录表](#8-结果记录表)
---
## 1. 评测体系
### 1.1 质量指标
| 指标 | 适用模态 | 计算方式 |
|------|---------|---------|
| **EM (Exact Match)** | 文本 QA | 标准化后精确字符串匹配，0/1 |
| **F1** | 文本 QA | token 级 precision/recall，`token_f1()` 已实现于 `answer_generator.py` |
| **Accuracy** | 视频 QA（多选） | 模型选项与标准答案选项完全匹配的比率 |
### 1.2 效率指标
| 指标 | 说明 |
|------|------|
| **Avg Rounds** | 所有样本的平均检索轮次（衡量 ACT Halt 效果，越少越好） |
| **Max Rounds Hit Rate** | 触达 `max_rounds` 上限而非主动停止的样本比率 |
### 1.3 诊断指标
| 指标 | 说明 |
|------|------|
| **Nav Accuracy @L1** | 第一轮检索选中正确 L1 节点的比率 |
| **Nav Accuracy @L2** | 在正确 L1 下选中正确 L2 节点的比率 |
| **Nav Accuracy @L3** | 在正确 L2 下选中正确 L3 节点的比率 |
---
## 2. 实验环境准备
### 2.1 安装依赖
```bash
conda activate Video-Tree-TRM
pip install -r requirements.txt
```
### 2.2 配置文件
复制并填写环境变量：
```bash
cp .env.example .env
# 填写 LLM_API_KEY / VLM_API_KEY / EMBED_API_KEY（若使用远程嵌入）
```
默认配置文件 `config/default.yaml` 已预置：
```yaml
embed.model_name: "BAAI/bge-base-zh-v1.5" # 嵌入模型
retriever.embed_dim: 2560
retriever.max_rounds: 5
train.max_epochs_phase1: 30
train.max_epochs_phase2: 20
```
> 消融实验时通过 `--set key=value` 覆盖，无需修改 yaml 文件。
---
## 3. 线路 A：长文本检索
> **优先级**: P0（LongBench）→ P1（NarrativeQA）
> **并行方式**: 与线路 B 同步推进，互不阻塞
### A.1 数据集准备
| 数据集 | 样本量 | 格式 | 优先级 |
|--------|--------|------|--------|
| **LongBench** | ~5K | JSONL `{"query", "answer", "source_path", "modality": "text"}` | P0 |
| **NarrativeQA** | ~30K | 同上 | P1 |
```bash
# 下载 LongBench（示例）
mkdir -p data/longbench data/narrativeqa
# 将转换后的 JSONL 文件放置于对应目录
```
配置覆盖：
```bash
# LongBench
--set train.dataset=longbench --set train.dataset_path=data/longbench
# NarrativeQA（P1，LongBench 完成后）
--set train.dataset=narrativeqa --set train.dataset_path=data/narrativeqa
```
### A.2 索引构建
```bash
# 对所有文本文件批量构建 TreeIndex（含磁盘缓存）
conda run -n Video-Tree-TRM python main.py index \
--source data/longbench/doc.txt \
--modality text \
--config config/default.yaml
```
> 缓存命中后重复运行零开销，构建结果保存至 `cache/trees/`。
### A.3 两阶段训练
```bash
# Phase 1：导航训练（单轮，~30 epoch）
conda run -n Video-Tree-TRM python train.py \
--config config/default.yaml \
--set train.dataset=longbench \
--set train.dataset_path=data/longbench
# Phase 2 权重加载自 Phase 1 最佳检查点
# 在 config/default.yaml 中设置:
# retriever.checkpoint: checkpoints/phase1_epoch30.pt
# train.max_epochs_phase2: 20
conda run -n Video-Tree-TRM python train.py \
--config config/default.yaml \
--set retriever.checkpoint=checkpoints/phase1_epoch30.pt
```
训练检查点保存至 `checkpoints/phase1_epoch{N}.pt` / `phase2_epoch{N}.pt`。
### A.4 推理评测
```bash
# 单样本问答验证
conda run -n Video-Tree-TRM python main.py query \
--source data/longbench/doc.txt \
--modality text \
--question "文档的核心观点是什么？" \
--config config/default.yaml
```
批量评测（自行实现评测脚本）：
- 遍历测试集 JSONL
- 调用 `Pipeline.query()` 获取预测答案
- 计算 EM / F1，汇总 Avg Rounds
### A.5 基线对比
见 [第 5 节](#5-基线方法)。
---
## 4. 线路 B：长视频检索
> **优先级**: P2（可与线路 A 并行推进）
### B.1 数据集准备
| 数据集 | 样本量 | 格式 | 评测指标 |
|--------|--------|------|---------|
| **VideoMME** | ~2K | JSONL `{"query", "answer", "source_path", "modality": "video", "timestamp": float}` | Accuracy（多选） |
```bash
mkdir -p data/videomme
# 下载 VideoMME 并转换为上述 JSONL 格式
```
### B.2 索引构建
```bash
conda run -n Video-Tree-TRM python main.py index \
--source data/videomme/video.mp4 \
--modality video \
--config config/default.yaml
```
关键配置：
```yaml
tree:
l1_segment_duration: 600.0 # L1 段时长（秒）
l2_clip_duration: 60.0 # L2 clip 时长（秒）
l3_fps: 1.0 # L3 帧提取频率
l2_representative_frames: 10 # VLM 描述用代表帧数
```
### B.3 两阶段训练
```bash
# Phase 1（视频模态）
conda run -n Video-Tree-TRM python train.py \
--config config/default.yaml \
--set train.dataset=videomme \
--set train.dataset_path=data/videomme
# Phase 2（视频模态，加载 Phase 1 检查点）
conda run -n Video-Tree-TRM python train.py \
--config config/default.yaml \
--set train.dataset=videomme \
--set train.dataset_path=data/videomme \
--set retriever.checkpoint=checkpoints/phase1_epoch30.pt
```
### B.4 推理评测
```bash
conda run -n Video-Tree-TRM python main.py query \
--source data/videomme/video.mp4 \
--modality video \
--question "视频中发生了什么事？" \
--config config/default.yaml
```
批量评测：遍历测试集 → `Pipeline.query()` → 与多选选项比对 → 计算 Accuracy。
### B.5 基线对比
见 [第 5 节](#5-基线方法)。
---
## 5. 基线方法
| 方法 | 描述 | 实现方式 |
|------|------|---------|
| **BM25 + LLM** | 传统稀疏检索 | `rank_bm25` 库检索 top-k 段落 → 拼接上下文 → `LLMClient.chat()` |
| **Dense Retrieval + LLM** | BGE 向量检索 + rerank | `EmbeddingModel.embed_tensor()` 全量检索 top-k → rerank → 生成 |
| **PageIndex（无 TRM）** | 树状导航，cosine 路由，无推理模块 | 替换 `CrossAttentionSelector` 为 cosine 相似度选节点 |
| **Tree-TRM（原论文）** | 原始实现 | 参考 `Reference/Tree-TRM/` 目录 |
| **Video-Tree-TRM（ours）** | 本项目实现 | `Pipeline.query()` |
评测时各方法使用相同数据集和评测脚本，确保公平对比。
---
## 6. 消融实验
在主实验（LongBench）最优配置基础上，逐一变更单一变量：
| 编号 | 变量 | 候选值 | 配置覆盖 | 预期观察 |
|------|------|--------|---------|---------|
| **A1** | 选择器类型 | Cross-Attention vs Cosine | 替换 `CrossAttentionSelector` 实现 | CA 路由是否带来 F1 提升 |
| **A2** | 推理深度 | L_cycles ∈ {1, 2, 4, 8} | `--set retriever.L_cycles=N` | 质量-计算量权衡 |
| **A3** | 推理模块层数 | L_layers ∈ {1, 2, 4} | `--set retriever.L_layers=N` | 网络深度的边际收益 |
| **A4** | 多轮检索上限 | max_rounds ∈ {1, 3, 5} | `--set retriever.max_rounds=N` | ACT 多轮边际收益 |
| **A5** | ACT Halt 机制 | 有 / 无 | `act_loss_weight=0.0` 禁用 | ACT 对效率和质量的贡献 |
| **A6** | 注意力头数 | num_heads ∈ {1, 4, 8} | `--set retriever.num_heads=N` | 多头注意力的容量影响 |
每组消融实验保持其余超参数为默认值（`config/default.yaml`）。
---
## 7. 里程碑检查表
### 线路 A（文本）
- [ ] 环境配置完成，`python main.py --help` 正常输出
- [ ] LongBench 数据集下载 & 转换为 JSONL 格式
- [ ] LongBench 索引构建完成（`cache/trees/` 缓存生成）
- [ ] Phase 1 训练完成（文本，LongBench）
- [ ] Phase 2 训练完成（文本，LongBench）
- [ ] LongBench 批量评测完成（EM / F1 / Avg Rounds）
- [ ] 4 条基线方法评测完成
- [ ] NarrativeQA 实验（P1，可选）
### 线路 B（视频）
- [ ] VideoMME 数据集下载 & 转换为 JSONL 格式
- [ ] VideoMME 索引构建完成（视频帧提取 + VLM 描述生成）
- [ ] Phase 1 训练完成（视频，VideoMME）
- [ ] Phase 2 训练完成（视频，VideoMME）
- [ ] VideoMME 批量评测完成（Accuracy / Avg Rounds）
### 消融实验
- [ ] A1: Cross-Attention vs Cosine 路由
- [ ] A2: L_cycles 扫描（1/2/4/8）
- [ ] A3: L_layers 扫描（1/2/4）
- [ ] A4: max_rounds 扫描（1/3/5）
- [ ] A5: 有/无 ACT Halt
- [ ] A6: num_heads 扫描（1/4/8）
---
## 8. 结果记录表
### 8.1 主实验结果
| 方法 | LongBench EM | LongBench F1 | VideoMME Acc | Avg Rounds |
|------|-------------|-------------|-------------|-----------|
| BM25 + LLM | | | | — |
| Dense Retrieval + LLM | | | | — |
| PageIndex（无 TRM） | | | | |
| Tree-TRM（原论文） | | | | |
| **Video-Tree-TRM（ours）** | | | | |
### 8.2 消融实验结果（LongBench F1）
| 变量 | 值 | F1 | Avg Rounds | 备注 |
|------|----|----|-----------|------|
| 选择器 | Cross-Attention | | | 默认 |
| 选择器 | Cosine | | | A1 |
| L_cycles | 1 | | | A2 |
| L_cycles | 2 | | | A2 |
| L_cycles | 4 | | | A2 默认 |
| L_cycles | 8 | | | A2 |
| L_layers | 1 | | | A3 |
| L_layers | 2 | | | A3 默认 |
| L_layers | 4 | | | A3 |
| max_rounds | 1 | | | A4 |
| max_rounds | 3 | | | A4 |
| max_rounds | 5 | | | A4 默认 |
| ACT | 有 | | | A5 默认 |
| ACT | 无 | | | A5 |
| num_heads | 1 | | | A6 |
| num_heads | 4 | | | A6 默认 |
| num_heads | 8 | | | A6 |
reference/docs/TD.md
+1414
View File
File diff suppressed because it is too large Load Diff
reference/docs/architecture.md
+557
View File
@@ -0,0 +1,557 @@
# Video-Tree-TRM 统一架构设计
## 1. 核心定位
结合 TRM 递归推理 + PageIndex 树状导航的 **模态无关 RAG 系统**。
```
设计原则:
- 检索阶段模态无关（全文本嵌入空间）
- 模态差异封装在两端（预处理 + 答案生成）
- TRM ACT 机制控制动态检索深度
- 树深度固定 3 层，检索轮次动态
```
---
## 2. 系统总览
```
┌──────────────────── 预处理（离线，一次性） ────────────────────┐
│ │
│ 原始输入 统一表示 │
│ ┌──────────┐ ┌──────────────────────┐ │
│ │ 长文本 │──TextTree──→│ │ │
│ └──────────┘ Builder │ TreeIndex │ │
│ ┌──────────┐ │ (全文本嵌入, [D]) │ │
│ │ 长视频 │──VideoTree─→│ │ │
│ └──────────┘ Builder └──────────┬───────────┘ │
│ │ │
└───────────────────────────────────────│───────────────────────┘
│
┌──────────────────── 检索（在线） ─────│───────────────────────┐
│ ↓ │
│ query ──text_embed──→ q ──→ RecursiveRetriever │
│ │ TRM 双循环推理 │
│ │ ACT halt 动态停止 │
│ │ 多轮 root-to-leaf 遍历 │
│ └──→ List[NodePath] │
│ │ │
└─────────────────────────────────────│─────────────────────────┘
│
┌──────────────────── 生成（在线） ────│─────────────────────────┐
│ ↓ │
│ 文本模式: LLM(query, raw_text_chunks) → answer │
│ 视频模式: VLM(query, frame_images) → answer │
│ │
└───────────────────────────────────────────────────────────────┘
```
---
## 3. TreeIndex：统一数据结构
三层固定深度，文本/视频同构。
```
TreeIndex
├── metadata: IndexMeta # 来源、嵌入器版本、创建时间
├── embed_dim: int # 嵌入维度 D
└── roots: List[L1Node] # L1 节点列表（树的根层）
L1Node (段/章)
├── id: str
├── summary: str # 聚合自子 L2 描述 (2-3句)
├── embedding: ndarray [D] # text_embed(summary)
├── time_range: (float, float) # 仅视频模式
└── children: List[L2Node]
L2Node (片段/小节)
├── id: str
├── description: str # 直接从原始内容生成 (1-2句)
├── embedding: ndarray [D] # text_embed(description)
├── time_range: (float, float) # 仅视频模式
└── children: List[L3Node]
L3Node (帧/文本块)
├── id: str
├── description: str # 视频=VLM帧描述(继承L2上下文), 文本=原始段落文本
├── embedding: ndarray [D] # text_embed(description)
├── raw_content: str # 原始文本块（文本模式）
├── frame_path: Optional[str] # 帧图像路径（视频模式）
└── timestamp: Optional[float] # 帧时间戳（视频模式）
```
**关键设计**：`embedding` 字段全部来自同一个 `text_embed()` 函数，不存在跨模态嵌入空间问题。`raw_content` / `frame_path` 仅供答案生成阶段使用，检索器不关心。
---
## 4. 预处理管线
### 4.1 摘要粒度规范
各层摘要的唯一职责是**生成好的路由嵌入**——让 `(q+z)·M^T/√D` 能选对节点。
| 层级 | 路由功能 | 摘要目标 | 粒度 |
|------|----------|----------|------|
| L1 | 粗路由：选主题区域 | "这个区域**关于什么**" | 2-3句：主题 + 覆盖范围 + 关键实体 |
| L2 | 细聚焦：选具体片段 | "这个片段**发生了什么**" | 1-2句：具体事件/内容 + 区分性细节 |
| L3(文本) | 精确定位：选文本块 | 段落原文 | 原始段落文本，不做摘要 |
| L3(视频) | 精确定位：选帧 | "这帧画面的**具体内容**" | 1-2句：继承L2上下文，聚焦区分性视觉信息 |
```
关键原则:
- L1 要宽: 涵盖其下所有 L2 的语义范围，让相关 query 能"进对门"
- L2 要窄: 只描述自己，与同级兄弟节点形成区分
- L3 要具体: 提供精确定位的细节信息
```
**文本示例**：
```
L1: "第三章讨论了用户认证系统的设计，涵盖OAuth流程、JWT令牌管理和权限控制。"
L2: "OAuth 2.0授权码流程的实现，包括重定向和回调处理。"
L2: "JWT令牌的生成、验证和刷新机制。"
L2: "基于角色的权限控制模型（RBAC）。"
```
**视频示例**：
```
L1: "厨师在厨房制作意大利面，从准备食材到最终装盘。"
L2: "切洋葱、蒜末和番茄，准备酱料食材。"
L3: "[准备酱料食材中] 厨师左手按住洋葱，右手快速切丁，砧板上已有切好的蒜末。"
L3: "[准备酱料食材中] 厨师将番茄放入沸水中烫皮，锅中冒出蒸汽。"
L2: "煮面条并制作番茄肉酱。"
L2: "装盘摆盘并撒上帕尔马干酪。"
```
### 4.2 构建顺序：L2 轴心策略
两种模态共享同一构建原则：**L2 为轴心，向下展开 L3，向上聚合 L1**。
```
构建依赖关系:
L3 需要 L2 上下文才能生成高质量描述（视频尤为关键）
L1 摘要应从 L2 聚合，确保覆盖面完整
→ 循环依赖 → 解法: L2 不依赖 L3，而是从原始内容独立生成
构建顺序:
Step 1: 结构切分 → 确定 L1/L2 边界
Step 2: L2 先行 → 从原始内容直接生成 L2 描述
Step 3: L3 向下 → 注入 L2 上下文，生成细粒度描述
Step 4: L1 向上 → 聚合 L2 描述，生成粗粒度摘要
```
```
Step 4: 聚合
↑
┌─────── L1 ───────┐ "第三章讨论了认证系统..."
│ │ ← LLM("总结这些片段: " + L2_descriptions)
│ Step 2: 先行 │
├── L2 ── L2 ── L2 ┤ "OAuth流程的实现..."
│ │ │ │ │ ← 文本: LLM 摘要段落组
│ ↓ ↓ ↓ │ 视频: VLM 看少量代表帧
│ L3s L3s L3s │
│ Step 3: 向下 │ "厨师左手按住洋葱..."
└────────────────────┘ ← 注入 L2 上下文后生成
```
### 4.3 TextTreeBuilder
```
输入: 长文本文档
输出: TreeIndex（全部 embedding=None）
Pipeline:
Step 1 — 结构切分
有 ToC → 解析章节层级 → L1/L2 边界
无 ToC → LLM 语义分段 (一次性调用)
Step 2 — L2 先行
L2: 段落组 → LLM 生成摘要 (1-2句)
摘要目标: 与兄弟 L2 形成区分
（batch_chat 并发生成所有 L2 摘要）
Step 3 — L3 向下
L3: 原始段落文本直接复用
存储 raw_content = description = 原始文本
（文本模式 L3 不需要 L2 上下文，不调用 LLM）
Step 4 — L1 向上聚合
L1: LLM("总结这些小节: " + 所有子 L2 摘要) → 摘要
摘要目标: 覆盖下属所有 L2 的语义范围 (2-3句)
Step 5 — 序列化 → TreeIndex（JSON，无 embedding）
⚠ 延迟 Embedding: 所有节点 embedding=None
首次检索时由 Pipeline._embed_tree() → tree.embed_all() 统一填充
```
### 4.4 VideoTreeBuilder
**状态**: ✅ 已实现（`video_tree_trm/video_tree_builder.py`）
```
输入: 长视频文件路径 或 YouTube URL
输出: TreeIndex（全部 embedding=None）
Pipeline（ThreadPoolExecutor 异步事件循环）:
Step 0 — 输入类型判断
本地文件: 直接使用 OpenCV 读取
YouTube URL: yt-dlp -g 获取 CDN 直链 + yt-dlp --dump-json 获取时长
Step 1 — 时间切分（固定步长）
本地: cv2 读取总时长
HTTP 流: 使用 yt-dlp 元数据时长（duration_hint，避免 cv2 流上不可靠）
固定步长 l1_segment_duration=600s → L1 区间列表
每个 L1 区间 → 等分 l2_clip_duration=60s → L2 clips
Step 2 — 预计算任务图
收集所有 L2 任务 + 记录每个 L1 的 L2 数量（l2_counts）
Step 3 — 一次性提交所有 L2 任务（非阻塞）
ThreadPoolExecutor(max_workers=concurrency)
每个 L2 clip: 均匀 seek l2_representative_frames=10 帧 → VLM → 1-2句描述
Step 4 — 事件循环（cfwait FIRST_COMPLETED）
L2 完成 → 立即提交 L3 任务（_build_l3_task，线程安全）
L3: 按 l3_fps 提取全量帧（持久化缓存），注入 L2 上下文，VLM 批量帧描述（JSON）
降级路径: JSON 解析失败 → 逐帧 VLM 调用
L3 完成 → 检查该 L1 的所有 L2 是否就绪 → 触发 L1 任务
L1: 拼接所有 L2 描述 → vlm.chat() → 2-3句摘要
主线程单线程操作 l1_l2_buckets，无竞争，无需 Lock
Step 5 — 有序重建 l1_nodes，组装 TreeIndex（写 IndexMeta + 日志）
⚠ 延迟 Embedding: 所有节点 embedding=None
首次检索时由 Pipeline._embed_tree() → tree.embed_all() 统一填充
```
**L2 代表帧采样**（均匀 seek，首尾均包含）：
```python
step = (end_sec - start_sec) / (n_rep - 1)
timestamps = [start_sec + i * step for i in range(n_rep)]
# → 直接 cap.set(CAP_PROP_POS_MSEC, ts * 1000) 提取，缓存为 l2_{ts:.3f}.jpg
```
**L3 帧描述 prompt**（继承 L2 上下文 + 要求 JSON 输出）：
```python
prompt = (
'该片段的整体内容: "{l2_description}"\n'
"以下是该片段中连续的 {n} 帧画面。\n"
"对每帧用一到两句话描述其具体画面内容。\n"
"重点关注: 动作、物体变化、文字信息、人物表情。\n"
"不要重复片段整体描述，聚焦每帧的区分性信息。\n"
'只返回 JSON 数组，格式: ["帧1描述", "帧2描述", ...]，不要其他内容。'
)
descs = VLM(frames_batch, prompt) # 主路径：一次调用，JSON 解析
# 降级：json.loads 失败 → 逐帧调用
```
**L1 摘要 prompt**：
```python
l2_texts = "\n".join(f"- {node.description}" for node in l2_children)
prompt = f"以下是一个视频段落中各片段的描述:\n{l2_texts}\n用2-3句话总结该段落的整体内容，涵盖所有片段的主题。"
l1_summary = vlm.chat(prompt) # 纯文本调用
```
---
## 5. RecursiveRetriever：TRM 式递归检索
### 5.1 核心算法
节点选择使用 **Cross-Attention**（学习 W_q/W_k/W_v/W_o 投影），替代简单 cosine 路由。
L_level 推理模块使用 **MLP-based**（RMSNorm + SwiGLU），操作对象为向量 `[B, D]`。
三个可学习组件（CrossAttentionSelector, ReasoningModule, q_head）**跨层级共享权重**。
```
┌─────────────────────────────────────────────────────────────┐
│ RecursiveRetriever │
│ │
│ 可训练组件（共享权重）: │
│ selector = CrossAttentionSelector(W_q, W_k, W_v, W_o) │
│ L_level = ReasoningModule(RMSNorm + SwiGLU, L_layers层) │
│ q_head = Linear(D, 1) │
│ │
│ 输入: query (str), tree (TreeIndex) │
│ 输出: List[RetrievalPath] │
│ │
│ q = text_embed(query) # 查询嵌入 [D]，冻结 │
│ z = q.clone() # 初始潜在状态 │
│ │
│ for round in range(max_rounds): ← ACT halt 控制 │
│ │ │
│ │ ┌─── Phase 1: L1 粗粒度路由 (Cross-Attention) ──┐ │
│ │ │ s1, w1, k1* = selector(q+z, M_L1) │ │
│ │ │ Q=W_q(q+z), K=W_k(M_L1), V=W_v(M_L1) │ │
│ │ │ s1 = W_o(MultiHeadAttn(Q,K,V)) # 软选择 │ │
│ │ │ k1* = argmax(mean_head_scores) # 硬索引 │ │
│ │ │ z = z + s1 │ │
│ │ │ for _ in L_cycles: │ │
│ │ │ z = L_level(z, s1 + q) │ │
│ │ └────────────────────────────────────────────────┘ │
│ │ │
│ │ ┌─── Phase 2: L2 细粒度聚焦 (Cross-Attention) ──┐ │
│ │ │ M_L2 = children_embeds(L1[k1*]) │ │
│ │ │ s2, w2, k2* = selector(q+z, M_L2) │ │
│ │ │ z = z + s2 │ │
│ │ │ for _ in L_cycles: │ │
│ │ │ z = L_level(z, s2 + q) │ │
│ │ └────────────────────────────────────────────────┘ │
│ │ │
│ │ ┌─── Phase 3: L3 精确定位 (Cross-Attention) ────┐ │
│ │ │ M_L3 = children_embeds(L2[k2*]) │ │
│ │ │ s3, w3, k3* = selector(q+z, M_L3) │ │
│ │ │ z = z + s3 │ │
│ │ └────────────────────────────────────────────────┘ │
│ │ │
│ │ collected.append(Path(k1*, k2*, k3*)) │
│ │ │
│ │ ┌─── ACT Halt Decision ───────────────────┐ │
│ │ │ halt_logit = q_head(z) │ │
│ │ │ if halt_logit > 0 and round_idx > 0: │ │
│ │ │ break # 至少跑 1 轮 │ │
│ │ └──────────────────────────────────────────┘ │
│ │ │
│ └── z 状态保留到下一轮（累积已检索信息） │
│ │
│ return collected │
└─────────────────────────────────────────────────────────────┘
```
### 5.2 算法伪代码
```python
class RecursiveRetriever:
"""TRM 式递归检索器（Cross-Attention 版本）。"""
def __init__(self, embed_dim: int, num_heads: int, L_layers: int, L_cycles: int, max_rounds: int):
self.selector = CrossAttentionSelector(embed_dim, num_heads) # 共享，跨层级
self.L_level = ReasoningModule(embed_dim, L_layers) # 共享，跨层级
self.q_head = Linear(embed_dim, 1) # ACT halt head
self.L_cycles = L_cycles
self.max_rounds = max_rounds
def forward(
self, q: Tensor, tree: TreeIndex, return_internals: bool = False
) -> Dict[str, Any]:
z = q.clone() # [B, D]，初始潜在状态 = 查询嵌入
paths = []
for round_idx in range(self.max_rounds):
# ── 三阶段树遍历（一次完整 root-to-leaf） ──
path, z, step_attns = self._traverse_one_path(q, z, tree)
paths.append(path)
# ── ACT halt（推理模式，至少走 1 轮） ──
halt_logit = self.q_head(z) # [B, 1]
if not self.training and halt_logit.item() > 0 and round_idx > 0:
break
return {"paths": paths, "num_rounds": len(paths), "z_final": z}
def _traverse_one_path(self, q, z, tree):
"""单次 root → L1 → L2 → L3 遍历。"""
# Phase 1
k1, z = self._select_and_reason(q, z, tree.l1_embeddings())
# Phase 2
k2, z = self._select_and_reason(q, z, tree.l2_embeddings_of(k1))
# Phase 3
k3, z = self._select_and_reason(q, z, tree.l3_embeddings_of(k1, k2))
return Path(k1, k2, k3), z
def _select_and_reason(self, q, z, M):
"""单层: Cross-Attention 选择 + L_cycles 内循环推理。"""
# Navigate: Cross-Attention
state = q + z # [B, D]
selected_info, attn_w, k_star = self.selector(state, M)
# Q = W_q(state) → [B, 1, D]
# K = W_k(M) → [B, N, D]
# V = W_v(M) → [B, N, D]
# selected_info = W_o(MultiHeadAttn(Q, K, V)) [B, D] ← 可微
# k_star = argmax(avg_head_scores) ← 路径索引
# Update: z += attention 加权信息
z = z + selected_info
# Reason: TRM L-level MLP 内循环
for _ in range(self.L_cycles):
z = self.L_level(z, selected_info + q)
return k_star, z
```
### 5.3 多轮检索的 z 状态流
```
Round 1: Round 2:
z₀ = q z₃ 来自 Round 1（包含已检索信息）
│ │
├─ Phase1: CA(q+z₀, M_L1) → s1 ├─ Phase1: CA(q+z₃, M_L1) → s4
│ z₁ = z₀ + s1 → L_level×L_cycles │ z₄ = z₃ + s4 → L_level×L_cycles
├─ Phase2: CA(q+z₁, M_L2[k1]) → s2 ├─ Phase2: CA(q+z₄, M_L2[k4]) → s5
│ z₂ = z₁ + s2 → L_level×L_cycles │ z₅ = z₄ + s5 → L_level×L_cycles
├─ Phase3: CA(q+z₂, M_L3[k2]) → s3 ├─ Phase3: CA(q+z₅, M_L3[k5]) → s6
│ z₃ = z₂ + s3 │ z₆ = z₅ + s6
│ │
ACT: q_head(z₃) < 0 → 继续 ACT: q_head(z₆) > 0 → 停止
→ 返回 [Path(k1,k2,k3), Path(k4,k5,k6)]
关键: s = W_o(MultiHeadAttn(W_q(q+z), W_k(M), W_v(M)))
比简单 cosine 路由更强:
- 学习的投影决定"关注什么"进行选择
- Multi-head 捕获多种相关性信号
- V 投影让"更新信息"与"匹配键"解耦
z₃ 累积了 Round 1 的 attention 信息
→ (q + z₃) 的投影方向自动偏离已选区域
```
---
## 6. ACT Halt 训练方案
### 6.1 训练目标
```
ACT halt head 学习: "已收集的信息是否足以回答 query"
┌──────────────────────────────┐
│ reward 定义 │
│ │
│ R = answer_quality - λ·rounds │
│ │
│ answer_quality: │
│ 文本QA: EM / F1 score │
│ 视频QA: 选项匹配准确率 │
│ │
│ λ: 步数惩罚系数 │
│ 鼓励用更少轮次达到同样质量 │
└──────────────────────────────┘
```
### 6.2 训练数据
| 数据集 | 模态 | 样本量 | reward 信号 |
|--------|------|--------|-------------|
| LongBench | 文本 | ~5K | ground truth → F1 |
| NarrativeQA | 文本 | ~30K | ground truth → ROUGE |
| VideoMME | 视频 | ~2K | 选项匹配 → 0/1 |
### 6.3 训练流程
```python
# 训练伪代码
for query, tree, ground_truth in dataset:
q = text_embed(query)
z = q.copy()
total_loss = 0
for round_idx in range(max_rounds):
path, z = retriever.traverse_one_path(q, z, tree)
# halt 决策
halt_logit = q_head(z)
# 用当前已收集的 context 生成答案，计算质量
context = collect_raw_content(paths_so_far)
answer = generator(query, context)
quality = compute_score(answer, ground_truth) # EM/F1
# Q-learning target
# 如果现在停 → reward = quality
# 如果继续 → reward = γ · future_quality - λ
target_q = quality if is_last else γ * next_quality - λ
loss += mse(sigmoid(halt_logit), target_q)
total_loss.backward()
optimizer.step()
```
### 6.4 可训练组件 vs 冻结组件
```
冻结 (不训练):
✗ text_embed() # 预训练嵌入器，冻结
✗ TreeIndex embeddings # 预计算的节点嵌入，冻结
可训练:
✓ CrossAttentionSelector (W_q, W_k, W_v, W_o) # 节点选择投影
✓ L_level (ReasoningModule: RMSNorm + SwiGLU) # MLP 推理模块
✓ q_head (ACT halt) # 停止决策头
```
---
## 7. 检索结果与答案生成的接口
```python
@dataclass
class RetrievalPath:
"""一条 root-to-leaf 路径。"""
k1: int # L1 索引
k2: int # L2 索引
k3: int # L3 索引
l1_summary: str # L1 摘要
l2_description: str # L2 描述
l3_description: str # L3 描述
raw_content: Optional[str] # 原始文本（文本模式）
frame_path: Optional[str] # 帧路径（视频模式）
timestamp: Optional[float] # 时间戳（视频模式）
@dataclass
class RetrievalResult:
"""检索器输出。"""
query: str
paths: List[RetrievalPath] # 多轮收集的路径
num_rounds: int # 实际检索轮次
z_final: ndarray # 最终潜在状态
# 答案生成器接口
def generate_answer(query: str, result: RetrievalResult) -> str:
if is_text_mode(result):
context = "\n".join(p.raw_content for p in result.paths)
return LLM(query=query, context=context)
else:
frames = [load_image(p.frame_path) for p in result.paths]
captions = [p.l3_description for p in result.paths]
return VLM(query=query, images=frames, captions=captions)
```
---
## 8. 与现有系统的关系
```
参考代码 新架构 变化
─────────────────────────────────────────────────────────────────
video_pyramid.py (HSP) → TreeIndex 重构为统一格式
video_tree_trm.py (cosine路由) → RecursiveRetriever Cross-Attention+ACT
select_next_node (CrossAttn) → CrossAttentionSelector 保留CA思路，简化为向量级
L_level (Transformer blocks) → ReasoningModule MLP-based (向量非序列)
visual_projection.py → 删除 L3 全文本化
video_indexer.py (CLIP encode) → embeddings.py 统一 text_embed()
pipeline.py → pipeline.py ✅ 已实现（含延迟 embed 策略）
answer_generator.py → answer_generator.py ✅ 已实现
config.py → config.py 全面重构
新增:
+ tree_index.py 统一数据结构 ✅ 已实现
+ embeddings.py 嵌入服务封装 ✅ 已实现
+ llm_client.py LLM/VLM 客户端 ✅ 已实现
+ text_tree_builder.py 文本模式预处理 ✅ 已实现
+ video_tree_builder.py 视频模式预处理 ✅ 已实现
+ recursive_retriever.py TRM 递归检索器 (CA+MLP+ACT) ✅ 已实现
+ losses.py NavigationLoss + ACTLoss ✅ 已实现
+ train.py 两阶段训练入口 ✅ 已实现
+ main.py 推理/演示入口 ✅ 已实现
```
reference/main.py
+235
View File
@@ -0,0 +1,235 @@
"""
推理/演示入口
=============
从原始文档（文本或视频）构建 TreeIndex，执行问答。
与 ``train.py``（训练入口）配对，是 Video-Tree-TRM 的端到端演示入口。
子命令
------
- ``index``: 仅构建并缓存 TreeIndex，不执行问答。
- ``query``: 构建索引（或复用缓存）后执行问答。
使用示例::
# 仅构建文本索引
python main.py index --source data/doc.txt --modality text
# 单次问答（视频）
python main.py query --source data/video.mp4 --modality video \\
--question "视频的主要内容是什么？"
# 交互式多轮问答（文本）
python main.py query --source data/doc.txt --modality text --interactive
# 自定义配置
python main.py query --source data/doc.txt --modality text \\
--config config/default.yaml --env .env \\
--question "文档结论是什么？"
"""
from __future__ import annotations
import argparse
import sys
from utils.logger_system import log_msg
from video_tree_trm.config import Config
from video_tree_trm.pipeline import Pipeline
# ---------------------------------------------------------------------------
# CLI 解析
# ---------------------------------------------------------------------------
def _parse_args() -> argparse.Namespace:
"""解析 CLI 参数，返回 Namespace。
返回:
包含子命令及所有选项的 Namespace 对象。
"""
parser = argparse.ArgumentParser(
prog="main.py",
description="Video-Tree-TRM 推理/演示入口",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
# 公共参数
common = argparse.ArgumentParser(add_help=False)
common.add_argument(
"--config",
default="config/default.yaml",
metavar="YAML",
help="配置文件路径（默认: config/default.yaml）",
)
common.add_argument(
"--env",
default=".env",
metavar="ENV",
help="环境变量文件路径（默认: .env）",
)
common.add_argument(
"--source",
required=True,
metavar="PATH",
help="原始文件路径（文本文件或视频文件）",
)
common.add_argument(
"--modality",
required=True,
choices=["text", "video"],
help="文件模态：text 或 video",
)
subparsers = parser.add_subparsers(dest="command", title="子命令")
subparsers.required = True
# ── index 子命令 ──
subparsers.add_parser(
"index",
parents=[common],
help="构建并缓存 TreeIndex（不执行问答）",
description="从原始文件构建三层树索引并保存到缓存目录。",
)
# ── query 子命令 ──
query_parser = subparsers.add_parser(
"query",
parents=[common],
help="构建索引并执行问答",
description="构建索引后执行单次问答或交互式多轮问答。",
)
mode_group = query_parser.add_mutually_exclusive_group(required=True)
mode_group.add_argument(
"--question",
metavar="TEXT",
help="单次问答：问题字符串",
)
mode_group.add_argument(
"--interactive",
action="store_true",
help="交互式多轮问答模式（输入 quit 退出）",
)
return parser.parse_args()
# ---------------------------------------------------------------------------
# 配置加载
# ---------------------------------------------------------------------------
def _load_config(args: argparse.Namespace) -> Config:
"""从 YAML + .env 加载配置。
参数:
args: CLI 解析结果，需含 config（YAML 路径）和 env（.env 路径）字段。
返回:
完整的 Config 实例。
异常:
FileNotFoundError: 配置文件不存在。
TypeError: YAML 缺少必需字段。
"""
log_msg("INFO", "加载配置", yaml_path=args.config, env_path=args.env)
return Config.load(
yaml_path=args.config,
env_path=args.env,
)
# ---------------------------------------------------------------------------
# 子命令实现
# ---------------------------------------------------------------------------
def cmd_index(args: argparse.Namespace) -> None:
"""index 子命令：构建并缓存 TreeIndex，打印缓存路径。
参数:
args: 含 source、modality、config、env 字段的 Namespace。
"""
config = _load_config(args)
pipeline = Pipeline(config)
log_msg("INFO", "开始构建索引", source=args.source, modality=args.modality)
tree = pipeline.build_index(args.source, args.modality)
print(f"索引构建完成。模态={args.modality}，来源={args.source}")
print(f"节点数量: L1={len(tree.roots)}")
log_msg("INFO", "索引构建完成", l1_count=len(tree.roots))
def cmd_query(args: argparse.Namespace) -> None:
"""query 子命令：构建索引 → 单次或交互式问答。
参数:
args: 含 source、modality、question/interactive、config、env 字段的 Namespace。
实现细节:
- 先调用 build_index（缓存命中时直接加载，无额外开销）。
- --question 模式：单次问答后退出。
- --interactive 模式：循环读取用户输入，输入 'quit' 退出。
"""
config = _load_config(args)
pipeline = Pipeline(config)
log_msg("INFO", "构建/加载索引", source=args.source, modality=args.modality)
tree = pipeline.build_index(args.source, args.modality)
if args.question:
# Phase 1: 单次问答
answer = pipeline.query(args.question, tree)
print(f"\n问题: {args.question}")
print(f"答案: {answer}\n")
log_msg("INFO", "问答完成", question=args.question[:50])
else:
# Phase 2: 交互式多轮问答
print("\n已进入交互式问答模式。输入 'quit' 退出。\n")
while True:
try:
question = input("问题 (输入 'quit' 退出): ").strip()
except (EOFError, KeyboardInterrupt):
print("\n再见。")
break
if question.lower() in ("quit", "exit", "q"):
print("再见。")
break
if not question:
continue
answer = pipeline.query(question, tree)
print(f"答案: {answer}\n")
log_msg("INFO", "交互式问答", question=question[:50])
# ---------------------------------------------------------------------------
# 主入口
# ---------------------------------------------------------------------------
def main() -> None:
"""主入口：解析参数 → 分发子命令。
退出码:
0: 正常退出。
1: 参数错误或运行时异常。
"""
args = _parse_args()
try:
if args.command == "index":
cmd_index(args)
elif args.command == "query":
cmd_query(args)
except (FileNotFoundError, TypeError, ValueError) as exc:
print(f"错误: {exc}", file=sys.stderr)
log_msg("ERROR", f"main.py 异常退出: {exc}")
sys.exit(1)
if __name__ == "__main__":
main()
reference/requirements.txt
+21
View File
@@ -0,0 +1,21 @@
# 核心
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
reference/scripts/_download_meta.py
+138
View File
@@ -0,0 +1,138 @@
"""
_download_meta.py — VideoMME 元数据下载与长视频列表提取
==========================================================
从 HuggingFace `lmms-lab/Video-MME` 下载数据集元数据，
过滤 duration_category == "long"（30-60 分钟）的视频，
输出两个文件：
- {meta_dir}/long_videos.jsonl 每行一条唯一长视频记录
- {meta_dir}/long_videos_qa.jsonl 每行一条 QA 对（含 video_id）
使用方式（由 build_videomme_trees.sh 调用）:
python _download_meta.py --meta-dir /data/videomme/metadata
依赖（在脚本中已安装）: datasets, huggingface_hub
"""
from __future__ import annotations
import argparse
import json
import os
import sys
from pathlib import Path
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description="下载 VideoMME 元数据并提取长视频列表")
p.add_argument("--meta-dir", required=True, help="元数据输出目录")
p.add_argument(
"--min-duration",
type=int,
default=1800,
help="最短视频时长（秒），默认 1800（30 分钟）",
)
p.add_argument(
"--max-duration",
type=int,
default=3600,
help="最长视频时长（秒），默认 3600（60 分钟）",
)
return p.parse_args()
def main() -> None:
args = parse_args()
meta_dir = Path(args.meta_dir)
meta_dir.mkdir(parents=True, exist_ok=True)
# Phase 1: 尝试加载 HuggingFace 数据集
print("[meta] 正在从 HuggingFace 加载 lmms-lab/Video-MME ...")
try:
from datasets import load_dataset # type: ignore
except ImportError:
print("[meta][ERROR] 未安装 datasets 库，请先运行: pip install datasets", file=sys.stderr)
sys.exit(1)
try:
# VideoMME 数据集，test split
ds = load_dataset("lmms-lab/Video-MME", split="test")
except Exception as e:
print(f"[meta][ERROR] 数据集加载失败: {e}", file=sys.stderr)
print("[meta] 请确认 HuggingFace 可访问，或配置 HF_ENDPOINT 镜像", file=sys.stderr)
sys.exit(1)
print(f"[meta] 数据集总条目数: {len(ds)}")
# Phase 2: 过滤长视频
# VideoMME 字段: video_id, youtube_id, url, duration, duration_category,
# domain, sub_category, question, answer, options
seen_video_ids: set[str] = set()
long_videos: list[dict] = []
long_qa: list[dict] = []
for row in ds:
# 实际字段结构（lmms-lab/Video-MME 真实格式）:
# video_id : "001", "002", ... （数据集内部序号）
# videoID : YouTube 视频 ID（如 "fFjv93ACGo8"）
# url : YouTube 完整链接
# duration : "short" | "medium" | "long"（字符串类别，非秒数）
# domain : 领域
# sub_category : 细分类别
# question_id : 问题序号
# question : 问题文本
# options : 选项列表（字符串）
# answer : 正确选项字母
duration_category = str(row.get("duration", "")).strip().lower()
if duration_category != "long":
continue
youtube_id = row.get("videoID") or row.get("video_id", "")
url = row.get("url", f"https://www.youtube.com/watch?v={youtube_id}")
# 唯一视频记录（以 youtube_id 去重）
if youtube_id not in seen_video_ids:
seen_video_ids.add(youtube_id)
long_videos.append(
{
"video_id": row.get("video_id", ""),
"youtube_id": youtube_id,
"url": url,
"duration_category": duration_category,
"domain": row.get("domain", ""),
"sub_category": row.get("sub_category", ""),
}
)
# QA 对记录
long_qa.append(
{
"video_id": row.get("video_id", ""),
"youtube_id": youtube_id,
"question_id": row.get("question_id", ""),
"question": row.get("question", ""),
"answer": row.get("answer", ""),
"options": row.get("options", []),
"duration_category": duration_category,
}
)
# Phase 3: 写出文件
videos_path = meta_dir / "long_videos.jsonl"
qa_path = meta_dir / "long_videos_qa.jsonl"
with open(videos_path, "w", encoding="utf-8") as f:
for v in long_videos:
f.write(json.dumps(v, ensure_ascii=False) + "\n")
with open(qa_path, "w", encoding="utf-8") as f:
for q in long_qa:
f.write(json.dumps(q, ensure_ascii=False) + "\n")
print(f"[meta] 长视频唯一数量: {len(long_videos)}")
print(f"[meta] 长视频 QA 对数: {len(long_qa)}")
print(f"[meta] 视频列表已保存: {videos_path}")
print(f"[meta] QA 列表已保存: {qa_path}")
if __name__ == "__main__":
main()
reference/scripts/build_tree_single.py
+60
View File
@@ -0,0 +1,60 @@
"""
单视频建树脚本（仅 VLM，不加载 EmbeddingModel）
================================================
直接调用 VideoTreeBuilder，跳过 Pipeline 的嵌入模型初始化。
结果保存为 JSON 到 cache/trees/ 目录。
用法::
conda run -n Video-Tree-TRM python scripts/build_tree_single.py \
--video data/videomme/videos/xKiRmesHWIA.mp4
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
# 项目根目录加入 sys.path
sys.path.insert(0, str(Path(__file__).parent.parent))
from video_tree_trm.config import Config
from video_tree_trm.llm_client import LLMClient
from video_tree_trm.video_tree_builder import VideoTreeBuilder
from utils.logger_system import log_msg
def main() -> None:
"""构建单个视频的 TreeIndex，仅使用 VLM，不加载 EmbeddingModel。"""
parser = argparse.ArgumentParser(description="单视频建树（仅 VLM）")
parser.add_argument("--video", required=True, help="视频文件路径")
parser.add_argument("--config", default="config/default.yaml", help="配置文件路径")
args = parser.parse_args()
# Phase 1: 加载配置 + 初始化 VLM
cfg = Config.load(args.config)
vlm = LLMClient(cfg.vlm)
# Phase 2: 构建树（纯 VLM 描述，embedding=None）
builder = VideoTreeBuilder(vlm, cfg.tree)
tree = builder.build(args.video)
# Phase 3: 保存 JSON
stem = Path(args.video).stem
cache_dir = Path(cfg.tree.cache_dir)
cache_dir.mkdir(parents=True, exist_ok=True)
out_path = str(cache_dir / f"{stem}_video.json")
tree.save_json(out_path)
log_msg("INFO", "建树完成，已保存", path=out_path)
print(f"\n[完成] TreeIndex 已保存到: {out_path}")
print(f" L1 节点数: {len(tree.roots)}")
total_l2 = sum(len(r.children) for r in tree.roots)
total_l3 = sum(len(l2.children) for r in tree.roots for l2 in r.children)
print(f" L2 节点数: {total_l2}")
print(f" L3 节点数: {total_l3}")
if __name__ == "__main__":
main()
reference/scripts/build_trees_batch.py
+224
View File
@@ -0,0 +1,224 @@
"""
批量视频建树脚本（视频间并行 + 视频内 asyncio 真并发）
==========================================================
扫描指定目录下所有 MP4，跳过已有 JSON 树的视频，
使用 ThreadPoolExecutor 进行视频间并行（--jobs，默认 1），
每个视频内部使用 asyncio + Semaphore(concurrency=16) 真并发 VLM 调用。
并发架构::
外层: ThreadPoolExecutor(max_workers=jobs=1) — 视频间并行（默认 1 路）
内层: asyncio.run(_build_async()) — 每视频独立事件循环
asyncio.Semaphore(concurrency=16) — 限制同时在途 VLM 请求数
总最大 VLM 并发: jobs × concurrency = 1 × 16 = 16
用法::
# 默认 1 路 16 并发
conda run -n Video-Tree-TRM python scripts/build_trees_batch.py
# 指定路数
conda run -n Video-Tree-TRM python scripts/build_trees_batch.py --jobs 2
# 指定目录和配置
conda run -n Video-Tree-TRM python scripts/build_trees_batch.py \\
--video-dir data/videomme/videos \\
--config config/videomme.yaml \\
--jobs 3
"""
from __future__ import annotations
import argparse
import sys
import time
from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor
from concurrent.futures import wait as cfwait
from pathlib import Path
from typing import Dict, List, Optional, Tuple
# 项目根目录加入 sys.path
sys.path.insert(0, str(Path(__file__).parent.parent))
from video_tree_trm.config import Config
from video_tree_trm.llm_client import LLMClient
from video_tree_trm.video_tree_builder import VideoTreeBuilder
from utils.logger_system import log_msg
def _build_one(
video_path: Path,
cfg: Config,
) -> Tuple[str, bool, str]:
"""构建单个视频的 TreeIndex 并保存 JSON。
参数:
video_path: 视频文件绝对路径。
cfg: 已加载的配置对象（由主进程共享，仅读取）。
返回:
(stem, success, message) 三元组。
实现细节:
每次调用独立初始化 LLMClient（避免多线程共享同一 httpx.Client 内部状态），
使用 VideoTreeBuilder.build() 内部的异步事件循环（L2→L3→L1 链式并发）。
"""
stem = video_path.stem
try:
progress_dir = Path(cfg.tree.cache_dir) / "progress"
progress_path = progress_dir / f"{stem}.json"
if progress_path.is_file():
log_msg("INFO", "检测到中间进度，启用断点续跑", stem=stem, progress_path=str(progress_path))
# 每线程独立 LLMClient（httpx.Client 线程安全，但独立更稳健）
vlm = LLMClient(cfg.vlm)
builder = VideoTreeBuilder(vlm, cfg.tree)
tree = builder.build(str(video_path))
# 保存 JSON
cache_dir = Path(cfg.tree.cache_dir)
cache_dir.mkdir(parents=True, exist_ok=True)
out_path = cache_dir / f"{stem}_video.json"
tree.save_json(str(out_path))
l1 = len(tree.roots)
l2 = sum(len(r.children) for r in tree.roots)
l3 = sum(len(l2n.children) for r in tree.roots for l2n in r.children)
msg = f"L1={l1} L2={l2} L3={l3} → {out_path}"
log_msg("INFO", "视频建树完成", stem=stem, l1=l1, l2=l2, l3=l3)
return stem, True, msg
except Exception as e: # noqa: BLE001
log_msg("ERROR", "视频建树失败", stem=stem, error=str(e))
return stem, False, str(e)
def main() -> None:
"""批量建树主函数：视频间 ThreadPoolExecutor 并行 + 视频内异步事件循环。"""
parser = argparse.ArgumentParser(description="批量视频建树（视频间并行）")
parser.add_argument(
"--video-dir",
default="data/videomme/videos",
help="MP4 视频目录（默认: data/videomme/videos）",
)
parser.add_argument(
"--config",
default="config/videomme.yaml",
help="配置文件路径（默认: config/videomme.yaml）",
)
parser.add_argument(
"--jobs",
type=int,
default=1,
help="视频间并行数（默认: 1，每路视频内 asyncio Semaphore(concurrency) 并发 VLM）",
)
args = parser.parse_args()
# Phase 1: 加载配置（所有线程共享，只读）
cfg = Config.load(args.config)
log_msg(
"INFO",
"批量建树配置",
video_dir=args.video_dir,
cache_dir=cfg.tree.cache_dir,
jobs=args.jobs,
intra_concurrency=cfg.tree.concurrency,
)
# Phase 2: 扫描视频 + 过滤已有树
video_dir = Path(args.video_dir)
assert video_dir.is_dir(), f"视频目录不存在: {video_dir}"
all_videos: List[Path] = sorted(video_dir.glob("*.mp4"))
cache_dir = Path(cfg.tree.cache_dir)
progress_dir = cache_dir / "progress"
with_progress: List[Path] = []
without_progress: List[Path] = []
for v in all_videos:
if (cache_dir / f"{v.stem}_video.json").exists():
continue
progress_path = progress_dir / f"{v.stem}.json"
if progress_path.is_file():
with_progress.append(v)
else:
without_progress.append(v)
pending: List[Path] = with_progress + without_progress
skipped = len(all_videos) - len(pending)
print(f"\n===== 批量建树 =====")
print(f" 总视频数: {len(all_videos)}")
print(f" 已跳过(已建): {skipped}")
print(f" 待处理: {len(pending)}")
print(f" 视频间并行: {args.jobs}")
print(f" 视频内并发: {cfg.tree.concurrency}")
print(f" 输出目录: {cache_dir}\n")
if not pending:
print("所有视频均已建树，无需处理。")
return
# Phase 3: 异步视频间并行（ThreadPoolExecutor + FIRST_COMPLETED 事件循环）
built: List[str] = []
failed: List[Tuple[str, str]] = []
start_total = time.time()
pending_futures: Dict[Future, Path] = {}
pending_queue = list(pending) # 待提交队列
pool = ThreadPoolExecutor(max_workers=args.jobs)
# 初始填满 jobs 个任务
while pending_queue and len(pending_futures) < args.jobs:
video = pending_queue.pop(0)
fut = pool.submit(_build_one, video, cfg)
pending_futures[fut] = video
print(f"[提交] {video.stem} ({len(built)+len(failed)+len(pending_futures)}/{len(pending)})")
# 事件循环：完成一个，补一个
done_count = 0
while pending_futures:
done, _ = cfwait(list(pending_futures), return_when=FIRST_COMPLETED)
for fut in done:
video = pending_futures.pop(fut)
stem, success, msg = fut.result()
done_count += 1
elapsed = time.time() - start_total
if success:
built.append(stem)
print(f"[OK {done_count}/{len(pending)}] {stem} {msg} (累计 {elapsed:.0f}s)")
else:
failed.append((stem, msg))
print(f"[FAIL {done_count}/{len(pending)}] {stem} {msg}")
# 补充下一个任务（非阻塞提交）
if pending_queue:
next_video = pending_queue.pop(0)
next_fut = pool.submit(_build_one, next_video, cfg)
pending_futures[next_fut] = next_video
print(f"[提交] {next_video.stem} ({done_count+len(pending_futures)}/{len(pending)})")
pool.shutdown(wait=False)
# Phase 4: 汇总
total_elapsed = time.time() - start_total
print(f"\n===== 汇总 =====")
print(f" 成功: {len(built)}")
print(f" 失败: {len(failed)}")
print(f" 跳过: {skipped}")
print(f" 总耗时: {total_elapsed:.1f}s")
if failed:
print("\n 失败列表:")
for stem, err in failed:
print(f" {stem}: {err}")
log_msg(
"INFO",
"批量建树完成",
built=len(built),
failed=len(failed),
skipped=skipped,
elapsed_s=round(total_elapsed, 1),
)
if __name__ == "__main__":
main()
reference/scripts/build_trees_from_mp4.sh
Executable
+87
View File
@@ -0,0 +1,87 @@
#!/usr/bin/env bash
# =============================================================================
# build_trees_from_mp4.sh — 从本地 MP4 批量建树（JSON 输出，无 embedding）
#
# 用法:
# bash scripts/build_trees_from_mp4.sh
# bash scripts/build_trees_from_mp4.sh data/videomme/videos/TGom0uiW130.mp4 # 单文件
#
# 环境变量:
# VIDEO_DIR — MP4 目录（默认: <PROJECT_ROOT>/data/videomme/videos）
# CONFIG — 配置文件（默认: config/videomme.yaml）
# =============================================================================
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
VIDEO_DIR="${VIDEO_DIR:-${PROJECT_ROOT}/data/videomme/videos}"
CONFIG="${CONFIG:-${PROJECT_ROOT}/config/videomme.yaml}"
TREE_DIR="${PROJECT_ROOT}/data/videomme/trees"
LOG_DIR="${PROJECT_ROOT}/data/videomme/logs"
LOG="${LOG_DIR}/mp4_build_$(date +%Y%m%d_%H%M%S).log"
FAILED="${LOG_DIR}/failed_mp4_builds.txt"
mkdir -p "${TREE_DIR}" "${LOG_DIR}"
# 绕过代理
export NO_PROXY="${NO_PROXY:+${NO_PROXY},}100.83.164.94"
export no_proxy="${no_proxy:+${no_proxy},}100.83.164.94"
# 激活 conda 环境
CONDA_BASE="$(conda info --base 2>/dev/null || echo "${HOME}/miniconda3")"
# shellcheck source=/dev/null
source "${CONDA_BASE}/etc/profile.d/conda.sh"
conda activate Video-Tree-TRM
cd "${PROJECT_ROOT}"
# 确定待处理文件列表
if [[ $# -gt 0 ]]; then
MP4_FILES=("$@")
else
mapfile -t MP4_FILES < <(find "${VIDEO_DIR}" -name "*.mp4" | sort)
fi
TOTAL=${#MP4_FILES[@]}
BUILT=0; SKIP=0; FAIL=0
echo "[$(date)] ===== MP4 本地建树开始 =====" | tee "${LOG}"
echo "[$(date)] 待处理: ${TOTAL} 个视频" | tee -a "${LOG}"
echo "[$(date)] 配置: ${CONFIG}" | tee -a "${LOG}"
for MP4 in "${MP4_FILES[@]}"; do
[[ -z "${MP4}" ]] && continue
STEM="$(basename "${MP4}" .mp4)"
CACHE="${TREE_DIR}/${STEM}_video.json"
# 缓存命中跳过
if [[ -f "${CACHE}" ]]; then
SKIP=$((SKIP+1))
echo "[$(date)] [SKIP] ${STEM}" | tee -a "${LOG}"
continue
fi
echo "[$(date)] [BUILD] ${STEM} file=${MP4}" | tee -a "${LOG}"
if python -u main.py index \
--source "${MP4}" \
--modality video \
--config "${CONFIG}" \
>> "${LOG}" 2>&1; then
BUILT=$((BUILT+1))
echo "[$(date)] [OK] ${STEM}" | tee -a "${LOG}"
else
FAIL=$((FAIL+1))
echo "[$(date)] [FAIL] ${STEM}" | tee -a "${LOG}"
echo "${STEM} ${MP4}" >> "${FAILED}"
fi
done
TREE_COUNT=$(find "${TREE_DIR}" -name "*_video.json" 2>/dev/null | wc -l)
echo "" | tee -a "${LOG}"
echo "[$(date)] ===== 汇总 =====" | tee -a "${LOG}"
echo "[$(date)] 本次新建: ${BUILT} 跳过: ${SKIP} 失败: ${FAIL}" | tee -a "${LOG}"
echo "[$(date)] 树索引总数: ${TREE_COUNT}" | tee -a "${LOG}"
[[ ${FAIL} -gt 0 ]] && echo "[$(date)] 失败列表: ${FAILED}" | tee -a "${LOG}"
echo "[$(date)] 日志: ${LOG}" | tee -a "${LOG}"
reference/scripts/build_trees_from_urls.sh
Executable
+125
View File
@@ -0,0 +1,125 @@
#!/usr/bin/env bash
# =============================================================================
# build_trees_from_urls.sh — 直接从 YouTube URL 批量建树（不下载视频）
#
# 用法:
# # 全量（300 个长视频）
# bash scripts/build_trees_from_urls.sh
#
# # 只处理前 N 个
# head -10 data/videomme/metadata/long_videos.jsonl \
# | bash scripts/build_trees_from_urls.sh --stdin
#
# 环境变量:
# DATA_DIR — 数据根目录（默认: <PROJECT_ROOT>/data/videomme）
# CONFIG — 配置文件路径（默认: config/videomme.yaml）
#
# 特性:
# - 自动激活 Video-Tree-TRM conda 环境
# - 缓存命中跳过（trees/{youtube_id}_video.pkl 已存在则跳过）
# - 断点续传（重复运行安全）
# - 失败记录写入 logs/failed_url_builds.txt
# =============================================================================
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
# ---------------------------------------------------------------------------
# 0. 全局配置
# ---------------------------------------------------------------------------
DATA_DIR="${DATA_DIR:-${PROJECT_ROOT}/data/videomme}"
CONFIG="${CONFIG:-${PROJECT_ROOT}/config/videomme.yaml}"
JSONL="${DATA_DIR}/metadata/long_videos.jsonl"
TREE_DIR="${DATA_DIR}/trees"
LOG_DIR="${DATA_DIR}/logs"
LOG="${LOG_DIR}/url_build_$(date +%Y%m%d_%H%M%S).log"
FAILED="${LOG_DIR}/failed_url_builds.txt"
mkdir -p "${TREE_DIR}" "${LOG_DIR}"
# ---------------------------------------------------------------------------
# 0.5 绕过代理：GPU 内网地址直连，不经过 http_proxy
# ---------------------------------------------------------------------------
export NO_PROXY="${NO_PROXY:+${NO_PROXY},}100.83.164.94"
export no_proxy="${no_proxy:+${no_proxy},}100.83.164.94"
# ---------------------------------------------------------------------------
# 1. 激活 conda 环境
# ---------------------------------------------------------------------------
CONDA_BASE="$(conda info --base 2>/dev/null || echo "${HOME}/miniconda3")"
# shellcheck source=/dev/null
source "${CONDA_BASE}/etc/profile.d/conda.sh"
conda activate Video-Tree-TRM
cd "${PROJECT_ROOT}"
# ---------------------------------------------------------------------------
# 2. 解析参数
# ---------------------------------------------------------------------------
STDIN_MODE=false
for arg in "$@"; do
[[ "${arg}" == "--stdin" ]] && STDIN_MODE=true
done
# ---------------------------------------------------------------------------
# 3. 读取输入
# ---------------------------------------------------------------------------
if [[ "${STDIN_MODE}" == true ]]; then
# 从 stdin 读取（支持 head -N | bash ... --stdin）
INPUT_DATA="$(cat)"
else
INPUT_DATA="$(cat "${JSONL}")"
fi
TOTAL=$(echo "${INPUT_DATA}" | wc -l)
BUILT=0; SKIP=0; FAIL=0
echo "[$(date)] ===== URL 流式建树开始 =====" | tee "${LOG}"
echo "[$(date)] 待处理: ${TOTAL} 个长视频" | tee -a "${LOG}"
echo "[$(date)] 配置: ${CONFIG}" | tee -a "${LOG}"
# ---------------------------------------------------------------------------
# 4. 主循环
# ---------------------------------------------------------------------------
while IFS= read -r line; do
[[ -z "${line}" ]] && continue
YID=$(python -c "import sys,json; print(json.loads(sys.argv[1])['youtube_id'])" "${line}")
URL=$(python -c "import sys,json; print(json.loads(sys.argv[1])['url'])" "${line}")
CACHE="${TREE_DIR}/${YID}_video.json"
# 缓存命中跳过
if [[ -f "${CACHE}" ]]; then
SKIP=$((SKIP+1))
echo "[$(date)] [SKIP] ${YID}" | tee -a "${LOG}"
continue
fi
echo "[$(date)] [BUILD] ${YID} url=${URL}" | tee -a "${LOG}"
if python main.py index \
--source "${URL}" \
--modality video \
--config "${CONFIG}" \
>> "${LOG}" 2>&1; then
BUILT=$((BUILT+1))
echo "[$(date)] [OK] ${YID}" | tee -a "${LOG}"
else
FAIL=$((FAIL+1))
echo "[$(date)] [FAIL] ${YID}" | tee -a "${LOG}"
echo "${YID} ${URL}" >> "${FAILED}"
fi
done <<< "${INPUT_DATA}"
# ---------------------------------------------------------------------------
# 5. 汇总
# ---------------------------------------------------------------------------
TREE_COUNT=$(find "${TREE_DIR}" -name "*_video.json" 2>/dev/null | wc -l)
echo "" | tee -a "${LOG}"
echo "[$(date)] ===== 汇总 =====" | tee -a "${LOG}"
echo "[$(date)] 本次新建: ${BUILT} 跳过: ${SKIP} 失败: ${FAIL}" | tee -a "${LOG}"
echo "[$(date)] 树索引总数: ${TREE_COUNT} / ${TOTAL}" | tee -a "${LOG}"
[[ ${FAIL} -gt 0 ]] && echo "[$(date)] 失败列表: ${FAILED}" | tee -a "${LOG}"
echo "[$(date)] 日志: ${LOG}" | tee -a "${LOG}"
reference/scripts/build_videomme_trees.sh
Executable
+320
View File
@@ -0,0 +1,320 @@
#!/usr/bin/env bash
# =============================================================================
# build_videomme_trees.sh — VideoMME 长视频数据预处理：下载 + 建树
# =============================================================================
# 功能:
# 1. 初始化目录结构 (/data/videomme/...)
# 2. 激活 Conda 环境 (Video-Tree-TRM)
# 3. 安装必要工具 (yt-dlp, datasets)
# 4. 从 HuggingFace 下载 VideoMME 元数据，提取长视频列表
# 5. 用 yt-dlp 批量下载长视频（断点续传，跳过已下载）
# 6. 为每个视频调用 main.py index 建树（跳过已缓存）
# 7. 汇总日志
#
# 使用方式:
# cd /home/undergraduate/Video-Tree-TRM
# bash scripts/build_videomme_trees.sh
#
# 可选环境变量覆盖:
# DATA_DIR=/other/path bash scripts/build_videomme_trees.sh
# WORKERS=4 bash scripts/build_videomme_trees.sh # 并行建树进程数
#
# 断点续传:
# 重复运行完全安全 —— yt-dlp 跳过已下载文件，main.py 跳过缓存命中的树。
# =============================================================================
set -euo pipefail
# ---------------------------------------------------------------------------
# 0. 全局配置
# ---------------------------------------------------------------------------
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
CONDA_ENV="${CONDA_ENV:-Video-Tree-TRM}"
DATA_DIR="${DATA_DIR:-${PROJECT_ROOT}/data/videomme}"
VIDEO_DIR="${DATA_DIR}/videos"
META_DIR="${DATA_DIR}/metadata"
TREE_DIR="${DATA_DIR}/trees"
LOG_DIR="${DATA_DIR}/logs"
CKPT_DIR="${DATA_DIR}/checkpoints"
CONFIG_YAML="${PROJECT_ROOT}/config/videomme.yaml"
ENV_FILE="${PROJECT_ROOT}/.env"
META_SCRIPT="${SCRIPT_DIR}/_download_meta.py"
WORKERS="${WORKERS:-1}" # 并行建树进程数（默认串行，保护 API 速率）
MIN_DURATION="${MIN_DURATION:-1800}" # 长视频最短时长（秒）
MAX_DURATION="${MAX_DURATION:-3600}" # 长视频最长时长（秒）
YTDLP_RATE="${YTDLP_RATE:-500K}" # yt-dlp 下载限速（防封）
YTDLP_RETRIES="${YTDLP_RETRIES:-5}" # yt-dlp 重试次数
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
MAIN_LOG="${LOG_DIR}/build_${TIMESTAMP}.log"
# 颜色输出
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m'
info() { echo -e "${GREEN}[INFO]${NC} $(date '+%H:%M:%S') $*" | tee -a "${MAIN_LOG}"; }
warn() { echo -e "${YELLOW}[WARN]${NC} $(date '+%H:%M:%S') $*" | tee -a "${MAIN_LOG}"; }
error() { echo -e "${RED}[ERROR]${NC} $(date '+%H:%M:%S') $*" | tee -a "${MAIN_LOG}"; }
# ---------------------------------------------------------------------------
# 1. 创建目录结构
# ---------------------------------------------------------------------------
info "=== Step 1: 初始化目录结构 ==="
mkdir -p "${VIDEO_DIR}" "${META_DIR}" "${TREE_DIR}" "${LOG_DIR}" "${CKPT_DIR}"
info "数据目录已就绪: ${DATA_DIR}"
info " videos/ → ${VIDEO_DIR}"
info " metadata/ → ${META_DIR}"
info " trees/ → ${TREE_DIR}"
info " logs/ → ${LOG_DIR}"
info " checkpoints/ → ${CKPT_DIR}"
# ---------------------------------------------------------------------------
# 2. 激活 Conda 环境
# ---------------------------------------------------------------------------
info "=== Step 2: 激活 Conda 环境 (${CONDA_ENV}) ==="
# 找到 conda 初始化脚本
CONDA_BASE="$(conda info --base 2>/dev/null || echo "")"
if [[ -z "${CONDA_BASE}" ]]; then
error "未找到 conda，请确保 conda 已安装并在 PATH 中"
exit 1
fi
# shellcheck source=/dev/null
source "${CONDA_BASE}/etc/profile.d/conda.sh"
conda activate "${CONDA_ENV}"
info "已激活环境: $(conda info --envs | grep '*' | awk '{print $1}')"
info "Python 路径: $(which python)"
# ---------------------------------------------------------------------------
# 3. 安装必要工具
# ---------------------------------------------------------------------------
info "=== Step 3: 安装必要工具 ==="
pip install --quiet --upgrade yt-dlp datasets
info "yt-dlp 版本: $(yt-dlp --version)"
python -c "import datasets; print(f'datasets 版本: {datasets.__version__}')"
# ---------------------------------------------------------------------------
# 4. 下载 VideoMME 元数据，提取长视频列表
# ---------------------------------------------------------------------------
info "=== Step 4: 下载 VideoMME 元数据 ==="
LONG_VIDEOS_JSONL="${META_DIR}/long_videos.jsonl"
LONG_QA_JSONL="${META_DIR}/long_videos_qa.jsonl"
if [[ -f "${LONG_VIDEOS_JSONL}" ]]; then
EXISTING_COUNT=$(wc -l < "${LONG_VIDEOS_JSONL}")
warn "元数据已存在 (${EXISTING_COUNT} 条长视频)，跳过下载。如需重新下载，删除 ${LONG_VIDEOS_JSONL}"
else
# 配置 HuggingFace 镜像（国内加速，如已能直连可注释掉）
export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
info "HuggingFace 端点: ${HF_ENDPOINT}"
python "${META_SCRIPT}" \
--meta-dir "${META_DIR}" \
--min-duration "${MIN_DURATION}" \
--max-duration "${MAX_DURATION}" \
2>&1 | tee -a "${MAIN_LOG}"
fi
# 确认文件存在
if [[ ! -f "${LONG_VIDEOS_JSONL}" ]]; then
error "元数据文件不存在，元数据下载失败: ${LONG_VIDEOS_JSONL}"
exit 1
fi
TOTAL_VIDEOS=$(wc -l < "${LONG_VIDEOS_JSONL}")
info "长视频总数: ${TOTAL_VIDEOS}"
# ---------------------------------------------------------------------------
# 5. 批量下载长视频（yt-dlp，断点续传）
# ---------------------------------------------------------------------------
info "=== Step 5: 批量下载长视频 ==="
info "下载目录: ${VIDEO_DIR}"
info "限速: ${YTDLP_RATE}，重试次数: ${YTDLP_RETRIES}"
DOWNLOAD_LOG="${LOG_DIR}/download_${TIMESTAMP}.log"
FAILED_DOWNLOADS="${LOG_DIR}/failed_downloads_${TIMESTAMP}.txt"
DOWNLOAD_COUNT=0
SKIP_COUNT=0
FAIL_COUNT=0
while IFS= read -r line; do
VIDEO_ID=$(echo "${line}" | python -c "import sys,json; d=json.load(sys.stdin); print(d.get('youtube_id') or d.get('video_id',''))")
URL=$(echo "${line}" | python -c "import sys,json; d=json.load(sys.stdin); print(d.get('url',''))")
if [[ -z "${VIDEO_ID}" || -z "${URL}" ]]; then
warn "跳过无效记录: ${line:0:80}"
continue
fi
# 检查是否已下载（任意格式均算）
EXISTING_FILE=$(find "${VIDEO_DIR}" -name "${VIDEO_ID}.*" -type f 2>/dev/null | head -1)
if [[ -n "${EXISTING_FILE}" ]]; then
SKIP_COUNT=$((SKIP_COUNT + 1))
continue
fi
info "[下载 ${DOWNLOAD_COUNT}/${TOTAL_VIDEOS}] ${VIDEO_ID}"
# yt-dlp 下载：优先 mp4，最高 720p（节省空间），单文件断点续传
if yt-dlp \
--output "${VIDEO_DIR}/%(id)s.%(ext)s" \
--format "bestvideo[ext=mp4][height<=720]+bestaudio[ext=m4a]/best[ext=mp4][height<=720]/best" \
--merge-output-format mp4 \
--retries "${YTDLP_RETRIES}" \
--rate-limit "${YTDLP_RATE}" \
--no-playlist \
--continue \
--no-overwrites \
--write-info-json \
--quiet \
"${URL}" \
>> "${DOWNLOAD_LOG}" 2>&1; then
DOWNLOAD_COUNT=$((DOWNLOAD_COUNT + 1))
info " ✓ 下载成功: ${VIDEO_ID}"
else
FAIL_COUNT=$((FAIL_COUNT + 1))
warn " ✗ 下载失败: ${VIDEO_ID} (${URL})"
echo "${VIDEO_ID} ${URL}" >> "${FAILED_DOWNLOADS}"
fi
done < "${LONG_VIDEOS_JSONL}"
info "下载汇总: 新下载=${DOWNLOAD_COUNT}, 跳过=${SKIP_COUNT}, 失败=${FAIL_COUNT}"
if [[ ${FAIL_COUNT} -gt 0 ]]; then
warn "失败列表已保存至: ${FAILED_DOWNLOADS}"
fi
# ---------------------------------------------------------------------------
# 6. 批量建树（main.py index，跳过缓存命中）
# ---------------------------------------------------------------------------
info "=== Step 6: 批量建树 ==="
info "项目根目录: ${PROJECT_ROOT}"
info "配置文件: ${CONFIG_YAML}"
info "并行进程数: ${WORKERS}"
BUILD_LOG="${LOG_DIR}/build_trees_${TIMESTAMP}.log"
FAILED_BUILDS="${LOG_DIR}/failed_builds_${TIMESTAMP}.txt"
BUILD_COUNT=0
BUILD_SKIP=0
BUILD_FAIL=0
# 构建函数（单个视频）
build_one_video() {
local video_path="$1"
local video_stem
video_stem="$(basename "${video_path%.*}")"
local cache_file="${TREE_DIR}/${video_stem}_video.pkl"
# 缓存命中则跳过（pipeline.py 内部也会检查，此处提前判断减少日志噪声）
if [[ -f "${cache_file}" ]]; then
echo "[SKIP] ${video_stem}"
return 0
fi
echo "[BUILD] ${video_stem} ← ${video_path}"
if conda run -n "${CONDA_ENV}" python "${PROJECT_ROOT}/main.py" \
index \
--source "${video_path}" \
--modality video \
--config "${CONFIG_YAML}" \
--env "${ENV_FILE}" \
>> "${BUILD_LOG}" 2>&1; then
echo "[OK] ${video_stem}"
return 0
else
echo "[FAIL] ${video_stem}"
echo "${video_path}" >> "${FAILED_BUILDS}"
return 1
fi
}
export -f build_one_video
export CONDA_ENV PROJECT_ROOT CONFIG_YAML ENV_FILE TREE_DIR BUILD_LOG FAILED_BUILDS
if [[ "${WORKERS}" -gt 1 ]]; then
# 并行模式：使用 GNU parallel
if ! command -v parallel &> /dev/null; then
warn "未找到 GNU parallel，降级为串行模式"
WORKERS=1
fi
fi
if [[ "${WORKERS}" -gt 1 ]]; then
info "并行建树 (jobs=${WORKERS})..."
find "${VIDEO_DIR}" -type f \( -name "*.mp4" -o -name "*.avi" -o -name "*.mkv" -o -name "*.webm" \) \
| parallel -j "${WORKERS}" --bar build_one_video {} \
2>&1 | tee -a "${MAIN_LOG}"
else
# 串行模式
while IFS= read -r -d '' video_path; do
video_stem="$(basename "${video_path%.*}")"
cache_file="${TREE_DIR}/${video_stem}_video.pkl"
if [[ -f "${cache_file}" ]]; then
BUILD_SKIP=$((BUILD_SKIP + 1))
continue
fi
info "[建树 $((BUILD_COUNT + BUILD_SKIP + 1))/${TOTAL_VIDEOS}] ${video_stem}"
if conda run -n "${CONDA_ENV}" python "${PROJECT_ROOT}/main.py" \
index \
--source "${video_path}" \
--modality video \
--config "${CONFIG_YAML}" \
--env "${ENV_FILE}" \
>> "${BUILD_LOG}" 2>&1; then
BUILD_COUNT=$((BUILD_COUNT + 1))
info " ✓ 建树成功: ${video_stem}"
else
BUILD_FAIL=$((BUILD_FAIL + 1))
warn " ✗ 建树失败: ${video_stem}"
echo "${video_path}" >> "${FAILED_BUILDS}"
fi
done < <(find "${VIDEO_DIR}" -type f \( -name "*.mp4" -o -name "*.avi" -o -name "*.mkv" -o -name "*.webm" \) -print0)
fi
# ---------------------------------------------------------------------------
# 7. 最终汇总
# ---------------------------------------------------------------------------
info "=== Step 7: 汇总 ==="
TREE_COUNT=$(find "${TREE_DIR}" -name "*_video.pkl" -type f 2>/dev/null | wc -l)
VIDEO_COUNT=$(find "${VIDEO_DIR}" -type f \( -name "*.mp4" -o -name "*.avi" -o -name "*.mkv" -o -name "*.webm" \) 2>/dev/null | wc -l)
info "======================================"
info " 长视频元数据数量: ${TOTAL_VIDEOS}"
info " 已下载视频数量: ${VIDEO_COUNT}"
info " 已完成树索引数量: ${TREE_COUNT}"
info " 本次新建树: ${BUILD_COUNT}"
info " 跳过（缓存命中）: ${BUILD_SKIP}"
info " 建树失败: ${BUILD_FAIL}"
info "======================================"
info "主日志: ${MAIN_LOG}"
info "下载日志: ${DOWNLOAD_LOG}"
info "建树日志: ${BUILD_LOG}"
if [[ ${BUILD_FAIL} -gt 0 ]]; then
warn "有 ${BUILD_FAIL} 个视频建树失败，详见: ${FAILED_BUILDS}"
warn "可重新运行脚本以续建失败项（自动跳过已缓存）"
fi
if [[ "${TREE_COUNT}" -ge "${TOTAL_VIDEOS}" ]]; then
info "✅ 所有长视频树索引已构建完成！"
else
warn "⚠ 树索引尚未全部完成 (${TREE_COUNT}/${TOTAL_VIDEOS})，可重新运行以续建"
fi
info "脚本完成。"
reference/scripts/download_videos.sh
Executable
+95
View File
@@ -0,0 +1,95 @@
#!/usr/bin/env bash
# =============================================================================
# download_videos.sh — 批量下载 VideoMME 长视频（MP4 含音轨）
#
# 用法:
# bash scripts/download_videos.sh # 下载全部未完成的视频
#
# 特性:
# - 已存在的 mp4 自动跳过（断点续传）
# - bestvideo+bestaudio 合并，保证有音轨
# - 并发数: 3（避免被 YouTube 限流）
# - 失败记录写入 logs/failed_downloads.txt
# =============================================================================
set -uo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
JSONL="${PROJECT_ROOT}/data/videomme/metadata/long_videos.jsonl"
VIDEO_DIR="${PROJECT_ROOT}/data/videomme/videos"
LOG_DIR="${PROJECT_ROOT}/data/videomme/logs"
LOG="${LOG_DIR}/download_$(date +%Y%m%d_%H%M%S).log"
FAILED_FILE="${LOG_DIR}/failed_downloads.txt"
CONCURRENT=3
mkdir -p "${VIDEO_DIR}" "${LOG_DIR}"
# 激活 conda 环境（yt-dlp、ffmpeg 在此环境中）
CONDA_BASE="$(conda info --base 2>/dev/null || echo "${HOME}/miniconda3")"
# shellcheck source=/dev/null
source "${CONDA_BASE}/etc/profile.d/conda.sh"
conda activate Video-Tree-TRM
echo "[$(date)] ===== 批量下载开始 =====" | tee "${LOG}"
echo "[$(date)] 视频目录: ${VIDEO_DIR}" | tee -a "${LOG}"
# 读取所有 youtube_id，写入临时文件供 xargs 读取
TMP_IDS=$(mktemp)
python3 -c "
import json
with open('${JSONL}') as f:
for line in f:
d = json.loads(line.strip())
print(d['youtube_id'])
" > "${TMP_IDS}"
TOTAL=$(wc -l < "${TMP_IDS}")
echo "[$(date)] 总计: ${TOTAL} 个视频，并发数: ${CONCURRENT}" | tee -a "${LOG}"
# 3 2.5Mb/s并发下载：每个 youtube_id 单独调用 yt-dlp
# 注意：直接在 xargs 中内联参数，不依赖 bash 数组导出
xargs -P "${CONCURRENT}" -I {} bash -c '
VIDEO_DIR="'"${VIDEO_DIR}"'"
LOG="'"${LOG}"'"
FAILED_FILE="'"${FAILED_FILE}"'"
YID="{}"
OUT="${VIDEO_DIR}/${YID}.mp4"
if [[ -f "${OUT}" && -s "${OUT}" ]]; then
echo "[$(date)] [SKIP] ${YID}" >> "${LOG}"
exit 0
fi
URL="https://www.youtube.com/watch?v=${YID}"
echo "[$(date)] [START] ${YID}" >> "${LOG}"
if yt-dlp \
--format "bestvideo[vcodec^=avc][ext=mp4]+bestaudio[ext=m4a]/bestvideo[ext=mp4]+bestaudio[ext=m4a]/bestvideo+bestaudio/best[ext=mp4]/best" \
--merge-output-format mp4 \
--output "${OUT}" \
--no-playlist \
--retries 5 \
--fragment-retries 5 \
--socket-timeout 60 \
--no-warnings \
--rate-limit 833K \
"${URL}" >> "${LOG}" 2>&1; then
echo "[$(date)] [OK] ${YID} size=$(du -sh "${OUT}" 2>/dev/null | cut -f1)" >> "${LOG}"
else
echo "[$(date)] [FAIL] ${YID}" >> "${LOG}"
echo "${YID}" >> "${FAILED_FILE}"
fi
' < "${TMP_IDS}"
rm -f "${TMP_IDS}"
TOTAL_MP4=$(find "${VIDEO_DIR}" -name "*.mp4" -size +0c | wc -l)
FAILED_COUNT=$(wc -l < "${FAILED_FILE}" 2>/dev/null || echo 0)
echo "" | tee -a "${LOG}"
echo "[$(date)] ===== 汇总 =====" | tee -a "${LOG}"
echo "[$(date)] 目录共 MP4: ${TOTAL_MP4} / ${TOTAL}" | tee -a "${LOG}"
[[ "${FAILED_COUNT}" -gt 0 ]] && echo "[$(date)] 失败数: ${FAILED_COUNT}，列表: ${FAILED_FILE}" | tee -a "${LOG}"
echo "[$(date)] 日志: ${LOG}" | tee -a "${LOG}"
reference/tests/conftest.py
+89
View File
@@ -0,0 +1,89 @@
"""
全局测试配置
============
- 代理修复: 从 .env 读取 EMBED_API_URL，将其中的 CGNAT 段 IP 加入 no_proxy。
- real_config fixture: 从真实 config/default.yaml + .env 加载 Config，供远程测试复用。
"""
from __future__ import annotations
import ipaddress
import os
import re
from pathlib import Path
import pytest
from dotenv import dotenv_values
# ---------------------------------------------------------------------------
# 代理修复: 将内网 API 地址加入 no_proxy（模块加载时立即执行）
# ---------------------------------------------------------------------------
_CGNAT_NETWORK = ipaddress.IPv4Network("100.64.0.0/10")
def _fix_no_proxy() -> None:
"""从 .env 读取 API URL，将 CGNAT 段 IP 加入 no_proxy / NO_PROXY。
httpx/urllib 的 no_proxy 不支持 CIDR 记法，必须逐个添加具体 IP。
"""
if not (os.environ.get("http_proxy") or os.environ.get("HTTP_PROXY")):
return
project_root = Path(__file__).resolve().parent.parent
env_file = project_root / ".env"
if not env_file.exists():
return
env_vars = dotenv_values(str(env_file))
# 收集所有需要绕过代理的 IP
hosts_to_add: list[str] = []
for url_key in ("EMBED_API_URL", "LLM_API_URL", "VLM_API_URL"):
url = env_vars.get(url_key, "")
if not url:
continue
match = re.match(r"https?://([^:/]+)", url)
if not match:
continue
host = match.group(1)
try:
if ipaddress.IPv4Address(host) in _CGNAT_NETWORK:
hosts_to_add.append(host)
except (ipaddress.AddressValueError, ValueError):
pass
if not hosts_to_add:
return
for var in ("no_proxy", "NO_PROXY"):
current = os.environ.get(var, "")
for host in hosts_to_add:
if host not in current:
separator = "," if current else ""
current = f"{current}{separator}{host}"
os.environ[var] = current
# 模块加载时执行代理修复
_fix_no_proxy()
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture(scope="session")
def real_config():
"""从真实 config/default.yaml + .env 加载 Config（session 级缓存）。
用于远程 API 测试，需要 .env 中配置有效的 API 密钥。
"""
from video_tree_trm.config import Config
project_root = Path(__file__).resolve().parent.parent
yaml_path = project_root / "config" / "default.yaml"
env_path = project_root / ".env"
return Config.load(str(yaml_path), env_path=str(env_path))
reference/tests/unit/test_config.py
+276
View File
@@ -0,0 +1,276 @@
"""
配置管理模块单元测试
====================
覆盖: YAML 加载、.env 覆盖、CLI 覆盖、缺字段报错、优先级验证、embed_dim 一致性。
"""
from __future__ import annotations
from pathlib import Path
import pytest
import yaml
from video_tree_trm.config import Config, TreeConfig
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
_FULL_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": {
"backend": "local",
"model_name": "test-model",
"embed_dim": 2560,
"device": "cpu",
"api_key": "",
"api_url": "",
},
"llm": {
"backend": "qwen",
"api_key": "yaml-llm-key",
"model": "qwen-plus",
"api_url": "https://example.com/llm",
"max_tokens": 256,
"temperature": 0.1,
},
"vlm": {
"backend": "qwen",
"api_key": "yaml-vlm-key",
"model": "qwen-vl-plus",
"api_url": "https://example.com/vlm",
"max_tokens": 256,
"temperature": 0.1,
},
"retriever": {
"embed_dim": 2560,
"num_heads": 4,
"L_layers": 2,
"L_cycles": 4,
"max_rounds": 5,
"ffn_expansion": 2.0,
"checkpoint": None,
},
"train": {
"lr": 1e-4,
"weight_decay": 1e-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",
},
}
@pytest.fixture()
def yaml_path(tmp_path: Path) -> Path:
"""创建完整配置的临时 YAML 文件。"""
p = tmp_path / "config" / "default.yaml"
p.parent.mkdir(parents=True, exist_ok=True)
with open(p, "w", encoding="utf-8") as f:
yaml.dump(_FULL_YAML, f, allow_unicode=True)
return p
@pytest.fixture()
def env_path(tmp_path: Path) -> Path:
"""创建临时 .env 文件。"""
p = tmp_path / ".env"
p.write_text(
"LLM_API_KEY=env-llm-key\n"
"LLM_MODEL=env-llm-model\n"
"LLM_API_URL=https://env.example.com/llm\n"
"VLM_API_KEY=env-vlm-key\n"
"VLM_MODEL=env-vlm-model\n"
"VLM_API_URL=https://env.example.com/vlm\n"
"EMBED_BACKEND=remote\n"
"EMBED_MODEL=env-embed-model\n"
"EMBED_API_KEY=env-embed-key\n"
"EMBED_API_URL=https://env.example.com/embed\n"
)
return p
# ---------------------------------------------------------------------------
# 测试: YAML 加载
# ---------------------------------------------------------------------------
class TestYAMLLoad:
"""YAML 基础加载测试。"""
def test_load_full_yaml(self, yaml_path: Path) -> None:
"""完整 YAML 应成功加载所有字段。"""
cfg = Config.load(
str(yaml_path), env_path=str(yaml_path.parent / ".env.nonexist")
)
assert isinstance(cfg.tree, TreeConfig)
assert cfg.tree.max_paragraphs_per_l2 == 5
assert cfg.tree.l1_segment_duration == 600.0
assert cfg.embed.embed_dim == 2560
assert cfg.retriever.checkpoint is None
assert cfg.train.dataset == "longbench"
def test_file_not_found(self, tmp_path: Path) -> None:
"""不存在的 YAML 应抛出 FileNotFoundError。"""
with pytest.raises(FileNotFoundError, match="配置文件不存在"):
Config.load(str(tmp_path / "nonexist.yaml"))
# ---------------------------------------------------------------------------
# 测试: 缺字段报错
# ---------------------------------------------------------------------------
class TestMissingField:
"""缺少必需字段时应抛出 TypeError。"""
def test_missing_tree_field(self, tmp_path: Path) -> None:
"""tree 节缺少字段应报 TypeError。"""
bad_yaml = _FULL_YAML.copy()
bad_yaml = {
k: (v.copy() if isinstance(v, dict) else v) for k, v in _FULL_YAML.items()
}
del bad_yaml["tree"]["cache_dir"]
p = tmp_path / "bad.yaml"
with open(p, "w") as f:
yaml.dump(bad_yaml, f)
with pytest.raises(TypeError):
Config.load(str(p), env_path=str(tmp_path / ".env.nonexist"))
def test_missing_section(self, tmp_path: Path) -> None:
"""缺少整个配置节应报 TypeError。"""
bad_yaml = {k: v for k, v in _FULL_YAML.items() if k != "train"}
p = tmp_path / "bad2.yaml"
with open(p, "w") as f:
yaml.dump(bad_yaml, f)
with pytest.raises(TypeError):
Config.load(str(p), env_path=str(tmp_path / ".env.nonexist"))
# ---------------------------------------------------------------------------
# 测试: .env 覆盖
# ---------------------------------------------------------------------------
class TestEnvOverride:
""".env 文件应覆盖 YAML 中的 api_key。"""
def test_env_overrides_api_keys(self, yaml_path: Path, env_path: Path) -> None:
"""api_key/model/api_url 应优先使用 .env 中的值。"""
cfg = Config.load(str(yaml_path), env_path=str(env_path))
assert cfg.llm.api_key == "env-llm-key"
assert cfg.llm.model == "env-llm-model"
assert cfg.llm.api_url == "https://env.example.com/llm"
assert cfg.vlm.api_key == "env-vlm-key"
assert cfg.vlm.model == "env-vlm-model"
assert cfg.vlm.api_url == "https://env.example.com/vlm"
def test_env_overrides_embed(self, yaml_path: Path, env_path: Path) -> None:
"""embed 相关字段应优先使用 .env 中的值。"""
cfg = Config.load(str(yaml_path), env_path=str(env_path))
assert cfg.embed.backend == "remote"
assert cfg.embed.model_name == "env-embed-model"
assert cfg.embed.api_key == "env-embed-key"
assert cfg.embed.api_url == "https://env.example.com/embed"
def test_yaml_fallback_when_no_env(self, yaml_path: Path) -> None:
"""无 .env 时应使用 YAML 中的值。"""
cfg = Config.load(
str(yaml_path), env_path=str(yaml_path.parent / ".env.nonexist")
)
assert cfg.llm.api_key == "yaml-llm-key"
assert cfg.vlm.api_key == "yaml-vlm-key"
# ---------------------------------------------------------------------------
# 测试: CLI 覆盖
# ---------------------------------------------------------------------------
class TestCLIOverride:
"""CLI args 应覆盖 YAML 和 .env 的值。"""
def test_cli_overrides_yaml(self, yaml_path: Path) -> None:
"""CLI 点路径覆盖应生效。"""
cfg = Config.load(
str(yaml_path),
cli_args={"retriever.num_heads": 8, "train.lr": 0.001},
env_path=str(yaml_path.parent / ".env.nonexist"),
)
assert cfg.retriever.num_heads == 8
assert cfg.train.lr == 0.001
def test_cli_overrides_env(self, yaml_path: Path, env_path: Path) -> None:
"""CLI 应覆盖 .env 中的 api_key。"""
cfg = Config.load(
str(yaml_path),
cli_args={"llm.api_key": "cli-key"},
env_path=str(env_path),
)
assert cfg.llm.api_key == "cli-key"
# ---------------------------------------------------------------------------
# 测试: 优先级
# ---------------------------------------------------------------------------
class TestPriority:
"""三层优先级: CLI > .env > YAML。"""
def test_full_priority_chain(self, yaml_path: Path, env_path: Path) -> None:
"""CLI > .env > YAML 的完整优先级链。"""
cfg = Config.load(
str(yaml_path),
cli_args={"llm.api_key": "cli-key"},
env_path=str(env_path),
)
# CLI 覆盖 .env
assert cfg.llm.api_key == "cli-key"
# .env 覆盖 YAML（vlm 未被 CLI 覆盖）
assert cfg.vlm.api_key == "env-vlm-key"
# ---------------------------------------------------------------------------
# 测试: embed_dim 一致性校验
# ---------------------------------------------------------------------------
class TestEmbedDimConsistency:
"""embed.embed_dim 与 retriever.embed_dim 必须一致。"""
def test_inconsistent_embed_dim(self, tmp_path: Path) -> None:
"""embed_dim 不一致应抛出 ValueError。"""
bad_yaml = {
k: (v.copy() if isinstance(v, dict) else v) for k, v in _FULL_YAML.items()
}
bad_yaml["retriever"] = bad_yaml["retriever"].copy()
bad_yaml["retriever"]["embed_dim"] = 512 # 与 embed.embed_dim=768 不一致
p = tmp_path / "bad_dim.yaml"
with open(p, "w") as f:
yaml.dump(bad_yaml, f)
with pytest.raises(ValueError, match="不一致"):
Config.load(str(p), env_path=str(tmp_path / ".env.nonexist"))
reference/tests/unit/test_embedding_persistence.py
+193
View File
@@ -0,0 +1,193 @@
"""
Embedding 持久化单元测试
=======================
测试 TreeIndex 的 embedding 序列化/反序列化功能。
"""
import json
import os
import tempfile
from pathlib import Path
import numpy as np
import pytest
from video_tree_trm.tree_index import (
L1Node,
L2Node,
L3Node,
IndexMeta,
TreeIndex,
_embed_to_str,
_embed_from_str,
)
class TestEmbeddingSerialization:
"""测试 embedding 序列化辅助函数。"""
def test_embed_to_str_and_back(self):
"""测试 base64 序列化/反序列化往返正确。"""
arr = np.random.randn(768).astype(np.float32)
s = _embed_to_str(arr)
recovered = _embed_from_str(s)
assert recovered is not None
assert recovered.dtype == np.float32
assert recovered.shape == (768,)
np.testing.assert_array_almost_equal(arr, recovered, decimal=6)
def test_embed_to_str_handles_none(self):
"""测试 None 输入返回 None。"""
assert _embed_to_str(None) is None
assert _embed_from_str(None) is None
assert _embed_from_str("") is None
class TestL3NodeEmbedding:
"""测试 L3Node embedding 序列化。"""
def test_l3_to_dict_without_embedding(self):
"""测试不带 embedding 序列化。"""
node = L3Node(
id="l3_0",
description="测试描述",
timestamp=1.0,
frame_path="frame.jpg",
raw_content="原始内容",
)
d = {
"id": node.id,
"description": node.description,
"timestamp": node.timestamp,
"frame_path": node.frame_path,
"raw_content": node.raw_content,
}
# 无 embedding 字段
assert "embedding" not in d
def test_l3_embedding_roundtrip(self):
"""测试 L3 embedding 序列化往返。"""
embed = np.random.randn(768).astype(np.float32)
s = _embed_to_str(embed)
recovered = _embed_from_str(s)
np.testing.assert_array_almost_equal(embed, recovered, decimal=6)
class TestTreeIndexEmbedding:
"""测试 TreeIndex 完整序列化/反序列化。"""
def test_save_load_without_embedding(self):
"""测试不带 embedding 保存/加载（向后兼容）。"""
with tempfile.TemporaryDirectory() as tmpdir:
path = os.path.join(tmpdir, "test.json")
# 创建简单树
l3 = L3Node(id="l3_0", description="L3 描述")
l2 = L2Node(id="l2_0", description="L2 描述", children=[l3])
l1 = L1Node(id="l1_0", summary="L1 摘要", children=[l2])
meta = IndexMeta(source_path="test.mp4", modality="video")
tree = TreeIndex(metadata=meta, roots=[l1])
# 保存（不含 embedding）
tree.save_json(path, include_embedding=False)
# 加载
loaded = TreeIndex.load_json(path)
assert len(loaded.roots) == 1
assert loaded.roots[0].summary == "L1 摘要"
assert not loaded.is_embedded # 无 embedding
def test_save_load_with_embedding(self):
"""测试带 embedding 保存/加载。"""
with tempfile.TemporaryDirectory() as tmpdir:
path = os.path.join(tmpdir, "test_embed.json")
# 创建带 embedding 的树
embed_l1 = np.random.randn(768).astype(np.float32)
embed_l2 = np.random.randn(768).astype(np.float32)
embed_l3 = np.random.randn(768).astype(np.float32)
l3 = L3Node(id="l3_0", description="L3 描述", embedding=embed_l3)
l2 = L2Node(id="l2_0", description="L2 描述", embedding=embed_l2, children=[l3])
l1 = L1Node(id="l1_0", summary="L1 摘要", embedding=embed_l1, children=[l2])
meta = IndexMeta(
source_path="test.mp4",
modality="video",
embed_model="test-model",
embed_dim=768,
)
tree = TreeIndex(metadata=meta, roots=[l1])
# 保存（含 embedding）
tree.save_json(path, include_embedding=True)
# 验证 JSON 中有 embedding 字段
with open(path, encoding="utf-8") as f:
data = json.load(f)
assert "embedding" in data["roots"][0]
assert data["metadata"]["embed_model"] == "test-model"
# 加载
loaded = TreeIndex.load_json(path)
assert loaded.is_embedded
np.testing.assert_array_almost_equal(
loaded.roots[0].embedding, embed_l1, decimal=6
)
np.testing.assert_array_almost_equal(
loaded.roots[0].children[0].embedding, embed_l2, decimal=6
)
np.testing.assert_array_almost_equal(
loaded.roots[0].children[0].children[0].embedding, embed_l3, decimal=6
)
def test_load_old_format_compatible(self):
"""测试加载旧格式（无 embedding 字段）JSON 兼容。"""
with tempfile.TemporaryDirectory() as tmpdir:
path = os.path.join(tmpdir, "old_format.json")
# 手动创建旧格式 JSON（无 embedding 字段）
old_data = {
"metadata": {
"source_path": "test.mp4",
"modality": "video",
"created_at": "2024-01-01T00:00:00",
},
"roots": [
{
"id": "l1_0",
"summary": "L1 摘要",
"time_range": None,
"children": [
{
"id": "l2_0",
"description": "L2 描述",
"time_range": None,
"children": [
{
"id": "l3_0",
"description": "L3 描述",
"timestamp": 1.0,
"frame_path": None,
"raw_content": None,
}
],
}
],
}
],
}
with open(path, "w", encoding="utf-8") as f:
json.dump(old_data, f)
# 加载
loaded = TreeIndex.load_json(path)
assert len(loaded.roots) == 1
assert not loaded.is_embedded # 无 embedding，向后兼容
if __name__ == "__main__":
pytest.main([__file__, "-v"])
reference/tests/unit/test_embeddings.py
+199
View File
@@ -0,0 +1,199 @@
"""
嵌入服务模块单元测试
====================
覆盖: 本地模式（embed/embed_tensor/归一化/dim/冻结）、远程模式（真实 API 调用）。
本地测试使用轻量模型 all-MiniLM-L6-v2 (dim=384) 加速。
远程测试使用 .env 中配置的真实 API（需有效密钥）。
"""
from __future__ import annotations
import numpy as np
import pytest
import torch
from video_tree_trm.config import Config, EmbedConfig
from video_tree_trm.embeddings import EmbeddingModel
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
_LOCAL_CONFIG = EmbedConfig(
backend="local",
model_name="all-MiniLM-L6-v2",
embed_dim=384,
device="cpu",
api_key="",
api_url="",
)
@pytest.fixture(scope="module")
def local_model() -> EmbeddingModel:
"""本地嵌入模型（模块级缓存，避免重复加载）。"""
return EmbeddingModel(_LOCAL_CONFIG)
@pytest.fixture(scope="module")
def remote_model(real_config: Config) -> EmbeddingModel:
"""远程嵌入模型（模块级缓存），使用真实 API。"""
return EmbeddingModel(real_config.embed)
# ---------------------------------------------------------------------------
# 测试: 本地模式
# ---------------------------------------------------------------------------
class TestLocalEmbed:
"""本地 sentence-transformers 后端测试。"""
def test_embed_single_text(self, local_model: EmbeddingModel) -> None:
"""单条文本，验证形状 [1, D]。"""
result = local_model.embed("hello world")
assert isinstance(result, np.ndarray)
assert result.shape == (1, 384)
def test_embed_batch(self, local_model: EmbeddingModel) -> None:
"""批量文本，验证形状 [N, D]。"""
texts = ["hello", "world", "test"]
result = local_model.embed(texts)
assert result.shape == (3, 384)
def test_embed_tensor_type(self, local_model: EmbeddingModel) -> None:
"""验证 embed_tensor() 返回 Tensor。"""
result = local_model.embed_tensor("hello")
assert isinstance(result, torch.Tensor)
assert result.shape == (1, 384)
assert result.dtype == torch.float32
def test_embeddings_normalized(self, local_model: EmbeddingModel) -> None:
"""验证 L2 范数 ≈ 1.0。"""
result = local_model.embed(["hello", "world"])
norms = np.linalg.norm(result, axis=1)
np.testing.assert_allclose(norms, 1.0, atol=1e-5)
def test_dim_property(self, local_model: EmbeddingModel) -> None:
"""验证 dim 属性一致。"""
assert local_model.dim == 384
def test_local_model_frozen(self, local_model: EmbeddingModel) -> None:
"""本地模式参数 requires_grad=False。"""
for param in local_model._model.parameters():
assert not param.requires_grad
# ---------------------------------------------------------------------------
# 测试: 远程模式（真实 API）
# ---------------------------------------------------------------------------
class TestRemoteEmbed:
"""远程 OpenAI 兼容 API 后端测试（真实 API 调用）。"""
def test_remote_embed_single(self, remote_model: EmbeddingModel) -> None:
"""单条中文文本，验证形状 [1, D]。"""
result = remote_model.embed("你好世界")
assert isinstance(result, np.ndarray)
assert result.shape == (1, remote_model.dim)
def test_remote_embed_batch(self, remote_model: EmbeddingModel) -> None:
"""5 条文本，验证形状 [5, D]。"""
texts = ["机器学习", "深度学习", "自然语言处理", "计算机视觉", "强化学习"]
result = remote_model.embed(texts)
assert result.shape == (5, remote_model.dim)
def test_remote_embed_normalized(self, remote_model: EmbeddingModel) -> None:
"""验证 L2 范数 ≈ 1.0。"""
result = remote_model.embed(["归一化测试文本", "另一段测试文本"])
norms = np.linalg.norm(result, axis=1)
np.testing.assert_allclose(norms, 1.0, atol=1e-5)
def test_remote_embed_tensor(self, remote_model: EmbeddingModel) -> None:
"""验证返回 torch.Tensor + float32。"""
result = remote_model.embed_tensor("张量测试")
assert isinstance(result, torch.Tensor)
assert result.dtype == torch.float32
assert result.shape == (1, remote_model.dim)
def test_remote_dim(self, remote_model: EmbeddingModel, real_config: Config) -> None:
"""验证 dim 属性与配置一致。"""
assert remote_model.dim == real_config.embed.embed_dim
def test_remote_semantic_similarity(self, remote_model: EmbeddingModel) -> None:
"""语义相近文本相似度 > 语义无关文本。"""
# 语义相近对
vec_cat = remote_model.embed("猫咪在沙发上睡觉")
vec_kitten = remote_model.embed("小猫趴在沙发上打盹")
# 语义无关
vec_math = remote_model.embed("二次方程的求解公式")
sim_close = float(np.dot(vec_cat[0], vec_kitten[0]))
sim_far = float(np.dot(vec_cat[0], vec_math[0]))
assert sim_close > sim_far, (
f"语义相近对相似度 ({sim_close:.4f}) 应大于语义无关对 ({sim_far:.4f})"
)
# ---------------------------------------------------------------------------
# 测试: 配置校验
# ---------------------------------------------------------------------------
class TestConfigValidation:
"""配置校验测试。"""
def test_invalid_backend(self) -> None:
"""无效 backend 应抛出 ValueError。"""
config = EmbedConfig(
backend="invalid",
model_name="test",
embed_dim=384,
device="cpu",
api_key="",
api_url="",
)
with pytest.raises(ValueError, match="backend"):
EmbeddingModel(config)
def test_remote_missing_api_key(self) -> None:
"""远程模式缺少 api_key 应抛出 ValueError。"""
config = EmbedConfig(
backend="remote",
model_name="test",
embed_dim=384,
device="cpu",
api_key="",
api_url="http://localhost:8080/v1",
)
with pytest.raises(ValueError, match="api_key"):
EmbeddingModel(config)
def test_remote_missing_api_url(self) -> None:
"""远程模式缺少 api_url 应抛出 ValueError。"""
config = EmbedConfig(
backend="remote",
model_name="test",
embed_dim=384,
device="cpu",
api_key="test-key",
api_url="",
)
with pytest.raises(ValueError, match="api_url"):
EmbeddingModel(config)
def test_local_dim_mismatch(self) -> None:
"""本地模式维度不一致应抛出 ValueError。"""
config = EmbedConfig(
backend="local",
model_name="all-MiniLM-L6-v2",
embed_dim=999, # 实际是 384
device="cpu",
api_key="",
api_url="",
)
with pytest.raises(ValueError, match="维度"):
EmbeddingModel(config)
reference/tests/unit/test_llm_client.py
+321
View File
@@ -0,0 +1,321 @@
"""
LLMClient 单元测试
==================
测试覆盖:
- TestLLMChatMock: mock OpenAI 客户端的纯文本对话功能
- TestVLMChatMock: mock 环境下的多模态图像对话功能
- TestConfigValidation: 配置校验（api_key/api_url 缺失）
- TestRealLLMChat: 真实 API 集成测试（需 .env 配置）
运行::
conda run -n Video-Tree-TRM python -m pytest tests/unit/test_llm_client.py -v \\
--cov=video_tree_trm/llm_client --cov-report=term-missing
"""
from __future__ import annotations
import base64
import os
import tempfile
from dataclasses import dataclass
from typing import Optional
from unittest.mock import MagicMock, patch
import pytest
from video_tree_trm.config import LLMConfig, VLMConfig
from video_tree_trm.llm_client import LLMClient
# ---------------------------------------------------------------------------
# 辅助：构造最小配置对象（避免加载真实 YAML）
# ---------------------------------------------------------------------------
def _make_llm_config(
api_key: str = "sk-test",
api_url: str = "https://api.example.com/v1",
model: str = "test-model",
max_tokens: int = 128,
temperature: float = 0.1,
) -> LLMConfig:
"""构造测试用 LLMConfig，所有字段可覆盖。"""
return LLMConfig(
backend="openai",
api_key=api_key,
api_url=api_url,
model=model,
max_tokens=max_tokens,
temperature=temperature,
)
def _make_vlm_config(
api_key: str = "sk-test",
api_url: str = "https://api.example.com/v1",
model: str = "test-vlm",
max_tokens: int = 128,
temperature: float = 0.1,
) -> VLMConfig:
"""构造测试用 VLMConfig。"""
return VLMConfig(
backend="openai",
api_key=api_key,
api_url=api_url,
model=model,
max_tokens=max_tokens,
temperature=temperature,
)
def _mock_completion(content: str) -> MagicMock:
"""构造 openai.ChatCompletion 返回值的 Mock。"""
choice = MagicMock()
choice.message.content = content
response = MagicMock()
response.choices = [choice]
return response
# ---------------------------------------------------------------------------
# TestLLMChatMock — 纯文本对话（Mock）
# ---------------------------------------------------------------------------
class TestLLMChatMock:
"""使用 mock openai.OpenAI 测试 chat() 和 batch_chat()。"""
@patch("video_tree_trm.llm_client.openai.OpenAI")
def test_chat_returns_string(self, mock_openai_cls: MagicMock) -> None:
"""chat() 应返回 API 返回内容的字符串。"""
mock_client = MagicMock()
mock_openai_cls.return_value = mock_client
mock_client.chat.completions.create.return_value = _mock_completion("你好！")
llm = LLMClient(_make_llm_config())
result = llm.chat("你好")
assert result == "你好！"
assert isinstance(result, str)
@patch("video_tree_trm.llm_client.openai.OpenAI")
def test_chat_uses_config_max_tokens(self, mock_openai_cls: MagicMock) -> None:
"""未传 max_tokens 时，应使用 config.max_tokens 的值。"""
mock_client = MagicMock()
mock_openai_cls.return_value = mock_client
mock_client.chat.completions.create.return_value = _mock_completion("ok")
cfg = _make_llm_config(max_tokens=256)
llm = LLMClient(cfg)
llm.chat("test")
call_kwargs = mock_client.chat.completions.create.call_args[1]
assert call_kwargs["max_tokens"] == 256
@patch("video_tree_trm.llm_client.openai.OpenAI")
def test_chat_overrides_max_tokens(self, mock_openai_cls: MagicMock) -> None:
"""显式传入 max_tokens 时，应覆盖 config.max_tokens。"""
mock_client = MagicMock()
mock_openai_cls.return_value = mock_client
mock_client.chat.completions.create.return_value = _mock_completion("ok")
cfg = _make_llm_config(max_tokens=256)
llm = LLMClient(cfg)
llm.chat("test", max_tokens=64)
call_kwargs = mock_client.chat.completions.create.call_args[1]
assert call_kwargs["max_tokens"] == 64
@patch("video_tree_trm.llm_client.openai.OpenAI")
def test_batch_chat_order_preserved(self, mock_openai_cls: MagicMock) -> None:
"""batch_chat() 应按输入顺序返回结果，即使并发完成顺序不同。"""
mock_client = MagicMock()
mock_openai_cls.return_value = mock_client
# 每次调用返回不同内容
responses = ["结果0", "结果1", "结果2"]
mock_client.chat.completions.create.side_effect = [
_mock_completion(r) for r in responses
]
llm = LLMClient(_make_llm_config())
results = llm.batch_chat(["prompt0", "prompt1", "prompt2"])
assert len(results) == 3
assert results == responses
@patch("video_tree_trm.llm_client.openai.OpenAI")
def test_batch_chat_empty_list(self, mock_openai_cls: MagicMock) -> None:
"""batch_chat() 传入空列表时，应返回空列表。"""
mock_openai_cls.return_value = MagicMock()
llm = LLMClient(_make_llm_config())
assert llm.batch_chat([]) == []
# ---------------------------------------------------------------------------
# TestVLMChatMock — 多模态对话（Mock）
# ---------------------------------------------------------------------------
class TestVLMChatMock:
"""使用 mock openai.OpenAI 测试 chat_with_images()。"""
@patch("video_tree_trm.llm_client.openai.OpenAI")
def test_chat_with_images_encodes_local_path(self, mock_openai_cls: MagicMock) -> None:
"""传入本地文件路径时，消息中应包含 data:image/jpeg;base64, 前缀。"""
mock_client = MagicMock()
mock_openai_cls.return_value = mock_client
mock_client.chat.completions.create.return_value = _mock_completion("图中有猫")
# 创建临时 JPEG 文件
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
f.write(b"\xff\xd8\xff\xe0" + b"\x00" * 20) # 最小 JPEG header
tmp_path = f.name
try:
vlm = LLMClient(_make_vlm_config())
result = vlm.chat_with_images("图中有什么？", images=[tmp_path])
assert result == "图中有猫"
# 验证消息结构包含 base64 编码图像
call_kwargs = mock_client.chat.completions.create.call_args[1]
content = call_kwargs["messages"][0]["content"]
assert isinstance(content, list)
image_items = [c for c in content if c.get("type") == "image_url"]
assert len(image_items) == 1
assert image_items[0]["image_url"]["url"].startswith("data:image/jpeg;base64,")
finally:
os.unlink(tmp_path)
@patch("video_tree_trm.llm_client.openai.OpenAI")
def test_chat_with_images_accepts_b64(self, mock_openai_cls: MagicMock) -> None:
"""传入已有 base64 字符串时，不应重复编码，直接透传。"""
mock_client = MagicMock()
mock_openai_cls.return_value = mock_client
mock_client.chat.completions.create.return_value = _mock_completion("ok")
b64_str = "data:image/jpeg;base64," + base64.b64encode(b"fake").decode()
vlm = LLMClient(_make_vlm_config())
vlm.chat_with_images("描述图片", images=[b64_str])
call_kwargs = mock_client.chat.completions.create.call_args[1]
content = call_kwargs["messages"][0]["content"]
image_items = [c for c in content if c.get("type") == "image_url"]
assert image_items[0]["image_url"]["url"] == b64_str
@patch("video_tree_trm.llm_client.openai.OpenAI")
def test_chat_with_images_png_mime(self, mock_openai_cls: MagicMock) -> None:
"""PNG 文件应编码为 data:image/png;base64, 格式。"""
mock_client = MagicMock()
mock_openai_cls.return_value = mock_client
mock_client.chat.completions.create.return_value = _mock_completion("ok")
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
f.write(b"\x89PNG\r\n\x1a\n" + b"\x00" * 20)
tmp_path = f.name
try:
vlm = LLMClient(_make_vlm_config())
vlm.chat_with_images("描述图片", images=[tmp_path])
call_kwargs = mock_client.chat.completions.create.call_args[1]
content = call_kwargs["messages"][0]["content"]
image_items = [c for c in content if c.get("type") == "image_url"]
assert image_items[0]["image_url"]["url"].startswith("data:image/png;base64,")
finally:
os.unlink(tmp_path)
@patch("video_tree_trm.llm_client.openai.OpenAI")
def test_chat_with_images_message_structure(self, mock_openai_cls: MagicMock) -> None:
"""多模态消息中图像应在文本之前。"""
mock_client = MagicMock()
mock_openai_cls.return_value = mock_client
mock_client.chat.completions.create.return_value = _mock_completion("ok")
b64_str = "data:image/jpeg;base64," + base64.b64encode(b"img").decode()
vlm = LLMClient(_make_vlm_config())
vlm.chat_with_images("提问", images=[b64_str])
call_kwargs = mock_client.chat.completions.create.call_args[1]
content = call_kwargs["messages"][0]["content"]
# 最后一项为 text
assert content[-1]["type"] == "text"
assert content[-1]["text"] == "提问"
# 前面各项为 image_url
for item in content[:-1]:
assert item["type"] == "image_url"
def test_encode_image_file_not_found(self) -> None:
"""_encode_image 传入不存在的路径时，应抛出 FileNotFoundError。"""
with patch("video_tree_trm.llm_client.openai.OpenAI"):
llm = LLMClient(_make_llm_config())
with pytest.raises(FileNotFoundError):
llm._encode_image("/nonexistent/path/image.jpg")
# ---------------------------------------------------------------------------
# TestConfigValidation — 配置校验
# ---------------------------------------------------------------------------
class TestConfigValidation:
"""测试 LLMClient 初始化时的配置校验逻辑。"""
def test_missing_api_key_raises(self) -> None:
"""api_key 为空时应抛出 ValueError。"""
cfg = _make_llm_config(api_key="")
with pytest.raises(ValueError, match="api_key"):
LLMClient(cfg)
def test_missing_api_url_raises(self) -> None:
"""api_url 为空时应抛出 ValueError。"""
cfg = _make_llm_config(api_url="")
with pytest.raises(ValueError, match="api_url"):
LLMClient(cfg)
@patch("video_tree_trm.llm_client.openai.OpenAI")
def test_valid_config_initializes(self, mock_openai_cls: MagicMock) -> None:
"""有效配置应正常初始化，不抛出异常。"""
mock_openai_cls.return_value = MagicMock()
cfg = _make_llm_config()
client = LLMClient(cfg)
assert client is not None
@patch("video_tree_trm.llm_client.openai.OpenAI")
def test_vlm_config_accepted(self, mock_openai_cls: MagicMock) -> None:
"""VLMConfig 也应被正常接受。"""
mock_openai_cls.return_value = MagicMock()
cfg = _make_vlm_config()
client = LLMClient(cfg)
assert client is not None
# ---------------------------------------------------------------------------
# TestRealLLMChat — 真实 API 集成测试（需 .env）
# ---------------------------------------------------------------------------
class TestRealLLMChat:
"""调用真实 LLM API 进行集成测试。
需要 .env 中配置有效的 LLM_API_KEY / LLM_API_URL / LLM_MODEL。
"""
def test_real_chat(self, real_config) -> None: # noqa: ANN001
"""真实 API 单轮对话，应返回非空字符串。"""
llm = LLMClient(real_config.llm)
result = llm.chat("请用一句话回答：天空是什么颜色？", max_tokens=32)
assert isinstance(result, str)
assert len(result) > 0
def test_real_batch_chat(self, real_config) -> None: # noqa: ANN001
"""真实 API 批量对话，应返回与输入等长的非空字符串列表。"""
llm = LLMClient(real_config.llm)
prompts = ["1+1等于几？请只回答数字。", "2+2等于几？请只回答数字。"]
results = llm.batch_chat(prompts, max_tokens=16)
assert len(results) == 2
for r in results:
assert isinstance(r, str)
assert len(r) > 0
reference/tests/unit/test_pipeline.py
+357
View File
@@ -0,0 +1,357 @@
"""
test_pipeline.py — Pipeline 单元测试
======================================
使用 unittest.mock.MagicMock + patch 隔离所有外部依赖（无真实 API / 文件 IO）。
"""
from __future__ import annotations
import os
from pathlib import Path
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
import torch
from video_tree_trm.pipeline import Pipeline
from video_tree_trm.tree_index import IndexMeta, L1Node, L2Node, L3Node, TreeIndex
# ---------------------------------------------------------------------------
# 辅助：构造最小 Config Mock
# ---------------------------------------------------------------------------
D = 8 # 嵌入维度（与 RetrieverConfig 一致）
def _make_config(checkpoint: str | None = None) -> MagicMock:
"""返回一个 Mock Config，字段值与实际 dataclass 对齐。"""
cfg = MagicMock()
cfg.embed.model_name = "test-embed"
cfg.embed.embed_dim = D
cfg.retriever.checkpoint = checkpoint
cfg.retriever.embed_dim = D
cfg.retriever.num_heads = 2
cfg.retriever.L_layers = 2
cfg.retriever.L_cycles = 2
cfg.retriever.max_rounds = 2
cfg.retriever.ffn_expansion = 2.0
cfg.tree.cache_dir = "/tmp/test_pipeline_cache"
return cfg
def _make_small_tree() -> TreeIndex:
"""构造最小 1×1×1 TreeIndex，用于 query() 测试。"""
meta = IndexMeta(
source_path="dummy",
modality="text",
embed_model="test",
embed_dim=D,
)
l3 = L3Node(
id="l3_0",
description="节点描述",
embedding=np.zeros(D, dtype=np.float32),
raw_content="节点内容",
)
l2 = L2Node(
id="l2_0",
description="L2",
embedding=np.zeros(D, dtype=np.float32),
children=[l3],
)
l1 = L1Node(
id="l1_0",
summary="L1",
embedding=np.zeros(D, dtype=np.float32),
children=[l2],
)
return TreeIndex(metadata=meta, roots=[l1])
# ---------------------------------------------------------------------------
# Patch 工厂：将所有子模块构造函数替换为 MagicMock
# ---------------------------------------------------------------------------
_PATCHES = [
"video_tree_trm.pipeline.EmbeddingModel",
"video_tree_trm.pipeline.LLMClient",
"video_tree_trm.pipeline.RecursiveRetriever",
"video_tree_trm.pipeline.AnswerGenerator",
"video_tree_trm.pipeline.TextTreeBuilder",
"video_tree_trm.pipeline.VideoTreeBuilder",
]
# ---------------------------------------------------------------------------
# Pipeline.__init__ 测试
# ---------------------------------------------------------------------------
def test_pipeline_init_components() -> None:
"""__init__ 后各属性（embed_model/llm/vlm/retriever/generator）均存在。"""
cfg = _make_config()
with patch.multiple(
"video_tree_trm.pipeline",
EmbeddingModel=MagicMock(),
LLMClient=MagicMock(),
RecursiveRetriever=MagicMock(),
AnswerGenerator=MagicMock(),
):
p = Pipeline(cfg)
assert hasattr(p, "embed_model"), "缺少 embed_model 属性"
assert hasattr(p, "llm"), "缺少 llm 属性"
assert hasattr(p, "vlm"), "缺少 vlm 属性"
assert hasattr(p, "retriever"), "缺少 retriever 属性"
assert hasattr(p, "generator"), "缺少 generator 属性"
def test_pipeline_init_no_checkpoint() -> None:
"""checkpoint=None 时 load_state_dict 不被调用。"""
cfg = _make_config(checkpoint=None)
mock_retriever_instance = MagicMock()
MockRetriever = MagicMock(return_value=mock_retriever_instance)
with patch.multiple(
"video_tree_trm.pipeline",
EmbeddingModel=MagicMock(),
LLMClient=MagicMock(),
RecursiveRetriever=MockRetriever,
AnswerGenerator=MagicMock(),
):
Pipeline(cfg)
mock_retriever_instance.load_state_dict.assert_not_called()
def test_pipeline_init_with_checkpoint(tmp_path: Path) -> None:
"""checkpoint 非 None 时 load_state_dict 被调用一次。"""
ckpt_file = tmp_path / "model.pt"
ckpt_file.write_bytes(b"") # 创建空文件，使 os.path.isfile 返回 True
cfg = _make_config(checkpoint=str(ckpt_file))
mock_retriever_instance = MagicMock()
MockRetriever = MagicMock(return_value=mock_retriever_instance)
fake_state_dict = {"weight": torch.zeros(1)}
with patch.multiple(
"video_tree_trm.pipeline",
EmbeddingModel=MagicMock(),
LLMClient=MagicMock(),
RecursiveRetriever=MockRetriever,
AnswerGenerator=MagicMock(),
), patch("video_tree_trm.pipeline.torch.load", return_value=fake_state_dict):
Pipeline(cfg)
mock_retriever_instance.load_state_dict.assert_called_once_with(fake_state_dict)
# ---------------------------------------------------------------------------
# Pipeline.build_index 测试
# ---------------------------------------------------------------------------
def test_build_index_text_calls_builder(tmp_path: Path) -> None:
"""文本模式调用 TextTreeBuilder.build，参数含文件内容。"""
src = tmp_path / "doc.txt"
src.write_text("文档内容", encoding="utf-8")
cfg = _make_config()
cfg.tree.cache_dir = str(tmp_path / "cache")
mock_tree = MagicMock(spec=TreeIndex)
mock_builder_instance = MagicMock()
mock_builder_instance.build.return_value = mock_tree
MockTextBuilder = MagicMock(return_value=mock_builder_instance)
with patch.multiple(
"video_tree_trm.pipeline",
EmbeddingModel=MagicMock(),
LLMClient=MagicMock(),
RecursiveRetriever=MagicMock(),
AnswerGenerator=MagicMock(),
TextTreeBuilder=MockTextBuilder,
):
p = Pipeline(cfg)
result = p.build_index(str(src), modality="text")
mock_builder_instance.build.assert_called_once()
call_args = mock_builder_instance.build.call_args
assert "文档内容" in call_args[0][0], "TextTreeBuilder.build 应传入文件内容"
assert result is mock_tree
def test_build_index_video_calls_builder(tmp_path: Path) -> None:
"""视频模式调用 VideoTreeBuilder.build，参数为 source_path。"""
cfg = _make_config()
cfg.tree.cache_dir = str(tmp_path / "cache")
mock_tree = MagicMock(spec=TreeIndex)
mock_builder_instance = MagicMock()
mock_builder_instance.build.return_value = mock_tree
MockVideoBuilder = MagicMock(return_value=mock_builder_instance)
video_path = "/fake/video.mp4"
with patch.multiple(
"video_tree_trm.pipeline",
EmbeddingModel=MagicMock(),
LLMClient=MagicMock(),
RecursiveRetriever=MagicMock(),
AnswerGenerator=MagicMock(),
VideoTreeBuilder=MockVideoBuilder,
):
p = Pipeline(cfg)
result = p.build_index(video_path, modality="video")
mock_builder_instance.build.assert_called_once_with(video_path)
assert result is mock_tree
def test_build_index_cache_hit(tmp_path: Path) -> None:
"""缓存文件存在时直接 TreeIndex.load，不重新构建。"""
cfg = _make_config()
cache_dir = tmp_path / "cache"
cache_dir.mkdir()
cfg.tree.cache_dir = str(cache_dir)
# 手动创建缓存文件（空文件即可让 isfile 返回 True）
cache_file = cache_dir / "doc_text.pkl"
cache_file.write_bytes(b"")
mock_tree = MagicMock(spec=TreeIndex)
mock_text_builder = MagicMock()
with patch.multiple(
"video_tree_trm.pipeline",
EmbeddingModel=MagicMock(),
LLMClient=MagicMock(),
RecursiveRetriever=MagicMock(),
AnswerGenerator=MagicMock(),
TextTreeBuilder=mock_text_builder,
), patch("video_tree_trm.pipeline.TreeIndex.load", return_value=mock_tree) as mock_load:
p = Pipeline(cfg)
result = p.build_index(str(tmp_path / "doc.txt"), modality="text")
mock_load.assert_called_once_with(str(cache_file))
mock_text_builder.return_value.build.assert_not_called()
assert result is mock_tree
def test_build_index_saves_cache(tmp_path: Path) -> None:
"""缓存不存在时构建后调用 tree.save。"""
cfg = _make_config()
cfg.tree.cache_dir = str(tmp_path / "cache")
src = tmp_path / "doc.txt"
src.write_text("内容", encoding="utf-8")
mock_tree = MagicMock(spec=TreeIndex)
mock_builder_instance = MagicMock()
mock_builder_instance.build.return_value = mock_tree
with patch.multiple(
"video_tree_trm.pipeline",
EmbeddingModel=MagicMock(),
LLMClient=MagicMock(),
RecursiveRetriever=MagicMock(),
AnswerGenerator=MagicMock(),
TextTreeBuilder=MagicMock(return_value=mock_builder_instance),
):
p = Pipeline(cfg)
p.build_index(str(src), modality="text")
mock_tree.save.assert_called_once()
saved_path: str = mock_tree.save.call_args[0][0]
assert "doc_text.pkl" in saved_path, f"保存路径应含 'doc_text.pkl'，实际={saved_path}"
# ---------------------------------------------------------------------------
# Pipeline.query 测试
# ---------------------------------------------------------------------------
def test_query_embeds_question() -> None:
"""query() 调用 embed_model.embed_tensor(question)。"""
cfg = _make_config()
tree = _make_small_tree()
mock_embed = MagicMock()
mock_embed.embed_tensor.return_value = torch.zeros(1, D)
MockEmbed = MagicMock(return_value=mock_embed)
mock_retriever_instance = MagicMock()
mock_retriever_instance.return_value = {"paths": [(0, 0, 0)], "num_rounds": 1}
with patch.multiple(
"video_tree_trm.pipeline",
EmbeddingModel=MockEmbed,
LLMClient=MagicMock(),
RecursiveRetriever=MagicMock(return_value=mock_retriever_instance),
AnswerGenerator=MagicMock(),
):
p = Pipeline(cfg)
p.query("测试问题", tree)
mock_embed.embed_tensor.assert_called_once_with("测试问题")
def test_query_calls_retriever() -> None:
"""query() 调用 retriever(q, tree)。"""
cfg = _make_config()
tree = _make_small_tree()
q_tensor = torch.zeros(1, D)
mock_embed = MagicMock()
mock_embed.embed_tensor.return_value = q_tensor
mock_retriever_instance = MagicMock()
mock_retriever_instance.return_value = {"paths": [(0, 0, 0)], "num_rounds": 1}
with patch.multiple(
"video_tree_trm.pipeline",
EmbeddingModel=MagicMock(return_value=mock_embed),
LLMClient=MagicMock(),
RecursiveRetriever=MagicMock(return_value=mock_retriever_instance),
AnswerGenerator=MagicMock(),
):
p = Pipeline(cfg)
p.query("测试问题", tree)
mock_retriever_instance.assert_called_once()
call_args = mock_retriever_instance.call_args
# 第一个位置参数应为嵌入 Tensor，第二个为 tree
assert call_args[0][1] is tree, "retriever 第二个参数应为 tree"
def test_query_returns_answer() -> None:
"""query() 返回 generator.generate 的返回值。"""
cfg = _make_config()
tree = _make_small_tree()
mock_embed = MagicMock()
mock_embed.embed_tensor.return_value = torch.zeros(1, D)
mock_retriever_instance = MagicMock()
mock_retriever_instance.return_value = {"paths": [(0, 0, 0)], "num_rounds": 1}
mock_generator_instance = MagicMock()
mock_generator_instance.generate.return_value = "生成的答案"
with patch.multiple(
"video_tree_trm.pipeline",
EmbeddingModel=MagicMock(return_value=mock_embed),
LLMClient=MagicMock(),
RecursiveRetriever=MagicMock(return_value=mock_retriever_instance),
AnswerGenerator=MagicMock(return_value=mock_generator_instance),
):
p = Pipeline(cfg)
answer = p.query("问题", tree)
assert answer == "生成的答案", f"query() 应返回 generator 的结果，实际='{answer}'"
mock_generator_instance.generate.assert_called_once_with(
"问题", [(0, 0, 0)], tree
)
reference/tests/unit/test_text_tree_builder.py
+546
View File
@@ -0,0 +1,546 @@
"""
TextTreeBuilder 单元测试
========================
使用 MagicMock 替代真实 LLM 和 Embed，测试文本树构建各阶段的正确性。
覆盖范围:
- _detect_toc：检测 Markdown 标题
- _segment_with_regex：正则切分 L1/L2 边界、超限分块
- _segment_with_llm：LLM JSON 分段解析、异常处理
- _build_l3_from_paragraphs：L3 节点字段验证
- _build_l2：L2 节点字段验证（通过 build() 间接调用）
- _build_l1：L1 节点字段验证（通过 build() 间接调用）
- build()：完整树结构与 IndexMeta 验证、MD 输出文件生成
MD 输出位置: tests/outputs/text_tree_builder/build_<timestamp>.md
"""
from __future__ import annotations
import json
import os
from datetime import datetime
from pathlib import Path
from typing import List
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
from video_tree_trm.config import TreeConfig
from video_tree_trm.text_tree_builder import TextTreeBuilder, _chunk
from video_tree_trm.tree_index import L1Node, L2Node, L3Node, TreeIndex
# ---------------------------------------------------------------------------
# 常量
# ---------------------------------------------------------------------------
_EMBED_DIM = 16 # 轻量维度，加速测试
_SAMPLE_TOC_TEXT = """\
# 第一章 引言
本章介绍研究背景。
信息检索系统的发展历程悠久，早期以关键词匹配为主。
## 1.1 研究背景
随着互联网数据量急剧增长，传统检索方法面临挑战。
语义理解成为关键技术突破口。
## 1.2 研究意义
本研究具有重要的理论和实践价值。
# 第二章 相关工作
本章回顾相关研究成果。
## 2.1 稠密检索
DPR 等模型实现了端到端稠密向量检索。
## 2.2 树状索引
PageIndex 引入了层次化树状检索结构。
"""
_SAMPLE_PLAIN_TEXT = (
"信息检索是计算机科学的重要分支，研究如何从大规模数据中找到相关信息。"
"\n\n"
"传统方法以 TF-IDF 和 BM25 为代表，基于词频统计进行相关性计算。"
"\n\n"
"近年来，基于深度学习的稠密检索方法取得了显著进展，DPR 是代表性工作之一。"
)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def tree_config() -> TreeConfig:
"""树构建配置（小 max_paragraphs_per_l2 便于测试分块）。"""
return TreeConfig(
max_paragraphs_per_l2=2,
l1_segment_duration=600.0,
l2_clip_duration=20.0,
l3_fps=1.0,
l2_representative_frames=3,
cache_dir="cache/trees",
)
@pytest.fixture
def mock_embed() -> MagicMock:
"""Mock EmbeddingModel，embed() 返回固定维度的随机向量。"""
m = MagicMock()
m.dim = _EMBED_DIM
m._model_name = "mock-embed-model"
def _embed(texts):
if isinstance(texts, str):
texts = [texts]
n = len(texts)
return np.ones((n, _EMBED_DIM), dtype=np.float32)
m.embed.side_effect = _embed
return m
@pytest.fixture
def mock_llm() -> MagicMock:
"""Mock LLMClient，chat() 返回固定字符串，batch_chat() 逐条映射。"""
m = MagicMock()
m.chat.return_value = "这是一段模拟的摘要描述。"
m.batch_chat.side_effect = lambda prompts: [
f"模拟摘要_{i}" for i in range(len(prompts))
]
return m
@pytest.fixture
def builder(mock_embed, mock_llm, tree_config) -> TextTreeBuilder:
"""标准 TextTreeBuilder（使用 mock 依赖）。"""
return TextTreeBuilder(
embed_model=mock_embed,
llm=mock_llm,
config=tree_config,
)
@pytest.fixture(scope="session", autouse=True)
def ensure_output_dir():
"""确保 MD 输出目录存在。"""
Path("tests/outputs/text_tree_builder").mkdir(parents=True, exist_ok=True)
# ---------------------------------------------------------------------------
# 辅助函数：保存 MD 输出
# ---------------------------------------------------------------------------
def _save_md_output(test_name: str, content: str) -> str:
"""将测试执行过程保存为 Markdown 文件。
参数:
test_name: 测试名称（用于文件名）。
content: Markdown 内容。
返回:
保存的文件路径。
"""
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
path = Path(f"tests/outputs/text_tree_builder/{test_name}_{ts}.md")
path.write_text(content, encoding="utf-8")
return str(path)
# ---------------------------------------------------------------------------
# 辅助函数：_chunk
# ---------------------------------------------------------------------------
class TestChunk:
"""测试 _chunk 等长分块辅助函数。"""
def test_exact_multiple(self):
"""整除时每组大小相等。"""
result = _chunk(["a", "b", "c", "d"], 2)
assert result == [["a", "b"], ["c", "d"]]
def test_remainder(self):
"""不整除时最后一组为余数。"""
result = _chunk(["a", "b", "c", "d", "e"], 2)
assert result == [["a", "b"], ["c", "d"], ["e"]]
def test_single_chunk(self):
"""列表长度 <= size 时只有一组。"""
result = _chunk(["a", "b"], 5)
assert result == [["a", "b"]]
def test_empty_list(self):
"""空列表返回空列表。"""
result = _chunk([], 3)
assert result == []
# ---------------------------------------------------------------------------
# _detect_toc
# ---------------------------------------------------------------------------
class TestDetectToc:
"""测试 Markdown 标题检测。"""
def test_with_h1_header(self, builder):
"""一级标题返回 True。"""
assert builder._detect_toc("# 第一章\n\n内容") is True
def test_with_h2_header(self, builder):
"""二级标题返回 True。"""
assert builder._detect_toc("## 1.1 小节\n\n内容") is True
def test_with_both_headers(self, builder):
"""同时含 # 和 ## 返回 True。"""
assert builder._detect_toc(_SAMPLE_TOC_TEXT) is True
def test_without_headers(self, builder):
"""纯段落文本返回 False。"""
assert builder._detect_toc(_SAMPLE_PLAIN_TEXT) is False
def test_hash_in_content_not_header(self, builder):
"""行中间的 # 不算标题。"""
text = "颜色代码 #FF0000 是红色。\n\n这是普通段落。"
assert builder._detect_toc(text) is False
def test_h3_not_counted(self, builder):
"""三级标题 ### 也被检测为有 ToC（属于 Markdown 结构文本）。
注：当前实现只检测 # 和 ##，所以 ### 开头应返回 False。
"""
text = "### 三级标题\n\n内容"
# _detect_toc 只匹配 #{1,2}，### 不匹配
assert builder._detect_toc(text) is False
# ---------------------------------------------------------------------------
# _segment_with_regex
# ---------------------------------------------------------------------------
class TestSegmentWithRegex:
"""测试正则切分 L1/L2 边界。"""
def test_basic_structure(self, builder):
"""含 2 个 L1 章节，正确返回 2 个外层元素。"""
sections = builder._segment_with_regex(_SAMPLE_TOC_TEXT)
assert len(sections) == 2
def test_all_sections_nonempty(self, builder):
"""每个 section 至少包含一个段落。"""
sections = builder._segment_with_regex(_SAMPLE_TOC_TEXT)
for s in sections:
assert len(s) > 0
def test_overflow_chunking_via_build(self, builder):
"""当段落数超过 max_paragraphs_per_l2=2 时，build() 会等长分块为多个 L2。"""
# 第一章有 4 个段落（含标题），超过 max=2 → 应有 2 个 L2
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
# 每个 L1 的 L2 数 >= 1
for l1 in index.roots:
assert len(l1.children) >= 1
def test_paragraphs_are_strings(self, builder):
"""所有段落应为非空字符串。"""
sections = builder._segment_with_regex(_SAMPLE_TOC_TEXT)
for section in sections:
for para in section:
assert isinstance(para, str)
assert para.strip() != ""
def test_single_chapter_text(self, builder):
"""只含一个 L1 章节的文本，返回 1 个外层元素。"""
text = "# 唯一章节\n\n第一段内容。\n\n第二段内容。"
sections = builder._segment_with_regex(text)
assert len(sections) == 1
def test_no_l2_header(self, builder):
"""只有 L1 无 L2 标题时，段落直接收集到 L1 组。"""
text = "# 第一章\n\n段落一。\n\n段落二。\n\n段落三。"
sections = builder._segment_with_regex(text)
assert len(sections) == 1
assert len(sections[0]) >= 2
# ---------------------------------------------------------------------------
# _segment_with_llm
# ---------------------------------------------------------------------------
class TestSegmentWithLLM:
"""测试 LLM 语义分段。"""
def test_json_array_parsing(self, builder):
"""mock LLM 返回合法 JSON 数组，应正确解析为段落列表。"""
paragraphs = ["第一段描述内容。", "第二段描述内容。", "第三段描述内容。"]
builder.llm.chat.return_value = json.dumps(paragraphs, ensure_ascii=False)
sections = builder._segment_with_llm(_SAMPLE_PLAIN_TEXT)
assert len(sections) == 1 # 单个 L1
assert len(sections[0]) == 3
assert sections[0][0] == "第一段描述内容。"
def test_json_wrapped_in_markdown(self, builder):
"""JSON 数组被 markdown 代码块包裹时也能正确解析。"""
paragraphs = ["段落A", "段落B"]
json_str = json.dumps(paragraphs, ensure_ascii=False)
builder.llm.chat.return_value = f"```json\n{json_str}\n```"
sections = builder._segment_with_llm(_SAMPLE_PLAIN_TEXT)
assert sections[0] == ["段落A", "段落B"]
def test_invalid_json_raises(self, builder):
"""LLM 返回非法 JSON 时应抛出 ValueError。"""
builder.llm.chat.return_value = "这不是 JSON 格式的内容"
with pytest.raises((ValueError, AssertionError)):
builder._segment_with_llm(_SAMPLE_PLAIN_TEXT)
def test_empty_paragraphs_filtered(self, builder):
"""空段落应被过滤掉。"""
paragraphs = ["有效段落", "", " ", "另一有效段落"]
builder.llm.chat.return_value = json.dumps(paragraphs, ensure_ascii=False)
sections = builder._segment_with_llm(_SAMPLE_PLAIN_TEXT)
assert len(sections[0]) == 2
# ---------------------------------------------------------------------------
# _build_l3_from_paragraphs
# ---------------------------------------------------------------------------
class TestBuildL3:
"""测试 L3 节点构建。"""
def test_description_equals_raw_content(self, builder):
"""L3 节点的 description 应等于 raw_content，等于原始段落。"""
paragraphs = ["段落一内容", "段落二内容"]
nodes = builder._build_l3_from_paragraphs(paragraphs, l1_i=0, l2_j=0)
for node, para in zip(nodes, paragraphs):
assert node.description == para
assert node.raw_content == para
def test_embedding_shape(self, builder):
"""每个 L3 节点的 embedding 形状应为 (dim,)。"""
paragraphs = ["A", "B", "C"]
nodes = builder._build_l3_from_paragraphs(paragraphs, l1_i=0, l2_j=0)
for node in nodes:
assert node.embedding.shape == (_EMBED_DIM,)
def test_node_count(self, builder):
"""返回的 L3 节点数应与输入段落数相等。"""
paragraphs = ["A", "B", "C"]
nodes = builder._build_l3_from_paragraphs(paragraphs, l1_i=0, l2_j=0)
assert len(nodes) == 3
def test_node_id_format(self, builder):
"""节点 ID 格式应为 l1_{i}_l2_{j}_l3_{k}。"""
nodes = builder._build_l3_from_paragraphs(["A", "B"], l1_i=1, l2_j=2)
assert nodes[0].id == "l1_1_l2_2_l3_0"
assert nodes[1].id == "l1_1_l2_2_l3_1"
def test_video_fields_are_none(self, builder):
"""文本模式下 frame_path 和 timestamp 应为 None。"""
nodes = builder._build_l3_from_paragraphs(["测试段落"], l1_i=0, l2_j=0)
assert nodes[0].frame_path is None
assert nodes[0].timestamp is None
def test_embedding_dtype(self, builder):
"""嵌入向量应为 float32。"""
nodes = builder._build_l3_from_paragraphs(["A"], l1_i=0, l2_j=0)
assert nodes[0].embedding.dtype == np.float32
# ---------------------------------------------------------------------------
# build() — 完整树结构验证
# ---------------------------------------------------------------------------
class TestBuild:
"""测试 build() 完整流程与 TreeIndex 结构。"""
def test_returns_tree_index(self, builder):
"""build() 应返回 TreeIndex 实例。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
assert isinstance(index, TreeIndex)
def test_roots_nonempty(self, builder):
"""roots 列表不为空。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
assert len(index.roots) > 0
def test_l1_has_l2_children(self, builder):
"""每个 L1 节点至少有一个 L2 子节点。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
for l1 in index.roots:
assert isinstance(l1, L1Node)
assert len(l1.children) > 0
def test_l2_has_l3_children(self, builder):
"""每个 L2 节点至少有一个 L3 子节点。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
for l1 in index.roots:
for l2 in l1.children:
assert isinstance(l2, L2Node)
assert len(l2.children) > 0
def test_l3_are_leaf_nodes(self, builder):
"""L3 节点是叶子层，类型为 L3Node。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
for l1 in index.roots:
for l2 in l1.children:
for l3 in l2.children:
assert isinstance(l3, L3Node)
def test_index_meta_modality(self, builder):
"""IndexMeta.modality 应为 'text'。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="doc.txt")
assert index.metadata.modality == "text"
def test_index_meta_source_path(self, builder):
"""IndexMeta.source_path 应与传入参数一致。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="my_doc.txt")
assert index.metadata.source_path == "my_doc.txt"
def test_index_meta_embed_dim(self, builder):
"""IndexMeta.embed_dim 应与 embed.dim 一致。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
assert index.metadata.embed_dim == _EMBED_DIM
def test_embedding_shapes(self, builder):
"""所有节点 embedding 形状应为 (dim,)。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
for l1 in index.roots:
assert l1.embedding.shape == (_EMBED_DIM,)
for l2 in l1.children:
assert l2.embedding.shape == (_EMBED_DIM,)
for l3 in l2.children:
assert l3.embedding.shape == (_EMBED_DIM,)
def test_l1_summary_nonempty(self, builder):
"""L1 summary 不为空。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
for l1 in index.roots:
assert l1.summary.strip() != ""
def test_l2_description_nonempty(self, builder):
"""L2 description 不为空。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
for l1 in index.roots:
for l2 in l1.children:
assert l2.description.strip() != ""
def test_plain_text_build(self, builder):
"""无 ToC 的纯段落文本也能正确构建。"""
# 修正 mock：返回合法 JSON
paragraphs = ["信息检索是计算机科学的重要分支。", "传统方法以 TF-IDF 为代表。", "近年来稠密检索方法兴起。"]
builder.llm.chat.return_value = json.dumps(paragraphs, ensure_ascii=False)
builder.llm.batch_chat.side_effect = lambda prompts: [
f"摘要_{i}" for i in range(len(prompts))
]
index = builder.build(_SAMPLE_PLAIN_TEXT, source_path="plain.txt")
assert len(index.roots) > 0
def test_empty_text_raises(self, builder):
"""空文本应抛出 ValueError（通过 ensure 检查）。"""
with pytest.raises((ValueError, AssertionError)):
builder.build(" ", source_path="empty.txt")
def test_batch_chat_called_once(self, builder):
"""build() 应调用 batch_chat() 一次（所有 L2 并发处理）。"""
builder.llm.batch_chat.reset_mock()
builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
builder.llm.batch_chat.assert_called_once()
def test_l1_node_ids_unique(self, builder):
"""所有 L1 节点 ID 唯一。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
ids = [l1.id for l1 in index.roots]
assert len(ids) == len(set(ids))
def test_l2_node_ids_unique(self, builder):
"""所有 L2 节点 ID 全局唯一。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
ids = [l2.id for l1 in index.roots for l2 in l1.children]
assert len(ids) == len(set(ids))
# ---------------------------------------------------------------------------
# MD 输出文件
# ---------------------------------------------------------------------------
class TestMDOutput:
"""测试 MD 输出文件生成（Agent 测试规范）。"""
def test_md_output_saved(self, builder):
"""build() 后应能保存 Markdown 执行记录文件。"""
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
# 统计树结构信息
total_l2 = sum(len(r.children) for r in index.roots)
total_l3 = sum(len(l2.children) for r in index.roots for l2 in r.children)
# 构造 MD 内容
l2_details = []
for l1 in index.roots:
for l2 in l1.children:
l2_details.append(f" - {l2.id}: {l2.description[:40]}...")
md_content = f"""\
# Agent 测试: TextTreeBuilder.build
## 任务: 长文本 → TreeIndex
## 输入信息
- **文本长度**: {len(_SAMPLE_TOC_TEXT)} 字符
- **是否含 ToC**: {builder._detect_toc(_SAMPLE_TOC_TEXT)}
- **source_path**: test.txt
- **max_paragraphs_per_l2**: {builder.config.max_paragraphs_per_l2}
- **embed_dim**: {_EMBED_DIM}
## Step 1: 结构切分
- **策略**: {'ToC 正则切分' if builder._detect_toc(_SAMPLE_TOC_TEXT) else 'LLM 语义分段'}
- **L1 数量**: {len(index.roots)}
- **各 L1 的 L2 数**: {[len(r.children) for r in index.roots]}
## Step 2: L2 先行（批量 LLM）
- **L2 节点总数**: {total_l2}
- **调用方式**: batch_chat()（并发生成所有 L2 摘要）
- **L2 描述示例**:
{chr(10).join(l2_details[:5])}
## Step 3: L3 向下（原始段落直接复用）
- **L3 节点总数**: {total_l3}
- **L3 特性**: description == raw_content（无 LLM 调用）
## Step 4: L1 向上（聚合摘要）
- **L1 摘要示例**:
{chr(10).join(f' - {r.id}: {r.summary[:60]}...' for r in index.roots)}
## 最终结果
- **roots 数量**: {len(index.roots)}
- **总 L2 节点**: {total_l2}
- **总 L3 节点**: {total_l3}
- **modality**: {index.metadata.modality}
- **embed_dim**: {index.metadata.embed_dim}
- **embedding shape 检查**: {'PASS' if all(r.embedding.shape == (_EMBED_DIM,) for r in index.roots) else 'FAIL'}
- **L3 description==raw_content 检查**: {'PASS' if all(l3.description == l3.raw_content for r in index.roots for l2 in r.children for l3 in l2.children) else 'FAIL'}
"""
out_path = _save_md_output("build_toc", md_content)
assert os.path.exists(out_path)
print(f"\n[MD 输出] 已保存到: {out_path}")
reference/tests/unit/test_tree_index.py
+212
View File
@@ -0,0 +1,212 @@
"""
tree_index 单元测试
===================
覆盖: 嵌入矩阵提取、节点访问、边界检查、序列化往返、空树处理。
"""
from __future__ import annotations
import os
import tempfile
import numpy as np
import pytest
from video_tree_trm.tree_index import (
IndexMeta,
L1Node,
L2Node,
L3Node,
TreeIndex,
)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
EMBED_DIM = 64
def _make_embed(seed: int = 0) -> np.ndarray:
"""生成固定种子的随机嵌入向量 [D]。"""
rng = np.random.RandomState(seed)
return rng.randn(EMBED_DIM).astype(np.float32)
def _make_meta() -> IndexMeta:
return IndexMeta(
source_path="test.mp4",
modality="video",
embed_model="test-model",
embed_dim=EMBED_DIM,
)
def _make_tree() -> TreeIndex:
"""构建一棵 2 x 2 x 3 的测试树。"""
meta = _make_meta()
roots = []
seed = 0
for i in range(2):
l2_nodes = []
for j in range(2):
l3_nodes = [
L3Node(
id=f"l3_{i}_{j}_{k}",
description=f"帧描述 {i}-{j}-{k}",
embedding=_make_embed(seed := seed + 1),
)
for k in range(3)
]
l2_nodes.append(
L2Node(
id=f"l2_{i}_{j}",
description=f"片段描述 {i}-{j}",
embedding=_make_embed(seed := seed + 1),
children=l3_nodes,
)
)
roots.append(
L1Node(
id=f"l1_{i}",
summary=f"摘要 {i}",
embedding=_make_embed(seed := seed + 1),
children=l2_nodes,
)
)
return TreeIndex(metadata=meta, roots=roots)
# ---------------------------------------------------------------------------
# 测试: 嵌入矩阵提取
# ---------------------------------------------------------------------------
class TestEmbeddings:
"""嵌入矩阵提取方法测试。"""
def test_l1_embeddings_shape(self) -> None:
"""l1_embeddings() 返回 [N1, D]。"""
tree = _make_tree()
emb = tree.l1_embeddings()
assert emb.shape == (2, EMBED_DIM)
assert emb.dtype == np.float32
def test_l2_embeddings_of_shape(self) -> None:
"""l2_embeddings_of(idx) 返回 [N2, D]。"""
tree = _make_tree()
emb = tree.l2_embeddings_of(0)
assert emb.shape == (2, EMBED_DIM)
assert emb.dtype == np.float32
def test_l3_embeddings_of_shape(self) -> None:
"""l3_embeddings_of(l1, l2) 返回 [N3, D]。"""
tree = _make_tree()
emb = tree.l3_embeddings_of(0, 1)
assert emb.shape == (3, EMBED_DIM)
assert emb.dtype == np.float32
# ---------------------------------------------------------------------------
# 测试: 节点访问
# ---------------------------------------------------------------------------
class TestGetNode:
"""节点访问方法测试。"""
def test_get_node(self) -> None:
"""正确返回目标 L3Node。"""
tree = _make_tree()
node = tree.get_node(1, 0, 2)
assert isinstance(node, L3Node)
assert node.id == "l3_1_0_2"
assert node.description == "帧描述 1-0-2"
def test_get_node_boundary_error(self) -> None:
"""越界索引抛出 IndexError。"""
tree = _make_tree()
with pytest.raises(IndexError):
tree.get_node(5, 0, 0)
with pytest.raises(IndexError):
tree.get_node(0, 5, 0)
with pytest.raises(IndexError):
tree.get_node(0, 0, 5)
def test_get_node_negative_index_error(self) -> None:
"""负数索引抛出 IndexError。"""
tree = _make_tree()
with pytest.raises(IndexError):
tree.get_node(-1, 0, 0)
# ---------------------------------------------------------------------------
# 测试: 序列化
# ---------------------------------------------------------------------------
class TestSerialization:
"""pickle 序列化测试。"""
def test_save_load_roundtrip(self) -> None:
"""pickle 序列化后反序列化，数据完整一致。"""
tree = _make_tree()
with tempfile.TemporaryDirectory() as tmpdir:
path = os.path.join(tmpdir, "test.pkl")
tree.save(path)
loaded = TreeIndex.load(path)
# 元数据一致
assert loaded.metadata.source_path == tree.metadata.source_path
assert loaded.metadata.embed_dim == tree.metadata.embed_dim
# 结构一致
assert len(loaded.roots) == len(tree.roots)
for orig_l1, load_l1 in zip(tree.roots, loaded.roots):
assert orig_l1.id == load_l1.id
assert len(orig_l1.children) == len(load_l1.children)
# 嵌入一致
np.testing.assert_array_equal(loaded.l1_embeddings(), tree.l1_embeddings())
np.testing.assert_array_equal(
loaded.l3_embeddings_of(0, 1), tree.l3_embeddings_of(0, 1)
)
def test_load_nonexistent_file(self) -> None:
"""加载不存在的文件抛出 FileNotFoundError。"""
with pytest.raises(FileNotFoundError):
TreeIndex.load("/tmp/nonexistent_tree_index_abc123.pkl")
# ---------------------------------------------------------------------------
# 测试: 空树边界
# ---------------------------------------------------------------------------
class TestEmptyTree:
"""空树边界情况测试。"""
def test_empty_tree_l1_embeddings(self) -> None:
"""空树的 l1_embeddings() 返回 [0, D]。"""
tree = TreeIndex(metadata=_make_meta(), roots=[])
emb = tree.l1_embeddings()
assert emb.shape == (0, EMBED_DIM)
assert emb.dtype == np.float32
def test_empty_tree_get_node_raises(self) -> None:
"""空树访问节点抛出 IndexError。"""
tree = TreeIndex(metadata=_make_meta(), roots=[])
with pytest.raises(IndexError):
tree.get_node(0, 0, 0)
def test_l2_embeddings_of_boundary(self) -> None:
"""l2_embeddings_of 越界抛出 ValueError。"""
tree = _make_tree()
with pytest.raises(ValueError):
tree.l2_embeddings_of(10)
def test_l3_embeddings_of_boundary(self) -> None:
"""l3_embeddings_of 越界抛出 ValueError。"""
tree = _make_tree()
with pytest.raises(ValueError):
tree.l3_embeddings_of(0, 10)
reference/tests/unit/test_video_tree_builder.py
+504
View File
@@ -0,0 +1,504 @@
"""
VideoTreeBuilder 单元测试
=========================
覆盖视频树构建的各个子方法和完整流程。
测试策略:
- mock cv2.VideoCapture 避免依赖真实视频（_segment_video 等）
- 使用 cv2 合成小视频进行帧提取和集成测试（tiny_video fixture）
- mock VLM/embed 依赖，隔离外部 API
- 测试 L3 批量降级路径（JSON 解析失败时退回逐帧调用）
"""
from __future__ import annotations
import json
import os
from datetime import datetime
from pathlib import Path
from typing import List
from unittest.mock import MagicMock, patch
import cv2
import numpy as np
import pytest
from video_tree_trm.config import TreeConfig
from video_tree_trm.tree_index import L1Node, L2Node, L3Node
from video_tree_trm.video_tree_builder import VideoTreeBuilder
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def tree_config(tmp_path: Path) -> TreeConfig:
"""构建测试用 TreeConfig，cache_dir 指向临时目录。"""
return TreeConfig(
max_paragraphs_per_l2=5,
l1_segment_duration=30.0,
l2_clip_duration=10.0,
l3_fps=1.0,
l2_representative_frames=3,
cache_dir=str(tmp_path / "cache"),
)
@pytest.fixture
def mock_embed() -> MagicMock:
"""返回 mock 嵌入模型，embed() 返回全 1 向量。"""
embed = MagicMock()
embed.dim = 4
embed._model_name = "mock-embed"
def _embed(texts):
"""根据输入类型返回 [1, D] 或 [N, D] 的 float32 数组。"""
if isinstance(texts, str):
return np.ones((1, 4), dtype=np.float32)
n = len(texts)
return np.ones((n, 4), dtype=np.float32)
embed.embed.side_effect = _embed
return embed
@pytest.fixture
def mock_vlm() -> MagicMock:
"""返回 mock VLM 客户端。"""
vlm = MagicMock()
vlm.chat.return_value = "这是一段精彩的视频内容摘要。"
vlm.chat_with_images.return_value = "这帧画面中有人物在移动。"
return vlm
@pytest.fixture
def builder(
mock_embed: MagicMock, mock_vlm: MagicMock, tree_config: TreeConfig
) -> VideoTreeBuilder:
"""构建测试用 VideoTreeBuilder 实例。"""
return VideoTreeBuilder(
embed_model=mock_embed,
vlm=mock_vlm,
config=tree_config,
)
@pytest.fixture
def tiny_video(tmp_path: Path) -> str:
"""用 cv2 生成 30 帧合成彩色视频（10fps，时长 3 秒），返回路径。
视频规格:
- 分辨率: 64×48
- 帧率: 10fps
- 时长: 3 秒（30 帧）
- 内容: 随机彩色帧
"""
video_path = str(tmp_path / "tiny.mp4")
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
writer = cv2.VideoWriter(video_path, fourcc, 10.0, (64, 48))
for _ in range(30):
frame = np.random.randint(0, 256, (48, 64, 3), dtype=np.uint8)
writer.write(frame)
writer.release()
return video_path
# ---------------------------------------------------------------------------
# 测试：_segment_video — 时间切分
# ---------------------------------------------------------------------------
def test_segment_video_fixed_step(builder: VideoTreeBuilder, tmp_path: Path) -> None:
"""mock cv2.VideoCapture（总时长=60s，l1_segment_duration=30s），
验证切分出 2 个均等 L1 区间：(0,30),(30,60)。"""
mock_cap = MagicMock()
mock_cap.isOpened.return_value = True
mock_cap.get.side_effect = lambda prop: (
10.0 if prop == cv2.CAP_PROP_FPS else 600.0 # 600帧/10fps = 60s
)
with patch("video_tree_trm.video_tree_builder.cv2.VideoCapture", return_value=mock_cap):
ranges = builder._segment_video("fake.mp4")
assert len(ranges) == 2
assert ranges[0] == (0.0, 30.0)
assert ranges[1] == (30.0, 60.0)
def test_segment_video_uneven(builder: VideoTreeBuilder) -> None:
"""总时长不能被 l1_segment_duration 整除时，最后一段应短于步长。"""
mock_cap = MagicMock()
mock_cap.isOpened.return_value = True
# 75s = 750帧 / 10fps
mock_cap.get.side_effect = lambda prop: (
10.0 if prop == cv2.CAP_PROP_FPS else 750.0
)
with patch("video_tree_trm.video_tree_builder.cv2.VideoCapture", return_value=mock_cap):
ranges = builder._segment_video("fake.mp4")
assert len(ranges) == 3
assert ranges[0] == (0.0, 30.0)
assert ranges[1] == (30.0, 60.0)
assert abs(ranges[2][0] - 60.0) < 1e-6
assert abs(ranges[2][1] - 75.0) < 1e-6
# ---------------------------------------------------------------------------
# 测试：_get_l2_clips — L2 切分
# ---------------------------------------------------------------------------
def test_get_l2_clips_even(builder: VideoTreeBuilder) -> None:
"""l1=(0,30)，l2_duration=10 → 3 clips 均等：(0,10),(10,20),(20,30)。"""
clips = builder._get_l2_clips((0.0, 30.0))
assert len(clips) == 3
assert clips[0] == (0.0, 10.0)
assert clips[1] == (10.0, 20.0)
assert clips[2] == (20.0, 30.0)
def test_get_l2_clips_uneven(builder: VideoTreeBuilder) -> None:
"""l1=(0,25)，l2_duration=10 → 3 clips，最后一段为 5s。"""
clips = builder._get_l2_clips((0.0, 25.0))
assert len(clips) == 3
assert clips[2] == (20.0, 25.0)
def test_get_l2_clips_shorter_than_step(builder: VideoTreeBuilder) -> None:
"""L1 区间短于 l2_clip_duration 时，返回 1 个 clip。"""
clips = builder._get_l2_clips((0.0, 5.0))
assert len(clips) == 1
assert clips[0] == (0.0, 5.0)
# ---------------------------------------------------------------------------
# 测试：_extract_frames — 帧提取
# ---------------------------------------------------------------------------
def test_extract_frames_saves_files(
builder: VideoTreeBuilder, tiny_video: str, tmp_path: Path
) -> None:
"""使用真实合成视频（3s），提取 1fps 帧，验证返回路径和文件存在。"""
frames = builder._extract_frames(tiny_video, (0.0, 3.0), fps=1.0)
assert len(frames) >= 1
for frame_path, ts in frames:
assert os.path.isfile(frame_path), f"帧文件不存在: {frame_path}"
assert ts >= 0.0
def test_extract_frames_cache_reuse(
builder: VideoTreeBuilder, tiny_video: str
) -> None:
"""第二次提取同一区间时，帧文件应直接复用（不重复写磁盘）。"""
frames1 = builder._extract_frames(tiny_video, (0.0, 2.0), fps=1.0)
assert len(frames1) >= 1
# 记录文件修改时间
mtimes_before = [os.path.getmtime(fp) for fp, _ in frames1]
frames2 = builder._extract_frames(tiny_video, (0.0, 2.0), fps=1.0)
mtimes_after = [os.path.getmtime(fp) for fp, _ in frames2]
assert frames1 == frames2
assert mtimes_before == mtimes_after, "缓存帧文件被重复写入"
def test_extract_frames_empty_range(
builder: VideoTreeBuilder, tiny_video: str
) -> None:
"""时间范围内无有效时间戳时，返回空列表。"""
# 起始=结束，无时间戳
frames = builder._extract_frames(tiny_video, (1.0, 1.0), fps=1.0)
assert frames == []
# ---------------------------------------------------------------------------
# 测试：_build_l2_video — L2 节点构建
# ---------------------------------------------------------------------------
def test_build_l2_video_node_structure(
builder: VideoTreeBuilder, tiny_video: str, mock_vlm: MagicMock
) -> None:
"""验证 L2Node 字段：description 非空、embedding shape 正确、time_range 正确。"""
mock_vlm.chat_with_images.return_value = "片段展示了室内场景的变化。"
l2_node = builder._build_l2_video(tiny_video, (0.0, 2.0), "l1_0_l2_0")
assert isinstance(l2_node, L2Node)
assert l2_node.id == "l1_0_l2_0"
assert len(l2_node.description) > 0
assert l2_node.embedding.shape == (4,)
assert l2_node.embedding.dtype == np.float32
assert l2_node.time_range == (0.0, 2.0)
assert l2_node.children == [] # 调用方填充
def test_build_l2_video_representative_frames_count(
builder: VideoTreeBuilder, tiny_video: str, mock_vlm: MagicMock
) -> None:
"""验证 VLM 被调用时传入的图像数不超过 l2_representative_frames。"""
mock_vlm.chat_with_images.return_value = "描述内容。"
builder._build_l2_video(tiny_video, (0.0, 3.0), "l1_0_l2_0")
call_args = mock_vlm.chat_with_images.call_args
images_passed = call_args.kwargs.get("images", call_args.args[1] if len(call_args.args) > 1 else [])
assert len(images_passed) <= builder.config.l2_representative_frames
# ---------------------------------------------------------------------------
# 测试：_build_l3_video — L3 节点构建
# ---------------------------------------------------------------------------
def _make_frames(n: int, tmp_path: Path) -> List[tuple]:
"""创建 n 个临时 JPEG 帧文件，返回 [(path, ts), ...]。"""
frames = []
frame_dir = tmp_path / "frames"
frame_dir.mkdir(exist_ok=True)
for i in range(n):
frame_path = str(frame_dir / f"frame_{i}.jpg")
img = np.zeros((48, 64, 3), dtype=np.uint8)
cv2.imwrite(frame_path, img)
frames.append((frame_path, float(i)))
return frames
def test_build_l3_video_batch_success(
builder: VideoTreeBuilder, mock_vlm: MagicMock, tmp_path: Path
) -> None:
"""mock VLM 返回合法 JSON 数组，验证 L3Node 列表结构正确。"""
frames = _make_frames(2, tmp_path)
mock_vlm.chat_with_images.return_value = json.dumps(["帧1描述内容", "帧2描述内容"])
nodes = builder._build_l3_video(frames, "片段整体描述", l1_i=0, l2_j=0)
assert len(nodes) == 2
for k, node in enumerate(nodes):
assert isinstance(node, L3Node)
assert node.id == f"l1_0_l2_0_l3_{k}"
assert len(node.description) > 0
assert node.embedding.shape == (4,)
assert node.embedding.dtype == np.float32
assert node.frame_path == frames[k][0]
assert node.timestamp == float(k)
assert node.raw_content is None
def test_build_l3_video_batch_fallback(
builder: VideoTreeBuilder, mock_vlm: MagicMock, tmp_path: Path
) -> None:
"""mock VLM 第一次返回非 JSON 字符串，验证降级逐帧调用（call_count == n+1）。
第一次 = 批量调用（失败），后 n 次 = 逐帧调用。
"""
n = 3
frames = _make_frames(n, tmp_path)
# 第一次返回无效 JSON，后续逐帧返回正常描述
mock_vlm.chat_with_images.side_effect = (
["这不是一个JSON数组，无法解析"] + [f"第{i}帧描述" for i in range(n)]
)
nodes = builder._build_l3_video(frames, "片段整体描述", l1_i=0, l2_j=0)
assert len(nodes) == n
# 1次批量 + n次逐帧
assert mock_vlm.chat_with_images.call_count == n + 1
for node in nodes:
assert len(node.description) > 0
def test_build_l3_video_json_length_mismatch_fallback(
builder: VideoTreeBuilder, mock_vlm: MagicMock, tmp_path: Path
) -> None:
"""VLM 返回 JSON 但长度不匹配时，也应降级逐帧调用。"""
n = 3
frames = _make_frames(n, tmp_path)
# 只返回 2 项，但期望 3 项
mock_vlm.chat_with_images.side_effect = (
[json.dumps(["描述1", "描述2"])] + [f"帧{i}" for i in range(n)]
)
nodes = builder._build_l3_video(frames, "片段描述", l1_i=0, l2_j=0)
assert len(nodes) == n
assert mock_vlm.chat_with_images.call_count == n + 1
# ---------------------------------------------------------------------------
# 测试：_build_l1_video — L1 节点构建
# ---------------------------------------------------------------------------
def test_build_l1_video_node_structure(
builder: VideoTreeBuilder, mock_vlm: MagicMock, mock_embed: MagicMock
) -> None:
"""验证 L1Node 字段：summary 非空、time_range 正确、children 已赋值。"""
mock_vlm.chat.return_value = "这段视频涵盖了户外活动和室内场景的切换。"
l2_children = [
L2Node(
id=f"l1_0_l2_{j}",
description=f"L2描述{j}",
embedding=np.ones(4, dtype=np.float32),
time_range=(j * 10.0, (j + 1) * 10.0),
)
for j in range(3)
]
l1_node = builder._build_l1_video(l2_children, "l1_0", (0.0, 30.0))
assert isinstance(l1_node, L1Node)
assert l1_node.id == "l1_0"
assert len(l1_node.summary) > 0
assert l1_node.time_range == (0.0, 30.0)
assert l1_node.children is l2_children
assert l1_node.embedding.shape == (4,)
assert l1_node.embedding.dtype == np.float32
def test_build_l1_video_prompt_contains_l2_descriptions(
builder: VideoTreeBuilder, mock_vlm: MagicMock
) -> None:
"""验证 L1 摘要的 prompt 包含所有 L2 描述文本。"""
mock_vlm.chat.return_value = "综合摘要内容。"
l2_descriptions = ["片段A描述", "片段B描述", "片段C描述"]
l2_children = [
L2Node(
id=f"l1_0_l2_{j}",
description=desc,
embedding=np.ones(4, dtype=np.float32),
time_range=(j * 10.0, (j + 1) * 10.0),
)
for j, desc in enumerate(l2_descriptions)
]
builder._build_l1_video(l2_children, "l1_0", (0.0, 30.0))
call_prompt = mock_vlm.chat.call_args.args[0]
for desc in l2_descriptions:
assert desc in call_prompt, f"L2 描述 '{desc}' 未出现在 L1 prompt 中"
# ---------------------------------------------------------------------------
# 测试：build 完整流程（集成测试，mock VLM）
# ---------------------------------------------------------------------------
def test_build_full_integration(
builder: VideoTreeBuilder,
tiny_video: str,
mock_vlm: MagicMock,
mock_embed: MagicMock,
tmp_path: Path,
) -> None:
"""用合成视频（3s）+ mock VLM 验证完整 TreeIndex 三层结构。
配置：l1_segment_duration=2s，l2_clip_duration=1s
预期：至少 1 个 L1，每 L1 至少 1 个 L2，每 L2 至少 1 个 L3。
"""
# 调整 config 使 3s 视频能切出多个节点
builder.config.l1_segment_duration = 2.0
builder.config.l2_clip_duration = 1.0
# VLM 批量调用返回 JSON 数组（按帧数动态生成）
def vlm_side_effect(prompt, images=None):
if images and len(images) > 1:
# L3 批量调用：返回 JSON
return json.dumps([f"帧{i}描述" for i in range(len(images))])
return "这是 VLM 的描述文本。"
mock_vlm.chat_with_images.side_effect = vlm_side_effect
mock_vlm.chat.return_value = "L1 整体摘要内容。"
index = builder.build(tiny_video)
# 验证元数据
assert index.metadata.modality == "video"
assert index.metadata.source_path == tiny_video
assert index.metadata.embed_dim == 4
# 验证三层结构非空
assert len(index.roots) >= 1, "应有至少 1 个 L1 节点"
for l1 in index.roots:
assert len(l1.children) >= 1, f"L1 {l1.id} 应有至少 1 个 L2 子节点"
assert l1.time_range is not None
for l2 in l1.children:
assert len(l2.children) >= 1, f"L2 {l2.id} 应有至少 1 个 L3 子节点"
assert l2.time_range is not None
for l3 in l2.children:
assert l3.frame_path is not None
assert l3.timestamp is not None
assert l3.embedding.shape == (4,)
def test_build_saves_output_md(
builder: VideoTreeBuilder,
tiny_video: str,
mock_vlm: MagicMock,
tmp_path: Path,
) -> None:
"""构建完成后，将执行摘要保存为 Markdown（CLAUDE.md 规范）。"""
builder.config.l1_segment_duration = 2.0
builder.config.l2_clip_duration = 1.0
def vlm_side_effect(prompt, images=None):
if images and len(images) > 1:
return json.dumps([f"帧{i}描述" for i in range(len(images))])
return "VLM 描述内容。"
mock_vlm.chat_with_images.side_effect = vlm_side_effect
mock_vlm.chat.return_value = "L1 摘要。"
index = builder.build(tiny_video)
# 保存 Markdown 输出
output_dir = Path(__file__).resolve().parent.parent / "outputs" / "video_tree_builder"
output_dir.mkdir(parents=True, exist_ok=True)
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
output_path = output_dir / f"build_video_{ts}.md"
total_l2 = sum(len(r.children) for r in index.roots)
total_l3 = sum(len(l2.children) for r in index.roots for l2 in r.children)
lines = [
f"# Agent 测试: test_build_saves_output_md",
f"## 任务: VideoTreeBuilder.build() 完整流程验证",
f"",
f"## 输入",
f"- 视频路径: `{tiny_video}`",
f"- l1_segment_duration: {builder.config.l1_segment_duration}s",
f"- l2_clip_duration: {builder.config.l2_clip_duration}s",
f"- l3_fps: {builder.config.l3_fps}",
f"",
f"## 输出结构",
f"- L1 节点数: {len(index.roots)}",
f"- L2 节点数: {total_l2}",
f"- L3 节点数: {total_l3}",
f"- embed_dim: {index.metadata.embed_dim}",
f"",
f"## L1 详情",
]
for l1 in index.roots:
lines.append(f"### {l1.id} (time_range={l1.time_range})")
lines.append(f"- summary: {l1.summary[:80]}...")
for l2 in l1.children:
lines.append(
f" - {l2.id} [{l2.time_range}]: {l2.description[:60]}... "
f"({len(l2.children)} L3)"
)
lines += [
"",
"## 最终结果",
"✅ TreeIndex 构建成功，三层结构完整。",
f"",
f"输出文件: `{output_path}`",
]
output_path.write_text("\n".join(lines), encoding="utf-8")
print(f"\n[测试输出] {output_path}")
assert output_path.is_file()
reference/utils/__init__.py
View File
reference/utils/logger_system.py
+164
View File
@@ -0,0 +1,164 @@
"""
日志系统模块
============
提供双通道日志（system.log 文本 + metrics.json 结构化），
以及全局便捷函数 log_msg / log_json / ensure / log_exception。
使用方式::
from utils.logger_system import log_msg, log_json, ensure, log_exception
log_msg("INFO", "模型加载完成", model="bert-base")
log_json("train_loss", {"epoch": 1, "loss": 0.35})
ensure(tensor.shape[0] > 0, "batch 不能为空")
"""
from __future__ import annotations
import json
import logging
import os
import threading
import traceback
from datetime import datetime
from pathlib import Path
from typing import Any, Optional
# ---------------------------------------------------------------------------
# 日志目录
# ---------------------------------------------------------------------------
_LOG_DIR = Path(os.getenv("LOG_DIR", "logs"))
_LOG_DIR.mkdir(parents=True, exist_ok=True)
_SYSTEM_LOG = _LOG_DIR / "system.log"
_METRICS_LOG = _LOG_DIR / "metrics.json"
# ---------------------------------------------------------------------------
# LoggerSystem
# ---------------------------------------------------------------------------
class LoggerSystem:
"""双通道日志系统。
通道 1: system.log — 文本格式，记录所有级别的运行日志。
通道 2: metrics.json — JSON Lines 格式，记录结构化指标数据。
Attributes:
_logger: 标准库 Logger 实例，输出到 system.log。
_metrics_path: metrics.json 文件路径。
"""
_instance: Optional["LoggerSystem"] = None
_lock: threading.Lock = threading.Lock()
def __init__(self, log_dir: Optional[str] = None) -> None:
"""初始化日志系统。
参数:
log_dir: 日志输出目录，默认为 'logs/'。
"""
log_path = Path(log_dir) if log_dir else _LOG_DIR
log_path.mkdir(parents=True, exist_ok=True)
self._metrics_path = log_path / "metrics.json"
# 文本日志
self._logger = logging.getLogger("video_tree_trm")
if not self._logger.handlers:
self._logger.setLevel(logging.DEBUG)
fh = logging.FileHandler(str(log_path / "system.log"), encoding="utf-8")
fh.setLevel(logging.DEBUG)
fmt = logging.Formatter(
"%(asctime)s | %(levelname)-7s | %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
fh.setFormatter(fmt)
self._logger.addHandler(fh)
@classmethod
def get(cls) -> "LoggerSystem":
"""获取全局单例（线程安全，双重检查锁定）。"""
if cls._instance is None:
with cls._lock:
if cls._instance is None:
cls._instance = cls()
return cls._instance
# ---- 文本日志 ----
def msg(self, level: str, message: str, **kwargs: Any) -> None:
"""写入文本日志。
参数:
level: 日志级别，如 "INFO", "ERROR", "DEBUG", "WARNING"。
message: 日志消息。
**kwargs: 附加键值对，追加到消息末尾。
"""
extra = " | ".join(f"{k}={v}" for k, v in kwargs.items()) if kwargs else ""
text = f"{message} | {extra}" if extra else message
log_fn = getattr(self._logger, level.lower(), self._logger.info)
log_fn(text)
# ---- 结构化指标 ----
def json(self, tag: str, data: dict[str, Any]) -> None:
"""写入结构化 JSON 指标。
参数:
tag: 指标标签，如 "train_loss"。
data: 指标数据字典。
"""
record = {"ts": datetime.now().isoformat(), "tag": tag, **data}
with open(self._metrics_path, "a", encoding="utf-8") as f:
f.write(json.dumps(record, ensure_ascii=False) + "\n")
# ---- 断言 ----
@staticmethod
def ensure(condition: bool, message: str) -> None:
"""运行时断言，失败时抛出 ValueError。
参数:
condition: 断言条件。
message: 失败时的错误消息。
"""
if not condition:
raise ValueError(message)
# ---- 异常记录 ----
def exception(self, message: str, exc: BaseException) -> None:
"""记录异常信息到日志。
参数:
message: 上下文描述。
exc: 异常实例。
"""
tb = traceback.format_exception(type(exc), exc, exc.__traceback__)
self._logger.error(f"{message} | {''.join(tb)}")
# ---------------------------------------------------------------------------
# 全局便捷函数
# ---------------------------------------------------------------------------
def log_msg(level: str, message: str, **kwargs: Any) -> None:
"""写入文本日志（全局便捷函数）。"""
LoggerSystem.get().msg(level, message, **kwargs)
def log_json(tag: str, data: dict[str, Any]) -> None:
"""写入结构化 JSON 指标（全局便捷函数）。"""
LoggerSystem.get().json(tag, data)
def ensure(condition: bool, message: str) -> None:
"""运行时断言（全局便捷函数）。"""
LoggerSystem.ensure(condition, message)
def log_exception(message: str, exc: BaseException) -> None:
"""记录异常（全局便捷函数）。"""
LoggerSystem.get().exception(message, exc)
reference/video_tree_trm/__init__.py
+29
View File
@@ -0,0 +1,29 @@
"""
Video-Tree-TRM 核心包
=====================
结合 TRM 多层对话探索能力与 PageIndex 树状检索能力的新型 Video RAG。
"""
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.text_tree_builder import TextTreeBuilder
from video_tree_trm.tree_index import (
IndexMeta,
L1Node,
L2Node,
L3Node,
TreeIndex,
)
__all__ = [
"Config",
"EmbeddingModel",
"LLMClient",
"TextTreeBuilder",
"IndexMeta",
"L1Node",
"L2Node",
"L3Node",
"TreeIndex",
]
reference/video_tree_trm/config.py
+440
View File
@@ -0,0 +1,440 @@
"""
配置管理模块
============
定义所有超参数的 dataclass 类型（无默认值）+ 多源加载。
三层优先级: CLI args > .env > YAML，统一归口到 Config dataclass。
使用方式::
from video_tree_trm.config import Config
cfg = Config.load("config/default.yaml")
cfg = Config.load("config/default.yaml", cli_args={"retriever.num_heads": "8"})
"""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional
import yaml
from dotenv import dotenv_values
from utils.logger_system import ensure, log_msg
# ---------------------------------------------------------------------------
# 子配置 Dataclass（全部无默认值，YAML 必须写全）
# ---------------------------------------------------------------------------
@dataclass
class TreeConfig:
"""树索引构建参数。
属性:
max_paragraphs_per_l2: 每个 L2 节点包含的最大段落数（文本模式）。
l1_segment_duration: L1 段时长，秒（视频模式）。
l2_clip_duration: L2 clip 时长，秒（视频模式）。
l3_fps: L3 帧提取频率（视频模式）。
l2_representative_frames: L2 VLM 描述用的代表帧数。
cache_dir: TreeIndex 缓存目录。
"""
max_paragraphs_per_l2: int
l1_segment_duration: float
l2_clip_duration: float
l3_fps: float
l2_representative_frames: int
cache_dir: str
concurrency: int
@dataclass
class EmbedConfig:
"""嵌入模型参数。
属性:
backend: 嵌入后端类型，"local"（sentence-transformers）或 "remote"（OpenAI 兼容 API）。
model_name: 本地模式为 HuggingFace 模型名，远程模式为 API 模型名。
embed_dim: 嵌入维度 D。
device: 推理设备，"cuda" 或 "cpu"（仅本地模式使用）。
api_key: 远程模式 API 密钥，从 .env 加载。本地模式为空串。
api_url: 远程模式 API 端点。本地模式为空串。
"""
backend: str
model_name: str
embed_dim: int
device: str
api_key: str
api_url: str
@dataclass
class LLMConfig:
"""大语言模型参数。
属性:
backend: 后端类型，"qwen" | "openai" | "ollama"。
api_key: API 密钥，从 .env 加载。
model: 模型名称。
api_url: API 端点 URL。
max_tokens: 最大生成 token 数。
temperature: 采样温度。
"""
backend: str
api_key: str
model: str
api_url: str
max_tokens: int
temperature: float
@dataclass
class VLMConfig:
"""视觉语言模型参数。
属性:
backend: 后端类型，"qwen" | "openai" | "ollama"。
api_key: API 密钥，从 .env 加载。
model: 模型名称。
api_url: API 端点 URL。
max_tokens: 最大生成 token 数。
temperature: 采样温度。
"""
backend: str
api_key: str
model: str
api_url: str
max_tokens: int
temperature: float
@dataclass
class RetrieverConfig:
"""TRM 检索器参数。
属性:
embed_dim: 嵌入维度，须与 EmbedConfig.embed_dim 一致。
num_heads: Cross-Attention 头数。
L_layers: ReasoningModule 层数。
L_cycles: 每级推理迭代次数。
max_rounds: ACT 最大遍历轮次。
ffn_expansion: SwiGLU 扩展比。
checkpoint: 训练好的模型权重路径，推理时必填。
"""
embed_dim: int
num_heads: int
L_layers: int
L_cycles: int
max_rounds: int
ffn_expansion: float
checkpoint: Optional[str]
# 多路径检索配置 (Top-k)
k_l1: int = 1
k_l2: int = 1
k_l3: int = 1
max_paths: int = 5
@dataclass
class TrainConfig:
"""训练参数。
属性:
lr: 学习率。
weight_decay: 权重衰减。
batch_size: 批大小。
max_epochs_phase1: Phase 1 导航训练轮数。
max_epochs_phase2: Phase 2 ACT 训练轮数。
nav_loss_weight: 导航损失权重。
act_loss_weight: ACT 损失权重。
act_lambda_step: ACT 步数惩罚系数。
act_gamma: ACT 折扣因子。
eval_interval: 每 N epoch 评估一次。
save_dir: 模型权重保存目录。
dataset: 数据集名称，"longbench" | "narrativeqa" | "videomme"。
dataset_path: 数据集路径。
"""
lr: float
weight_decay: float
batch_size: int
max_epochs_phase1: int
max_epochs_phase2: int
nav_loss_weight: float
act_loss_weight: float
margin_loss_weight: float
act_lambda_step: float
act_gamma: float
eval_interval: int
save_dir: str
dataset: str
dataset_path: str
@dataclass
class HierRetrieverConfig:
"""Hierarchical Cross-Encoder 检索器参数。
属性:
backbone_model: 预训练语言模型名称。
num_heads: Cross-Attention 头数。
hidden_dim: 隐藏层维度。
dropout: Dropout 概率。
use_query_dep_weights: 是否使用 Query-dependent 层级权重。
level_weight_type: 层级权重类型，"fixed" | "query_dependent" | "hybrid"。
# Stage 1 参数
stage1_epochs: Stage 1 训练轮数。
stage1_lr: Stage 1 学习率。
stage1_batch_size: Stage 1 批大小。
stage1_num_negatives: Stage 1 每样本的负例数量。
stage1_temperature: Stage 1 温度参数。
# Stage 2 参数
stage2_epochs: Stage 2 训练轮数。
stage2_lr: Stage 2 学习率。
stage2_batch_size: Stage 2 批大小。
stage2_num_negatives: Stage 2 每样本的负例数量。
stage2_temperature: Stage 2 温度参数。
stage2_hard_neg_update_freq: Stage 2 硬负例更新频率。
stage2_hier_loss_weight: Stage 2 层级一致性损失权重。
# Stage 3 参数
stage3_enabled: 是否启用 Stage 3。
stage3_epochs: Stage 3 训练轮数。
stage3_lr: Stage 3 学习率。
# 推理参数
coarse_top_k: 粗排候选数量。
fine_top_k: 精排返回数量。
use_bm25: 是否使用 BM25 辅助粗排。
"""
# 模型参数
backbone_model: str
num_heads: int
hidden_dim: int
dropout: float
use_query_dep_weights: bool
level_weight_type: str
# Stage 1
stage1_epochs: int
stage1_lr: float
stage1_batch_size: int
stage1_num_negatives: int
stage1_temperature: float
# Stage 2
stage2_epochs: int
stage2_lr: float
stage2_batch_size: int
stage2_num_negatives: int
stage2_temperature: float
stage2_hard_neg_update_freq: int
stage2_hier_loss_weight: float
# Stage 3
stage3_enabled: bool
stage3_epochs: int
stage3_lr: float
# 推理
coarse_top_k: int
fine_top_k: int
use_bm25: bool
# ---------------------------------------------------------------------------
# 辅助函数
# ---------------------------------------------------------------------------
_SECTION_TO_CLASS: dict[str, type] = {
"tree": TreeConfig,
"embed": EmbedConfig,
"llm": LLMConfig,
"vlm": VLMConfig,
"retriever": RetrieverConfig,
"train": TrainConfig,
}
def _deep_merge(base: dict, override: dict) -> dict:
"""递归合并字典，override 优先覆盖 base。
参数:
base: 基础字典。
override: 覆盖字典。
返回:
合并后的新字典。
"""
merged = base.copy()
for key, value in override.items():
if key in merged and isinstance(merged[key], dict) and isinstance(value, dict):
merged[key] = _deep_merge(merged[key], value)
else:
merged[key] = value
return merged
def _apply_dotpath(d: dict, key: str, value: Any) -> None:
"""通过点路径设置嵌套字典的值。
支持 "retriever.num_heads" 风格的路径，自动拆分并逐级写入。
参数:
d: 目标字典。
key: 点分隔的路径，如 "retriever.num_heads"。
value: 要设置的值。
"""
parts = key.split(".")
current = d
for part in parts[:-1]:
if part not in current:
current[part] = {}
current = current[part]
current[parts[-1]] = value
def _coerce_value(raw: str, target_type: type) -> Any:
"""将 CLI 字符串值转换为目标类型。
参数:
raw: 原始字符串值。
target_type: 目标 Python 类型。
返回:
转换后的值。
"""
if target_type is bool:
return raw.lower() in ("true", "1", "yes")
if target_type is type(None):
return None if raw.lower() in ("none", "null", "") else raw
return target_type(raw)
# ---------------------------------------------------------------------------
# 顶层配置
# ---------------------------------------------------------------------------
@dataclass
class Config:
"""全局配置容器。
统一管理所有子模块配置，通过 ``Config.load()`` 加载。
属性:
tree: 树索引构建参数。
embed: 嵌入模型参数。
llm: 大语言模型参数。
vlm: 视觉语言模型参数。
retriever: TRM 检索器参数。
train: 训练参数。
hier_retriever: Hierarchical Cross-Encoder 检索器参数。
"""
tree: TreeConfig
embed: EmbedConfig
llm: LLMConfig
vlm: VLMConfig
retriever: RetrieverConfig
train: TrainConfig
hier_retriever: Optional[HierRetrieverConfig] = None
@classmethod
def load(
cls,
yaml_path: str,
cli_args: Optional[dict[str, str]] = None,
env_path: Optional[str] = None,
) -> "Config":
"""三层合并加载配置。
优先级: CLI args > .env > YAML。
参数:
yaml_path: YAML 配置文件路径。
cli_args: CLI 覆盖参数，键为点路径（如 "retriever.num_heads"），值为字符串。
env_path: .env 文件路径，默认为项目根目录的 .env。
返回:
完整的 Config 实例。
异常:
FileNotFoundError: YAML 文件不存在。
TypeError: YAML 中缺少必需字段。
"""
# Phase 1: 读取 YAML
yaml_file = Path(yaml_path)
if not yaml_file.exists():
raise FileNotFoundError(f"配置文件不存在: {yaml_path}")
with open(yaml_file, encoding="utf-8") as f:
base_dict: dict = yaml.safe_load(f)
# Phase 2: 读取 .env 覆盖敏感字段
if env_path is None:
env_file = Path(yaml_path).parent.parent / ".env"
else:
env_file = Path(env_path)
if env_file.exists():
env_vars = dotenv_values(str(env_file))
env_overrides: dict[str, dict[str, str]] = {}
# .env 变量名 → (配置节, 字段名) 的映射
_ENV_MAP: dict[str, tuple[str, str]] = {
"LLM_API_KEY": ("llm", "api_key"),
"LLM_MODEL": ("llm", "model"),
"LLM_API_URL": ("llm", "api_url"),
"VLM_API_KEY": ("vlm", "api_key"),
"VLM_MODEL": ("vlm", "model"),
"VLM_API_URL": ("vlm", "api_url"),
"EMBED_BACKEND": ("embed", "backend"),
"EMBED_MODEL": ("embed", "model_name"),
"EMBED_API_KEY": ("embed", "api_key"),
"EMBED_API_URL": ("embed", "api_url"),
}
for env_name, (section, field) in _ENV_MAP.items():
if env_vars.get(env_name):
env_overrides.setdefault(section, {})[field] = env_vars[env_name]
base_dict = _deep_merge(base_dict, env_overrides)
# Phase 3: CLI args 覆盖
if cli_args:
for dotpath, value in cli_args.items():
_apply_dotpath(base_dict, dotpath, value)
# Phase 4: 构造 dataclass（缺字段自动抛 TypeError）
sections = {}
for section_name, dc_class in _SECTION_TO_CLASS.items():
section_data = base_dict.get(section_name, {})
if not isinstance(section_data, dict):
# 对于 Optional 的 section，跳过
if section_name == "hier_retriever":
continue
raise TypeError(
f"配置节 '{section_name}' 必须是字典，实际为 {type(section_data)}"
)
sections[section_name] = dc_class(**section_data)
config = cls(**sections)
# 校验: embed_dim 一致性
ensure(
config.embed.embed_dim == config.retriever.embed_dim,
f"embed.embed_dim ({config.embed.embed_dim}) 与 "
f"retriever.embed_dim ({config.retriever.embed_dim}) 不一致",
)
log_msg("INFO", "配置加载完成", yaml=yaml_path)
return config
reference/video_tree_trm/embeddings.py
+190
View File
@@ -0,0 +1,190 @@
"""
嵌入服务模块
============
封装文本嵌入器，支持本地 sentence-transformers 和远程 OpenAI 兼容 API 两种后端。
提供统一的 ``embed()`` / ``embed_tensor()`` 接口，冻结不训练。
使用方式::
from video_tree_trm.embeddings import EmbeddingModel
from video_tree_trm.config import Config
cfg = Config.load("config/default.yaml")
model = EmbeddingModel(cfg.embed)
vecs = model.embed(["你好世界"]) # ndarray [1, D]
tens = model.embed_tensor(["你好"]) # Tensor [1, D]
"""
from __future__ import annotations
from typing import List, Union
import numpy as np
import torch
from numpy import ndarray
from torch import Tensor
from utils.logger_system import ensure, log_msg
from video_tree_trm.config import EmbedConfig
class EmbeddingModel:
"""文本嵌入器封装（冻结），支持本地和远程双后端。
本地模式: 使用 sentence-transformers 加载 HuggingFace 模型，本地推理。
远程模式: 调用 OpenAI 兼容 API（如 GPUStack 上的 qwen3-embedding）。
属性:
dim: 嵌入维度 D。
"""
def __init__(self, config: EmbedConfig) -> None:
"""初始化嵌入模型。
参数:
config: 嵌入配置，包含 backend、model_name、embed_dim 等。
异常:
ValueError: backend 不是 "local" 或 "remote"。
ValueError: 远程模式缺少 api_key 或 api_url。
"""
ensure(
config.backend in ("local", "remote"),
f"embed.backend 必须为 'local' 或 'remote'，实际为 '{config.backend}'",
)
self._backend = config.backend
self._dim = config.embed_dim
if self._backend == "local":
self._init_local(config)
else:
self._init_remote(config)
log_msg(
"INFO", "嵌入模型初始化完成", backend=self._backend, model=config.model_name
)
# ------------------------------------------------------------------
# 初始化
# ------------------------------------------------------------------
def _init_local(self, config: EmbedConfig) -> None:
"""初始化本地 sentence-transformers 模型。
参数:
config: 嵌入配置。
"""
from sentence_transformers import SentenceTransformer
self._model = SentenceTransformer(config.model_name, device=config.device)
self._model.eval()
# 冻结所有参数
for param in self._model.parameters():
param.requires_grad = False
actual_dim = self._model.get_sentence_embedding_dimension()
ensure(
actual_dim == self._dim,
f"模型实际维度 ({actual_dim}) 与配置 embed_dim ({self._dim}) 不一致",
)
def _init_remote(self, config: EmbedConfig) -> None:
"""初始化远程 OpenAI 兼容 API 客户端。
参数:
config: 嵌入配置。
"""
ensure(bool(config.api_key), "远程模式必须提供 embed.api_key")
ensure(bool(config.api_url), "远程模式必须提供 embed.api_url")
from openai import OpenAI
self._client = OpenAI(base_url=config.api_url, api_key=config.api_key)
self._model_name = config.model_name
# ------------------------------------------------------------------
# 公共接口
# ------------------------------------------------------------------
@property
def dim(self) -> int:
"""嵌入维度 D。"""
return self._dim
def embed(self, texts: Union[str, List[str]]) -> ndarray:
"""文本 → 嵌入向量（L2 归一化）。
参数:
texts: 单条文本或文本列表。
返回:
[N, D] ndarray，每行 L2 范数为 1.0。单条文本时 N=1。
"""
if isinstance(texts, str):
texts = [texts]
if self._backend == "local":
return self._embed_local(texts)
return self._embed_remote(texts)
def embed_tensor(self, texts: Union[str, List[str]]) -> Tensor:
"""文本 → 嵌入 Tensor（L2 归一化）。
参数:
texts: 单条文本或文本列表。
返回:
[N, D] torch.Tensor（float32）。
"""
arr = self.embed(texts)
return torch.from_numpy(arr).float()
# ------------------------------------------------------------------
# 后端实现
# ------------------------------------------------------------------
def _embed_local(self, texts: List[str]) -> ndarray:
"""本地 sentence-transformers 推理。
参数:
texts: 文本列表。
返回:
[N, D] ndarray，L2 归一化。
"""
with torch.no_grad():
embeddings = self._model.encode(
texts,
normalize_embeddings=True,
convert_to_numpy=True,
)
# sentence-transformers encode 返回 ndarray [N, D]
if embeddings.ndim == 1:
embeddings = embeddings.reshape(1, -1)
return embeddings
def _embed_remote(self, texts: List[str]) -> ndarray:
"""远程 OpenAI 兼容 API 调用。
参数:
texts: 文本列表。
返回:
[N, D] ndarray，L2 归一化。
"""
response = self._client.embeddings.create(
model=self._model_name,
input=texts,
)
# 按 index 排序，确保顺序一致
sorted_data = sorted(response.data, key=lambda x: x.index)
embeddings = np.array(
[item.embedding for item in sorted_data], dtype=np.float32
)
# L2 归一化
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
norms = np.maximum(norms, 1e-12) # 避免除零
embeddings = embeddings / norms
return embeddings
reference/video_tree_trm/llm_client.py
+421
View File
@@ -0,0 +1,421 @@
"""
LLM/VLM 客户端模块
==================
统一封装 LLM（纯文本）和 VLM（多模态）API 调用。
仅支持 OpenAI-compatible 接口，通过配置 api_url + model 适配不同服务商
（如 Qwen DashScope、OpenAI、本地推理服务等）。
提供同步版（chat / chat_with_images）和异步版（chat_async / chat_with_images_async）。
异步版基于 openai.AsyncOpenAI，适配 asyncio 事件循环，零线程阻塞。
使用方式::
from video_tree_trm.llm_client import LLMClient
from video_tree_trm.config import Config
cfg = Config.load("config/default.yaml")
vlm = LLMClient(cfg.vlm)
# 同步
answer = vlm.chat_with_images("图中有什么？", images=["frame.jpg"])
# 异步
import asyncio
answer = asyncio.run(vlm.chat_with_images_async("图中有什么？", images=["frame.jpg"]))
"""
from __future__ import annotations
import asyncio
import base64
import os
import re
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, List, Optional, Union
import httpx
import openai
from utils.logger_system import log_exception, log_msg
from video_tree_trm.config import LLMConfig, VLMConfig
# 502/503 时的重试参数
_RETRY_STATUS_CODES = {502, 503}
_MAX_RETRIES = 20 # 最多重试次数（约等待 20+ 分钟）
_RETRY_BASE_WAIT = 60 # 首次等待 60 秒
_RETRY_MAX_WAIT = 300 # 单次等待上限 5 分钟
def _call_with_retry(fn, label: str):
"""对 fn() 调用执行指数退避重试（重试 502/503 及超时）。
参数:
fn: 无参调用的函数，返回 API response。
label: 日志标识（如方法名）。
返回:
fn() 的返回值。
异常:
openai.OpenAIError: 超过最大重试次数或非可重试错误时抛出。
"""
wait = _RETRY_BASE_WAIT
for attempt in range(1, _MAX_RETRIES + 1):
try:
return fn()
except openai.APITimeoutError:
log_msg(
"WARNING",
f"{label} 请求超时，等待 {wait}s 后重试",
attempt=attempt,
max_retries=_MAX_RETRIES,
)
time.sleep(wait)
wait = min(wait * 2, _RETRY_MAX_WAIT)
except openai.InternalServerError as exc:
status = getattr(exc, "status_code", None)
if status not in _RETRY_STATUS_CODES:
raise
log_msg(
"WARNING",
f"{label} 遇到 {status}，等待 {wait}s 后重试",
attempt=attempt,
max_retries=_MAX_RETRIES,
)
time.sleep(wait)
wait = min(wait * 2, _RETRY_MAX_WAIT)
raise RuntimeError(f"{label} 已重试 {_MAX_RETRIES} 次仍失败")
async def _async_call_with_retry(coro_fn, label: str):
"""异步版指数退避重试，适配 asyncio 事件循环。
参数:
coro_fn: 无参调用的协程工厂函数（每次调用返回新协程）。
label: 日志标识（如方法名）。
返回:
coro_fn() 的返回值。
实现细节:
使用 await asyncio.sleep() 替代 time.sleep()，不阻塞事件循环。
每次重试需重新调用 coro_fn() 构造新协程（协程不可复用）。
"""
wait = _RETRY_BASE_WAIT
for attempt in range(1, _MAX_RETRIES + 1):
try:
return await coro_fn()
except openai.APITimeoutError:
log_msg(
"WARNING",
f"{label} 请求超时，等待 {wait}s 后重试",
attempt=attempt,
max_retries=_MAX_RETRIES,
)
await asyncio.sleep(wait)
wait = min(wait * 2, _RETRY_MAX_WAIT)
except openai.InternalServerError as exc:
status = getattr(exc, "status_code", None)
if status not in _RETRY_STATUS_CODES:
raise
log_msg(
"WARNING",
f"{label} 遇到 {status}，等待 {wait}s 后重试",
attempt=attempt,
max_retries=_MAX_RETRIES,
)
await asyncio.sleep(wait)
wait = min(wait * 2, _RETRY_MAX_WAIT)
raise RuntimeError(f"{label} 已重试 {_MAX_RETRIES} 次仍失败")
class LLMClient:
"""OpenAI-compatible LLM/VLM 统一客户端。
同时提供同步接口（chat / chat_with_images）和异步接口（chat_async / chat_with_images_async）。
异步接口使用独立的 AsyncOpenAI 实例，零线程阻塞，与 asyncio.Semaphore 配合实现真并发。
属性:
_config: LLMConfig 或 VLMConfig 配置对象。
_client: openai.OpenAI 同步客户端。
_async_client: openai.AsyncOpenAI 异步客户端。
_extra_body: 关闭 Qwen3 thinking 模式的额外参数。
"""
def __init__(self, config: Union[LLMConfig, VLMConfig]) -> None:
"""初始化 LLM/VLM 客户端（同步 + 异步双客户端）。
参数:
config: LLMConfig 或 VLMConfig，包含 api_key、api_url、model 等参数。
异常:
ValueError: api_key 或 api_url 为空时抛出。
"""
if not config.api_key:
raise ValueError(
"LLMClient 初始化失败: config.api_key 不能为空，请在 .env 中设置"
)
if not config.api_url:
raise ValueError(
"LLMClient 初始化失败: config.api_url 不能为空，请在 config/default.yaml 中设置"
)
self._config = config
# 同步客户端（向后兼容）
self._client = openai.OpenAI(
api_key=config.api_key,
base_url=config.api_url,
http_client=httpx.Client(proxy=None),
)
# 异步客户端（asyncio 场景，零阻塞）
self._async_client = openai.AsyncOpenAI(
api_key=config.api_key,
base_url=config.api_url,
http_client=httpx.AsyncClient(proxy=None),
)
# 关闭 Qwen3 thinking 模式（vLLM 正确格式）
self._extra_body: Dict = {"chat_template_kwargs": {"enable_thinking": False}}
log_msg(
"INFO", "LLMClient 初始化完成", model=config.model, api_url=config.api_url
)
# ── 同步接口（向后兼容）─────────────────────────────────────────────────
def chat(self, prompt: str, max_tokens: Optional[int] = None) -> str:
"""纯文本单轮对话（同步）。
参数:
prompt: 用户输入文本。
max_tokens: 最大生成 token 数，为 None 时使用 config.max_tokens。
返回:
生成的文本字符串。
"""
messages = self._build_messages(prompt)
tokens = max_tokens if max_tokens is not None else self._config.max_tokens
try:
response = _call_with_retry(
lambda: self._client.chat.completions.create(
model=self._config.model,
messages=messages,
max_tokens=tokens,
temperature=self._config.temperature,
extra_body=self._extra_body,
),
label="LLMClient.chat",
)
return self._strip_thinking(response.choices[0].message.content)
except Exception as exc:
log_exception("LLMClient.chat 调用失败", exc)
raise
def chat_with_images(
self,
prompt: str,
images: List[str],
max_tokens: Optional[int] = None,
) -> str:
"""多模态单轮对话（VLM，同步）。
参数:
prompt: 文本指令。
images: 图像列表，每项可为本地文件路径或已编码的 base64 字符串。
max_tokens: 最大生成 token 数，为 None 时使用 config.max_tokens。
返回:
生成的文本字符串。
"""
encoded = [self._encode_image(img) for img in images]
messages = self._build_messages(prompt, images=encoded)
tokens = max_tokens if max_tokens is not None else self._config.max_tokens
try:
response = _call_with_retry(
lambda: self._client.chat.completions.create(
model=self._config.model,
messages=messages,
max_tokens=tokens,
temperature=self._config.temperature,
extra_body=self._extra_body,
),
label="LLMClient.chat_with_images",
)
return self._strip_thinking(response.choices[0].message.content)
except Exception as exc:
log_exception("LLMClient.chat_with_images 调用失败", exc)
raise
def batch_chat(
self,
prompts: List[str],
max_tokens: Optional[int] = None,
) -> List[str]:
"""批量纯文本并发对话，保序返回（同步）。
参数:
prompts: 文本输入列表。
max_tokens: 最大生成 token 数。
返回:
与 prompts 等长的生成文本列表，顺序与输入对应。
"""
results: List[str] = [""] * len(prompts)
with ThreadPoolExecutor(max_workers=8) as executor:
future_to_idx = {
executor.submit(self.chat, prompt, max_tokens): idx
for idx, prompt in enumerate(prompts)
}
for future in as_completed(future_to_idx):
idx = future_to_idx[future]
results[idx] = future.result()
return results
# ── 异步接口（asyncio 事件循环，零阻塞）──────────────────────────────────
async def chat_async(self, prompt: str, max_tokens: Optional[int] = None) -> str:
"""纯文本单轮对话（异步，零线程阻塞）。
参数:
prompt: 用户输入文本。
max_tokens: 最大生成 token 数，为 None 时使用 config.max_tokens。
返回:
生成的文本字符串。
实现细节:
使用 AsyncOpenAI 客户端，await 期间事件循环可处理其他协程，
配合 asyncio.Semaphore 实现受控并发。
"""
messages = self._build_messages(prompt)
tokens = max_tokens if max_tokens is not None else self._config.max_tokens
try:
response = await _async_call_with_retry(
lambda: self._async_client.chat.completions.create(
model=self._config.model,
messages=messages,
max_tokens=tokens,
temperature=self._config.temperature,
extra_body=self._extra_body,
),
label="LLMClient.chat_async",
)
return self._strip_thinking(response.choices[0].message.content)
except Exception as exc:
log_exception("LLMClient.chat_async 调用失败", exc)
raise
async def chat_with_images_async(
self,
prompt: str,
images: List[str],
max_tokens: Optional[int] = None,
) -> str:
"""多模态单轮对话（VLM，异步，零线程阻塞）。
参数:
prompt: 文本指令。
images: 图像列表，每项可为本地文件路径或已编码的 base64 字符串。
max_tokens: 最大生成 token 数，为 None 时使用 config.max_tokens。
返回:
生成的文本字符串。
实现细节:
图像编码（磁盘读取 + base64）在默认线程池执行器中并行执行，
避免阻塞事件循环；VLM API 调用通过 AsyncOpenAI 零阻塞。
"""
loop = asyncio.get_event_loop()
# 并行编码所有图像（I/O 密集，交给线程池）
encoded: List[str] = await asyncio.gather(
*[loop.run_in_executor(None, self._encode_image, img) for img in images]
)
messages = self._build_messages(prompt, images=list(encoded))
tokens = max_tokens if max_tokens is not None else self._config.max_tokens
try:
response = await _async_call_with_retry(
lambda: self._async_client.chat.completions.create(
model=self._config.model,
messages=messages,
max_tokens=tokens,
temperature=self._config.temperature,
extra_body=self._extra_body,
),
label="LLMClient.chat_with_images_async",
)
return self._strip_thinking(response.choices[0].message.content)
except Exception as exc:
log_exception("LLMClient.chat_with_images_async 调用失败", exc)
raise
# ── 私有辅助方法 ──────────────────────────────────────────────────────────
@staticmethod
def _strip_thinking(content: str) -> str:
"""剥离 Qwen3 thinking 模式生成的 <think>...</think> 块。
参数:
content: VLM/LLM 原始返回文本（可能含 <think> 块）。
返回:
去除 think 块后的纯净文本。
实现细节:
当 API 参数无法完全禁用 thinking 时作为兜底保障。
<think> 块可能跨多行，使用 DOTALL 模式匹配。
"""
cleaned = re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL)
return cleaned.strip()
def _encode_image(self, path_or_b64: str) -> str:
"""将图像转换为 data URI 格式的 base64 字符串。
参数:
path_or_b64: 本地文件路径，或已是 "data:image/...;base64,..." 格式的字符串。
返回:
"data:image/jpeg;base64,<base64数据>" 格式字符串。
异常:
FileNotFoundError: 指定路径文件不存在时抛出。
"""
if "base64," in path_or_b64:
return path_or_b64
if not os.path.exists(path_or_b64):
raise FileNotFoundError(f"图像文件不存在: {path_or_b64}")
with open(path_or_b64, "rb") as f:
raw = f.read()
b64_data = base64.b64encode(raw).decode("utf-8")
ext = os.path.splitext(path_or_b64)[1].lower()
mime = "image/png" if ext == ".png" else "image/jpeg"
return f"data:{mime};base64,{b64_data}"
def _build_messages(
self,
prompt: str,
images: Optional[List[str]] = None,
) -> List[Dict]:
"""拼装 OpenAI-compatible 消息结构。
参数:
prompt: 文本指令。
images: 已编码的 base64 data URI 列表（可为 None）。
返回:
OpenAI messages 格式的列表。
实现细节:
- 无图像：content 为纯字符串。
- 有图像：content 为列表，图像在前，文本在后。
"""
if not images:
return [{"role": "user", "content": prompt}]
content: List[Dict] = [
{"type": "image_url", "image_url": {"url": img}} for img in images
]
content.append({"type": "text", "text": prompt})
return [{"role": "user", "content": content}]
reference/video_tree_trm/pipeline.py
+266
View File
@@ -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")
)
reference/video_tree_trm/text_tree_builder.py
+466
View File
@@ -0,0 +1,466 @@
"""
文本树构建模块
==============
将长文本通过 L2 轴心策略转化为三层 TreeIndex。
构建策略::
Step 1: _segment_text — 结构切分，确定 L1/L2 边界
Step 2: L2 先行 — 从原始内容独立生成 L2 摘要（batch_chat 并发）
Step 3: L3 向下 — 原始段落文本直接作为 L3，无需二次生成
Step 4: L1 向上 — 聚合 L2 描述，生成 L1 粗粒度摘要
Step 5: 组装 TreeIndex
L2 轴心策略解决了循环依赖：
- L2 描述不依赖 L3，从原始段落直接生成
- L3 直接使用原始段落文本，不调用 LLM
- L1 聚合 L2 描述，保证完整覆盖
使用方式::
builder = TextTreeBuilder(embed_model, llm_client, config.tree)
index = builder.build(text, source_path="docs/my_doc.txt")
index.save("cache/my_doc.pkl")
"""
from __future__ import annotations
import json
import re
from datetime import datetime
from typing import List, Tuple
import numpy as np
from utils.logger_system import ensure, log_json, log_msg
from video_tree_trm.config import TreeConfig
from video_tree_trm.llm_client import LLMClient
from video_tree_trm.tree_index import (
IndexMeta,
L1Node,
L2Node,
L3Node,
TreeIndex,
)
# ---------------------------------------------------------------------------
# Prompt 常量
# ---------------------------------------------------------------------------
_L2_PROMPT = (
"用1-2句话描述以下段落的核心内容，与同级小节形成区分:\n\n{text}"
)
_L1_PROMPT = (
"用2-3句话总结以下小节的核心内容:\n\n{l2_descriptions}"
)
_SEG_PROMPT = (
"将以下文本分成若干语义段落，每段为完整语义单元。\n"
"只返回 JSON 数组，格式: [\"段落1\", \"段落2\", ...]，不要其他内容。\n"
"文本:\n\n{text}"
)
# ---------------------------------------------------------------------------
# 辅助函数
# ---------------------------------------------------------------------------
def _chunk(lst: List[str], size: int) -> List[List[str]]:
"""将列表等长分块（固定步长，无重叠）。
参数:
lst: 待分块的列表。
size: 每块的最大长度。
返回:
分块后的列表，每个元素为一个子列表。
"""
return [lst[i : i + size] for i in range(0, len(lst), size)]
# ---------------------------------------------------------------------------
# 主类
# ---------------------------------------------------------------------------
class TextTreeBuilder:
"""文本模态树构建器。
将长文本通过 L2 轴心策略（先构建 L2，再向下扩展 L3，向上聚合 L1）
转化为三层 TreeIndex。节点 embedding 均为 None（由 Pipeline.embed_all 延迟填充）。
属性:
llm: LLM 客户端。
config: 树构建配置。
"""
def __init__(
self,
llm: LLMClient,
config: TreeConfig,
) -> None:
"""初始化文本树构建器。
参数:
llm: 已初始化的 LLM 客户端（LLMClient）。
config: 树构建配置（TreeConfig），关键字段 max_paragraphs_per_l2。
实现细节:
构建器不持有 EmbeddingModel，所有 embedding 延迟到检索阶段由 Pipeline 统一计算。
"""
self.llm = llm
self.config = config
# ------------------------------------------------------------------
# 公共接口
# ------------------------------------------------------------------
def build(self, text: str, source_path: str) -> TreeIndex:
"""将长文本构建为三层 TreeIndex。
参数:
text: 输入长文本（UTF-8 字符串）。
source_path: 原始文件路径，写入 IndexMeta。
返回:
三层 TreeIndex 对象。
实现细节:
1. _segment_text 切分文本 → List[List[str]]（外层=L1，内层=L2段落组）
2. 将所有 L2 段落组的 prompt 批量送入 llm.batch_chat()，并发获取摘要
3. 逐层组装 L3→L2→L1 节点
4. 构建 TreeIndex 并写入日志
"""
ensure(bool(text.strip()), "输入文本不能为空")
log_msg("INFO", "开始构建文本树索引", source_path=source_path)
# Phase 1: 结构切分
sections = self._segment_text(text)
ensure(len(sections) > 0, "文本切分结果为空")
log_msg(
"INFO",
"文本切分完成",
l1_count=len(sections),
l2_groups=[len(s) for s in sections],
)
# Phase 2: 收集所有 L2 段落组，批量生成摘要（L2 先行）
all_groups: List[Tuple[int, int, List[str]]] = []
for i, section_paragraphs in enumerate(sections):
for j, group in enumerate(
_chunk(section_paragraphs, self.config.max_paragraphs_per_l2)
):
all_groups.append((i, j, group))
l2_prompts = [
_L2_PROMPT.format(text="\n\n".join(group))
for _, _, group in all_groups
]
l2_descs = self.llm.batch_chat(l2_prompts)
log_msg("INFO", "L2 摘要生成完成", total_l2=len(l2_descs))
# Phase 3-4: 按 L1 组装三层节点
# 构建索引映射：(i, j) → 在 all_groups / l2_descs 中的位置
group_index: dict = {
(i, j): idx for idx, (i, j, _) in enumerate(all_groups)
}
l1_nodes: List[L1Node] = []
for i, section_paragraphs in enumerate(sections):
groups = _chunk(section_paragraphs, self.config.max_paragraphs_per_l2)
l2_nodes: List[L2Node] = []
for j, group in enumerate(groups):
idx = group_index[(i, j)]
desc = l2_descs[idx]
l3_nodes = self._build_l3_from_paragraphs(group, i, j)
l2_node = L2Node(
id=f"l1_{i}_l2_{j}",
description=desc,
embedding=None,
time_range=None,
children=l3_nodes,
)
l2_nodes.append(l2_node)
l1_node = self._build_l1(l2_nodes, f"l1_{i}")
l1_nodes.append(l1_node)
# Phase 5: 组装 TreeIndex（embedding 延迟到 Pipeline.embed_all，此处为 None）
metadata = IndexMeta(
source_path=source_path,
modality="text",
created_at=datetime.now().isoformat(),
)
index = TreeIndex(metadata=metadata, roots=l1_nodes)
total_l2 = sum(len(r.children) for r in l1_nodes)
total_l3 = sum(
len(l2.children) for r in l1_nodes for l2 in r.children
)
log_json(
"text_tree_build",
{
"source_path": source_path,
"l1_count": len(l1_nodes),
"l2_count": total_l2,
"l3_count": total_l3,
"embedded": False,
},
)
log_msg(
"INFO",
"文本树索引构建完成",
l1=len(l1_nodes),
l2=total_l2,
l3=total_l3,
)
return index
# ------------------------------------------------------------------
# 内部方法：切分策略
# ------------------------------------------------------------------
def _segment_text(self, text: str) -> List[List[str]]:
"""结构切分长文本为 L1/L2 层次。
参数:
text: 输入文本。
返回:
sections[i] = [paragraph_1, paragraph_2, ...]
外层列表 = L1 段（章节），内层列表 = L2 单元（段落组内段落）。
策略:
有 Markdown 标题 → 正则解析 #/## 边界
无 Markdown 标题 → LLM 单次调用语义分段
"""
if self._detect_toc(text):
log_msg("INFO", "检测到 Markdown 标题，使用正则切分")
return self._segment_with_regex(text)
else:
log_msg("INFO", "未检测到 Markdown 标题，使用 LLM 语义分段")
return self._segment_with_llm(text)
def _detect_toc(self, text: str) -> bool:
"""检测文本是否包含 Markdown 标题（有 ToC 结构）。
参数:
text: 输入文本。
返回:
True 表示有 # 或 ## 开头的标题行，False 表示无。
"""
return bool(re.search(r"^#{1,2}\s+\S", text, re.MULTILINE))
def _segment_with_regex(self, text: str) -> List[List[str]]:
"""通过正则解析 Markdown 标题边界进行结构切分。
参数:
text: 含 Markdown 标题的文本。
返回:
List[List[str]]，外层=L1章节，内层=该章节下的段落列表。
若二级标题下段落数超过 max_paragraphs_per_l2，则进一步等长分块。
实现细节:
- # 标题 → L1 边界
- ## 标题 → L2 边界
- ### 及以下标题视为段落内容，收集到最近 L2 段落组
- 空段落过滤掉
"""
lines = text.split("\n")
sections: List[List[str]] = [] # 外层=L1
current_section: List[str] = [] # 当前 L1 下的段落（扁平）
current_para_lines: List[str] = [] # 积累段落文本行
def _flush_para() -> None:
"""将当前积累的行合并为一个段落加入 current_section。"""
para = "\n".join(current_para_lines).strip()
if para:
current_section.append(para)
current_para_lines.clear()
def _flush_section() -> None:
"""将当前 section 保存，重置。"""
_flush_para()
if current_section:
sections.append(list(current_section))
current_section.clear()
for line in lines:
h1_match = re.match(r"^#\s+(.+)", line)
h2_match = re.match(r"^##\s+(.+)", line)
if h1_match:
# L1 边界：保存当前 section
_flush_section()
# 将 H1 标题本身作为第一段落（可选：也可忽略标题行）
title = h1_match.group(1).strip()
if title:
current_section.append(title)
elif h2_match:
# L2 边界：只冲刷当前段落，不切换 section
_flush_para()
title = h2_match.group(1).strip()
if title:
current_section.append(title)
else:
# 普通内容行（含 ###、####、正文段落）
if line.strip() == "":
# 空行触发段落分隔
_flush_para()
else:
current_para_lines.append(line)
_flush_section()
# 若没有产生任何 section（如文本只有一个 L1），保底处理
if not sections:
sections = [self._collect_paragraphs(text)]
# 对超出 max_paragraphs_per_l2 的段落组不做处理（由 build() 负责分块）
return sections
def _segment_with_llm(self, text: str) -> List[List[str]]:
"""通过 LLM 单次调用语义分段无结构文本。
参数:
text: 无 Markdown 标题的纯文本。
返回:
List[List[str]]，只有一个外层元素（整篇视为单个 L1）。
内层为 LLM 返回的语义段落列表。
异常:
ValueError: LLM 返回的内容无法解析为 JSON 数组。
"""
prompt = _SEG_PROMPT.format(text=text)
raw = self.llm.chat(prompt)
# 尝试从返回结果中提取 JSON 数组
raw = raw.strip()
# 提取可能被 markdown 代码块包裹的 JSON
code_match = re.search(r"```(?:json)?\s*(\[.*?\])\s*```", raw, re.DOTALL)
if code_match:
raw = code_match.group(1)
ensure(
raw.startswith("["),
f"LLM 语义分段返回格式错误，期望 JSON 数组，实际: {raw[:100]}",
)
try:
paragraphs: List[str] = json.loads(raw)
except json.JSONDecodeError as e:
raise ValueError(f"LLM 语义分段 JSON 解析失败: {e}\n原始输出: {raw}") from e
ensure(isinstance(paragraphs, list), "LLM 返回值不是列表")
ensure(len(paragraphs) > 0, "LLM 语义分段返回空列表")
# 过滤空段落
paragraphs = [p.strip() for p in paragraphs if p.strip()]
log_msg("INFO", "LLM 语义分段完成", paragraph_count=len(paragraphs))
return [paragraphs] # 整篇视为单个 L1
def _collect_paragraphs(self, text: str) -> List[str]:
"""按双换行符切分段落（保底策略）。
参数:
text: 输入文本。
返回:
非空段落列表。
"""
paras = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
return paras if paras else [text.strip()]
# ------------------------------------------------------------------
# 内部方法：节点构建
# ------------------------------------------------------------------
def _build_l2(self, paragraphs: List[str], l2_id: str) -> L2Node:
"""将段落组构建为 L2 节点（含 LLM 摘要和嵌入）。
参数:
paragraphs: 该 L2 节点下的段落文本列表。
l2_id: 节点 ID。
返回:
L2Node（children 为空，由调用方填充）。
"""
ensure(len(paragraphs) > 0, f"L2 节点 {l2_id} 的段落列表为空")
prompt = _L2_PROMPT.format(text="\n\n".join(paragraphs))
description = self.llm.chat(prompt)
return L2Node(
id=l2_id,
description=description,
embedding=None,
time_range=None,
)
def _build_l3_from_paragraphs(
self,
paragraphs: List[str],
l1_i: int,
l2_j: int,
) -> List[L3Node]:
"""将段落列表批量构建为 L3 节点（原始文本直接复用，不调用 LLM）。
参数:
paragraphs: 段落文本列表。
l1_i: 父 L1 索引（用于生成 ID）。
l2_j: 父 L2 索引（用于生成 ID）。
返回:
L3Node 列表，description == raw_content == 原始段落文本。
实现细节:
使用 embed.embed(paragraphs) 批量嵌入，一次调用获取全部向量。
"""
ensure(len(paragraphs) > 0, f"L3 段落列表为空 (l1={l1_i}, l2={l2_j})")
nodes: List[L3Node] = []
for k, para in enumerate(paragraphs):
nodes.append(
L3Node(
id=f"l1_{l1_i}_l2_{l2_j}_l3_{k}",
description=para,
embedding=None,
raw_content=para,
frame_path=None,
timestamp=None,
)
)
return nodes
def _build_l1(self, l2_children: List[L2Node], l1_id: str) -> L1Node:
"""聚合 L2 描述，构建 L1 节点（含 LLM 摘要和嵌入）。
参数:
l2_children: 该 L1 节点下的所有 L2 节点。
l1_id: 节点 ID。
返回:
L1Node（children 已由调用方赋值，或在此赋值）。
实现细节:
将所有 L2 描述拼接，用序号标注后送入 LLM 生成 2-3 句摘要。
"""
ensure(len(l2_children) > 0, f"L1 节点 {l1_id} 没有 L2 子节点")
l2_descriptions = "\n".join(
f"{idx + 1}. {node.description}"
for idx, node in enumerate(l2_children)
)
prompt = _L1_PROMPT.format(l2_descriptions=l2_descriptions)
summary = self.llm.chat(prompt)
return L1Node(
id=l1_id,
summary=summary,
embedding=None,
time_range=None,
children=l2_children,
)
reference/video_tree_trm/tree_index.py
+642
View File
@@ -0,0 +1,642 @@
"""
三层树索引核心数据结构
======================
定义 Video-Tree-TRM 的三层树状索引结构，是所有后续模块
（builder、retriever、losses、pipeline）的基础依赖。
数据结构层次::
TreeIndex
└─ List[L1Node] 全局叙事节点
└─ List[L2Node] 片段级语义节点
└─ List[L3Node] 帧/细节级节点
与参考项目 (Tree-TRM/video_pyramid.py) 的关键区别：
- 统一嵌入空间：所有 embedding 均来自 text_embed()，无跨模态问题
- 序列化方式：pickle 整体序列化（而非 JSON + NPY 分文件存储）
- L3 全文本化：无需 VisualProjectionLayer
"""
from __future__ import annotations
import base64
import json
import pickle
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import numpy as np
from utils.logger_system import ensure, log_msg
# ---------------------------------------------------------------------------
# Embedding 序列化辅助函数
# ---------------------------------------------------------------------------
def _embed_to_str(arr: Optional[np.ndarray]) -> Optional[str]:
"""float32 ndarray → base64 字符串（用于 JSON 序列化）。
参数:
arr: float32 数组，形状任意。
返回:
base64 编码字符串，或 None（输入为 None 时）。
"""
if arr is None:
return None
return base64.b64encode(arr.astype(np.float32).tobytes()).decode()
def _embed_from_str(s: Optional[str]) -> Optional[np.ndarray]:
"""base64 字符串 → float32 ndarray（用于 JSON 反序列化）。
参数:
s: base64 编码字符串。
返回:
float32 数组，或 None（输入为 None/空时）。
"""
if s is None or s == "":
return None
return np.frombuffer(base64.b64decode(s), dtype=np.float32)
# ---------------------------------------------------------------------------
# 元数据
# ---------------------------------------------------------------------------
@dataclass
class IndexMeta:
"""树索引元数据。
Attributes:
source_path: 原始数据路径（视频文件或文本文件）。
modality: 数据模态，"text" 或 "video"。
embed_model: 嵌入模型名称（建树时为 None，embed_all 后填充）。
embed_dim: 嵌入向量维度（建树时为 None，embed_all 后填充）。
created_at: 创建时间（ISO 格式字符串）。
"""
source_path: str
modality: str
embed_model: Optional[str] = None
embed_dim: Optional[int] = None
created_at: str = field(default_factory=lambda: datetime.now().isoformat())
# ---------------------------------------------------------------------------
# 节点数据结构
# ---------------------------------------------------------------------------
@dataclass
class L3Node:
"""L3 帧/细节级节点（叶子层）。
代表最细粒度的语义单元，对应一个具体的描述片段。
Attributes:
id: 节点唯一标识。
description: 文本描述。
embedding: 文本嵌入向量，形状 [D]，float32。
raw_content: 原始文本内容（可选）。
frame_path: 关联的帧图像路径（可选，仅视频模态）。
timestamp: 对应的时间戳（秒，可选）。
"""
id: str
description: str
embedding: Optional[np.ndarray] = None # [D]，build 时为 None，embed_all 后填充
raw_content: Optional[str] = None
frame_path: Optional[str] = None
timestamp: Optional[float] = None
@dataclass
class L2Node:
"""L2 片段级语义节点（中间层）。
连接 L1 宏观叙事与 L3 细节描述。
Attributes:
id: 节点唯一标识。
description: 片段文本描述。
embedding: 文本嵌入向量，形状 [D]，float32。
time_range: 时间范围 (start, end)（秒，可选）。
children: 所属的 L3 子节点列表。
"""
id: str
description: str
embedding: Optional[np.ndarray] = None # [D]，build 时为 None，embed_all 后填充
time_range: Optional[Tuple[float, float]] = None
children: List[L3Node] = field(default_factory=list)
@dataclass
class L1Node:
"""L1 全局叙事节点（根层）。
代表最粗粒度的语义单元，包含宏观事件摘要。
Attributes:
id: 节点唯一标识。
summary: 高层叙事摘要。
embedding: 文本嵌入向量，形状 [D]，float32。
time_range: 时间范围 (start, end)（秒，可选）。
children: 所属的 L2 子节点列表。
"""
id: str
summary: str
embedding: Optional[np.ndarray] = None # [D]，build 时为 None，embed_all 后填充
time_range: Optional[Tuple[float, float]] = None
children: List[L2Node] = field(default_factory=list)
# ------------------------------------------------------------------
# JSON 辅助方法（单个 L1 段的轻量序列化）
# ------------------------------------------------------------------
def to_dict(self, include_embedding: bool = False) -> Dict[str, Any]:
"""将当前 L1 节点（及其全部 L2/L3 子树）序列化为纯 dict。
参数:
include_embedding: 若 True，将 embedding 向量序列化为 base64 字符串。
返回:
包含 id/summary/time_range/children 的字典，可选包含 embedding。
"""
def l3_to_dict(n: L3Node) -> Dict[str, Any]:
d = {
"id": n.id,
"description": n.description,
"timestamp": n.timestamp,
"frame_path": n.frame_path,
"raw_content": n.raw_content,
}
if include_embedding:
d["embedding"] = _embed_to_str(n.embedding)
return d
def l2_to_dict(n: L2Node) -> Dict[str, Any]:
d = {
"id": n.id,
"description": n.description,
"time_range": list(n.time_range) if n.time_range else None,
"children": [l3_to_dict(c) for c in n.children],
}
if include_embedding:
d["embedding"] = _embed_to_str(n.embedding)
return d
d = {
"id": self.id,
"summary": self.summary,
"time_range": list(self.time_range) if self.time_range else None,
"children": [l2_to_dict(c) for c in self.children],
}
if include_embedding:
d["embedding"] = _embed_to_str(self.embedding)
return d
@staticmethod
def from_dict(d: Dict[str, Any]) -> "L1Node":
"""从 dict 反序列化单个 L1 节点（支持 embedding 恢复）。
参数:
d: to_dict() 输出的字典，可包含 embedding 字段。
返回:
L1Node 实例（embedding 自动从 base64 恢复，若无则为 None）。
"""
l2_nodes: List[L2Node] = []
for l2d in d.get("children", []):
l3_nodes: List[L3Node] = []
for l3d in l2d.get("children", []):
l3_nodes.append(
L3Node(
id=l3d["id"],
description=l3d["description"],
embedding=_embed_from_str(l3d.get("embedding")),
timestamp=l3d.get("timestamp"),
frame_path=l3d.get("frame_path"),
raw_content=l3d.get("raw_content"),
)
)
tr2 = l2d.get("time_range")
l2_nodes.append(
L2Node(
id=l2d["id"],
description=l2d["description"],
embedding=_embed_from_str(l2d.get("embedding")),
time_range=tuple(tr2) if tr2 else None,
children=l3_nodes,
)
)
tr1 = d.get("time_range")
return L1Node(
id=d["id"],
summary=d["summary"],
embedding=_embed_from_str(d.get("embedding")),
time_range=tuple(tr1) if tr1 else None,
children=l2_nodes,
)
# ---------------------------------------------------------------------------
# 树索引容器
# ---------------------------------------------------------------------------
@dataclass
class TreeIndex:
"""三层树索引容器。
组织和管理三层节点结构，提供嵌入矩阵提取、节点访问、
以及 pickle 序列化/反序列化接口。
典型工作流::
# 1. 构建索引
index = TreeIndex(metadata=meta, roots=[l1_node_1, l1_node_2])
# 2. 提取嵌入矩阵（用于 Tree-TRM 检索）
M_L1 = index.l1_embeddings() # [N1, D]
M_L2 = index.l2_embeddings_of(l1_idx=0) # [N2, D]
M_L3 = index.l3_embeddings_of(0, 1) # [N3, D]
# 3. 序列化
index.save("cache/my_index.pkl")
loaded = TreeIndex.load("cache/my_index.pkl")
Attributes:
metadata: 索引元数据。
roots: L1 节点列表。
"""
metadata: IndexMeta
roots: List[L1Node] = field(default_factory=list)
# ------------------------------------------------------------------ #
# 嵌入矩阵提取
# ------------------------------------------------------------------ #
# ------------------------------------------------------------------ #
# 懒加载嵌入支持
# ------------------------------------------------------------------ #
@property
def is_embedded(self) -> bool:
"""检查所有节点是否已填充嵌入向量。
返回:
True 表示所有 L1/L2/L3 节点的 embedding 均非 None；False 表示尚未 embed。
"""
for l1 in self.roots:
if l1.embedding is None:
return False
for l2 in l1.children:
if l2.embedding is None:
return False
for l3 in l2.children:
if l3.embedding is None:
return False
return True
def embed_all(
self,
embed_fn: Callable[[Union[str, List[str]]], np.ndarray],
model_name: str,
embed_dim: int,
) -> None:
"""对所有节点批量执行 embedding，更新 metadata。
建树阶段不调用此方法（embedding=None）。
首次检索前由 Pipeline 调用，结果缓存在节点上。
参数:
embed_fn: EmbeddingModel.embed 方法，接受 str 或 List[str]，返回 [N, D] ndarray。
model_name: 嵌入模型名称，写入 metadata。
embed_dim: 嵌入维度，写入 metadata。
实现细节:
- L3 节点按 L2 分组批量 embed（一次调用），减少 API 开销。
- L1/L2 各单独 embed（数量少，不值得合并）。
- 仅对 embedding 为 None 的节点执行（支持增量更新）。
"""
ensure(len(self.roots) > 0, "embed_all: 树为空，无节点可 embed")
for l1 in self.roots:
if l1.embedding is None:
l1.embedding = embed_fn(l1.summary)[0].astype(np.float32)
for l2 in l1.children:
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):
l3.embedding = emb
self.metadata.embed_model = model_name
self.metadata.embed_dim = embed_dim
log_msg("INFO", "embed_all 完成", model=model_name, embed_dim=embed_dim)
def l1_embeddings(self) -> np.ndarray:
"""返回所有 L1 节点的嵌入矩阵。
返回:
形状 [N1, D] 的 float32 矩阵。空树返回 [0, D]。
异常:
RuntimeError: 节点 embedding 尚未计算（请先调用 embed_all）。
"""
ensure(self.is_embedded, "L1 embedding 尚未计算，请先调用 tree.embed_all()")
if not self.roots:
return np.zeros((0, self.metadata.embed_dim), dtype=np.float32)
return np.stack([r.embedding for r in self.roots], axis=0).astype(np.float32)
def l2_embeddings_of(self, l1_idx: int) -> np.ndarray:
"""返回指定 L1 节点下所有 L2 子节点的嵌入矩阵。
参数:
l1_idx: L1 节点索引。
返回:
形状 [N2, D] 的 float32 矩阵。
异常:
IndexError: l1_idx 越界。
RuntimeError: embedding 尚未计算。
"""
ensure(self.is_embedded, "L2 embedding 尚未计算，请先调用 tree.embed_all()")
if not (0 <= l1_idx < len(self.roots)):
raise IndexError(f"l1_idx={l1_idx} 越界，L1 节点数={len(self.roots)}")
children = self.roots[l1_idx].children
if not children:
return np.zeros((0, self.metadata.embed_dim), dtype=np.float32)
return np.stack([c.embedding for c in children], axis=0).astype(np.float32)
def l3_embeddings_of(self, l1_idx: int, l2_idx: int) -> np.ndarray:
"""返回指定 L2 节点下所有 L3 子节点的嵌入矩阵。
参数:
l1_idx: L1 节点索引。
l2_idx: L2 节点索引（相对于 L1）。
返回:
形状 [N3, D] 的 float32 矩阵。
异常:
IndexError: 索引越界。
RuntimeError: embedding 尚未计算。
"""
ensure(self.is_embedded, "L3 embedding 尚未计算，请先调用 tree.embed_all()")
if not (0 <= l1_idx < len(self.roots)):
raise IndexError(f"l1_idx={l1_idx} 越界，L1 节点数={len(self.roots)}")
l2_children = self.roots[l1_idx].children
if not (0 <= l2_idx < len(l2_children)):
raise IndexError(f"l2_idx={l2_idx} 越界，L2 节点数={len(l2_children)}")
l3_children = l2_children[l2_idx].children
if not l3_children:
return np.zeros((0, self.metadata.embed_dim), dtype=np.float32)
return np.stack([c.embedding for c in l3_children], axis=0).astype(np.float32)
# ------------------------------------------------------------------ #
# 节点访问
# ------------------------------------------------------------------ #
def get_node(self, l1: int, l2: int, l3: int) -> L3Node:
"""按三级路径索引获取 L3 节点。
参数:
l1: L1 节点索引。
l2: L2 节点索引。
l3: L3 节点索引。
返回:
目标 L3Node。
异常:
IndexError: 任意层级索引越界。
"""
if l1 < 0 or l1 >= len(self.roots):
raise IndexError(f"l1={l1} 越界，L1 节点数={len(self.roots)}")
l2_children = self.roots[l1].children
if l2 < 0 or l2 >= len(l2_children):
raise IndexError(f"l2={l2} 越界，L2 节点数={len(l2_children)}")
l3_children = l2_children[l2].children
if l3 < 0 or l3 >= len(l3_children):
raise IndexError(f"l3={l3} 越界，L3 节点数={len(l3_children)}")
return l3_children[l3]
# ------------------------------------------------------------------ #
# JSON 序列化（主格式，无 embedding）
# ------------------------------------------------------------------ #
def to_dict(self, include_embedding: bool = False) -> Dict[str, Any]:
"""将树索引序列化为纯 Python dict。
参数:
include_embedding: 若 True，将所有节点的 embedding 向量序列化为 base64。
返回:
可直接 json.dump 的字典，结构为 {metadata, roots[...]}。
当 include_embedding=True 时，每个节点包含 embedding 字段。
"""
def l3_to_dict(n: L3Node) -> Dict[str, Any]:
d = {
"id": n.id,
"description": n.description,
"timestamp": n.timestamp,
"frame_path": n.frame_path,
"raw_content": n.raw_content,
}
if include_embedding:
d["embedding"] = _embed_to_str(n.embedding)
return d
def l2_to_dict(n: L2Node) -> Dict[str, Any]:
d = {
"id": n.id,
"description": n.description,
"time_range": list(n.time_range) if n.time_range else None,
"children": [l3_to_dict(c) for c in n.children],
}
if include_embedding:
d["embedding"] = _embed_to_str(n.embedding)
return d
metadata_dict = {
"source_path": self.metadata.source_path,
"modality": self.metadata.modality,
"created_at": self.metadata.created_at,
}
if include_embedding:
metadata_dict["embed_model"] = self.metadata.embed_model
metadata_dict["embed_dim"] = self.metadata.embed_dim
return {
"metadata": metadata_dict,
"roots": [r.to_dict(include_embedding=include_embedding) for r in self.roots],
}
@classmethod
def from_dict(cls, d: Dict[str, Any]) -> "TreeIndex":
"""从 dict 反序列化为 TreeIndex（支持 embedding 恢复）。
参数:
d: to_dict() 的输出或等价结构，可包含 embedding 字段。
返回:
TreeIndex 实例（若 JSON 中有 embedding 字段，自动反序列化填充）。
"""
meta = IndexMeta(
source_path=d["metadata"]["source_path"],
modality=d["metadata"]["modality"],
embed_model=d["metadata"].get("embed_model"),
embed_dim=d["metadata"].get("embed_dim"),
created_at=d["metadata"].get("created_at", datetime.now().isoformat()),
)
roots: List[L1Node] = []
for r in d["roots"]:
l2_nodes: List[L2Node] = []
for l2d in r.get("children", []):
l3_nodes: List[L3Node] = []
for l3d in l2d.get("children", []):
l3_nodes.append(L3Node(
id=l3d["id"],
description=l3d["description"],
embedding=_embed_from_str(l3d.get("embedding")),
timestamp=l3d.get("timestamp"),
frame_path=l3d.get("frame_path"),
raw_content=l3d.get("raw_content"),
))
tr2 = l2d.get("time_range")
l2_nodes.append(L2Node(
id=l2d["id"],
description=l2d["description"],
embedding=_embed_from_str(l2d.get("embedding")),
time_range=tuple(tr2) if tr2 else None,
children=l3_nodes,
))
tr1 = r.get("time_range")
roots.append(L1Node(
id=r["id"],
summary=r["summary"],
embedding=_embed_from_str(r.get("embedding")),
time_range=tuple(tr1) if tr1 else None,
children=l2_nodes,
))
return cls(metadata=meta, roots=roots)
def save_json(self, path: str, include_embedding: bool = False) -> None:
"""将树索引以 JSON 格式保存到磁盘。
参数:
path: 保存文件路径（推荐 .json 后缀）。
include_embedding: 若 True，将所有节点的 embedding 向量保存到 JSON。
"""
with open(path, "w", encoding="utf-8") as f:
json.dump(self.to_dict(include_embedding=include_embedding), f, ensure_ascii=False, indent=2)
log_msg(
"INFO",
f"树索引（JSON）已保存至 {path}",
n_l1=len(self.roots),
include_embedding=include_embedding,
)
@classmethod
def load_json(cls, path: str) -> "TreeIndex":
"""从 JSON 文件加载树索引（自动检测并恢复 embedding）。
参数:
path: JSON 文件路径。
返回:
TreeIndex 实例。若 JSON 中包含 embedding 字段，自动反序列化填充；
否则 embedding=None（向后兼容旧格式）。
异常:
FileNotFoundError: 文件不存在。
"""
with open(path, "r", encoding="utf-8") as f:
d = json.load(f)
obj = cls.from_dict(d)
log_msg(
"INFO",
f"树索引（JSON）已从 {path} 加载",
n_l1=len(obj.roots),
is_embedded=obj.is_embedded,
)
return obj
def save_l1_json(path: str, l1_node: L1Node) -> None:
"""将单个 L1 节点（及其子树）以 JSON 形式保存到磁盘。
参数:
path: 目标文件路径。
l1_node: 待序列化的 L1 节点。
"""
with open(path, "w", encoding="utf-8") as f:
json.dump(l1_node.to_dict(), f, ensure_ascii=False, indent=2)
log_msg("INFO", "L1 中间结果已保存", path=path, l1_id=l1_node.id)
def load_l1_json(path: str) -> L1Node:
"""从 JSON 文件加载单个 L1 节点（embedding=None）。"""
with open(path, "r", encoding="utf-8") as f:
data = json.load(f)
node = L1Node.from_dict(data)
log_msg("INFO", "L1 中间结果已加载", path=path, l1_id=node.id)
return node
# ------------------------------------------------------------------ #
# 序列化（pickle，保留供向后兼容）
# ------------------------------------------------------------------ #
def save(self, path: str) -> None:
"""将整个树索引序列化到磁盘（pickle 格式）。
参数:
path: 保存文件路径（推荐 .pkl 后缀）。
"""
with open(path, "wb") as f:
pickle.dump(self, f, protocol=pickle.HIGHEST_PROTOCOL)
log_msg("INFO", f"树索引已保存至 {path}", n_l1=len(self.roots))
@classmethod
def load(cls, path: str) -> "TreeIndex":
"""从磁盘加载树索引。
.. warning::
pickle 反序列化可执行任意代码，切勿加载不受信任的文件。
如需安全替代方案，请考虑 JSON + NPY 分文件存储。
参数:
path: pickle 文件路径。
返回:
TreeIndex 实例。
异常:
FileNotFoundError: 文件不存在。
TypeError: 文件内容不是 TreeIndex 实例。
"""
with open(path, "rb") as f:
obj = pickle.load(f) # noqa: S301
if not isinstance(obj, cls):
msg = f"文件内容不是 TreeIndex 实例: {type(obj)}"
log_msg("ERROR", msg, path=path)
raise TypeError(msg)
log_msg("INFO", f"树索引已从 {path} 加载", n_l1=len(obj.roots))
return obj
reference/video_tree_trm/video_tree_builder.py
+993
View File
@@ -0,0 +1,993 @@
"""
视频树构建模块
==============
将长视频通过 L2 轴心策略 + VLM 帧描述转化为三层 TreeIndex。
构建策略::
Step 1: _segment_video — 固定步长切分，确定 L1 时间边界
Step 2: L2 先行 — 每个 L2 clip 均匀 seek l2_representative_frames 帧（稀疏），VLM 生成片段描述（1-2句）
Step 3: L3 向下 — 注入 L2 上下文，VLM 批量帧描述（每帧1-2句）
Step 4: L1 向上 — 聚合 L2 描述，LLM 生成 L1 摘要（2-3句）
Step 5: 组装 TreeIndex
并发模型（异步版）::
build() → asyncio.run(_build_async())
_build_async():
asyncio.Semaphore(concurrency=16) 控制最大 VLM 并发数
Phase 1: asyncio.gather(所有L2任务) — 16路同时 VLM
Phase 2: asyncio.gather(所有L3任务) — 每个L3任务内的12批次同时并发
Phase 3: asyncio.gather(各L1摘要) — L1收齐后并发
ffmpeg 提帧在 ThreadPoolExecutor(max_workers=8) 中并行执行
L2 轴心策略解决了循环依赖：
- L2 描述不依赖 L3，从代表帧直接生成
- L3 注入 L2 上下文后逐帧描述
- L1 聚合 L2 描述，保证完整覆盖
帧持久化：
- 帧图像保存到 {cache_dir}/frames/{video_stem}/，长期有效
- 已提取的帧自动跳过（缓存复用）
使用方式::
builder = VideoTreeBuilder(vlm_client, config.tree)
index = builder.build("path/to/video.mp4") # 同步壳，内部 asyncio.run()
index.save("cache/my_video.pkl")
"""
from __future__ import annotations
import asyncio
import json
import os
import re
import subprocess
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import cv2
from utils.logger_system import ensure, log_json, log_msg
from video_tree_trm.config import TreeConfig
from video_tree_trm.llm_client import LLMClient
from video_tree_trm.tree_index import (
IndexMeta,
L1Node,
L2Node,
L3Node,
TreeIndex,
load_l1_json,
save_l1_json,
)
# ---------------------------------------------------------------------------
# Prompt 常量
# ---------------------------------------------------------------------------
_L2_VIDEO_PROMPT = "用1-2句话描述以下视频片段的核心内容，与同级片段形成区分。"
_L3_VIDEO_PROMPT = (
'该片段的整体内容: "{l2_description}"\n'
"以下是该片段中连续的 {n} 帧画面。\n"
"对每帧用一到两句话描述其具体画面内容。\n"
"重点关注: 动作、物体变化、文字信息、人物表情。\n"
"不要重复片段整体描述，聚焦每帧的区分性信息。\n"
'只返回 JSON 数组，格式: ["帧1描述", "帧2描述", ...]，不要其他内容。'
)
_L1_VIDEO_PROMPT = (
"以下是一个视频段落中各片段的描述:\n{l2_texts}\n"
"用2-3句话总结该段落的整体内容，涵盖所有片段的主题。"
)
# 每次 VLM 调用携带的最大帧数：5 帧 payload 小、JSON 解析成功率高
_L3_BATCH_SIZE = 5
_L3_SINGLE_PROMPT = (
'该片段的整体内容: "{l2_description}"\n'
"用一到两句话描述这帧画面的具体内容。"
"重点关注: 动作、物体变化、文字信息、人物表情。"
)
# ffmpeg 并发提帧的线程池大小（CPU 密集型，避免过度并发）
_FFMPEG_MAX_WORKERS = 8
# ---------------------------------------------------------------------------
# 主类
# ---------------------------------------------------------------------------
class VideoTreeBuilder:
"""视频模态树构建器（asyncio 真并发版）。
将长视频通过 L2 轴心策略（先构建 L2，再向下扩展 L3，向上聚合 L1）
转化为三层 TreeIndex。
并发架构:
build() 为同步壳，内部调用 asyncio.run(_build_async())。
_build_async() 使用 asyncio.Semaphore(concurrency) 控制并发 VLM 数量。
所有 VLM 调用通过 LLMClient 的异步接口（AsyncOpenAI）发起，零线程阻塞。
ffmpeg 提帧在独立 ThreadPoolExecutor 中并行，不阻塞事件循环。
属性:
vlm: VLM/LLM 客户端（用于图文和纯文本调用）。
config: 树构建配置。
_ffmpeg_pool: ffmpeg 专用线程池（max_workers=_FFMPEG_MAX_WORKERS）。
"""
def __init__(
self,
vlm: LLMClient,
config: TreeConfig,
) -> None:
"""初始化视频树构建器。
参数:
vlm: 已初始化的 VLM/LLM 客户端（LLMClient），需支持异步接口。
config: 树构建配置（TreeConfig），关键字段：
l1_segment_duration, l2_clip_duration, l3_fps,
l2_representative_frames, cache_dir, concurrency。
实现细节:
ffmpeg 线程池在构建器级别创建，所有异步协程共用，
避免每次提帧都重建线程池的开销。
"""
self.vlm = vlm
self.config = config
self._ffmpeg_pool = ThreadPoolExecutor(max_workers=_FFMPEG_MAX_WORKERS)
# 进度与中间结果目录均挂在 cache_dir 下，避免散落其它位置
self._cache_root = Path(self.config.cache_dir)
# ------------------------------------------------------------------
# URL 流式辅助方法
# ------------------------------------------------------------------
@staticmethod
def _is_url(path_or_url: str) -> bool:
"""判断输入是否为网络 URL（而非本地路径）。
参数:
path_or_url: 文件路径或 URL 字符串。
返回:
True 表示 URL，False 表示本地路径。
"""
return path_or_url.startswith(("http://", "https://"))
@staticmethod
def _source_stem(video_path: str) -> str:
"""从视频路径或 YouTube URL 中提取短标识符，用于帧缓存目录命名。
参数:
video_path: 本地文件路径或 YouTube 视频页面 URL。
返回:
短字符串标识符（本地文件取 stem，YouTube URL 取 v= 后的视频 ID）。
"""
if "youtube.com/watch" in video_path or "youtu.be/" in video_path:
match = re.search(r"(?:v=|youtu\.be/)([A-Za-z0-9_-]{8,15})", video_path)
if match:
return match.group(1)
stem = Path(video_path).stem
return stem[:64] if len(stem) > 64 else stem
@staticmethod
def _resolve_stream(url: str) -> str:
"""通过 yt-dlp 获取 YouTube 视频的 CDN 直链，供 cv2.VideoCapture 直接使用。
参数:
url: YouTube 视频页面 URL。
返回:
CDN HTTPS 直链（ffmpeg/OpenCV 可直接流式读取）。
"""
log_msg("INFO", "获取 YouTube CDN 直链", url=url)
result = subprocess.run(
[
"yt-dlp",
"-g",
"--format",
"best[ext=mp4][height<=720]/best[ext=mp4]/best",
url,
],
capture_output=True,
text=True,
timeout=30,
)
ensure(result.returncode == 0, f"yt-dlp 获取直链失败: {result.stderr.strip()}")
stream_url = result.stdout.strip().splitlines()[0]
ensure(stream_url.startswith("http"), f"yt-dlp 返回非 URL: {stream_url[:100]}")
log_msg("INFO", "CDN 直链获取成功", stream_url=stream_url[:80])
return stream_url
@staticmethod
def _get_video_duration(url: str) -> float:
"""通过 yt-dlp --dump-json 获取视频时长（秒）。
参数:
url: YouTube 视频页面 URL。
返回:
视频总时长（秒，浮点数）。
"""
log_msg("INFO", "获取视频时长元数据", url=url)
result = subprocess.run(
["yt-dlp", "--dump-json", "--no-playlist", url],
capture_output=True,
text=True,
timeout=30,
)
ensure(result.returncode == 0, f"yt-dlp 元数据获取失败: {result.stderr.strip()}")
meta = json.loads(result.stdout)
duration = float(meta.get("duration", 0))
ensure(duration > 0, f"视频时长读取异常: {duration}")
log_msg("INFO", "视频时长确认", duration_sec=round(duration, 1))
return duration
# ------------------------------------------------------------------
# 公共接口
# ------------------------------------------------------------------
def build(self, video_path: str) -> TreeIndex:
"""将长视频构建为三层 TreeIndex（同步壳，内部 asyncio.run 驱动）。
参数:
video_path: 视频文件路径（.mp4/.avi/.mkv 等）或 YouTube URL。
返回:
三层 TreeIndex 对象。
实现细节:
同步壳设计保持与 build_trees_batch.py 的接口兼容性。
每次调用 asyncio.run() 创建独立事件循环，多线程安全（各线程独立循环）。
"""
return asyncio.run(self._build_async(video_path))
# ------------------------------------------------------------------
# 核心异步构建逻辑
# ------------------------------------------------------------------
async def _build_async(self, video_path: str) -> TreeIndex:
"""异步构建三层 TreeIndex（真并发核心，L2→L3 链式触发）。
参数:
video_path: 视频文件路径或 YouTube URL。
返回:
三层 TreeIndex 对象。
实现细节:
并发架构：每个 L1 段内启动一组"L2→L3 链式协程"，
L2 完成后立即触发 L3（不等待其他 L2），L3 完成后触发 L1 摘要。
各 L1 段独立并发，彼此不阻塞。
Semaphore(concurrency=16) 全局限制同时在途 VLM 调用数量。
关键调用链（每个 L2 clip 独立）::
_build_segment(i) → asyncio.gather(
_chain(i,0): _build_l2_video_async → _build_l3_task_async
_chain(i,1): _build_l2_video_async → _build_l3_task_async
...
) → _build_l1_video_async(i)
"""
# Phase 0: URL vs 本地文件处理
if self._is_url(video_path):
stream_url = self._resolve_stream(video_path)
duration_hint: Optional[float] = self._get_video_duration(video_path)
log_msg("INFO", "开始构建视频树索引（URL 流式模式）", source_url=video_path)
else:
ensure(os.path.isfile(video_path), f"视频文件不存在: {video_path}")
stream_url = video_path
duration_hint = None
log_msg("INFO", "开始构建视频树索引", video_path=video_path)
source_id = self._source_stem(video_path)
# Phase 1: 时间切分（同步，仅一次）
l1_ranges = self._segment_video(stream_url, duration_hint=duration_hint)
ensure(len(l1_ranges) > 0, "视频时间切分结果为空")
log_msg("INFO", "视频切分完成", l1_count=len(l1_ranges))
total_l1 = len(l1_ranges)
# Phase 1.1: 读取已有进度（支持断点续跑）
finished_l1_ids: set[int] = set()
progress = self._load_progress(source_id)
if progress is not None and progress.get("total_l1") == total_l1:
finished_l1_ids = set(progress.get("finished_l1_ids", []))
if finished_l1_ids:
log_msg(
"INFO",
"检测到中间进度，启用断点续跑",
stem=source_id,
finished_l1=list(sorted(finished_l1_ids)),
)
else:
# 进度不存在或形状不匹配时，从零开始，旧进度视为无效
if progress is not None:
log_msg(
"WARNING",
"进度文件与当前 L1 段数不一致，忽略旧进度",
stem=source_id,
recorded_total_l1=progress.get("total_l1"),
current_total_l1=total_l1,
)
# 创建 VLM 并发控制信号量（每视频独立，限制同时在途 VLM 请求数）
vlm_sem = asyncio.Semaphore(self.config.concurrency)
# Phase 2-5: 按 L1 段并发，段内 L2→L3 链式触发，L3 收齐后触发 L1
async def _build_segment(i: int, l1_range: Tuple[float, float]) -> L1Node:
"""单个 L1 段的完整构建：L2+L3 并发链式 → L1 摘要。
参数:
i: L1 段索引。
l1_range: L1 时间区间 (start, end)。
返回:
完整的 L1Node（含所有 L2 和 L3 子节点）。
实现细节:
段内所有 L2 clip 同时启动（asyncio.gather），
每个 clip 的 L2 VLM 完成后立即触发 L3，不等待其他 clip 的 L2。
所有 clip 的 L2+L3 完成后，触发 L1 文本摘要。
"""
clips = self._get_l2_clips(l1_range)
async def _chain(j: int, clip_range: Tuple[float, float]) -> Tuple[int, L2Node]:
"""L2→L3 链：L2 完成立即触发 L3，返回 (j, 含children的L2Node)。"""
l2_id = f"l1_{i}_l2_{j}"
l2_node = await self._build_l2_video_async(
stream_url, clip_range, l2_id, source_id, vlm_sem
)
log_msg("INFO", "L2 VLM 完成，已触发 L3 任务", l2_id=l2_id)
completed_l2 = await self._build_l3_task_async(
stream_url, l2_node, clip_range, source_id, i, j, vlm_sem
)
log_msg(
"INFO", "L3 完成",
l2_id=l2_id,
l3_count=len(completed_l2.children),
)
return (j, completed_l2)
# 所有 clip 同时启动（不等 L2 全部结束再开 L3）
pairs = await asyncio.gather(*[_chain(j, clip) for j, clip in enumerate(clips)])
ordered_l2 = [p[1] for p in sorted(pairs, key=lambda x: x[0])]
log_msg("INFO", "L1 触发", l1_id=f"l1_{i}")
l1_node = await self._build_l1_video_async(
ordered_l2, f"l1_{i}", l1_range, vlm_sem
)
log_msg(
"INFO", "L1 节点构建完成",
l1_id=f"l1_{i}",
l2_count=len(ordered_l2),
)
return l1_node
total_clips = sum(len(self._get_l2_clips(r)) for r in l1_ranges)
log_msg(
"INFO",
"开始并发构建（L2→L3链式，L1段间并发，支持断点续跑）",
total_l2=total_clips,
concurrency=self.config.concurrency,
)
# Phase 2: 并发构建尚未完成的 L1 段（段内 L2+L3 链式并发）
tasks: List[asyncio.Task[L1Node]] = []
task_indices: List[int] = []
for i, r in enumerate(l1_ranges):
# 已完成且中间 JSON 存在 → 直接复用
if i in finished_l1_ids and self._has_l1_intermediate(source_id, i):
continue
tasks.append(asyncio.create_task(_build_segment(i, r)))
task_indices.append(i)
new_l1_nodes: Dict[int, L1Node] = {}
if tasks:
results = await asyncio.gather(*tasks)
for idx, node in zip(task_indices, results):
# 每完成一个 L1 段就写入中间 JSON，并刷新进度文件
self._save_l1_intermediate(source_id, node, idx)
finished_l1_ids.add(idx)
new_l1_nodes[idx] = node
self._save_progress(source_id, total_l1, finished_l1_ids)
# Phase 3: 汇总所有 L1 段（中间 + 新生成）
l1_nodes: List[L1Node] = []
for i in range(total_l1):
if i in new_l1_nodes:
l1_nodes.append(new_l1_nodes[i])
continue
node = self._load_l1_intermediate(source_id, i)
ensure(node is not None, f"L1 段 {i} 缺失中间结果，无法恢复")
l1_nodes.append(node)
# Phase 6: 组装 TreeIndex
metadata = IndexMeta(
source_path=video_path,
modality="video",
created_at=datetime.now().isoformat(),
)
index = TreeIndex(metadata=metadata, roots=l1_nodes)
total_l2_count = sum(len(r.children) for r in l1_nodes)
total_l3_count = sum(len(l2.children) for r in l1_nodes for l2 in r.children)
log_json(
"video_tree_build",
{
"source_path": video_path,
"l1_count": len(l1_nodes),
"l2_count": total_l2_count,
"l3_count": total_l3_count,
"embedded": False,
},
)
log_msg(
"INFO",
"视频树索引构建完成",
l1=len(l1_nodes),
l2=total_l2_count,
l3=total_l3_count,
)
# 最终 JSON 写入成功后，清理由断点机制生成的中间文件
self._cleanup_intermediate_and_progress(source_id)
return index
# ------------------------------------------------------------------
# 内部方法：时间切分（同步，仅执行一次）
# ------------------------------------------------------------------
def _segment_video(
self,
video_path: str,
duration_hint: Optional[float] = None,
) -> List[Tuple[float, float]]:
"""读取视频总时长，按固定步长切分为 L1 时间区间列表。
参数:
video_path: 视频文件路径或 CDN 流式 URL。
duration_hint: 已知视频时长（秒）。传入时跳过 cv2 读取。
返回:
L1 时间区间列表，每项为 (start_sec, end_sec)。
"""
if duration_hint is not None:
total_duration = duration_hint
else:
cap = cv2.VideoCapture(video_path)
ensure(cap.isOpened(), f"无法打开视频文件: {video_path}")
fps = cap.get(cv2.CAP_PROP_FPS)
total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
cap.release()
ensure(fps > 0, f"视频 FPS 读取异常: {fps}")
ensure(total_frames > 0, f"视频总帧数读取异常: {total_frames}")
total_duration = total_frames / fps
step = self.config.l1_segment_duration
ranges: List[Tuple[float, float]] = []
start = 0.0
while start < total_duration:
end = min(start + step, total_duration)
ranges.append((start, end))
start = end
log_msg(
"INFO",
"L1 时间切分",
total_duration=round(total_duration, 2),
l1_count=len(ranges),
)
return ranges
def _get_l2_clips(self, l1_range: Tuple[float, float]) -> List[Tuple[float, float]]:
"""将 L1 时间区间等分为 L2 clips。
参数:
l1_range: L1 时间区间 (start, end)，单位秒。
返回:
L2 clip 时间区间列表。
"""
start, end = l1_range
step = self.config.l2_clip_duration
clips: List[Tuple[float, float]] = []
t = start
while t < end:
clip_end = min(t + step, end)
clips.append((t, clip_end))
t = clip_end
return clips
# ------------------------------------------------------------------
# 内部方法：帧提取（ffmpeg subprocess，在线程池执行）
# ------------------------------------------------------------------
def _ffmpeg_extract_frame(self, video_path: str, ts: float, out_path: str) -> bool:
"""用 ffmpeg subprocess 提取单帧图像，兼容 AV1/H.264 等所有编码格式。
参数:
video_path: 视频文件路径（本地 MP4 或 CDN URL）。
ts: 目标时间戳（秒）。
out_path: 输出 JPEG 文件路径。
返回:
True 表示提取成功，False 表示失败。
"""
cmd = [
"ffmpeg", "-hide_banner", "-loglevel", "error",
"-ss", f"{ts:.3f}",
"-i", video_path,
"-frames:v", "1",
"-q:v", "2",
"-y", out_path,
]
result = subprocess.run(cmd, capture_output=True)
return result.returncode == 0 and os.path.isfile(out_path)
async def _extract_frames_async(
self,
video_path: str,
time_range: Tuple[float, float],
fps: float,
source_id: Optional[str] = None,
) -> List[Tuple[str, float]]:
"""异步并发提取时间范围内的帧，保存到 cache 目录。
参数:
video_path: 视频文件路径或 CDN 流式 URL。
time_range: 提取时间区间 (start_sec, end_sec)。
fps: 提取帧率（帧/秒）。
source_id: 帧缓存目录名。
返回:
[(frame_path, timestamp_sec), ...]，按时间顺序排列。
实现细节:
所有 ffmpeg 提取任务通过 run_in_executor(self._ffmpeg_pool, ...) 并发执行，
已缓存的帧直接跳过（无需调用 ffmpeg）。
ffmpeg 线程池 max_workers=_FFMPEG_MAX_WORKERS 防止过度并发占满 CPU。
"""
video_stem = source_id if source_id is not None else self._source_stem(video_path)
frame_dir = Path(self.config.cache_dir) / "frames" / video_stem
frame_dir.mkdir(parents=True, exist_ok=True)
start_sec, end_sec = time_range
step = 1.0 / fps
timestamps: List[float] = []
t = start_sec
while t < end_sec:
timestamps.append(t)
t += step
if not timestamps:
log_msg("WARNING", "帧提取时间区间内无有效时间戳", time_range=time_range, fps=fps)
return []
loop = asyncio.get_event_loop()
async def _extract_one(ts: float) -> Optional[Tuple[str, float]]:
"""提取单帧：缓存命中直接返回，否则在线程池中调用 ffmpeg。"""
frame_name = f"{start_sec:.1f}_{ts:.3f}.jpg"
frame_path = str(frame_dir / frame_name)
if os.path.isfile(frame_path):
return (frame_path, ts)
success = await loop.run_in_executor(
self._ffmpeg_pool,
self._ffmpeg_extract_frame, video_path, ts, frame_path,
)
if not success:
log_msg("WARNING", "帧读取失败，跳过", timestamp=ts, video_path=video_path)
return None
return (frame_path, ts)
# 并发提取所有帧（受 ffmpeg 线程池限制，不会无限并发）
results = await asyncio.gather(*[_extract_one(ts) for ts in timestamps])
return [r for r in results if r is not None]
# ------------------------------------------------------------------
# 内部方法：L1 中间结果与进度管理（断点续跑）
# ------------------------------------------------------------------
def _intermediate_dir(self, stem: str) -> Path:
"""获取某视频的中间结果目录路径。"""
return self._cache_root / "intermediate" / stem
def _progress_path(self, stem: str) -> Path:
"""获取某视频的进度文件路径。"""
return self._cache_root / "progress" / f"{stem}.json"
def _has_l1_intermediate(self, stem: str, l1_idx: int) -> bool:
"""检查某 L1 段的中间 JSON 是否存在。"""
path = self._intermediate_dir(stem) / f"l1_{l1_idx}.json"
return path.is_file()
def _save_l1_intermediate(self, stem: str, l1_node: L1Node, l1_idx: int) -> None:
"""将单个 L1 段的中间结果保存到 JSON 文件。"""
dir_path = self._intermediate_dir(stem)
dir_path.mkdir(parents=True, exist_ok=True)
out_path = dir_path / f"l1_{l1_idx}.json"
save_l1_json(str(out_path), l1_node)
def _load_l1_intermediate(self, stem: str, l1_idx: int) -> Optional[L1Node]:
"""从中间 JSON 加载单个 L1 段，若不存在则返回 None。"""
path = self._intermediate_dir(stem) / f"l1_{l1_idx}.json"
if not path.is_file():
return None
return load_l1_json(str(path))
def _load_progress(self, stem: str) -> Optional[Dict[str, object]]:
"""加载某视频的进度文件（若不存在则返回 None）。"""
path = self._progress_path(stem)
if not path.is_file():
return None
with open(path, "r", encoding="utf-8") as f:
try:
data: Dict[str, object] = json.load(f)
except json.JSONDecodeError:
log_msg("WARNING", "进度文件 JSON 解析失败，忽略", path=str(path))
return None
return data
def _save_progress(self, stem: str, total_l1: int, finished_l1_ids: set[int]) -> None:
"""将最新进度写回磁盘，支持断点续跑。"""
path = self._progress_path(stem)
path.parent.mkdir(parents=True, exist_ok=True)
payload = {
"video_id": stem,
"total_l1": total_l1,
"finished_l1_ids": sorted(finished_l1_ids),
"updated_at": datetime.now().isoformat(),
}
if not path.is_file():
payload["created_at"] = payload["updated_at"]
else:
# 尝试保留旧 created_at
old = self._load_progress(stem)
if old and isinstance(old.get("created_at"), str):
payload["created_at"] = old["created_at"]
with open(path, "w", encoding="utf-8") as f:
json.dump(payload, f, ensure_ascii=False, indent=2)
log_msg(
"INFO",
"进度文件已更新",
path=str(path),
total_l1=total_l1,
finished_l1=list(sorted(finished_l1_ids)),
)
def _cleanup_intermediate_and_progress(self, stem: str) -> None:
"""在最终 JSON 写入成功后清理中间结果与进度文件。"""
# 清理 progress
progress_path = self._progress_path(stem)
if progress_path.is_file():
try:
progress_path.unlink()
except OSError:
log_msg("WARNING", "删除进度文件失败", path=str(progress_path))
# 清理 intermediate 目录
inter_dir = self._intermediate_dir(stem)
if inter_dir.is_dir():
for child in inter_dir.glob("l1_*.json"):
try:
child.unlink()
except OSError:
log_msg("WARNING", "删除 L1 中间 JSON 失败", path=str(child))
try:
# 目录可能仍有其它调试文件，忽略删除异常
inter_dir.rmdir()
except OSError:
pass
# ------------------------------------------------------------------
# 内部方法：异步节点构建
# ------------------------------------------------------------------
async def _build_l2_video_async(
self,
video_path: str,
clip_range: Tuple[float, float],
l2_id: str,
source_id: Optional[str],
vlm_sem: asyncio.Semaphore,
) -> L2Node:
"""异步构建 L2 视频节点（代表帧 VLM 描述）。
参数:
video_path: 视频文件路径或 CDN 流式 URL。
clip_range: L2 clip 时间区间 (start, end)，单位秒。
l2_id: 节点 ID。
source_id: 帧缓存目录名。
vlm_sem: VLM 并发控制信号量。
返回:
L2Node（children 为空，由后续 L3 阶段填充）。
实现细节:
均匀采样 l2_representative_frames 帧，并行 ffmpeg 提取，
async with vlm_sem 限制 VLM 并发量。
"""
start_sec, end_sec = clip_range
n_rep = self.config.l2_representative_frames
if n_rep == 1:
timestamps = [(start_sec + end_sec) / 2.0]
else:
step = (end_sec - start_sec) / (n_rep - 1)
timestamps = [start_sec + i * step for i in range(n_rep)]
video_stem = source_id if source_id is not None else self._source_stem(video_path)
frame_dir = Path(self.config.cache_dir) / "frames" / video_stem
frame_dir.mkdir(parents=True, exist_ok=True)
loop = asyncio.get_event_loop()
async def _extract_rep(ts: float) -> Optional[str]:
frame_name = f"l2_{ts:.3f}.jpg"
frame_path = str(frame_dir / frame_name)
if os.path.isfile(frame_path):
return frame_path
success = await loop.run_in_executor(
self._ffmpeg_pool,
self._ffmpeg_extract_frame, video_path, ts, frame_path,
)
if not success:
log_msg("WARNING", "L2 代表帧读取失败，跳过", timestamp=ts)
return None
return frame_path
# 并发提取所有代表帧
rep_results = await asyncio.gather(*[_extract_rep(ts) for ts in timestamps])
rep_frames = [p for p in rep_results if p is not None]
ensure(len(rep_frames) > 0, f"L2 节点 {l2_id} 代表帧提取结果为空")
# VLM 调用受信号量保护
async with vlm_sem:
description = await self.vlm.chat_with_images_async(
_L2_VIDEO_PROMPT, images=rep_frames
)
return L2Node(
id=l2_id,
description=description,
embedding=None,
time_range=clip_range,
)
async def _build_l3_task_async(
self,
video_path: str,
l2_node: L2Node,
clip_range: Tuple[float, float],
source_id: str,
l1_i: int,
l2_j: int,
vlm_sem: asyncio.Semaphore,
) -> L2Node:
"""异步 L3 任务：并发提帧 + 批次级并发 VLM 帧描述。
参数:
video_path: 视频文件路径或 CDN 流式 URL。
l2_node: 已构建的 L2 节点。
clip_range: L2 clip 时间区间。
source_id: 帧缓存目录名。
l1_i: 父 L1 索引。
l2_j: 父 L2 索引。
vlm_sem: VLM 并发控制信号量。
返回:
已填充 children 的 L2Node。
实现细节:
提帧阶段完全并行（受 ffmpeg 线程池限制）；
VLM 调用阶段：12个批次同时提交（asyncio.gather），受信号量限流。
"""
all_frames = await self._extract_frames_async(
video_path, clip_range, self.config.l3_fps, source_id=source_id
)
l3_nodes = await self._build_l3_video_async(
all_frames, l2_node.description, l1_i, l2_j, vlm_sem
)
l2_node.children = l3_nodes
return l2_node
async def _build_l3_video_async(
self,
frames: List[Tuple[str, float]],
l2_description: str,
l1_i: int,
l2_j: int,
vlm_sem: asyncio.Semaphore,
) -> List[L3Node]:
"""异步批次级并发构建 L3 节点（核心加速点）。
参数:
frames: [(frame_path, timestamp), ...]。
l2_description: L2 节点描述，注入 prompt 上下文。
l1_i: 父 L1 索引（用于节点 ID 生成）。
l2_j: 父 L2 索引（用于节点 ID 生成）。
vlm_sem: VLM 并发控制信号量。
返回:
L3Node 列表，每项对应一帧。
实现细节:
将全部帧按 _L3_BATCH_SIZE 分批，所有批次同时提交（asyncio.gather），
每批通过 vlm_sem 限流，实现批次级真并发。
对比旧版串行：12批 × 6s = 72s → 现在 ~6s（受信号量限流取最慢批次）。
"""
ensure(len(frames) > 0, f"L3 帧列表为空 (l1={l1_i}, l2={l2_j})")
# Phase 1: 构建所有批次的协程（同时提交，asyncio.gather 并发执行）
batches: List[List[Tuple[str, float]]] = []
for batch_start in range(0, len(frames), _L3_BATCH_SIZE):
batches.append(frames[batch_start : batch_start + _L3_BATCH_SIZE])
batch_results: List[List[str]] = list(
await asyncio.gather(
*[
self._call_vlm_batch_async(
batch, l2_description, l1_i, l2_j, vlm_sem
)
for batch in batches
]
)
)
# Phase 2: 展平所有批次描述，构建 L3 节点
all_descriptions: List[str] = [
desc for batch_descs in batch_results for desc in batch_descs
]
nodes: List[L3Node] = []
for k, (desc, (frame_path, ts)) in enumerate(zip(all_descriptions, frames)):
nodes.append(
L3Node(
id=f"l1_{l1_i}_l2_{l2_j}_l3_{k}",
description=desc,
embedding=None,
raw_content=None,
frame_path=frame_path,
timestamp=ts,
)
)
return nodes
async def _call_vlm_batch_async(
self,
batch: List[Tuple[str, float]],
l2_description: str,
l1_i: int,
l2_j: int,
vlm_sem: asyncio.Semaphore,
) -> List[str]:
"""异步单批次 VLM 调用（≤ _L3_BATCH_SIZE 帧），解析失败时逐帧 fallback。
参数:
batch: 本批帧列表 [(frame_path, ts), ...]。
l2_description: L2 描述，用于 prompt 和 fallback prompt。
l1_i: 父 L1 索引（日志用）。
l2_j: 父 L2 索引（日志用）。
vlm_sem: VLM 并发控制信号量。
返回:
与 batch 等长的描述文本列表。
实现细节:
async with vlm_sem 确保并发量不超过 config.concurrency。
fallback 时逐帧并发（asyncio.gather），同样受信号量保护。
"""
batch_paths = [fp for fp, _ in batch]
n = len(batch_paths)
prompt = _L3_VIDEO_PROMPT.format(l2_description=l2_description, n=n)
# Phase 1: 尝试批量调用
try:
async with vlm_sem:
raw = await self.vlm.chat_with_images_async(prompt, images=batch_paths)
descriptions = self._parse_json_descriptions(raw, n)
if descriptions is not None:
return descriptions
log_msg(
"WARNING",
"L3 小批量 VLM JSON 解析失败，对本批逐帧 fallback",
l1=l1_i,
l2=l2_j,
batch_n=n,
raw_preview=raw[:100],
)
except Exception as exc:
log_msg(
"WARNING",
f"L3 小批量 VLM 调用异常，对本批逐帧 fallback: {exc}",
l1=l1_i,
l2=l2_j,
batch_n=n,
)
# Phase 2: 逐帧 fallback（并发，受信号量保护）
single_prompt = _L3_SINGLE_PROMPT.format(l2_description=l2_description)
async def _single_frame(fp: str) -> str:
async with vlm_sem:
return await self.vlm.chat_with_images_async(single_prompt, images=[fp])
return list(await asyncio.gather(*[_single_frame(fp) for fp in batch_paths]))
async def _build_l1_video_async(
self,
l2_children: List[L2Node],
l1_id: str,
l1_range: Tuple[float, float],
vlm_sem: asyncio.Semaphore,
) -> L1Node:
"""异步构建 L1 节点（LLM 文本摘要）。
参数:
l2_children: 该 L1 节点下的所有 L2 节点。
l1_id: 节点 ID。
l1_range: L1 时间区间 (start, end)，单位秒。
vlm_sem: VLM 并发控制信号量。
返回:
L1Node（children 已赋值）。
"""
ensure(len(l2_children) > 0, f"L1 节点 {l1_id} 没有 L2 子节点")
l2_texts = "\n".join(f"- {node.description}" for node in l2_children)
prompt = _L1_VIDEO_PROMPT.format(l2_texts=l2_texts)
async with vlm_sem:
summary = await self.vlm.chat_async(prompt)
log_msg("INFO", "L1 触发", l1_id=l1_id)
return L1Node(
id=l1_id,
summary=summary,
embedding=None,
time_range=l1_range,
children=l2_children,
)
# ------------------------------------------------------------------
# 内部方法：JSON 解析（同步，纯 CPU）
# ------------------------------------------------------------------
def _parse_json_descriptions(
self, raw: str, expected_n: int
) -> Optional[List[str]]:
"""从 VLM 输出中解析 JSON 描述数组。
参数:
raw: VLM 原始返回字符串。
expected_n: 期望的描述条数。
返回:
成功解析且长度匹配时返回 List[str]，否则返回 None。
"""
raw = raw.strip()
code_match = re.search(r"```(?:json)?\s*(\[.*?\])\s*```", raw, re.DOTALL)
if code_match:
raw = code_match.group(1)
if not raw.startswith("["):
return None
try:
items: List[str] = json.loads(raw)
except json.JSONDecodeError:
return None
if not isinstance(items, list) or len(items) != expected_n:
return None
return [str(item).strip() for item in items]