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.
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
4REPO = "aitraineracc/intent-classification-multilabel"
5tokenizer = AutoTokenizer.from_pretrained(REPO)
6model = AutoModelForSequenceClassification.from_pretrained(REPO)
7model.eval()
8
9thresholds = model.config.thresholds # {'web_search': 0.35, 'diagram_enabled': 0.6}
10
11def predict(text: str) -> dict:
12 inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
13 with torch.no_grad():
14 logits = model(**inputs).logits
15 probs = torch.sigmoid(logits).squeeze().tolist()
16 return {
17 "web_search" : int(probs[0] >= thresholds["web_search"]),
18 "diagram_enabled" : int(probs[1] >= thresholds["diagram_enabled"]),
19 "probs" : {"web_search": round(probs[0], 4),
20 "diagram_enabled": round(probs[1], 4)},
21 }
22
23print(predict("What is the weather today in Singapore?"))
24# {'web_search': 1, 'diagram_enabled': 0, 'probs': ...}
25
26print(predict("Draw me a diagram of how TCP/IP works"))
27# {'web_search': 0, 'diagram_enabled': 1, 'probs': ...}
