"""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()