Add shared-AV multiwindow Qwen feature extraction
This commit is contained in:
Xin Wang
@@ -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
|
||||
Reference in New Issue
Block a user