cc72b78715f3499bb4e0abbe41648d06

Preparing CAD Data for HOOPS Embeddings Training

This notebook is Part 1 of 2 in the custom HOOPS Embeddings workflow. It encodes a raw CAD dataset into the tensor format that the EmbeddingFlowModel expects, producing a reusable .dataset store and an .infoset metadata file.

The workflow is split into two notebooks because data preparation typically runs once, while training runs many times as you tune hyperparameters and iterate toward convergence:

``demo_HOOPS_EMBEDDINGS_DataPrep.ipynb`` (this notebook) — gather and encode CAD files.
`demo_HOOPS_EMBEDDINGS_training.ipynb <./demo_HOOPS_EMBEDDINGS_training.ipynb>`__ — split the encoded dataset, then train and save your model.

Once the encoded dataset is written to disk, you can re-run the training notebook against it as often as you like without repeating this preprocessing step.

Set the HOOPS AI license

A valid license must be set before any HOOPS AI functionality is used.

[1]:

import hoops_ai
import os
import sys

license_key = os.environ.get("HOOPS_AI_LICENSE")
if not license_key:
    sys.exit("HOOPS_AI_LICENSE environment variable is required.")

hoops_ai.set_license(license_key, validate=True)

------------------------------------------------------------
HOOPS AI
------------------------------------------------------------
  Platform      : Windows 11
  Architecture  : AMD64
  Python        : 3.12.10
------------------------------------------------------------
  Core          : hoops-ai             1.1.0  (build: ed23c844 2026-06-12T14:14:13Z)
  CAD Access    : hoops-exchange       26.2.0  (build: 1e11169 2026-06-12T10:38:16Z)
  Conversion    : hoops-converter      26.1.1  (build: 00dc9f6 2026-06-12T10:22:46Z)
  Insights      : hoops-web-viewer     26.1.1  (build: d30058f 2026-06-12T10:22:25Z)
------------------------------------------------------------
======================================================================
[OK] HOOPS AI License: Valid
======================================================================

Dataset

This notebook uses the TMCAD dataset (Truly Mechanical CAD Dataset), introduced in:

Zou, Q., & Zhu, L. (2025). Bringing attention to CAD: Boundary representation learning via transformer. Computer-Aided Design, 103940. Elsevier.

Download links:

TMCAD v2 (2025-11-02, recommended): https://pan.zju.edu.cn/share/218d10a88e8c18f5b96e94a7e0
Dataset documentation: https://github.com/Qiang-Zou/BRT/blob/main/DATASET.md

TMCAD v2 is a cleaned and verified collection of ~10,000 CAD models in STEP (.stp) format across 10 mechanical part categories (bearing, bolt-screw, bracket, flange, gear, nut, shaft, coupling, pulley, spring). Released under the GPL-3.0 license.

Step 1: Prepare your CAD dataset

Training EmbeddingFlowModel requires a preprocessing pipeline that gathers CAD files and encodes them into the tensor format the model expects. The pipeline runs in parallel across workers, producing a .dataset store and an .infoset metadata file that the DatasetLoader (in the training notebook) will consume.

[2]:

import pathlib
import sys

path_tmcad_v2_dataset = os.environ.get("PATH_DATASET_TMCAD")
if not path_tmcad_v2_dataset:
    sys.exit("PATH_DATASET_TMCAD environment variable is required.")

datasources_dir = pathlib.Path(path_tmcad_v2_dataset)

Parallel encoding pipeline

The flow processes each CAD file through two tasks defined in scripts/cad_tasks_embeddings.py:

``gather_cad_files`` — walks datasources_dir and emits one record per file.
``encode_data_for_ml_training`` — converts each file to the B-rep tensor representation the embedding model expects (face types, edge connectivity, etc.).

max_workers controls parallelism. Set it to the number of physical CPU cores available (or slightly below to leave headroom). Results are written to flows_outputdir as a .dataset store and an .infoset metadata file.

[3]:

# We define our tasks in a separate file for multiprocessing compatibility.
from scripts.cad_tasks_embeddings import EmbeddingModel, flows_outputdir, get_flow_name, gather_cad_files, encode_data_for_ml_training

# Create and run the data flow
flow_name = get_flow_name()

data_prep = hoops_ai.create_flow(
    name=flow_name,
    tasks=[gather_cad_files, encode_data_for_ml_training],
    max_workers=12,
    flows_outputdir=str(flows_outputdir),
    ml_task="Private HOOPS Embedings Model",
    export_visualization=False
)

# Run the flow to process all files
flow_output, output_dict, flow_file = data_prep.process(inputs={'cad_datasources': [str(datasources_dir)]}, clean_ouput_dir=True)

c:\Users\LuisSalazar.LY-LS-LEGION\Documents\hoops-ai-packages-test\.venv\Lib\site-packages\pytorch_lightning\utilities\parsing.py:136: pytorch_lightning save_hyperparameters: frame introspection failed. Nuitka-compiled frames have empty locals, which breaks PL's collect_init_args. Pass an explicit dict to save_hyperparameters() for correct hparams under Nuitka.
Cleaning up existing flow directory: c:\Users\LuisSalazar.LY-LS-LEGION\Documents\repos\HOOPS-AI-tutorials\embeddings_pipeline\out\flows\HOOPS_Embedding_Training
Removing all previous outputs for flow 'HOOPS_Embedding_Training' to avoid build conflicts.

Total number of items with errors: 290 (2.66%)
Corrupted items are listed in 'c:\Users\LuisSalazar.LY-LS-LEGION\Documents\repos\HOOPS-AI-tutorials\embeddings_pipeline\out\flows\HOOPS_Embedding_Training\error_summary.json'.
  [218x] Data extraction error: division by zero
  [9x] Data extraction error: "Data key 'graph' does not exist in Zarr storage."
  [22x] Timeout (CUMULATIVE): item exceeded time_limit_s=120.0s (total_elapsed=120.1s). Worker restarted.
  [26x] Timeout (CUMULATIVE): item exceeded time_limit_s=120.0s (total_elapsed=120.2s). Worker restarted.
  [9x] Timeout (CUMULATIVE): item exceeded time_limit_s=120.0s (total_elapsed=120.0s). Worker restarted.
  [1x] Timeout (CUMULATIVE): item exceeded time_limit_s=120.0s (total_elapsed=120.8s). Worker restarted.
  [3x] Data extraction error: Exception: Failed
  [1x] Data extraction error: list index out of range
  [1x] Timeout (CUMULATIVE): item exceeded time_limit_s=120.0s (total_elapsed=120.3s). Worker restarted.

[DatasetInfo] Warning: 170 .json files have no matching .data
[DatasetInfo] Skipped 170 files without valid IDs
[DatasetMerger] Using streaming merge into temporary directory store for large dataset...

Sequential Task end=====================

Data preparation complete

The encoded dataset is now written to disk under flows_outputdir:

flows/{flow_name}/{flow_name}.dataset — the encoded B-rep tensor store.
flows/{flow_name}/{flow_name}.infoset — the metadata file used for splitting.

Because these artifacts are reusable, you only need to run this notebook once per dataset.

Next step — train your model: open `demo_HOOPS_EMBEDDINGS_training.ipynb <./demo_HOOPS_EMBEDDINGS_training.ipynb>`__. It loads the .dataset and .infoset produced here, splits them into train/validation/test sets, and runs the training loop. You can re-run the training notebook as many times as needed — tuning hyperparameters and iterating toward convergence — without repeating this preprocessing step.