Add encoder-only Ego4D comparison workflow

prepare_encoder_only_features.py

@@ -0,0 +1,137 @@
+"""Pack Ego4D local Qwen encoder vectors as standalone head inputs on CPU.
+
+No model, raw media, Thinker states or symlink index is loaded. All manifest
+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 load_ego4d_manifest_split, resolve_feature_path
+from local_turn_features import MODALITIES, load_local_features
+
+
+DEFAULT_DATA_ROOT = Path("/data1/yfzx/wangx/Datasets/ego4d")
+PACKING = "native-vectors-plus-validity-bits"
+
+
+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("--data-root", type=Path, default=DEFAULT_DATA_ROOT)
+    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.manifest = args.manifest or args.data_root / "ego4d_egospeak.tsv"
+    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 / "ego4d_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"):
+        setattr(args, key, getattr(args, key).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):
+    """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="ego4d", 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": "ego4d", "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 = load_ego4d_manifest_split(args.manifest, "all")
+    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="ego4d") 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="clip", disable=args.no_progress):
+        packed, contract = pack_sample(paths[name], targets[name], name, args.modalities, contract)
+        total += len(packed["states"])
+        width = packed["states"].shape[-1]
+        del packed
+    print(f"Validated {len(names)} clips / {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="clip", disable=args.no_progress):
+        packed, _ = pack_sample(paths[name], targets[name], name, args.modalities, contract)
+        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({
+        "manifest": str(args.manifest), "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, "clips": 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)} clips\n")
+    print(f"Complete: {args.output_root}", flush=True)
+
+
+if __name__ == "__main__":
+    main()