Auto-sync: 2026-06-25 23:02:08 (part 2)

scripts/merge_task_datasets.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+"""
+Merge multiple task-specific h=16 datasets into single training-ready dataset.
+Creates unified manifest and symlinks to shards.
+"""
+import json
+import shutil
+from pathlib import Path
+import argparse
+
+def merge_datasets(task_dirs, output_dir):
+    """Merge multiple CIL datasets into one."""
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Collect all manifests
+    all_groups = []
+    all_group_index = []
+    all_record_index = []
+    total_records = 0
+    shard_counter = 0
+
+    for task_dir in task_dirs:
+        task_path = Path(task_dir)
+        manifest_path = task_path / "manifest.json"
+
+        if not manifest_path.exists():
+            print(f"⚠️  Skipping {task_path.name} (no manifest)")
+            continue
+
+        print(f"Processing {task_path.name}...")
+        manifest = json.loads(manifest_path.read_text())
+
+        # Copy shard files (rename to .jsonl for compatibility)
+        for shard_info in manifest.get("shards", []):
+            shard_file = task_path / shard_info["path"]
+            if shard_file.exists():
+                dest = output_dir / f"shard_{shard_counter:04d}.jsonl"
+                shutil.copy2(shard_file, dest)
+                shard_counter += 1
+
+        # Fallback: glob for .shard files if shards key missing
+        if not manifest.get("shards"):
+            for shard_file in task_path.glob("*.shard"):
+                dest = output_dir / f"shard_{shard_counter:04d}.jsonl"
+                shutil.copy2(shard_file, dest)
+                shard_counter += 1
+
+        # Read and merge group index
+        group_index_path = task_path / manifest.get("group_index_path", "group_index.jsonl")
+        if group_index_path.exists():
+            with open(group_index_path) as f:
+                for line in f:
+                    if line.strip():
+                        entry = json.loads(line)
+                        all_groups.append(entry["group_id"])
+                        all_group_index.append(entry)
+
+        # Read and merge record index
+        record_index_path = task_path / manifest.get("record_index_path", "record_index.jsonl")
+        if record_index_path.exists():
+            with open(record_index_path) as f:
+                for line in f:
+                    if line.strip():
+                        all_record_index.append(json.loads(line))
+
+        # Accumulate metadata
+        total_records += manifest["record_count"]
+
+        print(f"  ✅ {manifest['group_count']} groups, {manifest['record_count']} records")
+
+    # Write merged indexes
+    group_index_out = output_dir / "group_index.jsonl"
+    with open(group_index_out, 'w') as f:
+        for entry in all_group_index:
+            f.write(json.dumps(entry) + '\n')
+
+    record_index_out = output_dir / "record_index.jsonl"
+    with open(record_index_out, 'w') as f:
+        for entry in all_record_index:
+            f.write(json.dumps(entry) + '\n')
+
+    # Create unified manifest
+    shard_list = [{"path": f"shard_{i:04d}.jsonl"} for i in range(shard_counter)]
+
+    merged_manifest = {
+        "backend": "maniskill",
+        "created_at": Path(task_dirs[0]).joinpath("manifest.json").stat().st_mtime if task_dirs else 0,
+        "dataset_name": "cil_h16_merged",
+        "format": "dovla_cil",
+        "group_count": len(all_groups),
+        "group_index_path": "group_index.jsonl",
+        "index_format": "jsonl",
+        "k": 16,
+        "num_groups": len(all_groups),
+        "num_records": total_records,
+        "record_count": total_records,
+        "record_index_path": "record_index.jsonl",
+        "schema_version": "0.1",
+        "shard_count": shard_counter,
+        "shard_format": "jsonl",
+        "shards": shard_list,
+        "horizon": 16,
+        "tasks": [Path(d).name for d in task_dirs if (Path(d) / "manifest.json").exists()]
+    }
+
+    manifest_out = output_dir / "manifest.json"
+    manifest_out.write_text(json.dumps(merged_manifest, indent=2))
+
+    print("\n" + "="*60)
+    print(f"✅ Merged dataset created: {output_dir}")
+    print(f"   Total groups: {merged_manifest['group_count']}")
+    print(f"   Total records: {merged_manifest['record_count']}")
+    print(f"   Shards: {merged_manifest['shard_count']}")
+    print(f"   Tasks: {', '.join(merged_manifest['tasks'])}")
+    print("="*60)
+
+    return merged_manifest
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", required=True, help="Parent dir containing task subdirs")
+    parser.add_argument("--output", required=True, help="Output merged dataset dir")
+    args = parser.parse_args()
+
+    input_path = Path(args.input)
+    task_dirs = [d for d in input_path.iterdir() if d.is_dir() and (d / "manifest.json").exists()]
+
+    if not task_dirs:
+        print(f"❌ No valid task directories found in {input_path}")
+        return 1
+
+    print(f"Found {len(task_dirs)} tasks to merge:")
+    for d in task_dirs:
+        print(f"  - {d.name}")
+    print()
+
+    merge_datasets(task_dirs, args.output)
+    return 0
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(main())