Add NeurIPS 2025 acceptance info
Browse files
README.md
CHANGED
|
@@ -17,6 +17,8 @@ library_name: transformers
|
|
| 17 |
|
| 18 |
**CAD-Coder: Text-to-CAD Generation with Chain-of-Thought and Geometric Reward**
|
| 19 |
|
|
|
|
|
|
|
| 20 |
This is the reinforcement learning (GRPO) fine-tuned model for generating CadQuery code from natural language descriptions.
|
| 21 |
|
| 22 |
## Model Description
|
|
@@ -47,58 +49,14 @@ pip install "numpy<2.0" cadquery==2.3.1 # Optional: for code execution
|
|
| 47 |
### Quick Start
|
| 48 |
|
| 49 |
```python
|
| 50 |
-
"""CAD-Coder Inference Example"""
|
| 51 |
-
import re
|
| 52 |
-
import os
|
| 53 |
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 54 |
|
| 55 |
-
|
| 56 |
-
|
| 57 |
-
|
| 58 |
-
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype="auto", device_map="auto")
|
| 59 |
-
|
| 60 |
-
# Test data (from the first record in cad_data_test_cot.json)
|
| 61 |
-
prompt = """Please create a CadQuery Python code to generate a model based on the following description. The reasoning process MUST BE enclosed within <think> </think> tags. The final CadQuery code MUST BE put in \\boxed{```python
|
| 62 |
-
code
|
| 63 |
-
```} with ONLY the executable code inside the box, nothing else.The final model is represented by r.
|
| 64 |
-
|
| 65 |
-
In the <think> section, simulate the thought process of an engineer converting text descriptions into a CAD model. Follow these steps:
|
| 66 |
-
|
| 67 |
-
1. **Description Analysis**:
|
| 68 |
-
- Break down the description into different parts or components
|
| 69 |
-
- Identify key parameters for each part (coordinate systems, Euler angles, translation vectors)
|
| 70 |
-
- Understand the spatial relationships and assembly sequence between parts
|
| 71 |
-
|
| 72 |
-
2. **Coordinate System Planning**:
|
| 73 |
-
- Determine the coordinate systems used for each part
|
| 74 |
-
- Parse how Euler angle rotations and translation vectors are applied
|
| 75 |
-
- Ensure understanding of local-to-global coordinate system transformations
|
| 76 |
-
|
| 77 |
-
3. **Sketch Construction Strategy**:
|
| 78 |
-
- Analyze how to create each 2D sketch (loops, lines, points)
|
| 79 |
-
- Determine scaling factors for each sketch
|
| 80 |
-
- Plan how to transform sketches into 3D space
|
| 81 |
-
|
| 82 |
-
4. **Extrusion Operation Planning**:
|
| 83 |
-
- Identify extrusion parameters for each part (direction, distance)
|
| 84 |
-
- Understand how to add extrusions to existing solids (new or merge)
|
| 85 |
-
- Verify dimensions after extrusion match the description
|
| 86 |
-
|
| 87 |
-
5. **Code Implementation Strategy**:
|
| 88 |
-
- Plan the sequence of CadQuery operations
|
| 89 |
-
- Determine necessary CadQuery functions and methods
|
| 90 |
-
- Consider how to organize code for clarity and readability
|
| 91 |
-
|
| 92 |
-
6. **When scaling in CadQuery:**:
|
| 93 |
-
- Directly scale the size, define a scaling factor variable, and apply the scaling factor directly to all coordinates and dimensions
|
| 94 |
-
- Don't try to use a non-existent `.scale()` method on Workplane objects
|
| 95 |
|
| 96 |
-
|
| 97 |
-
Think step by step, but only keep a minimum draft for each thinking step, with 50 words at most.
|
| 98 |
-
description:
|
| 99 |
-
Start by creating a new coordinate system with Euler angles set to zero and a translation vector also set to zero. Next, draw a two-dimensional sketch on the first face. This sketch consists of a single loop made up of four lines. The first line starts at the origin (0.0, 0.0) and ends at (0.6, 0.0). The second line starts at (0.6, 0.0) and ends at (0.6, 0.375). The third line starts at (0.6, 0.375) and ends at (0.0, 0.375). Finally, the fourth line completes the loop by starting at (0.0, 0.375) and ending at the origin (0.0, 0.0). After drawing the sketch, apply a scale factor of 0.6 to the entire sketch. To transform the scaled two-dimensional sketch into a three-dimensional model, extrude the sketch 0.075 units along the normal direction. The final dimensions of the rectangular block are a length of 0.6 units, a width of 0.375 units, and a height of 0.075 units."""
|
| 100 |
|
| 101 |
-
# Generate
|
| 102 |
text = tokenizer.apply_chat_template(
|
| 103 |
[{"role": "user", "content": prompt}],
|
| 104 |
tokenize=False,
|
|
@@ -106,8 +64,7 @@ text = tokenizer.apply_chat_template(
|
|
| 106 |
)
|
| 107 |
inputs = tokenizer([text], return_tensors="pt").to(model.device)
|
| 108 |
outputs = model.generate(**inputs, max_new_tokens=2048)
|
| 109 |
-
|
| 110 |
-
print(response)
|
| 111 |
```
|
| 112 |
|
| 113 |
## Performance
|
|
|
|
| 17 |
|
| 18 |
**CAD-Coder: Text-to-CAD Generation with Chain-of-Thought and Geometric Reward**
|
| 19 |
|
| 20 |
+
**Accepted at NeurIPS 2025 (Poster)**
|
| 21 |
+
|
| 22 |
This is the reinforcement learning (GRPO) fine-tuned model for generating CadQuery code from natural language descriptions.
|
| 23 |
|
| 24 |
## Model Description
|
|
|
|
| 49 |
### Quick Start
|
| 50 |
|
| 51 |
```python
|
|
|
|
|
|
|
|
|
|
| 52 |
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 53 |
|
| 54 |
+
model_name = "gudo7208/CAD-Coder"
|
| 55 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
| 56 |
+
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 57 |
|
| 58 |
+
prompt = "Create a cylinder with radius 10mm and height 20mm, with a central hole of radius 5mm."
|
|
|
|
|
|
|
|
|
|
| 59 |
|
|
|
|
| 60 |
text = tokenizer.apply_chat_template(
|
| 61 |
[{"role": "user", "content": prompt}],
|
| 62 |
tokenize=False,
|
|
|
|
| 64 |
)
|
| 65 |
inputs = tokenizer([text], return_tensors="pt").to(model.device)
|
| 66 |
outputs = model.generate(**inputs, max_new_tokens=2048)
|
| 67 |
+
print(tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True))
|
|
|
|
| 68 |
```
|
| 69 |
|
| 70 |
## Performance
|