Add Qwen tail and chunk ablations

This commit is contained in:
Xin Wang
2026-09-17 02:32:08 +08:00
parent a183fe5944
commit a9f748d6a6
7 changed files with 1219 additions and 32 deletions

View File

@@ -1025,6 +1025,8 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
choices=(
"causal-prefix",
"open-stream-replay",
"completed-block-hold",
"micro-chunk-200ms",
"overlap-working-memory",
),
default="open-stream-replay",
@@ -1414,8 +1416,17 @@ def main(
use_audio_in_video=True,
show_progress=not args.no_progress,
)
elif args.mode == "open-stream-replay":
extractor = core.OpenStreamReplayExtractor(
elif args.mode in {
"open-stream-replay",
"completed-block-hold",
"micro-chunk-200ms",
}:
extractor_class = {
"open-stream-replay": core.OpenStreamReplayExtractor,
"completed-block-hold": core.CompletedBlockHoldExtractor,
"micro-chunk-200ms": core.MicroChunkReplayExtractor,
}[args.mode]
extractor = extractor_class(
model, processor, config, fixed_grid_av=reader
)
result = extractor.extract(
@@ -1459,25 +1470,33 @@ def main(
)
persistent_memory = (
args.mode == "open-stream-replay" and args.memory_seconds == 0
args.mode
in {
"open-stream-replay",
"completed-block-hold",
"micro-chunk-200ms",
}
and args.memory_seconds == 0
)
payload: dict[str, Any] = {
"model_id": args.model_id,
"backbone_loader": args.backbone_loader,
"backbone_metadata": core.backbone_metadata(model),
"mode": args.mode,
"stream_representation": (
"single_causal_prefix"
if args.mode == "causal-prefix"
else (
"single_open_media_turn_with_2s_payload_commits"
if args.mode == "open-stream-replay"
else (
"serialized_cached_complete_blocks_then_overlapping_"
"global_history_current_frame_recent_recap"
)
)
),
"stream_representation": {
"causal-prefix": "single_causal_prefix",
"open-stream-replay": (
"native_2s_completed_blocks_plus_temporary_partial_tail"
),
"completed-block-hold": (
"native_2s_completed_blocks_forward_held_to_5hz"
),
"micro-chunk-200ms": "non_native_200ms_payload_commits",
"overlap-working-memory": (
"serialized_cached_complete_blocks_then_overlapping_"
"global_history_current_frame_recent_recap"
),
}[args.mode],
"config": asdict(config),
"requested_memory_seconds": args.memory_seconds,
"persistent_thinker_kv": persistent_memory,
@@ -1514,6 +1533,45 @@ def main(
if args.mode == "open-stream-replay"
else None
),
"chunk_ablation_policy": (
{
"kind": "native_completed_block_hold",
"native_chunk_seconds": config.native_block_seconds,
"output_step_seconds": config.step_seconds,
"tail_policy": "discard_incomplete_tail_and_hold_last_state",
"initial_state": "zeros_until_first_complete_native_block",
"fresh_at_every_5hz_step": False,
"memory": (
"persistent_full_extraction_thinker_kv"
if persistent_memory
else "rolling_native_completed_blocks"
),
"backbone": "frozen",
}
if args.mode == "completed-block-hold"
else (
{
"kind": "non_native_micro_chunk_stress_test",
"micro_chunk_seconds": config.step_seconds,
"native_chunk_seconds": config.native_block_seconds,
"tmrope_positions_per_micro_step": (
core._tmrope_positions_per_step(
model, config.step_seconds
)
),
"tail_policy": "commit_each_5hz_interval_as_a_payload",
"fresh_at_every_5hz_step": True,
"memory": (
"persistent_full_extraction_thinker_kv"
if persistent_memory
else "rolling_micro_payloads_on_native_boundaries"
),
"backbone": "frozen",
}
if args.mode == "micro-chunk-200ms"
else None
)
),
"overlap_working_memory_policy": (
{
"total_presentation_seconds": config.memory_seconds,
@@ -1556,6 +1614,7 @@ def main(
"times": result.times,
"states": result.states,
"verification": result.verification,
"extraction_diagnostics": result.diagnostics,
}
if args.head_checkpoint is not None:
head, head_loss_name = core._load_head(