Add dense-cache evaluation for Omni LoRA

README.md

@@ -36,6 +36,7 @@ classification head.
 | `TurnPredictionHead` | Trained and evaluated | Map `[2048]` to `[10, 3]` logits |
 | `train_turn_head.py` | Verified on EasyCom | Train the frozen-Omni anticipation head with the official EasyCom split |
 | `infer_turn_head.py` | Verified on EasyCom | Save EgoSpeak-format predictions and reproduce per-horizon AP/mAP |
+| `eval_omni_turn_lora.py` | Implemented, remote validation pending | Evaluate Thinker LoRA from dense AV caches with the same AP/mAP protocol |
 
 The full Thinker input path, embedded audio, direct video decoding,
 FlashAttention-2, 5 Hz streaming extraction, cache replay, target alignment,
@@ -1030,6 +1031,47 @@ for example, change 8 GPUs with batch 2 and accumulation 1 to 4 GPUs with batch
 post-hoc continuation because a one-epoch cosine schedule has already reached
 zero.
 
+### Evaluate a Thinker-LoRA checkpoint
+
+Use `eval_omni_turn_lora.py` for the formal EasyCom Table 2 evaluation.  It
+loads the trained Thinker LoRA and `10 x 3` turn head, reconstructs every causal
+state from dense AV caches, and reports both the EgoSpeak full-timeline and
+strict causal-valid protocols.  Sparse training caches are rejected because
+they do not contain predictions at every 200 ms timestamp.
+
+The following command evaluates the official EasyCom test sessions on three
+GPUs, with one complete session assigned to each rank:
+
+```bash
+CHECKPOINT=/data1/yfzx/wangx/Code/icassp2027-experiment/checkpoints/omni_turn_lora_easycom_sparse_8s_ddp8_epoch1
+
+CUDA_VISIBLE_DEVICES=0,1,2 \
+OMP_NUM_THREADS=1 \
+MKL_NUM_THREADS=1 \
+OPENBLAS_NUM_THREADS=1 \
+PYTORCH_ALLOC_CONF=expandable_segments:True \
+TORCH_NCCL_ASYNC_ERROR_HANDLING=1 \
+torchrun --standalone --nproc_per_node=3 eval_omni_turn_lora.py \
+  --model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_3 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_2 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_1 \
+  --adapter-checkpoint "${CHECKPOINT}/adapter" \
+  --head-checkpoint "${CHECKPOINT}/turn_head.pt" \
+  --output-dir results/omni_turn_lora_easycom_sparse_8s_ddp8_epoch1 \
+  --batch-size 2 \
+  --block-cache-size 32 \
+  --dtype bfloat16 \
+  --flash-attention-2 \
+  --ddp-timeout-minutes 60
+```
+
+The evaluator saves per-session probabilities and valid masks under `pred/`
+and `valid/`, and writes the aggregate metrics to `metrics.json`.  Add
+`--overwrite` only when intentionally replacing an existing result directory.
+Treat Sessions 1, 2, and 3 as the official test set: inspect their metrics only
+after the training recipe and hyperparameters have been fixed.
+
 For a speed-only alternative, `train_omni_turn_lora.py` supports a two-GPU
 producer-consumer benchmark.  The producer GPU runs the frozen AV towers and
 places current-tail embeddings in a bounded CPU queue while the consumer GPU

eval_omni_turn_lora.py

