diff --git a/PROJECT_CONTEXT.md b/PROJECT_CONTEXT.md index 3fd55a7..0ebfb51 100644 --- a/PROJECT_CONTEXT.md +++ b/PROJECT_CONTEXT.md @@ -81,6 +81,11 @@ All are pending in that log; Omni LoRA is deferred and no semantic gate is compl ### Implemented but not GPU-verified +- `extract_easycom_multiwindow.py` exports 2/4/6/8/10/12-second native-block + context banks in one shared AV preparation pass, with independent Thinker KV + per cap. CPU tests match separate single-window replay and both head/local + loaders; real-GPU replay and media smoke remain required. See + [`docs/easycom_multiwindow_features.md`](docs/easycom_multiwindow_features.md). - Qwen2.5-Omni now has two matched chunk-handling controls around the existing 5 Hz partial-tail path: native completed 2-second blocks held between boundaries, and non-native 200-ms payload commits. Both preserve dense causal diff --git a/docs/README.md b/docs/README.md index fe73777..354cd29 100644 --- a/docs/README.md +++ b/docs/README.md @@ -23,6 +23,7 @@ tests. | Document | Use when | |---|---| +| [`easycom_multiwindow_features.md`](easycom_multiwindow_features.md) | Shared-AV extraction of independent Qwen 2/4/6/8/10/12-second context banks, GPU smoke and matched head training | | [`local_feature_baseline.md`](local_feature_baseline.md) | Extracting/validating EasyCom pooled local Qwen visual and audio sidecars | | [`easycom_encoder_only.md`](easycom_encoder_only.md) | Comparing EasyCom Qwen encoder-only A/V/AV representations | | [`ego4d_encoder_only.md`](ego4d_encoder_only.md) | Preparing and running the Ego4D encoder-only comparison | diff --git a/docs/easycom_multiwindow_features.md b/docs/easycom_multiwindow_features.md new file mode 100644 index 0000000..25f92c9 --- /dev/null +++ b/docs/easycom_multiwindow_features.md @@ -0,0 +1,151 @@ +# EasyCom：一次共享 AV 提取，多种独立 context cap + +状态：实现和 CPU 合成测试已通过；**未跑真实 Qwen GPU smoke、全量提取或训练**。 +入口：[extract_easycom_multiwindow.py](../extract_easycom_multiwindow.py)； +缓存实现：[streaming_turn_multiwindow.py](../streaming_turn_multiwindow.py)。 +既有 [streaming_turn_model_directly.py](../streaming_turn_model_directly.py) +单窗口命令不变，也可直接传入 `--memory-seconds-list`。 + +## 实验契约 + +默认导出 2/4/6/8/10/12 s；2 s native block、5 Hz 截止时间、2 FPS 因果视觉网格、 +observed partial tail；默认保存 last4 stack `[T,4,D]`。必须使用现有 Thinker replay +要求的 `transformers==4.52.3`，不能用局部 encoder 对 4.57.6 的兼容性替代它。 + +- 一次 session reader，共享每个截止点的 preprocessing、AV block/tail embeddings。 + 模型在所选 sessions 间只加载一次，所有缓存在 session 边界清空。 +- **每个 cap 有独立 Thinker KV、滚动重建和 readout**。旧历史移出后，用缓存的 + encoder embeddings 从空 Thinker KV 重放留下的完整 blocks。 + 不裁剪长窗口 KV 或 pooled state 来假装短窗口。 +- 同一尾块仍随 200 ms 截止点增长并重编码；不是整个视频只编码一次。 + 六窗口生产 AV preparation 为 `T` 次，不是 `6T` 次；Thinker 成本仍按窗口增加。 + `--verify-every` 额外从原媒体独立重放，不计入生产调用计数。 +- Cap 是**最大呈现预算**，不是每行精确固定长度的滑窗。稳态 span 为 + `cap - 2s + 当前 tail phase`，phase 在 0.2--2.0 s；startup 少于 cap。 + 例如 cutoff=20.6 s：cap=2/4/8/12 s 分别读取 + `[20,20.6)` / `[18,20.6)` / `[14,20.6)` / `[10,20.6)`。 + 局部 A 都是 `[19.6,20.6)`；0.2 s horizon 都指向 20.8 s。 +- Sessions 4--12 训练、1--3 测试。当前测试已多次查看，曲线如实标 exploratory； + 不按测试涨点选择窗口、训练 epoch 或 early stopping。 +- 同一份 1 s local sidecar 可复用，但其 dense rows、times、label indices、源媒体和 + encoder contract 必须一致；加载失败不得 padding/clamp。 +- 联合提取 wall time/显存不是单窗口流式效率。后者需单独运行所选 cap，测整个流水线。 + +## 1. 先跑 Session 4 的 24 s GPU smoke + +从仓库根目录执行，路径自行设置，不使用某台机器的固定挂载。 +24 s 会触发最大 12 s cap 的滚动重建；4 s smoke 不足以覆盖全部窗口。 + +```bash +MODEL_ID=/path/to/Qwen2.5-Omni-3B +VIDEO_ROOT=/path/to/easycom/Main/Video_Compressed +TARGET_ROOT=/path/to/EasyCom/target_perframe +MULTI_ROOT=/path/to/omni_features/qwen_multiwindow_2fps_last4_stack + +CUDA_VISIBLE_DEVICES=0 python3 -u extract_easycom_multiwindow.py \ + --model-id "$MODEL_ID" \ + --easycom-video-root "$VIDEO_ROOT" \ + --sessions Session_4 \ + --memory-seconds-list 2 4 6 8 10 12 \ + --end 24 --verify-every 11 \ + --flash-attention-2 \ + --output-root features/easycom_multiwindow_smoke +``` + +也可用 `--dataset-root /path/to/easycom` 代替 `--easycom-video-root`。 +`--media-only` 只检验媒体/decoder、不加载模型、不写 features。 +`--end` 是 bounded smoke，故即便给了 target root 也不附加整段 target；输出标记 +`partial_extraction=true`，不能与整段 target/local sidecar 混用。 + +检查 shape、时间、freshness、预算、commit schedule 和独立 replay： + +```bash +python3 - <<'PY' +from pathlib import Path +import torch +from local_turn_features import reference_alignment + +root = Path("features/easycom_multiwindow_smoke") +for cap in (2, 4, 6, 8, 10, 12): + p = torch.load(root / f"context_{cap}s/Session_4.pt", map_location="cpu", weights_only=False) + times, indices = reference_alignment(p) + d = p["extraction_diagnostics"] + assert p["states"].ndim == 3 and p["states"].shape[:2] == (120, 4) + assert torch.isfinite(p["states"]).all() + assert torch.equal(times, d["state_source_times"]) + assert d["state_age_seconds"].eq(0).all() + assert (d["actual_context_seconds"] <= cap + 1e-6).all() + assert d["payload_commit_mask"].nonzero().flatten().tolist() == list(range(9,120,10)) + valid = indices[:, None] + torch.arange(1,11)[None, :] < len(times) + assert int(valid.all(dim=1).sum()) == p["valid_10_horizon_steps"] == 110 + checks = [c for c in p["verification"] if c["kind"] == "cache_replay"] + assert len(checks) == 10 + print(cap, tuple(p["states"].shape), "shared AV calls=", d["production_processor_calls_shared"], + "max_abs=", max(c["max_abs"] for c in checks), + "min_cosine=", min(c["cosine"] for c in checks)) +PY +``` + +Same-segmentation `cache_replay` 应 exact/effectively exact；显著差异先排查，不开始全量。 +若要核对既有单窗口生产路径，再运行相同媒体/参数的 `--memory-seconds 8 --end 24` +到另一目录，逐行比较其 states 与 `context_8s/Session_4.pt`。 +CPU 测试已覆盖 >12 s 的滚动重建、可变 KV、从头 replay、future-only 合成输入改动、 +跨调用 reset、last4/mean、CLI 单窗口兼容及 head/local 加载成功和失败路径； +这不能替代真实模型/媒体的因果 GPU 审计。 + +## 2. Smoke 审计通过后，一条命令提取所有 sessions + +去掉 `--end`；完整目标长度和原始 20 FPS 视频帧数必须匹配。 +默认 sessions 为 1--12，也可以先只提取训练侧。验证关闭以避免额外 reference 开销。 + +```bash +CUDA_VISIBLE_DEVICES=0 python3 -u extract_easycom_multiwindow.py \ + --model-id "$MODEL_ID" \ + --easycom-video-root "$VIDEO_ROOT" \ + --target-root "$TARGET_ROOT" \ + --memory-seconds-list 2 4 6 8 10 12 \ + --flash-attention-2 \ + --output-root "$MULTI_ROOT" +``` + +输出结构：每个 `context_2s/` ... `context_12s/` 含独立 `Session_1.pt` ... +`Session_12.pt`。每个 payload 仍是 `mode=open-stream-replay`，保留 model、source、 +target、readout 和 sampling provenance，并记录 actual history range、phase boundary +和共享生产计数。`extraction_summary.json` 每完成一个 session 原子更新，记录本次调用的 +sessions、输出 shape、cap、运行 commit/dirty 状态、依赖版本及联合耗时。 + +默认拒绝覆盖任一已有 bank。中断后检查文件；可用 `--sessions` 选择未完成 sessions， +或换新输出目录。**不实现 checkpoint resume**；显式 `--overwrite` 才会重跑/替换。 +Summary 只描述本次调用，不能将上次与本次不同配置的输出默认为同一个实验。 +联合保存每个文件原子替换，但六个文件不是跨文件事务；异常后仍需核查是否齐全。 + +## 3. 每种 cap 独立训练 G 或 G+A head + +复用 [run_easycom_local_baselines.sh](../run_easycom_local_baselines.sh) 的固定四 epoch +recipe；不要复用 8 s head 到其他 cap 当作训练对照。 +`audio` 表示 **G+A**，不是纯 audio-only。`LOCAL_ROOT` 指向已验证的 1 s local sidecar。 + +```bash +export TARGET_ROOT +export LOCAL_ROOT=/path/to/omni_features/local_qwen_1s +for cap in 2 4 6 8 10 12; do + CUDA_VISIBLE_DEVICES=0 \ + FEATURE_ROOT="$MULTI_ROOT/context_${cap}s" \ + CHECKPOINT_ROOT="checkpoints/easycom_context_${cap}s_2026-09-17" \ + RESULT_ROOT="results/easycom_context_${cap}s_2026-09-17" \ + bash run_easycom_local_baselines.sh audio 0 1 2 +done +``` + +要检验 local A 是否缓解 context 敏感性，在相同 roots 下另外运行 `baseline 0 1 2` +（G-only）；不要在已生成的 `audio_seedN` 路径重复启动。局部 V 是独立条件追加对照。 +现有 driver 训练 4--12 后立即推理 1--3，因此开始整组运行前冻结分析方案。 +报告完整预指定六-cap 曲线、dense 三类 AP/mAP、raw wearer endpoint onset/offset AP、 +prevalence 和 common valid pairs；semantic ENTRY/YIELD 仍等标注审计。 + +CPU 回归命令： + +```bash +python3 -m unittest test_streaming_multiwindow test_streaming_chunk_ablation \ + test_easycom_head_objectives test_local_turn_features test_encoder_only_workflow +``` diff --git a/docs/local_evidence_analysis_experiment_log_2026-09-17.md b/docs/local_evidence_analysis_experiment_log_2026-09-17.md index 55aef7a..e222259 100644 --- a/docs/local_evidence_analysis_experiment_log_2026-09-17.md +++ b/docs/local_evidence_analysis_experiment_log_2026-09-17.md @@ -1,7 +1,8 @@ # Local evidence analysis：待做实验日志 -日期：2026-09-17。状态：**设计已记录，以下实验均未在本轮执行**。 -只修改日志与手稿占位，不新增训练、提取或 evaluator 实现；未产生新的实测结果。 +日期：2026-09-17。状态：**设计已记录，以下 GPU 实验均未在本轮执行**。 +初始记录只修改日志与手稿占位；同日已追加多窗口提取实现和 CPU 测试，见文末。 +未产生新的模型性能或 GPU 效率实测结果。 当前稿件边界见 [四页稿结构决策](manuscript_local_evidence_structure_2026-09-17.md)。 ## 本轮决定与篇幅 @@ -47,7 +48,7 @@ MiniROADA stage-two continuation 不是 Omni LoRA，也不替代其适配对照 | L4 | 因果 current-activity classifier，当前预测概率重复至十个 horizon | Dense 收益是否主要是更好的当前状态识别？ | Main Comparison TODO | | T1 | 8-s cap 下 tail refresh × local A 的四组训练对照 | 两条最新证据路径互补，还是冗余？ | Streaming 表上半部 | | T2 | 同 block decomposition 从头 replay、future-only media intervention、cache reset；按十档 phase 分层 | cache 正确性与 tail-length dependence 是否成立？ | Streaming TODO | -| C1 | 4/8/16-s cap 的 tail-aware G+A；按 cap 重提取并分别训练 | 历史长度是否影响主方法，8 s 是否是合理折中？ | Streaming 表下半部 | +| C1 | 初始 4/8/16-s 计划已扩展为 2/4/6/8/10/12-s 共享提取、独立 tail-aware G+A 训练；16 s 暂不执行 | 历史长度是否影响主方法，8 s 是否是合理折中？ | Streaming 表下半部；完整曲线留日志 | | C2 | 仅当要声称 local 分支缓解窗口敏感性时，在各 cap 增加 G 对照 | 窗口能否替代近期证据？融合是否对 cap 更稳定？ | 条件追加，不新增正文小节 | | B0 | 审计各骨干的 causal readout、维度、刷新与实际历史范围；先短片 GPU smoke | 跨骨干样本和接口是否可比？ | Setup / Backbone transfer 说明 | | B1 | Qwen2.5-Omni 3B/7B 各跑 G 与 G+A，保持同一份 local A | 增量是否在同系列不同规模下保留？ | Local Evidence + compact 表末三行 | @@ -88,7 +89,8 @@ completed-only 未得到首个 state 的行在全部条件中统一处理并报 ## 窗口与运行代价 -C1 中 local A/V 固定 1 s，Omni cap 为预先指定的 4/8/16 s。各 cap 的 global state +C1 中 local A/V 固定 1 s；同日按用户确认将 Omni cap 改为 2/4/6/8/10/12 s， +暂不跑初始计划中的 16 s。六个 cap 均保留，不按测试 AP 再挑选长度。各 cap 的 global state 必须从对应因果输入提取，不能通过裁剪缓存向量伪造窗口。8-s cap 行与 T1-T-GA 复用同一个确切 run，不重复训练后择优填写两处。 固定 checkpoint 改 cap 若执行，单独标成 distribution-shift sensitivity；不能混入 @@ -197,3 +199,30 @@ metrics.csv，保留 provenance 与协议。正文替换 TODO/--，不要把 pen 同日追加 backbone-transfer 短段和三行增量空位后重新编译、渲染检查全部四页， 仍为四页含参考文献；链接、TeX/bib keys、单展示公式与 `git diff --check` 再次通过。 未修改旧运行或宣称跨骨干提取/对照已完成。 + +## 同日追加：C1 多窗口提取实现，GPU 实验仍 pending + +用户确认“一条提取命令得到六个 cap 的特征，再分别训练”，已实现 +[extract_easycom_multiwindow.py](../extract_easycom_multiwindow.py)， +底层 [streaming_turn_multiwindow.py](../streaming_turn_multiwindow.py)。 +共享 session reader、各 200 ms cutoff 的 AV preprocessing/encoder embeddings， +但每个 cap 独立 Thinker KV、滚动重建和 readout；绝不裁剪长窗口高层 KV 代替短窗口。 +每个 `context_Cs/Session_N.pt` 是正常 `open-stream-replay` dense bank， +可复用同一份已验证的 1 s local sidecar 和现有四 epoch 训练脚本。 + +CPU 合成测试覆盖 24 s（含所有 cap 的滚动重建），逐点匹配六次单窗口生产提取； +AV 生产调用从 `6T` 降为共享 `T`，Thinker 仍逐 cap 执行。验证调用不计入该生产统计。 +同时检查从原输入从头 replay、可变 KV 隔离、未来输入干预不回填过去、session reset、 +last4 stack/mean、单窗口 CLI 兼容、head/local 成功与失败加载、输出覆盖防护。 +既有 head objective、local feature、encoder-only 和 chunk-ablation CPU 回归通过。 + +实际时间轴 cutoff=20.6 s：cap=2/4/6/8/10/12 s 的 global history start 分别为 +20/18/16/14/12/10 s；所有状态 source time=20.6 s，local A=[19.6,20.6)， +0.2 s target=20.8 s。Cap 是 native-block 最大预算而非逐行精确固定长度滑窗； +必须同时记录 phase 和 actual context span。 + +尚未执行真实 Qwen/媒体 GPU smoke、全量提取、六-cap head 训练或吞吐/显存测量。 +先用 Session 4 的 24 s smoke 检查六 bank shape/time/valid 与 independent replay， +再按 [多窗口运行手册](easycom_multiwindow_features.md) 提取所有 sessions。 +联合多-bank 耗时不是单窗口在线效率；不能据 CPU 调用计数写性能增益或实时结论。 +本次不修改 manuscript 数值或将 C1/T2/S1 标为 measured。 diff --git a/extract_easycom_multiwindow.py b/extract_easycom_multiwindow.py new file mode 100644 index 0000000..e3f1f0d --- /dev/null +++ b/extract_easycom_multiwindow.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 +"""Extract EasyCom 2/4/6/8/10/12s context banks with one shared AV pass. + +One frozen model is loaded for all selected sessions. The current partial +native block is re-encoded at each 5-Hz cutoff, once rather than once per cap. +Existing direct-media validation and single-window CLI remain unchanged. +""" + +from __future__ import annotations + +import argparse +import importlib.metadata +import json +import math +import re +import subprocess +import time +from pathlib import Path +from typing import Sequence + +import torch + +import streaming_turn_model as core +import streaming_turn_model_directly as direct +from streaming_turn_multiwindow import DEFAULT_MEMORY_SECONDS, output_paths, validate_memory_seconds + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--dataset-root", type=Path, help="EasyCom root containing Main/Video_Compressed") + source.add_argument("--easycom-video-root", type=Path, help="Directory directly containing Session_N") + parser.add_argument("--target-root", type=Path, help="Required for full extraction; target_perframe [T,3]") + parser.add_argument("--sessions", nargs="+", default=[f"Session_{i}" for i in range(1, 13)]) + parser.add_argument("--memory-seconds-list", type=float, nargs="+", default=list(DEFAULT_MEMORY_SECONDS)) + parser.add_argument("--output-root", type=Path, required=True) + parser.add_argument("--model-id", default=core.MODEL_ID) + parser.add_argument("--backbone-loader", choices=("qwen", "roma"), default="qwen") + parser.add_argument("--dtype", choices=("auto", "bfloat16", "float16", "float32"), default="bfloat16") + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--flash-attention-2", action="store_true") + parser.add_argument("--end", type=float, help="Bounded zero-origin smoke only; does not attach full targets") + parser.add_argument("--readout-last-n-layers", type=int, default=4) + parser.add_argument("--readout-layer-fusion", choices=("mean", "stack"), default="stack") + parser.add_argument("--video-fps", type=float, default=2.0) + parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28) + parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28) + parser.add_argument("--decoder", choices=("decord", "torchcodec", "pynvcodec"), default="decord") + parser.add_argument("--decoder-gpu-id", type=int, default=0) + parser.add_argument("--decoder-threads", type=int, default=1) + parser.add_argument("--max-open-readers", type=int, default=2) + parser.add_argument("--frame-cache-size", type=int, default=8) + parser.add_argument("--verify-every", type=int, default=0) + parser.add_argument("--verify-continuous-every", type=int, default=0) + parser.add_argument("--no-progress", action="store_true") + parser.add_argument("--media-only", action="store_true") + parser.add_argument("--overwrite", action="store_true") + return parser.parse_args(argv) + + +def session_commands(args: argparse.Namespace) -> list[tuple[str, list[str]]]: + """Preflight all paths/settings before loading the backbone.""" + caps = validate_memory_seconds(args.memory_seconds_list, 2.0) + core.StreamConfig( + memory_seconds=max(caps), video_fps=args.video_fps, + video_min_pixels=args.video_min_pixels, video_max_pixels=args.video_max_pixels, + readout_last_n_layers=args.readout_last_n_layers, readout_layer_fusion=args.readout_layer_fusion, + ).validate() + if len(set(args.sessions)) != len(args.sessions) or not args.sessions: + raise ValueError("sessions must be nonempty and unique") + if args.end is None and args.target_root is None: + raise ValueError("full extraction requires --target-root; use --end only for smoke") + if args.end is not None: + if not math.isfinite(args.end) or args.end <= 0: + raise ValueError("smoke end must be finite and positive") + core._require_integer_ratio(args.end, 0.2, "duration / step") + if args.verify_every < 0 or args.verify_continuous_every < 0: + raise ValueError("verification intervals must be non-negative") + if args.decoder_threads <= 0 or args.max_open_readers <= 0 or args.frame_cache_size < 0: + raise ValueError("invalid reader/cache limits") + video_root = args.easycom_video_root or args.dataset_root / "Main" / "Video_Compressed" + commands = [] + for session in args.sessions: + if re.fullmatch(r"Session_(?:[1-9]|1[0-2])", session) is None: + raise ValueError(f"invalid EasyCom session: {session}") + session_dir = video_root / session + if not session_dir.is_dir(): + raise FileNotFoundError(session_dir) + direct._easycom_directory_specs(session_dir) + output = args.output_root / f"{session}.pt" + if not args.overwrite and not args.media_only: + for path in output_paths(output, caps).values(): + if path.exists(): + raise FileExistsError(f"{path}; use --overwrite explicitly or a fresh output root") + command = ["--easycom-session-dir", str(session_dir), "--output", str(output), + "--mode", "open-stream-replay", "--memory-seconds-list", *map(str, caps)] + for name in ( + "model_id", "backbone_loader", "dtype", "device", "readout_last_n_layers", + "readout_layer_fusion", "video_fps", "video_min_pixels", "video_max_pixels", + "decoder", "decoder_gpu_id", "decoder_threads", "max_open_readers", "frame_cache_size", + "verify_every", "verify_continuous_every", + ): + command.extend(["--" + name.replace("_", "-"), str(getattr(args, name))]) + for name in ("flash_attention_2", "no_progress", "media_only", "overwrite"): + if getattr(args, name): + command.append("--" + name.replace("_", "-")) + if args.end is not None: + command.extend(["--end", str(args.end)]) + else: + target = args.target_root / f"{session}.npy" + core.target_num_steps(target) # Missing/invalid targets fail before model loading. + command.extend(["--target", str(target)]) + commands.append((session, command)) + return commands + + +def main(argv: Sequence[str] | None = None) -> None: + args = parse_args(argv) + commands = session_commands(args) + if args.media_only: + for _, command in commands: + direct.main(command) + return + if args.end is not None: + print("SMOKE ONLY: partial banks are not a full-session training/evaluation dataset.", flush=True) + core._configure_qwen_logging() + load_started = time.perf_counter() + model, processor = core.load_qwen_omni( + args.model_id, dtype=args.dtype, device_map={"": args.device}, + use_flash_attention_2=args.flash_attention_2, backbone_loader=args.backbone_loader, + ) + model_loading_seconds = time.perf_counter() - load_started + repository = Path(__file__).resolve().parent + revision = subprocess.run(["git", "rev-parse", "HEAD"], cwd=repository, capture_output=True, text=True, check=False) + dirty = subprocess.run(["git", "status", "--porcelain"], cwd=repository, capture_output=True, text=True, check=False) + summary = { + "format": "easycom-multiwindow-extraction-v1", "arguments": { + key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items() + }, + "split": {"train": [f"Session_{i}" for i in range(4, 13)], + "test": ["Session_1", "Session_2", "Session_3"]}, + "backbone": "frozen", "partial_extraction": args.end is not None, + "code_revision": revision.stdout.strip() if revision.returncode == 0 else None, + "worktree_dirty": bool(dirty.stdout.strip()) if dirty.returncode == 0 else None, + "versions": {name: importlib.metadata.version(name) for name in ("torch", "transformers", "qwen-omni-utils")}, + "model_loading_seconds": model_loading_seconds, + "session_wall_scope": "direct_media_index_decode_extract_save_excluding_model_load_and_summary_scan", + "runtime_scope": "joint_multi_bank_with_shared_AV_not_single_window_online_runtime", + "sessions": {}, + } + for session, command in commands: + started = time.perf_counter() + direct.main(command, model=model, processor=processor) + elapsed = time.perf_counter() - started + records = {} + for cap, path in output_paths(args.output_root / f"{session}.pt", sorted(args.memory_seconds_list)).items(): + payload = torch.load(path, map_location="cpu", weights_only=False) + checks = [c for c in payload["verification"] if c.get("kind") == "cache_replay"] + spans = payload["extraction_diagnostics"]["actual_context_seconds"] + records[f"{cap:g}s"] = { + "path": str(path), "shape": list(payload["states"].shape), + "num_steps": payload["num_steps"], "max_actual_context_seconds": spans.max().item(), + "replay_checks": len(checks), + "worst_replay_max_abs": max((c["max_abs"] for c in checks), default=None), + "lowest_replay_cosine": min((c["cosine"] for c in checks), default=None), + "production_processor_calls_shared": payload["num_steps"], + } + del payload + summary["sessions"][session] = {"joint_wall_seconds": elapsed, "banks": records} + args.output_root.mkdir(parents=True, exist_ok=True) + summary_path = args.output_root / "extraction_summary.json" + temporary = summary_path.with_suffix(".json.tmp") + temporary.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") + temporary.replace(summary_path) + print(f"Summary: {args.output_root / 'extraction_summary.json'}") + + +if __name__ == "__main__": + main() diff --git a/streaming_turn_model_directly.py b/streaming_turn_model_directly.py index 4f0ba10..dbf1727 100644 --- a/streaming_turn_model_directly.py +++ b/streaming_turn_model_directly.py @@ -1061,6 +1061,14 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: help="0 keeps Thinker KV for the full extraction; positive values use a rolling window", ) parser.add_argument("--video-fps", type=float, default=2.0) + parser.add_argument( + "--memory-seconds-list", type=float, nargs="+", + help="Export independent open-stream banks in context_Cs/ with shared AV encoding", + ) + parser.add_argument( + "--overwrite", action="store_true", + help="Allow replacing existing multiwindow banks (single-window behavior is unchanged)", + ) parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28) parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28) parser.add_argument( @@ -1222,6 +1230,21 @@ def main( args = parse_args(argv) core._configure_qwen_logging() + multi_caps = None + if args.memory_seconds_list is not None: + from streaming_turn_multiwindow import output_paths, validate_memory_seconds + + if args.mode != "open-stream-replay" or args.head_checkpoint is not None: + raise SystemExit("multiwindow export requires open-stream-replay without --head-checkpoint") + if args.memory_seconds != 0: + raise SystemExit("use --memory-seconds-list instead of --memory-seconds, not both") + if not math.isclose(args.start, 0.0, abs_tol=1e-9) or not math.isclose(args.step_seconds, 0.2, abs_tol=1e-9): + raise SystemExit("multiwindow dense banks require --start 0 and --step-seconds 0.2") + multi_caps = validate_memory_seconds(args.memory_seconds_list, args.native_block_seconds) + if args.output is not None and not args.media_only and not args.overwrite: + existing = [str(path) for path in output_paths(args.output, multi_caps).values() if path.exists()] + if existing: + raise FileExistsError(f"multiwindow outputs already exist; use --overwrite explicitly: {existing}") if (model is None) != (processor is None): raise ValueError("model and processor must be supplied together") if args.verify_every < 0 or args.verify_continuous_every < 0: @@ -1340,7 +1363,7 @@ def main( "pass a positive --memory-seconds (recommended: 8)" ) effective_memory = _effective_memory_seconds( - args.memory_seconds, + max(multi_caps) if multi_caps else args.memory_seconds, duration=duration, block_seconds=args.native_block_seconds, ) @@ -1404,7 +1427,20 @@ def main( overlap_extractor = None working_grid_av = None - if args.mode == "causal-prefix": + multi_results = None + if multi_caps: + from streaming_turn_multiwindow import MultiWindowOpenStreamReplayExtractor + + extractor = MultiWindowOpenStreamReplayExtractor( + model, processor, config, memory_seconds_list=multi_caps, fixed_grid_av=reader, + ) + multi_results = extractor.extract( + video_path=media_path_for_api, audio_path=None, start=args.start, end=end, + use_audio_in_video=True, verify_every=args.verify_every, + verify_continuous_every=args.verify_continuous_every, show_progress=not args.no_progress, + ) + result = multi_results[max(multi_caps)] + elif args.mode == "causal-prefix": extractor = core.CausalPrefixExtractor( model, processor, config, fixed_grid_av=reader ) @@ -1616,6 +1652,9 @@ def main( "verification": result.verification, "extraction_diagnostics": result.diagnostics, } + if multi_results is not None: + save_multiwindow_banks(payload, multi_results, args.output, overwrite=args.overwrite) + return if args.head_checkpoint is not None: head, head_loss_name = core._load_head( args.head_checkpoint, core._hidden_size(model), model.device @@ -1647,5 +1686,56 @@ def main( ) +def save_multiwindow_banks( + payload: dict[str, Any], results: dict[float, core.ExtractionResult], output: Path, + *, overwrite: bool = False, +) -> dict[float, Path]: + """Write ordinary head-compatible payloads; never overwrite implicitly.""" + from local_turn_features import reference_alignment + from streaming_turn_multiwindow import output_paths + + paths = output_paths(output, sorted(results)) + banks = {} + for cap, result in results.items(): + if len(result.times) != payload["num_steps"]: + raise ValueError(f"{cap:g}s bank has a different number of dense rows") + bank = dict(payload) + bank.update( + config={**payload["config"], "memory_seconds": cap}, + requested_memory_seconds=cap, persistent_thinker_kv=False, + times=result.times, states=result.states, verification=result.verification, + extraction_diagnostics=result.diagnostics, + partial_extraction=payload["target_path"] is None, + open_stream_policy={**payload["open_stream_policy"], + "memory": "rolling_complete_blocks_plus_current_partial_block", + "rolling_replay_inputs": "reuse_cached_complete_block_av_embeddings"}, + multiwindow_policy={ + "memory_seconds_list": sorted(results), "context_cap_seconds": cap, + "shared": "media_reader_preprocessing_AV_block_and_tail_embeddings", + "independent": "per_context_Thinker_KV_replay_and_readout", + "long_context_KV_cropping": False, + "equivalence_reference": "single_window_open_stream_replay_same_cap", + "runtime_scope": "joint_multi_bank_extraction_not_single_bank_online_latency", + }, + ) + reference_alignment(bank) + if not torch.isfinite(result.states).all(): + raise ValueError(f"{cap:g}s bank contains non-finite states") + if paths[cap].exists() and not overwrite: + raise FileExistsError(paths[cap]) + banks[cap] = bank + for cap, result in results.items(): + paths[cap].parent.mkdir(parents=True, exist_ok=True) + temporary = paths[cap].with_suffix(".pt.tmp") + torch.save(banks[cap], temporary) + temporary.replace(paths[cap]) + print(f"Saved {cap:g}s bank: {paths[cap]} states={tuple(result.states.shape)}") + checks = [item for item in result.verification if item.get("kind") == "cache_replay"] + if checks: + print(f" independent replay n={len(checks)} max_abs={max(c['max_abs'] for c in checks):.6g} " + f"min_cosine={min(c['cosine'] for c in checks):.8f}") + return paths + + if __name__ == "__main__": main() diff --git a/streaming_turn_multiwindow.py b/streaming_turn_multiwindow.py new file mode 100644 index 0000000..42f2031 --- /dev/null +++ b/streaming_turn_multiwindow.py @@ -0,0 +1,211 @@ +"""Shared AV preparation, independent strict-window Thinker replay banks. + +Only encoder-native block embeddings are shared. High-layer KV is never +cropped from a longer context to stand in for a shorter-context extraction. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Sequence + +import torch +from tqdm.auto import tqdm + +import streaming_turn_model as core + +DEFAULT_MEMORY_SECONDS = (2.0, 4.0, 6.0, 8.0, 10.0, 12.0) + + +def validate_memory_seconds(values: Sequence[float], native_seconds: float) -> tuple[float, ...]: + """Reject ambiguous/non-native budgets before loading media or a model.""" + if not math.isfinite(native_seconds) or native_seconds <= 0: + raise ValueError("native block seconds must be finite and positive") + caps = tuple(float(value) for value in values) + if not caps or len(set(caps)) != len(caps): + raise ValueError("memory seconds list must be nonempty with no duplicates") + for cap in caps: + if not math.isfinite(cap) or cap < native_seconds: + raise ValueError("every memory cap must be finite and at least one native block") + blocks = core._require_integer_ratio(cap, native_seconds, "memory / native block") + if cap != blocks * native_seconds: + raise ValueError("memory caps must be exact native-block multiples, not rounded budgets") + return tuple(sorted(caps)) + + +def context_directory(cap: float) -> str: + number = str(int(cap)) if float(cap).is_integer() else repr(float(cap)) + return f"context_{number}s" + + +def output_paths(output: Path, caps: Sequence[float]) -> dict[float, Path]: + return {cap: output.parent / context_directory(cap) / output.name for cap in caps} + + +@dataclass +class _Bank: + helper: core.OpenStreamReplayExtractor + history_start: float + committed_cache: Any = None + payload_base: int | None = None + states: list[torch.Tensor] = field(default_factory=list) + checks: list[dict[str, Any]] = field(default_factory=list) + history_starts: list[float] = field(default_factory=list) + rebuilds: int = 0 + replayed_blocks: int = 0 + + +class MultiWindowOpenStreamReplayExtractor(core.OpenStreamReplayExtractor): + """Same tail-aware semantics as separate single-window extractions. + +Each bank has an independent KV history. All banks consume the same prepared +current tail and a bounded shared ring of completed AV block embeddings. +Temporary tails, chat closure and readout queries are not committed until a +native boundary; only the completed media payload is committed there. +""" + + def __init__(self, model, processor, config, *, memory_seconds_list, fixed_grid_av=None): + self.memory_seconds_list = validate_memory_seconds( + memory_seconds_list, config.native_block_seconds + ) + super().__init__( + model, processor, replace(config, memory_seconds=max(self.memory_seconds_list)), + fixed_grid_av=fixed_grid_av, + ) + + @torch.inference_mode() + def extract( + self, *, video_path: str | Path | None, audio_path: str | Path | None, + start: float, end: float, use_audio_in_video: bool, + verify_every: int = 0, verify_continuous_every: int = 0, + show_progress: bool = True, + ) -> dict[float, core.ExtractionResult]: + if verify_every < 0 or verify_continuous_every < 0: + raise ValueError("verification intervals must be non-negative") + if not math.isclose(start, 0.0, abs_tol=1e-9) or not math.isclose( + self.config.step_seconds, 0.2, abs_tol=1e-9 + ): + raise ValueError("multiwindow dense banks require start=0 and step_seconds=0.2") + if not math.isfinite(end) or end <= start: + raise ValueError("end must be finite and greater than start") + core._require_integer_ratio(end - start, self.config.step_seconds, "duration / step") + times = core._time_grid(start, end, self.config.step_seconds) + native = self.config.native_block_seconds + positions_per_chunk = core._tmrope_positions_per_chunk(self.model) + banks = { + cap: _Bank( + helper=core.OpenStreamReplayExtractor( + self.model, self.processor, replace(self.config, memory_seconds=cap), + fixed_grid_av=self.fixed_grid_av, + ), + history_start=start, + ) + for cap in self.memory_seconds_list + } + completed = core.CompletedBlockPieceCache( + origin=start, block_seconds=native, + max_blocks=round(max(self.memory_seconds_list) / native) - 1, + ) + block_start = start + prepared_seconds = 0.0 + with tqdm(total=len(times), desc="multiwindow-replay", unit="cutoff", + disable=not show_progress, dynamic_ncols=True) as progress: + for step, cutoff in enumerate(times, start=1): + for cap, bank in banks.items(): + if bank.committed_cache is not None and block_start - bank.history_start >= cap - 1e-6: + bank.history_start = block_start - (cap - native) + intervals = core._block_intervals(bank.history_start, block_start, native) + bank.rebuilds += 1 + bank.replayed_blocks += len(intervals) + if intervals: + bank.committed_cache, _, bank.payload_base, _, _ = bank.helper._cache_from_intervals( + video_path=video_path, audio_path=audio_path, intervals=intervals, + use_audio_in_video=use_audio_in_video, + prepared_blocks=completed.get_many(intervals), + ) + else: + bank.committed_cache = None + bank.payload_base = None + + # Exactly one production processor/tower invocation per cutoff. + pieces = self._prepare_pieces( + video_path=video_path, audio_path=audio_path, + interval=(block_start, cutoff), use_audio_in_video=use_audio_in_video, + ) + prepared_seconds += cutoff - block_start + boundary = math.isclose(cutoff, block_start + native, abs_tol=1e-6) + for cap, bank in banks.items(): + if bank.payload_base is None: + bank.payload_base = int(pieces.payload_position_ids.min().item()) + block_index = round((block_start - bank.history_start) / native) + media_cache, _, closure_positions = core._append_open_pieces( + self.model, pieces, past_key_values=core.clone_cache(bank.committed_cache), + rope_offset=bank.payload_base + block_index * positions_per_chunk, + include_prefix=bank.committed_cache is None, + ) + state = core._read_open_state( + self.model, media_cache, closure_ids=pieces.closure_ids, + closure_position_ids=closure_positions, query_ids=bank.helper.query_ids, + last_n_layers=self.config.readout_last_n_layers, + layer_fusion=self.config.readout_layer_fusion, + ) + bank.states.append(state.cpu()) + bank.history_starts.append(bank.history_start) + intervals = core._block_intervals(bank.history_start, cutoff, native) + if verify_every and step % verify_every == 0: + # Deliberately independent processor/tower calls, not the shared ring. + reference = bank.helper.state_from_scratch( + video_path=video_path, audio_path=audio_path, intervals=intervals, + use_audio_in_video=use_audio_in_video, + ) + bank.checks.append(core._compare_states(cutoff, state, reference, kind="cache_replay")) + if verify_continuous_every and step % verify_continuous_every == 0: + reference = bank.helper.continuous_reference.state_for_intervals( + video_path=video_path, audio_path=audio_path, + intervals=[(bank.history_start, cutoff)], use_audio_in_video=use_audio_in_video, + ) + bank.checks.append(core._compare_states( + cutoff, state, reference, kind="continuous_window_semantics" + )) + if boundary: + bank.committed_cache = media_cache + if boundary: + completed.put(core.CachedOpenMediaBlock(block_start, cutoff, pieces)) + block_start = cutoff + progress.set_postfix_str(f"media={cutoff:.1f}s banks={len(banks)}", refresh=False) + progress.update(1) + + time_tensor = torch.tensor(times, dtype=torch.float64) + boundary_mask = core._periodic_boundary_mask(times, origin=start, period=native) + results = {} + for cap, bank in banks.items(): + history = torch.tensor(bank.history_starts, dtype=torch.float64) + span = time_tensor - history + if (span <= 0).any() or (span > cap + 1e-6).any(): + raise RuntimeError(f"invalid context span in {cap:g}s bank") + results[cap] = core.ExtractionResult( + times=time_tensor.clone(), states=torch.cat(bank.states, dim=0), + verification=bank.checks, + diagnostics={ + "history_start_times": history, + "actual_context_seconds": span, + "state_source_times": time_tensor.clone(), + "state_age_seconds": torch.zeros_like(time_tensor), + "initialization_mask": torch.zeros(len(times), dtype=torch.bool), + "state_update_mask": torch.ones(len(times), dtype=torch.bool), + "payload_commit_mask": boundary_mask.clone(), + "native_boundary_mask": boundary_mask.clone(), + "production_processor_calls_shared": len(times), + "production_prepared_media_seconds_shared": prepared_seconds, + "production_thinker_tail_calls": len(times), + "rolling_rebuilds": bank.rebuilds, + "replayed_complete_blocks": bank.replayed_blocks, + "shared_completed_block_capacity": round(max(self.memory_seconds_list) / native) - 1, + "memory_seconds_list": list(self.memory_seconds_list), + "sharing": "AV_embeddings_only_independent_Thinker_KV", + "verification_excluded_from_production_counters": True, + }, + ) + return results diff --git a/test_streaming_multiwindow.py b/test_streaming_multiwindow.py new file mode 100644 index 0000000..08a5367 --- /dev/null +++ b/test_streaming_multiwindow.py @@ -0,0 +1,299 @@ +"""CPU cache semantics, output routing and both feature-loader contracts.""" + +from __future__ import annotations + +import copy +import json +import tempfile +import types +import unittest +from pathlib import Path +from unittest import mock + +import numpy as np +import torch + +import extract_easycom_multiwindow as driver +import streaming_turn_model as core +import streaming_turn_model_directly as direct +from easycom_head import load_session_data +from local_turn_features import LOCAL_FORMAT, load_local_features, reference_alignment +from streaming_turn_multiwindow import DEFAULT_MEMORY_SECONDS, MultiWindowOpenStreamReplayExtractor, validate_memory_seconds + + +class FakeModel: + config = types.SimpleNamespace(seconds_per_chunk=2.0, position_id_per_seconds=25.0, + text_config=types.SimpleNamespace(hidden_size=3)) + device = torch.device("cpu") + dtype = torch.float32 + + +class FakeCache: + def __init__(self): + self.tokens = [] + + +def append(model, pieces, *, past_key_values, rope_offset, include_prefix): + cache = FakeCache() if past_key_values is None else past_key_values + if include_prefix != (not cache.tokens): + raise AssertionError("prefix ownership mismatch or readout committed") + values = pieces.payload_embeddings.flatten().tolist() + # Mutation exercises clone_cache isolation, not a conveniently immutable fake. + cache.tokens.append((values, rope_offset)) + return cache, rope_offset + 1, pieces.closure_position_ids + + +def read(model, cache, *, last_n_layers=1, layer_fusion="mean", **kwargs): + value = 0.0 + for values, offset in cache.tokens: + value = value * 1.17 + sum(values) + offset / 100 + state = torch.tensor([[value, cache.tokens[0][0][0], cache.tokens[-1][0][1]]]) + if layer_fusion == "stack": + return torch.stack([state + layer for layer in range(last_n_layers)], dim=1) + return state + + +class MultiwindowTests(unittest.TestCase): + def setUp(self): + self.prepared = [] + self.future_change = False + + def prepare(extractor, *, interval, **kwargs): + start, end = interval + self.prepared.append(interval) + signal = 1000.0 if self.future_change and end > 13.0 else 0.0 + return core.OpenMediaPieces( + prefix_embeddings=torch.zeros(1, 1, 3), + prefix_position_ids=torch.zeros(3, 1, 1, dtype=torch.long), + payload_embeddings=torch.tensor([[[start, end, signal]]]), + payload_position_ids=torch.ones(3, 1, 1, dtype=torch.long), + closure_ids=torch.zeros(1, 1, dtype=torch.long), + closure_position_ids=torch.ones(3, 1, 1, dtype=torch.long), + ) + + for patch in ( + mock.patch.object(core, "_readout_query_ids", return_value=torch.zeros(1, 1, dtype=torch.long)), + mock.patch.object(core.OpenStreamReplayExtractor, "_prepare_pieces", new=prepare), + mock.patch.object(core, "_append_open_pieces", side_effect=append), + mock.patch.object(core, "_read_open_state", side_effect=read), + ): + patch.start() + self.addCleanup(patch.stop) + self.config = core.StreamConfig(memory_seconds=12, readout_last_n_layers=4, readout_layer_fusion="stack") + self.kwargs = dict(video_path="unused.mp4", audio_path=None, start=0.0, end=24.0, + use_audio_in_video=True, show_progress=False) + + def multi(self, **kwargs): + extractor = MultiWindowOpenStreamReplayExtractor( + FakeModel(), object(), self.config, memory_seconds_list=DEFAULT_MEMORY_SECONDS + ) + return extractor.extract(**{**self.kwargs, **kwargs}) + + def test_all_caps_match_separate_single_window_extractions_and_encode_once(self): + results = self.multi() + self.assertEqual(len(self.prepared), 120) + self.assertEqual(len(set(self.prepared)), 120) + for cap, result in results.items(): + reference = core.OpenStreamReplayExtractor( + FakeModel(), object(), core.StreamConfig(memory_seconds=cap, + readout_last_n_layers=4, readout_layer_fusion="stack") + ).extract(**self.kwargs) + self.assertTrue(torch.equal(result.times, reference.times)) + self.assertTrue(torch.equal(result.states, reference.states), f"mismatch cap={cap}") + self.assertEqual(result.states.shape, (120, 4, 3)) + self.assertEqual(result.diagnostics["production_processor_calls_shared"], 120) + self.assertTrue((result.diagnostics["actual_context_seconds"] <= cap + 1e-6).all()) + self.assertTrue(torch.equal(result.diagnostics["state_source_times"], result.times)) + self.assertEqual(result.diagnostics["payload_commit_mask"].sum().item(), 12) + # Representative physical timeline, independent of any learned scores. + index = 102 # cutoff=20.6 s + print("Timeline cutoff=20.6s, local A=[19.6,20.6), horizon0.2 target=20.8s") + for cap, result in results.items(): + start = result.diagnostics["history_start_times"][index].item() + self.assertAlmostEqual(start, 22.0 - cap) + print(f" cap={cap:g}s global=[{start:g},20.6) age=0s") + + def test_independent_replay_verification_and_no_future_backfill(self): + original = self.multi(verify_every=11) + for result in original.values(): + self.assertEqual(len(result.verification), 10) + self.assertTrue(all(check["max_abs"] == 0 for check in result.verification)) + self.future_change = True + changed = self.multi() + for cap in original: + self.assertTrue(torch.equal(original[cap].states[:65], changed[cap].states[:65])) + self.assertFalse(torch.equal(original[cap].states[65:], changed[cap].states[65:])) + + def test_repeated_extract_resets_history_and_supports_mean(self): + extractor = MultiWindowOpenStreamReplayExtractor( + FakeModel(), object(), core.StreamConfig(), memory_seconds_list=[2, 4] + ) + first = extractor.extract(**self.kwargs) + second = extractor.extract(**self.kwargs) + self.assertEqual(first[2].states.shape, (120, 3)) + self.assertTrue(torch.equal(first[4].states, second[4].states)) + for kwargs in ({"start": 0.2}, {"end": 0.1}, {"end": float("nan")}, {"verify_every": -1}): + with self.assertRaises(ValueError): + extractor.extract(**{**self.kwargs, **kwargs}) + + def test_invalid_caps(self): + for caps in ([], [2, 2], [2, 2.000000001], [0], [1], [3], [float("nan")], [float("inf")]): + with self.subTest(caps=caps), self.assertRaises(ValueError): + validate_memory_seconds(caps, 2) + + def test_export_loads_both_head_and_same_local_audio_and_rejects_bad_alignment(self): + results = self.multi(end=4) + n = 20 + payload = {"mode": "open-stream-replay", "model_id": "synthetic", "start": 0.0, + "num_steps": n, "target_path": "synthetic.npy", "config": {"step_seconds": .2}, + "label_indices": torch.arange(n), "open_stream_policy": {}} + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + target_root = root / "targets" + target_root.mkdir() + targets = np.tile(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 1]], dtype=np.float32), (5, 1)) + np.save(target_root / "Session_4.npy", targets) + paths = direct.save_multiwindow_banks(payload, results, root / "Session_4.pt") + side = {"format": LOCAL_FORMAT, "dataset": "easycom", "sample_id": "Session_4", + "partial": False, "times": results[2].times.clone(), "label_indices": torch.arange(n), + "audio": torch.ones(n, 5), "audio_valid": torch.ones(n, dtype=torch.bool), + "config": {"audio_dim": 5, "causal_policy": "half-open-trailing-window", + "uses_speech_labels": False, "encoder_id": "synthetic"}} + side_path = root / "audio.pt" + torch.save(side, side_path) + for cap, path in paths.items(): + bank = torch.load(path, weights_only=False) + loaded = load_session_data(path.parent, target_root, 4, expected_mode="open-stream-replay") + self.assertTrue(torch.equal(loaded.states, results[cap].states)) + local, _ = load_local_features(side_path, bank, sample_id="Session_4", dataset="easycom", modalities="audio") + self.assertEqual(local.shape, (n, 6)) # encoder vector + availability + self.assertFalse(bank["persistent_thinker_kv"]) + self.assertEqual(bank["config"]["memory_seconds"], cap) + self.assertFalse(bank["partial_extraction"]) + for indices in (torch.arange(n) + 1, torch.arange(n - 1)): + broken = {**bank, "label_indices": indices} + torch.save(broken, path) + with self.assertRaises(ValueError): + load_session_data(path.parent, target_root, 4, expected_mode="open-stream-replay") + with self.assertRaises(ValueError): + load_local_features(side_path, broken, sample_id="Session_4", dataset="easycom", modalities="audio") + torch.save(bank, path) + for broken in ({**bank, "times": bank["times"] + .2}, {**bank, "start": .2}): + with self.assertRaises(ValueError): + reference_alignment(broken) + with self.assertRaises(ValueError): + load_session_data(path.parent, target_root, 4) # Wrong mode. + with self.assertRaises(FileExistsError): + direct.save_multiwindow_banks(payload, results, root / "Session_4.pt") + bad = copy.deepcopy(results) + bad[12].states[0, 0, 0] = float("nan") + with self.assertRaises(ValueError): + direct.save_multiwindow_banks(payload, bad, root / "fresh" / "Session_4.pt") + self.assertFalse((root / "fresh" / "context_2s" / "Session_4.pt").exists()) + + +class DriverTests(unittest.TestCase): + def test_driver_loads_model_once_across_sessions_and_writes_summary(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + target_root = root / "targets" + target_root.mkdir() + for session in ("Session_4", "Session_5"): + (root / session).mkdir() + np.save(target_root / f"{session}.npy", np.tile(np.array([[1, 0, 0]], dtype=np.float32), (20, 1))) + model, processor = object(), object() + + def run(command, *, model=None, processor=None): + self.assertIs(model, loaded_model) + self.assertIs(processor, loaded_processor) + args = direct.parse_args(command) + times = (torch.arange(20, dtype=torch.float64) + 1) / 5 + results = {cap: core.ExtractionResult( + times, torch.ones(20, 8), [], {"actual_context_seconds": torch.full((20,), cap)} + ) for cap in args.memory_seconds_list} + payload = {"mode": "open-stream-replay", "start": 0, "num_steps": 20, + "target_path": str(args.target), "config": {"step_seconds": .2}, + "label_indices": torch.arange(20), "open_stream_policy": {}} + direct.save_multiwindow_banks(payload, results, args.output) + + loaded_model, loaded_processor = model, processor + with mock.patch.object(direct, "_easycom_directory_specs"), \ + mock.patch.object(core, "load_qwen_omni", return_value=(model, processor)) as loader, \ + mock.patch.object(direct, "main", side_effect=run) as runner: + driver.main(["--easycom-video-root", str(root), "--target-root", str(target_root), + "--sessions", "Session_4", "Session_5", "--memory-seconds-list", "2", "4", + "--output-root", str(root / "output")]) + loader.assert_called_once() + self.assertEqual(runner.call_count, 2) + summary = json.loads((root / "output" / "extraction_summary.json").read_text()) + self.assertEqual(set(summary["sessions"]), {"Session_4", "Session_5"}) + self.assertFalse(summary["partial_extraction"]) + self.assertEqual(summary["sessions"]["Session_5"]["banks"]["4s"]["shape"], [20, 8]) + + def test_direct_main_routes_multi_banks_and_preserves_scalar_output(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + times = (torch.arange(20, dtype=torch.float64) + 1) / 5 + result = core.ExtractionResult(times, torch.ones(20, 8), [], {}) + reader = types.SimpleNamespace( + clips=[types.SimpleNamespace(path=Path("unused.mp4"))], metadata=lambda: {"synthetic": True} + ) + model = FakeModel() + model.eval = lambda: model + model.requires_grad_ = lambda value: model + base = ["--video", "unused.mp4", "--end", "4", "--output", str(root / "Session_4.pt")] + with mock.patch.object(direct, "CausalSourceAV", return_value=reader), \ + mock.patch.object(direct, "_print_media_summary"), \ + mock.patch.object(core, "backbone_metadata", return_value={"loader": "unknown"}), \ + mock.patch.object(MultiWindowOpenStreamReplayExtractor, "__init__", return_value=None), \ + mock.patch.object(MultiWindowOpenStreamReplayExtractor, "extract", return_value={2: result, 4: result}): + direct.main(base + ["--memory-seconds-list", "2", "4"], model=model, processor=object()) + self.assertFalse((root / "Session_4.pt").exists()) + bank = torch.load(root / "context_2s" / "Session_4.pt", weights_only=False) + self.assertEqual(bank["config"]["memory_seconds"], 2) + self.assertFalse(bank["persistent_thinker_kv"]) + self.assertTrue(bank["partial_extraction"]) + with mock.patch.object(direct, "CausalSourceAV", return_value=reader), \ + mock.patch.object(direct, "_print_media_summary"), \ + mock.patch.object(core, "backbone_metadata", return_value={"loader": "unknown"}), \ + mock.patch.object(core.OpenStreamReplayExtractor, "__init__", return_value=None), \ + mock.patch.object(core.OpenStreamReplayExtractor, "extract", return_value=result): + direct.main(base + ["--memory-seconds", "4"], model=model, processor=object()) + self.assertTrue((root / "Session_4.pt").exists()) + + def test_cli_preflight_and_partial_target_policy(self): + with tempfile.TemporaryDirectory() as directory, mock.patch.object(direct, "_easycom_directory_specs"): + root = Path(directory) + (root / "Session_4").mkdir() + common = ["--easycom-video-root", str(root), "--sessions", "Session_4", + "--output-root", str(root / "output"), "--end", "24"] + args = driver.parse_args(common) + commands = driver.session_commands(args) + self.assertNotIn("--target", commands[0][1]) + self.assertIn("--memory-seconds-list", commands[0][1]) + self.assertEqual(direct.parse_args(commands[0][1]).readout_last_n_layers, 4) + with mock.patch.object(core, "load_qwen_omni") as loader, mock.patch.object(direct, "main") as main: + driver.main(common + ["--media-only"]) + loader.assert_not_called() + main.assert_called_once() + for extra in (["--sessions", "Session_99"], ["--sessions", "Session_4", "Session_4"], + ["--memory-seconds-list", "3"], ["--end", "0.1"]): + with self.assertRaises(ValueError): + driver.session_commands(driver.parse_args(common + extra)) + existing = root / "output" / "context_2s" / "Session_4.pt" + existing.parent.mkdir(parents=True) + existing.touch() + with self.assertRaises(FileExistsError): + driver.session_commands(args) + + def test_direct_legacy_parser_and_multi_conflicts(self): + base = ["--video", "unused.mp4", "--output", "unused.pt", "--end", "24"] + self.assertIsNone(direct.parse_args(base).memory_seconds_list) + for extra in (["--mode", "causal-prefix"], ["--memory-seconds", "8"], + ["--head-checkpoint", "head.pt"], ["--start", ".2"]): + with self.assertRaises(SystemExit): + direct.main(base + ["--memory-seconds-list", "2", "4"] + extra) + + +if __name__ == "__main__": + unittest.main()