优化流程

cadfs_to_cdsl/dataset.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import hashlib, json
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any
+
+MODALITIES = {
+    "featurescript": ("featurescript_rp", ".txt"),
+    "step": ("step_abc", ".step"),
+    "stl": ("stl_abc", ".stl"),
+    "image": ("multiview_images_abc", ".png"),
+    "annotation": ("text_annotations", ".txt"),
+}
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk)
+    return digest.hexdigest()
+
+
+@dataclass
+class Sample:
+    sample_id: str
+    files: dict[str, str] = field(default_factory=dict)
+    hashes: dict[str, str] = field(default_factory=dict)
+    metadata: dict[str, Any] = field(default_factory=dict)
+    diagnostics: list[dict[str, Any]] = field(default_factory=list)
+
+    def as_dict(self) -> dict[str, Any]: return asdict(self)
+
+
+def _files(root: Path, directory: str, suffix: str) -> dict[str, Path]:
+    base = root / directory; result: dict[str, Path] = {}
+    if not base.exists(): return result
+    for path in base.rglob(f"*{suffix}"):
+        if path.is_file() and path.stem.isdigit(): result[path.stem] = path
+    return result
+
+
+def _jsonl_records(path: Path) -> list[dict[str, Any]]:
+    if not path.exists(): return []
+    records = []
+    for index, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
+        try:
+            value = json.loads(line); value["_line"] = index; records.append(value)
+        except json.JSONDecodeError:
+            records.append({"_line": index, "_invalid": True})
+    return records
+
+
+def _assistant_script(record: dict[str, Any]) -> str | None:
+    for message in record.get("messages") or []:
+        if message.get("role") == "assistant" and isinstance(message.get("content"), str) and "FeatureScript" in message["content"]:
+            return message["content"]
+    return None
+
+
+def scan_dataset(root: Path, *, include_hashes: bool = True) -> list[Sample]:
+    if not root.is_dir(): raise FileNotFoundError(f"CADFS input directory does not exist: {root}")
+    by_modality = {name: _files(root, directory, suffix) for name, (directory, suffix) in MODALITIES.items()}
+    ids = sorted(set().union(*(set(values) for values in by_modality.values())))
+    text_records = _jsonl_records(root / "CADFS_text_test.jsonl")
+    image_records = _jsonl_records(root / "CADFS_image_test.jsonl")
+    text_by_hash = {hashlib.sha256(script.encode()).hexdigest(): record for record in text_records if (script := _assistant_script(record))}
+    samples: list[Sample] = []
+    for index, sample_id in enumerate(ids):
+        sample = Sample(sample_id)
+        for name, values in by_modality.items():
+            path = values.get(sample_id)
+            if path is None: sample.diagnostics.append({"code": "missing_modality", "modality": name}); continue
+            sample.files[name] = str(path)
+            if include_hashes: sample.hashes[name] = sha256(path)
+        fs_path = by_modality["featurescript"].get(sample_id)
+        if fs_path:
+            script_hash = hashlib.sha256(fs_path.read_text(encoding="utf-8").encode()).hexdigest()
+            record = text_by_hash.get(script_hash)
+            if record:
+                sample.metadata["text_jsonl"] = {"line": record.get("_line"), "cad_file_id": record.get("cad_file_id")}
+            elif index < len(text_records):
+                fallback = text_records[index]; sample.metadata["text_jsonl"] = {"line": fallback.get("_line"), "cad_file_id": fallback.get("cad_file_id"), "alignment": "fallback"}
+                sample.diagnostics.append({"code": "alignment_fallback", "modality": "text_jsonl"})
+        if index < len(image_records):
+            record = image_records[index]
+            sample.metadata["image_jsonl"] = {"line": record.get("_line"), "cad_file_id": record.get("cad_file_id")}
+        samples.append(sample)
+    return samples