CoolFace
Modelpublic

aitraineracc/intent-classification-multilabel

sourceHugging Faceapache-2.0updated 6mo agoView on Hugging Face
0likes9downloads
Model Card

Intent Classification — Multi-Label (web_search / diagram_enabled)

Fine-tuned from `Falconsai/intent_classification` (DistilBERT-base-uncased, apache-2.0) for multi-label binary intent classification. The original 15-class head was replaced with a 2-label sigmoid head trained with BCEWithLogitsLoss.

Labels

IndexLabelMeaning
0web_searchQuery requires a live web search
1diagram_enabledQuery benefits from a diagram / visualisation

Training details

SettingValue
Base modelFalconsai/intent_classification (DistilBERT)
Problem typemulti_label_classification
Frozen layersembeddings + transformer.layer[0-3]
Trainable params~7M / 67M total (~10%)
Classifier dropout0.3
Learning rate5e-6
Early stoppingpatience=3 on eval_loss
Threshold floor0.30
Max sequence length128
Split80 / 10 / 10 (train / val / test)
Seed42

Decision thresholds

LabelThreshold
web_search0.35
diagram_enabled0.6
Thresholds are stored in thresholds.json and embedded in config.json under config.thresholds — no separate download needed.

Usage

python
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

REPO = "aitraineracc/intent-classification-multilabel"
tokenizer = AutoTokenizer.from_pretrained(REPO)
model     = AutoModelForSequenceClassification.from_pretrained(REPO)
model.eval()

thresholds = model.config.thresholds  # {'web_search': 0.35, 'diagram_enabled': 0.6}

def predict(text: str) -> dict:
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
    with torch.no_grad():
        logits = model(**inputs).logits
    probs = torch.sigmoid(logits).squeeze().tolist()
    return {
        "web_search"      : int(probs[0] >= thresholds["web_search"]),
        "diagram_enabled" : int(probs[1] >= thresholds["diagram_enabled"]),
        "probs"           : {"web_search": round(probs[0], 4),
                              "diagram_enabled": round(probs[1], 4)},
    }

print(predict("What is the weather today in Singapore?"))
# {'web_search': 1, 'diagram_enabled': 0, 'probs': ...}

print(predict("Draw me a diagram of how TCP/IP works"))
# {'web_search': 0, 'diagram_enabled': 1, 'probs': ...}