@@ -0,0 +1,365 @@
+#!/usr/bin/env python3
+"""Evaluate Thinker LoRA + turn head from dense Omni AV caches.
+
+The script accepts one dense cache directory per EasyCom session.  With
+``torchrun``, complete sessions are assigned round-robin to ranks; rank zero
+then loads the saved per-session predictions and reports the exact
+EgoSpeak-compatible and strict-causal metrics used by ``infer_turn_head.py``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+from pathlib import Path
+from typing import Any, Sequence
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.utils.data import DataLoader, Dataset
+from tqdm.auto import tqdm
+
+from easycom_head import CLASS_NAMES, NUM_CLASSES, NUM_HORIZONS
+from infer_turn_head import (
+    _align_to_egospeak_timeline,
+    _evaluate,
+    _load_head,
+    _print_metrics,
+)
+from omni_cached_av import (
+    SCHEMA_VERSION,
+    CachedBlockStore,
+    collate_cached_sequences,
+    forward_last_state,
+    load_json,
+    reconstruct_open_stream_sequence,
+    torch_load,
+)
+from train_omni_turn_lora import (
+    DistributedRun,
+    _distributed_barrier,
+    _initialize_distributed,
+    _load_base_model,
+)
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model-id", required=True)
+    parser.add_argument(
+        "--cache-dir",
+        type=Path,
+        action="append",
+        required=True,
+        help="Dense cache directory from cache_omni_av_pieces.py; repeat per session",
+    )
+    parser.add_argument("--adapter-checkpoint", type=Path, required=True)
+    parser.add_argument("--head-checkpoint", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument("--batch-size", type=int, default=2)
+    parser.add_argument("--block-cache-size", type=int, default=32)
+    parser.add_argument(
+        "--dtype",
+        choices=("bfloat16", "float16"),
+        default="bfloat16",
+    )
+    parser.add_argument(
+        "--device",
+        default="cuda:0",
+        help="Single-GPU device; torchrun overrides this with LOCAL_RANK",
+    )
+    parser.add_argument("--flash-attention-2", action="store_true")
+    parser.add_argument("--ddp-timeout-minutes", type=float, default=60.0)
+    parser.add_argument("--overwrite", action="store_true")
+    parser.add_argument("--no-progress", action="store_true")
+    return parser.parse_args(argv)
+
+
+class DenseCacheSession(Dataset[dict[str, Any]]):
+    """Reconstruct every 5-Hz causal anchor for one cached session."""
+
+    def __init__(self, cache_dir: Path, *, block_cache_size: int) -> None:
+        self.cache_dir = cache_dir.expanduser().resolve()
+        self.manifest = load_json(self.cache_dir / "manifest.json")
+        if int(self.manifest.get("schema_version", -1)) != SCHEMA_VERSION:
+            raise ValueError(
+                f"Unsupported cache schema in {self.cache_dir}: "
+                f"{self.manifest.get('schema_version')}"
+            )
+        policy = self.manifest.get("anchor_selection", {}).get("policy")
+        if policy != "dense":
+            raise ValueError(
+                f"Evaluation requires --anchor-policy dense, got {policy!r}: "
+                f"{self.cache_dir}"
+            )
+        self.session = str(self.manifest["session"])
+        self.static = torch_load(self.cache_dir / self.manifest["static_path"])
+        self.store = CachedBlockStore(
+            self.cache_dir,
+            max_blocks=block_cache_size,
+        )
+        self.anchors = list(self.manifest["anchors"])
+        target_path = Path(self.manifest["target_path"])
+        targets = np.load(target_path, allow_pickle=False)
+        if targets.ndim != 2 or targets.shape[1] != NUM_CLASSES:
+            raise ValueError(
+                f"Expected target [T,{NUM_CLASSES}], got {targets.shape}: "
+                f"{target_path}"
+            )
+        if len(targets) != int(self.manifest["target_steps"]):
+            raise ValueError(
+                f"Target length changed for {self.session}: {len(targets)} vs "
+                f"{self.manifest['target_steps']}"
+            )
+        self.targets = np.asarray(targets, dtype=np.float32)
+        indices = np.asarray(
+            [int(anchor["target_index"]) for anchor in self.anchors],
+            dtype=np.int64,
+        )
+        expected = np.arange(len(self.targets), dtype=np.int64)
+        if not np.array_equal(indices, expected):
+            raise ValueError(
+                f"{self.session} dense cache does not contain exactly one ordered "
+                "anchor for every target timestep"
+            )
+        self.label_indices = torch.from_numpy(indices.copy()).long()
+
+    def __len__(self) -> int:
+        return len(self.anchors)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        anchor = self.anchors[index]
+        embeddings, positions = reconstruct_open_stream_sequence(
+            store=self.store,
+            static=self.static,
+            anchor=anchor,
+        )
+        return {
+            "inputs_embeds": embeddings,
+            "position_ids": positions,
+            # Reuse the verified padding/readout collator. Evaluation ignores
+            # this zero-length placeholder after stacking it.
+            "targets": torch.empty(0, dtype=torch.float32),
+            "step_index": int(anchor["target_index"]),
+        }
+
+
+@torch.inference_mode()
+def _predict_session(
+    *,
+    dataset: DenseCacheSession,
+    language_model: torch.nn.Module,
+    head: torch.nn.Module,
+    device: torch.device,
+    model_dtype: torch.dtype,
+    batch_size: int,
+    show_progress: bool,
+) -> tuple[np.ndarray, np.ndarray]:
+    loader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=0,
+        pin_memory=True,
+        collate_fn=collate_cached_sequences,
+    )
+    current_probabilities = np.zeros(
+        (len(dataset.targets), NUM_HORIZONS, NUM_CLASSES),
+        dtype=np.float32,
+    )
+    seen = np.zeros(len(dataset.targets), dtype=bool)
+    progress = tqdm(
+        loader,
+        total=len(loader),
+        desc=f"evaluate {dataset.session}",
+        unit="batch",
+        disable=not show_progress,
+    )
+    for batch in progress:
+        inputs_embeds = batch["inputs_embeds"].to(
+            device=device,
+            dtype=model_dtype,
+            non_blocking=True,
+        )
+        position_ids = batch["position_ids"].to(device=device, non_blocking=True)
+        attention_mask = batch["attention_mask"].to(device=device, non_blocking=True)
+        readout_indices = batch["readout_indices"].to(device=device, non_blocking=True)
+        states = forward_last_state(
+            language_model,
+            inputs_embeds=inputs_embeds,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            readout_indices=readout_indices,
+        )
+        probabilities = head(states.float()).softmax(dim=-1).cpu().numpy()
+        step_indices = batch["step_indices"].numpy()
+        current_probabilities[step_indices] = probabilities
+        seen[step_indices] = True
+    if not bool(seen.all()):
+        missing = np.flatnonzero(~seen)
+        raise RuntimeError(
+            f"{dataset.session} missed {len(missing)} dense anchors; "
+            f"first={missing[:10].tolist()}"
+        )
+    return _align_to_egospeak_timeline(
+        current_probabilities,
+        dataset.label_indices,
+        len(dataset.targets),
+    )
+
+
+def _load_session_specs(
+    cache_dirs: Sequence[Path], block_cache_size: int
+) -> list[DenseCacheSession]:
+    sessions = [
+        DenseCacheSession(path, block_cache_size=block_cache_size)
+        for path in cache_dirs
+    ]
+    names = [session.session for session in sessions]
+    if len(names) != len(set(names)):
+        raise ValueError(f"Duplicate cache sessions: {names}")
+    contracts = {
+        (
+            item.manifest["model_id"],
+            item.manifest["memory_seconds"],
+            item.manifest["video_fps"],
+            item.manifest["positions_per_chunk"],
+            item.manifest["hidden_size"],
+        )
+        for item in sessions
+    }
+    if len(contracts) != 1:
+        raise ValueError("Evaluation caches do not share one model/timing contract")
+    return sessions
+
+
+def _run(args: argparse.Namespace, run: DistributedRun) -> None:
+    if args.batch_size <= 0 or args.block_cache_size <= 0:
+        raise ValueError("Batch and block-cache sizes must be positive")
+    if run.world_size > len(args.cache_dir):
+        raise ValueError(
+            f"Requested {run.world_size} ranks for only {len(args.cache_dir)} sessions"
+        )
+    output_dir = args.output_dir.expanduser().resolve()
+    metrics_path = output_dir / "metrics.json"
+    if metrics_path.exists() and not args.overwrite:
+        raise FileExistsError(
+            f"Evaluation output exists: {metrics_path}; pass --overwrite"
+        )
+    if run.is_main:
+        (output_dir / "pred").mkdir(parents=True, exist_ok=True)
+        (output_dir / "valid").mkdir(parents=True, exist_ok=True)
+    _distributed_barrier(run)
+
+    sessions = _load_session_specs(args.cache_dir, args.block_cache_size)
+    assigned = [
+        session
+        for index, session in enumerate(sessions)
+        if index % run.world_size == run.rank
+    ]
+    print(
+        f"[rank {run.rank}/{run.world_size}] sessions="
+        + ",".join(item.session for item in assigned),
+        flush=True,
+    )
+
+    base_model, processor = _load_base_model(args, run)
+    base_model.requires_grad_(False)
+    try:
+        from peft import PeftModel
+    except ImportError as exc:
+        raise RuntimeError("PEFT is required to evaluate the LoRA adapter") from exc
+    adapter_path = args.adapter_checkpoint.expanduser().resolve()
+    if not (adapter_path / "adapter_config.json").is_file():
+        raise FileNotFoundError(f"Missing adapter config: {adapter_path}")
+    language_model = PeftModel.from_pretrained(
+        base_model.model,
+        str(adapter_path),
+        is_trainable=False,
+    ).eval()
+    language_model.config.use_cache = False
+    model_dtype = next(language_model.parameters()).dtype
+    head, head_payload = _load_head(args.head_checkpoint, run.device)
+    hidden_size = int(head.hidden_size)
+    for session in sessions:
+        if int(session.manifest["hidden_size"]) != hidden_size:
+            raise ValueError(
+                f"{session.session} hidden size {session.manifest['hidden_size']} "
+                f"!= head {hidden_size}"
+            )
+    del base_model
+    del processor
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    for session in assigned:
+        predictions, valid = _predict_session(
+            dataset=session,
+            language_model=language_model,
+            head=head,
+            device=run.device,
+            model_dtype=model_dtype,
+            batch_size=args.batch_size,
+            show_progress=not args.no_progress,
+        )
+        np.save(output_dir / "pred" / f"{session.session}.npy", predictions)
+        np.save(output_dir / "valid" / f"{session.session}.npy", valid)
+        print(
+            f"[rank {run.rank}] {session.session}: pred={predictions.shape}, "
+            f"valid={valid.sum(axis=0).tolist()}",
+            flush=True,
+        )
+
+    _distributed_barrier(run)
+    if run.is_main:
+        session_predictions = {
+            item.session: np.load(
+                output_dir / "pred" / f"{item.session}.npy",
+                allow_pickle=False,
+            )
+            for item in sessions
+        }
+        session_valid = {
+            item.session: np.load(
+                output_dir / "valid" / f"{item.session}.npy",
+                allow_pickle=False,
+            )
+            for item in sessions
+        }
+        session_targets = {item.session: item.targets for item in sessions}
+        metrics = _evaluate(session_predictions, session_valid, session_targets)
+        metrics.update(
+            {
+                "sessions": [item.session for item in sessions],
+                "classes": list(CLASS_NAMES),
+                "prediction_layout": "[target_time, class, future_horizon]",
+                "model_id": args.model_id,
+                "adapter_checkpoint": str(adapter_path),
+                "head_checkpoint": str(args.head_checkpoint.expanduser().resolve()),
+                "head_training_step": head_payload.get("training_step"),
+                "cache_dirs": [str(item.cache_dir) for item in sessions],
+                "world_size": run.world_size,
+                "batch_size_per_gpu": args.batch_size,
+            }
+        )
+        metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
+        _print_metrics(metrics)
+        print(f"\nmetrics: {metrics_path}")
+        print(f"predictions: {output_dir / 'pred'}")
+    _distributed_barrier(run)
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    args = parse_args(argv)
+    run = _initialize_distributed(args)
+    try:
+        _run(args, run)
+    finally:
+        if run.enabled and dist.is_initialized():
+            dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()