icassp2027-experiment/prepare_encoder_only_features.py

"""Pack EasyCom/Ego4D local Qwen encoder vectors as standalone head inputs on CPU.

No model, raw media, Thinker states or symlink index is loaded. All fixed-split
train/test sidecars are validated before creating a fresh output directory.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import tempfile

import numpy as np
import torch
from tqdm.auto import tqdm

from easycom_head import (
    DATASET_NAMES, OFFICIAL_TEST_SESSIONS, OFFICIAL_TRAIN_SESSIONS,
    load_ego4d_manifest_split, resolve_feature_path,
)
from local_turn_features import MODALITIES, load_local_features


DEFAULT_DATA_ROOTS = {dataset: Path("/data1/yfzx/wangx/Datasets") / dataset for dataset in DATASET_NAMES}
DEFAULT_DATA_ROOT = DEFAULT_DATA_ROOTS["ego4d"]  # Preserve the original Ego4D default/API.
PACKING = "native-vectors-plus-validity-bits"


def encoder_sample_split(dataset, manifest=None):
    """Fixed benchmark membership; EasyCom order matches the original B1 run."""
    if dataset == "easycom":
        if manifest is not None:
            raise ValueError("EasyCom uses fixed sessions, not an Ego4D manifest")
        number = lambda name: int(name.split("_")[-1])
        # B1 explicitly passes sessions 4..12. Preserve that concatenation
        # order too, so seeded shuffling sees the same sample ordering.
        train = sorted(OFFICIAL_TRAIN_SESSIONS, key=number)
        test = list(OFFICIAL_TEST_SESSIONS)
        return sorted(train + test, key=number), train, test
    if dataset != "ego4d" or manifest is None:
        raise ValueError("Ego4D requires its EgoSpeak manifest")
    return tuple(load_ego4d_manifest_split(manifest, part) for part in ("all", "train", "test"))


def sample_unit(dataset):
    return "session" if dataset == "easycom" else "clip"


def encoder_only_mode(modalities):
    if modalities not in MODALITIES:
        raise ValueError(f"Unknown encoder modalities: {modalities}")
    return f"encoder-only-{modalities}"


def parse_args(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset", choices=DATASET_NAMES, default="ego4d")
    parser.add_argument("--data-root", type=Path)
    parser.add_argument("--manifest", type=Path)
    parser.add_argument("--local-feature-root", type=Path)
    parser.add_argument("--target-root", type=Path)
    parser.add_argument("--output-root", type=Path)
    parser.add_argument("--modalities", choices=tuple(MODALITIES), default="av")
    parser.add_argument("--dry-run", action="store_true", help="Validate all inputs; no output files or directory")
    parser.add_argument("--no-progress", action="store_true")
    args = parser.parse_args(argv)
    args.data_root = args.data_root or DEFAULT_DATA_ROOTS[args.dataset]
    if args.dataset == "ego4d":
        args.manifest = args.manifest or args.data_root / "ego4d_egospeak.tsv"
    elif args.manifest is not None:
        parser.error("--manifest is only supported for Ego4D; EasyCom uses all 12 fixed sessions")
    args.local_feature_root = args.local_feature_root or args.data_root / "omni_features/local_qwen_1s"
    args.target_root = args.target_root or args.data_root / f"{args.dataset}_target_frame"
    args.output_root = args.output_root or args.data_root / f"omni_features/encoder_only_{args.modalities}_1s"
    for key in ("data_root", "manifest", "local_feature_root", "target_root", "output_root"):
        value = getattr(args, key)
        if value is not None:
            setattr(args, key, value.expanduser().resolve())
    if args.output_root == args.local_feature_root:
        parser.error("Encoder-only output must be separate from local source features")
    return args


def pack_sample(path, target_path, sample_id, modalities, expected_contract=None, *, dataset="ego4d"):
    """Read only selected encoder vectors; even a sidecar's reference path is unused."""
    source = torch.load(path, map_location="cpu", weights_only=False)
    if not isinstance(source, dict):
        raise ValueError(f"Invalid local sidecar: {path}")
    anchor = MODALITIES[modalities][0]
    reference = {
        "states": source.get(anchor), "times": source.get("times"),
        "label_indices": source.get("label_indices"), "start": 0.,
    }
    states, contract = load_local_features(
        path, reference, sample_id=sample_id, dataset=dataset, modalities=modalities,
        expected_contract=expected_contract,
    )
    target = np.load(target_path, allow_pickle=False, mmap_mode="r")
    if target.shape != (len(states), 3):
        raise ValueError(f"{sample_id}: feature/target mismatch: {tuple(states.shape)} vs {target.shape}")
    if len(states) <= 10:
        raise ValueError(f"{sample_id}: no examples with all ten future horizons")
    states = states.to(torch.bfloat16)
    if not torch.isfinite(states).all():
        raise ValueError(f"{sample_id}: encoder vectors overflow bfloat16 storage")
    return {
        "format": "qwen-encoder-only-v1", "dataset": dataset, "sample_id": sample_id,
        "states": states, "times": source["times"], "label_indices": source["label_indices"],
        "mode": encoder_only_mode(modalities), "start": 0., "end": len(states) / 5,
        "num_steps": len(states), "source_local_features": str(path),
        "config": {
            "input_source": "qwen-encoders-only", "uses_thinker_states": False,
            "packing": PACKING, "encoder_contract": contract,
        },
    }, contract


