Jerry
96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import random
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
OUTPUT = ROOT / "parser" / "output"
|
|
HERE = Path(__file__).resolve().parent
|
|
DATASET = Path("/Users/jerry/linkhand/cad/SLDPRT-NET/sub_datset/sub_datset")
|
|
TARGET_COUNT = 50
|
|
|
|
|
|
def main() -> None:
|
|
rows: list[dict] = []
|
|
for path in sorted(OUTPUT.glob("*.json")):
|
|
payload = json.loads(path.read_text())
|
|
stem = Path(payload["provenance"]["source_path"]).stem
|
|
image = DATASET / "image" / f"{stem}.png"
|
|
description = DATASET / "Des_Text" / f"{stem}.txt"
|
|
reference = DATASET / "model_step" / f"{stem}.step"
|
|
if image.is_file() and description.is_file() and reference.is_file():
|
|
rows.append(
|
|
{
|
|
"id": stem,
|
|
"family": payload["design_ir"]["part_family"],
|
|
"case_json": str(path.resolve()),
|
|
"image": str(image),
|
|
"description": str(description),
|
|
"reference_step": str(reference),
|
|
}
|
|
)
|
|
|
|
by_family: dict[str, list[dict]] = defaultdict(list)
|
|
for row in rows:
|
|
by_family[row["family"]].append(row)
|
|
rng = random.Random(20260723)
|
|
selected: list[dict] = []
|
|
total = len(rows)
|
|
for family, family_rows in sorted(by_family.items()):
|
|
rng.shuffle(family_rows)
|
|
quota = max(1, round(TARGET_COUNT * len(family_rows) / total))
|
|
selected.extend(family_rows[: min(quota, len(family_rows))])
|
|
selected_ids = {row["id"] for row in selected}
|
|
if len(selected) > TARGET_COUNT:
|
|
selected = selected[:TARGET_COUNT]
|
|
elif len(selected) < TARGET_COUNT:
|
|
remainder = [row for row in rows if row["id"] not in selected_ids]
|
|
rng.shuffle(remainder)
|
|
selected.extend(remainder[: TARGET_COUNT - len(selected)])
|
|
selected = sorted(selected, key=lambda row: row["id"])
|
|
selected_ids = {row["id"] for row in selected}
|
|
|
|
train_dir = HERE / "train_cases"
|
|
train_dir.mkdir(parents=True, exist_ok=True)
|
|
for path in train_dir.glob("*.json"):
|
|
path.unlink()
|
|
for row in rows:
|
|
if row["id"] in selected_ids:
|
|
continue
|
|
source = Path(row["case_json"])
|
|
(train_dir / source.name).symlink_to(source)
|
|
|
|
for row in selected:
|
|
for variant in ("baseline", "experience"):
|
|
output_dir = HERE / "generated" / variant / row["id"]
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
row[f"{variant}_step"] = str(output_dir / f"{row['id']}.step")
|
|
row[f"{variant}_log"] = str(output_dir / "codex.jsonl")
|
|
row[f"{variant}_message"] = str(output_dir / "final.txt")
|
|
|
|
manifest = {
|
|
"schema_version": "1.0",
|
|
"selection_seed": 20260723,
|
|
"holdout_count": len(selected),
|
|
"train_count": len(rows) - len(selected),
|
|
"cases": selected,
|
|
}
|
|
(HERE / "manifest.json").write_text(
|
|
json.dumps(manifest, indent=2, ensure_ascii=False) + "\n"
|
|
)
|
|
print(json.dumps({
|
|
"holdout_count": len(selected),
|
|
"train_count": len(rows) - len(selected),
|
|
"family_counts": {
|
|
family: sum(row["family"] == family for row in selected)
|
|
for family in sorted({row["family"] for row in selected})
|
|
},
|
|
}, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|