Add overlap working-memory ablation and LoRA epoch2 record

streaming_turn_model.py

@@ -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."""