def main(argv=None):
    args = parse_args(argv)
    names, train, test = encoder_sample_split(args.dataset, args.manifest)
    unit = sample_unit(args.dataset)
    if args.output_root.exists() and not args.dry_run:
        raise FileExistsError(f"Refusing to overwrite {args.output_root}; use a fresh output directory")
    paths = {name: resolve_feature_path(args.local_feature_root, name, dataset=args.dataset) for name in names}
    targets = {name: args.target_root / f"{name}.npy" for name in names}
    for path in targets.values():
        if not path.is_file():
            raise FileNotFoundError(path)
    contract, total, width = None, 0, None
    # Two passes keep memory bounded and prevent bad/partial test inputs from
    # leaving an apparently ready export. No full-dataset tensors are retained.
    for name in tqdm(names, desc="encoder-only preflight", unit=unit, disable=args.no_progress):
        packed, contract = pack_sample(paths[name], targets[name], name, args.modalities, contract, dataset=args.dataset)
        total += len(packed["states"])
        width = packed["states"].shape[-1]
        del packed
    print(f"Validated {len(names)} {unit}s / {total} states; {encoder_only_mode(args.modalities)}; "
          f"input width={width}; no Thinker features", flush=True)
    if args.dry_run:
        return
    args.output_root.mkdir(parents=True, exist_ok=False)
    rows = {}
    for name in tqdm(names, desc="pack encoder-only", unit=unit, disable=args.no_progress):
        packed, _ = pack_sample(paths[name], targets[name], name, args.modalities, contract, dataset=args.dataset)
        destination = args.output_root / f"{name}.pt"
        with tempfile.NamedTemporaryFile(dir=args.output_root, prefix=f".{name}.", suffix=".tmp", delete=False) as handle:
            temporary = Path(handle.name)
        try:
            torch.save(packed, temporary)
            os.replace(temporary, destination)
        finally:
            temporary.unlink(missing_ok=True)
        rows[name] = len(packed["states"])
        del packed
    (args.output_root / "export_manifest.json").write_text(json.dumps({
        "dataset": args.dataset, "manifest": str(args.manifest) if args.manifest else None,
        "train_samples": train, "test_samples": test, "mode": encoder_only_mode(args.modalities),
        "local_feature_root": str(args.local_feature_root), "target_root": str(args.target_root),
        "encoder_contract": contract, "input_width": width, "num_states": total, f"{unit}s": rows,
    }, indent=2) + "\n")
    # Written last. The runner also validates every train/test input instead of
    # trusting this marker alone. Compatible with the original inline command.
    (args.output_root / "complete.txt").write_text(f"{len(names)} {unit}s\n")
    print(f"Complete: {args.output_root}", flush=True)


if __name__ == "__main__":
    main()