yennik16/text-to-stl-parameter-extractor

Text-to-STL parameter extractor (LoRA fine-tune of Qwen2.5-1.5B-Instruct)

Stage 2 of a text-to-STL pipeline. Given the part type (from the classifier) and a normalized description, it writes the part's dimensions as compact JSON in one standard schema. The JSON is validated by rules and then built in CadQuery.

Model

• Base model: Qwen/Qwen2.5-1.5B-Instruct, with a LoRA adapter (rank 16, alpha 32, dropout 0.05) on all attention and MLP projections (q,k,v,o,gate,up,down_proj). This repo holds only the adapter.
• Prompt: Part type: <type>\nDescription: <normalized text>\nJSON:\n, answered with one JSON object.
• Output schema (decimal inches; hole_count integer, screw_size text):

Part type	Keys
pipe_gasket	od_in, id_in, thickness_in, hole_count
washer	screw_size, od_in, id_in, thickness_in
bracket	length_a_in, length_b_in, height_in, thickness_in
rect_container	length_a_in, length_b_in, height_in, thickness_in
cyl_container	id_in, height_in, thickness_in

Training

• Data: 823 rows: 323 real training descriptions and 500 augmented ones (values swapped between training rows, with typos and unit wording added). Validation (91) and test (87) sets are real, grouped by part size so test sizes are unseen.
• Objective: causal language-modelling loss on the JSON answer tokens only (the prompt is masked), maximum length 256 tokens.
• Optimiser: AdamW, learning rate 0.0002, 5 epochs, effective batch 8 (micro-batches of 2 with gradient accumulation), 5% warm-up then linear decay, gradient clipping 1.0, gradient checkpointing; full precision on a free Colab T4. The epoch with the lowest validation loss is kept.
• Decoding: greedy, so the same input always gives the same output.

Results (held-out test set)

A field counts as correct within a small tolerance of the label; hole_count and screw_size must match exactly. The baseline is the same base model with the adapter switched off, prompted with the schema and one example per part type.

model	JSON valid	od_in	id_in	thickness_in	hole_count	screw_size	length_a_in	length_b_in	height_in	ALL fields correct
fine-tuned (LoRA)	100.0%	100.0%	96.2%	100.0%	100.0%	100.0%	100.0%	100.0%	100.0%	97.7%
off-the-shelf, few-shot	100.0%	65.8%	69.2%	95.4%	100.0%	86.7%	100.0%	88.6%	93.9%	74.7%

Classifier then extractor on the test set: part type right 100.0%, part type and every field right 97.7%.

Usage

import sys, torch
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

path = snapshot_download("yennik16/text-to-stl-parameter-extractor"); sys.path.insert(0, path)
from normalizer import normalize_text

tok = AutoTokenizer.from_pretrained(path)
model = PeftModel.from_pretrained(AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct"), path).eval()
prompt = "Part type: washer\nDescription: " + normalize_text("washer for an M7 screw, 1/16 thick") + "\nJSON:\n"
ids = tok(prompt, return_tensors="pt", add_special_tokens=False)
out = model.generate(**ids, max_new_tokens=96, do_sample=False)
print(tok.decode(out[0, ids["input_ids"].shape[1]:], skip_special_tokens=True))

Intended use and limitations

• Only for the five part families and the schema above; it does not handle features such as holes or chamfers (that is the feature editor).
• It can produce values that are not in the description; the app flags those and checks every value before building.
• Descriptions should go through normalizer.py first, as during training.

Training and evaluation: Text_to_STL_Pipeline.ipynb (CMU Project 1).