File size: 5,100 Bytes
b56cdca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | #!/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())
|