yennik16/text-to-stl-part-classifier

Text-to-STL part-type classifier (trained from scratch)

Stage 1 of a text-to-STL pipeline. It reads a plain-English description of a simple mechanical part and decides which of five part families it is: pipe_gasket, washer, bracket, rect_container, cyl_container. The predicted type then selects the schema the parameter extractor fills in.

Model

A small bag-of-words neural classifier written in PyTorch, with all weights trained from random initialisation (nothing pretrained).

• Input features: the description is normalized (normalizer.py: fractions, number words and units become decimal inches), lowercased, every number is replaced by a <num> token, request-framing words ("I want it to be", "make", "please"...) are removed, and the remaining words plus adjacent word pairs are looked up in a 1,201-token vocabulary (tokens seen in at least 2 training texts).
• Architecture: EmbeddingBag (mean of 48-dimensional learned embeddings) -> dropout 0.3 -> 64-unit ReLU layer -> dropout 0.3 -> 5 outputs (softmax).
• Output: a part type and a probability for each type. The app treats under 85% as uncertain and shows the runner-up types.

Training

• Data: 1,323 rows: 323 real training descriptions, 500 augmented descriptions and 500 synthetic phrase variations (see the dataset). Validation and test sets are real data only.
• Split: grouped by part size, so no size appears in both training and test.
• Optimiser: AdamW, learning rate 5e-3, weight decay 1e-3, batch size 32, 80 epochs; the epoch with the best validation accuracy is kept.
• Why request-framing words are removed: in the raw data they identified who wrote each spreadsheet sheet (for example "it", "be" and "thats" appear in ~40% of gasket descriptions and almost no container descriptions), so the model learned the writer instead of the part.

Results

Metric	Value
Validation accuracy (91 real rows)	100.0%
Test accuracy (87 real rows, unseen sizes)	100.0%
Casual stress prompts (60, unusual wording, never trained on)	100.0%

Usage

import sys
from huggingface_hub import snapshot_download
path = snapshot_download("yennik16/text-to-stl-part-classifier"); sys.path.insert(0, path)
from model import PartTypeClassifier
from normalizer import normalize_text

clf = PartTypeClassifier(path)
print(clf.predict(normalize_text('a box 6 by 4 by 3 inches with 1/8" walls')))   # ('rect_container', {...})

Intended use and limitations

• Built for five part families only; anything else is forced into the closest of the five.
• Trained on English descriptions written by a small team plus synthetic variations, so unusual wording can still confuse it.
• It decides the part type only; dimensions come from the extractor, and every value is checked by rule-based validation before any CAD is built.

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