138 lines
6.5 KiB
Python
138 lines
6.5 KiB
Python
"""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()
|