Spaces:
Sleeping
Sleeping
File size: 3,152 Bytes
896a559 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """
Trains one multi-label classifier per attribute type on the labeled dataset.
Approach
--------
- Features: character n-gram TF-IDF (char_wb, 2-5) over the raw description.
Character n-grams generalize better than word n-grams on a 61-row dataset
because they share sub-word signal across morphological variants
("embroidered"/"embroidery", "beaded"/"beading", "off-shoulder"/
"off shoulder") without needing a stemmer.
- Label space: for each attribute type, the label space is exactly the
canonical vocabulary in lexicon.py (colors use COLOR_VOCAB). This keeps the
model closed-vocabulary and directly comparable to the rules baseline.
- Classifier: OneVsRestClassifier(LogisticRegression) per attribute type,
i.e. independent binary classifiers per label, combined into a multi-label
prediction per attribute (a description can have 0, 1, or many labels for
a given attribute, e.g. two colors).
- Given the small dataset, model artifacts are ensembled at inference time
with the deterministic rules extractor (see predict.py) rather than relied
on in isolation.
"""
import json
import joblib
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.multiclass import OneVsRestClassifier
from sklearn.preprocessing import MultiLabelBinarizer
from lexicon import ATTRIBUTE_LEXICON, COLOR_VOCAB
ATTR_TYPES = ["silhouette", "fabric", "neckline", "sleeve", "length",
"embellishment", "category", "color"]
def label_space(attr):
if attr == "color":
return COLOR_VOCAB
return list(ATTRIBUTE_LEXICON[attr].keys())
def load_dataset(path="../data/dataset.json"):
with open(path) as f:
return json.load(f)
def train_all(dataset_path="../data/dataset.json", out_path="model.joblib"):
data = load_dataset(dataset_path)
texts = [d["text"] for d in data]
vectorizer = TfidfVectorizer(analyzer="char_wb", ngram_range=(2, 5), min_df=1)
X = vectorizer.fit_transform(texts)
models = {}
binarizers = {}
for attr in ATTR_TYPES:
classes = label_space(attr)
mlb = MultiLabelBinarizer(classes=classes)
y_labels = [d["labels"][attr] for d in data]
Y = mlb.fit_transform(y_labels)
# Skip labels with zero positive examples (OneVsRest can't fit them);
# in practice every label in our lexicon appears >=1 time in the 61
# rows, but guard anyway for robustness on future dataset edits.
if Y.sum() == 0:
models[attr] = None
binarizers[attr] = mlb
continue
clf = OneVsRestClassifier(
LogisticRegression(max_iter=1000, class_weight="balanced")
)
clf.fit(X, Y)
models[attr] = clf
binarizers[attr] = mlb
artifact = {
"vectorizer": vectorizer,
"models": models,
"binarizers": binarizers,
"attr_types": ATTR_TYPES,
}
joblib.dump(artifact, out_path)
print(f"Trained models for attrs: {ATTR_TYPES}")
print(f"Saved artifact to {out_path}")
return artifact
if __name__ == "__main__":
train_all()
|