Add overlap working-memory ablation and LoRA epoch2 record
This commit is contained in:
Xin Wang
+298
-2
@@ -1,11 +1,13 @@
|
||||
"""Omni-only causal turn-taking features at 5 Hz.
|
||||
|
||||
This module implements two inference paths for Qwen2.5-Omni-3B:
|
||||
This module implements three inference paths for Qwen2.5-Omni-3B:
|
||||
|
||||
1. ``causal-prefix``: recompute a causal AV window ending at every 200 ms.
|
||||
This is deliberately slow, but is the simplest correctness baseline.
|
||||
2. ``open-stream-replay``: cache completed AV blocks inside one
|
||||
still-open user/media turn; closure and readout tokens remain temporary.
|
||||
3. ``overlap-working-memory``: a frozen ablation that serializes cached
|
||||
complete blocks and an overlapping ``[t-2,t]`` recent recap.
|
||||
|
||||
Neither path uses an RNN. A small head maps one Thinker state to the same
|
||||
``10 future horizons x 3 classes`` output used by EgoSpeak.
|
||||
@@ -54,7 +56,11 @@ from transformers import (
|
||||
)
|
||||
from transformers.cache_utils import DynamicCache
|
||||
|
||||
Mode = Literal["causal-prefix", "open-stream-replay"]
|
||||
Mode = Literal[
|
||||
"causal-prefix",
|
||||
"open-stream-replay",
|
||||
"overlap-working-memory",
|
||||
]
|
||||
MODEL_ID = "Qwen/Qwen2.5-Omni-3B"
|
||||
_KNOWN_QWEN_WARNING_TEXT = (
|
||||
"System prompt modified, audio output may not work",
|
||||
@@ -1642,6 +1648,296 @@ class OpenStreamReplayExtractor:
|
||||
)
|
||||
|
||||
|
||||
class OverlappingWorkingMemoryExtractor:
|
||||
"""Serialize cached long memory followed by an overlapping recent recap.
|
||||
|
||||
This is a frozen-representation ablation, not the native Qwen streaming
|
||||
path. At each prediction time, the total AV presentation budget is split
|
||||
between source-aligned complete blocks and one sliding working interval.
|
||||
For an 8-second budget and a 2-second working interval, the serialized
|
||||
Thinker input is three cached 2-second blocks followed by ``[t-2,t]``.
|
||||
The final fixed block and the recap intentionally overlap.
|
||||
|
||||
Fixed and working media readers may implement different visual grids. In
|
||||
the direct-media experiment the fixed blocks remain half-open/global-grid,
|
||||
while the recap contains a frame sampled exactly at current time ``t``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
||||
processor: Qwen2_5OmniProcessor,
|
||||
config: StreamConfig,
|
||||
*,
|
||||
working_memory_seconds: float,
|
||||
fixed_grid_av: FixedGridAV,
|
||||
working_grid_av: FixedGridAV,
|
||||
) -> None:
|
||||
config.validate()
|
||||
if working_memory_seconds <= 0:
|
||||
raise ValueError("working_memory_seconds must be positive")
|
||||
if not math.isclose(
|
||||
working_memory_seconds,
|
||||
config.native_block_seconds,
|
||||
rel_tol=0.0,
|
||||
abs_tol=1e-6,
|
||||
):
|
||||
raise ValueError(
|
||||
"The first overlap-working-memory experiment requires one "
|
||||
"native block of working memory: "
|
||||
f"{working_memory_seconds:g}s vs "
|
||||
f"{config.native_block_seconds:g}s"
|
||||
)
|
||||
total_blocks = _require_integer_ratio(
|
||||
config.memory_seconds,
|
||||
config.native_block_seconds,
|
||||
"total presentation memory / native block",
|
||||
)
|
||||
if total_blocks < 1:
|
||||
raise ValueError("The presentation budget must contain one block")
|
||||
|
||||
self.model = model
|
||||
self.processor = processor
|
||||
self.config = config
|
||||
self.working_memory_seconds = float(working_memory_seconds)
|
||||
# The working recap consumes one block from the matched token budget.
|
||||
self.fixed_block_budget = total_blocks - 1
|
||||
self.fixed_extractor = OpenStreamReplayExtractor(
|
||||
model,
|
||||
processor,
|
||||
config,
|
||||
fixed_grid_av=fixed_grid_av,
|
||||
)
|
||||
self.working_extractor = OpenStreamReplayExtractor(
|
||||
model,
|
||||
processor,
|
||||
config,
|
||||
fixed_grid_av=working_grid_av,
|
||||
)
|
||||
self.query_ids = self.fixed_extractor.query_ids
|
||||
self.layout_records: list[dict[str, Any]] = []
|
||||
|
||||
def _state_from_prepared(
|
||||
self,
|
||||
*,
|
||||
video_path: str | Path | None,
|
||||
audio_path: str | Path | None,
|
||||
fixed_intervals: list[tuple[float, float]],
|
||||
fixed_blocks: list[CachedOpenMediaBlock],
|
||||
working_pieces: OpenMediaPieces,
|
||||
use_audio_in_video: bool,
|
||||
cached_fixed_payload: tuple[Any, int] | None = None,
|
||||
) -> torch.Tensor:
|
||||
positions_per_chunk = _tmrope_positions_per_chunk(self.model)
|
||||
if fixed_intervals:
|
||||
if cached_fixed_payload is None:
|
||||
(
|
||||
fixed_cache,
|
||||
_,
|
||||
payload_base,
|
||||
_,
|
||||
_,
|
||||
) = self.fixed_extractor._cache_from_intervals(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
intervals=fixed_intervals,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
prepared_blocks=fixed_blocks,
|
||||
)
|
||||
else:
|
||||
fixed_cache, payload_base = cached_fixed_payload
|
||||
branch_cache = clone_cache(fixed_cache)
|
||||
working_offset = payload_base + len(fixed_intervals) * positions_per_chunk
|
||||
include_prefix = False
|
||||
else:
|
||||
branch_cache = None
|
||||
working_offset = int(working_pieces.payload_position_ids.min().item())
|
||||
include_prefix = True
|
||||
|
||||
(
|
||||
media_cache,
|
||||
_,
|
||||
closure_position_ids,
|
||||
) = _append_open_pieces(
|
||||
self.model,
|
||||
working_pieces,
|
||||
past_key_values=branch_cache,
|
||||
rope_offset=working_offset,
|
||||
include_prefix=include_prefix,
|
||||
)
|
||||
return _read_open_state(
|
||||
self.model,
|
||||
media_cache,
|
||||
closure_ids=working_pieces.closure_ids,
|
||||
closure_position_ids=closure_position_ids,
|
||||
query_ids=self.query_ids,
|
||||
)
|
||||
|
||||
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,
|
||||
show_progress: bool = True,
|
||||
) -> ExtractionResult:
|
||||
if verify_every < 0:
|
||||
raise ValueError("verify_every must be >= 0")
|
||||
times = _time_grid(start, end, self.config.step_seconds)
|
||||
block_seconds = self.config.native_block_seconds
|
||||
completed_blocks = CompletedBlockPieceCache(
|
||||
origin=start,
|
||||
block_seconds=block_seconds,
|
||||
max_blocks=max(1, self.fixed_block_budget),
|
||||
)
|
||||
next_fixed_index = 0
|
||||
cached_interval_key: tuple[tuple[float, float], ...] | None = None
|
||||
cached_fixed_payload: tuple[Any, int] | None = None
|
||||
states: list[torch.Tensor] = []
|
||||
checks: list[dict[str, Any]] = []
|
||||
self.layout_records = []
|
||||
|
||||
with tqdm(
|
||||
total=len(times),
|
||||
desc="overlap-working-memory",
|
||||
unit="state",
|
||||
dynamic_ncols=True,
|
||||
mininterval=0.0,
|
||||
miniters=1,
|
||||
disable=not show_progress,
|
||||
) as progress:
|
||||
for step_index, current_time in enumerate(times, start=1):
|
||||
completed_count = int(
|
||||
math.floor((current_time - start) / block_seconds + 1e-7)
|
||||
)
|
||||
if self.fixed_block_budget:
|
||||
while next_fixed_index < completed_count:
|
||||
block_start = start + next_fixed_index * block_seconds
|
||||
block_end = block_start + block_seconds
|
||||
pieces = self.fixed_extractor._prepare_pieces(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
interval=(block_start, block_end),
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
)
|
||||
completed_blocks.put(
|
||||
CachedOpenMediaBlock(
|
||||
start=block_start,
|
||||
end=block_end,
|
||||
pieces=pieces,
|
||||
)
|
||||
)
|
||||
next_fixed_index += 1
|
||||
|
||||
first_fixed_index = max(
|
||||
0, completed_count - self.fixed_block_budget
|
||||
)
|
||||
fixed_intervals = [
|
||||
(
|
||||
start + index * block_seconds,
|
||||
start + (index + 1) * block_seconds,
|
||||
)
|
||||
for index in range(first_fixed_index, completed_count)
|
||||
]
|
||||
fixed_blocks = (
|
||||
completed_blocks.get_many(fixed_intervals)
|
||||
if fixed_intervals
|
||||
else []
|
||||
)
|
||||
interval_key = tuple(fixed_intervals)
|
||||
if interval_key != cached_interval_key:
|
||||
if fixed_intervals:
|
||||
(
|
||||
fixed_cache,
|
||||
_,
|
||||
payload_base,
|
||||
_,
|
||||
_,
|
||||
) = self.fixed_extractor._cache_from_intervals(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
intervals=fixed_intervals,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
prepared_blocks=fixed_blocks,
|
||||
)
|
||||
cached_fixed_payload = (fixed_cache, payload_base)
|
||||
else:
|
||||
cached_fixed_payload = None
|
||||
cached_interval_key = interval_key
|
||||
|
||||
working_start = max(start, current_time - self.working_memory_seconds)
|
||||
working_pieces = self.working_extractor._prepare_pieces(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
interval=(working_start, current_time),
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
)
|
||||
state = self._state_from_prepared(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
fixed_intervals=fixed_intervals,
|
||||
fixed_blocks=fixed_blocks,
|
||||
working_pieces=working_pieces,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
cached_fixed_payload=cached_fixed_payload,
|
||||
)
|
||||
states.append(state.cpu())
|
||||
|
||||
fixed_start = (
|
||||
fixed_intervals[0][0] if fixed_intervals else current_time
|
||||
)
|
||||
fixed_end = fixed_intervals[-1][1] if fixed_intervals else current_time
|
||||
overlap_seconds = (
|
||||
max(0.0, min(fixed_end, current_time) - max(fixed_start, working_start))
|
||||
if fixed_intervals
|
||||
else 0.0
|
||||
)
|
||||
self.layout_records.append(
|
||||
{
|
||||
"time": float(current_time),
|
||||
"fixed_start": float(fixed_start),
|
||||
"fixed_end": float(fixed_end),
|
||||
"fixed_blocks": len(fixed_intervals),
|
||||
"working_start": float(working_start),
|
||||
"working_end": float(current_time),
|
||||
"overlap_seconds": float(overlap_seconds),
|
||||
}
|
||||
)
|
||||
|
||||
if verify_every and step_index % verify_every == 0:
|
||||
exact = self._state_from_prepared(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
fixed_intervals=fixed_intervals,
|
||||
fixed_blocks=fixed_blocks,
|
||||
working_pieces=working_pieces,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
)
|
||||
checks.append(
|
||||
_compare_states(
|
||||
current_time,
|
||||
state,
|
||||
exact,
|
||||
kind="overlap_cache_replay",
|
||||
)
|
||||
)
|
||||
|
||||
progress.set_postfix_str(
|
||||
f"media={current_time:.1f}s overlap={overlap_seconds:.1f}s",
|
||||
refresh=False,
|
||||
)
|
||||
progress.update(1)
|
||||
|
||||
return ExtractionResult(
|
||||
times=torch.tensor(times, dtype=torch.float64),
|
||||
states=torch.cat(states, dim=0),
|
||||
verification=checks,
|
||||
)
|
||||
|
||||
|
||||
def _block_intervals(start: float, end: float, block_seconds: float) -> list[tuple[float, float]]:
|
||||
"""Exact AV payload intervals represented by an open-stream cache."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user