Views
No views yet
distilbert-base-uncased1import torch
2import json
3import re
4import numpy as np
5from transformers import AutoTokenizer, AutoConfig, AutoModel
6
7# Repo name on Hugging Face
8model_name = "openchlsystem/CHS_tz_classifier_distilbert"
9# Load tokenizer
10tokenizer = AutoTokenizer.from_pretrained(model_name)
11
12# Load config
13config = AutoConfig.from_pretrained(model_name)
14
15# We re-define the distilbert Custom model class used in training
16import torch.nn as nn
17from transformers import DistilBertModel, DistilBertPreTrainedModel
18
19class MultiTaskDistilBert(DistilBertPreTrainedModel):
20 def __init__(self, config, num_main, num_sub, num_interv, num_priority):
21 super().__init__(config)
22 self.distilbert = DistilBertModel(config)
23 self.pre_classifier = nn.Linear(config.dim, config.dim)
24 self.classifier_main = nn.Linear(config.dim, num_main)
25 self.classifier_sub = nn.Linear(config.dim, num_sub)
26 self.classifier_interv = nn.Linear(config.dim, num_interv)
27 self.classifier_priority = nn.Linear(config.dim, num_priority)
28 self.dropout = nn.Dropout(config.dropout)
29 self.init_weights()
30
31 def forward(self, input_ids=None, attention_mask=None, **kwargs):
32 distilbert_output = self.distilbert(
33 input_ids=input_ids,
34 attention_mask=attention_mask,
35 return_dict=True
36 )
37 hidden_state = distilbert_output.last_hidden_state
38 pooled_output = hidden_state[:, 0]
39 pooled_output = self.pre_classifier(pooled_output)
40 pooled_output = nn.ReLU()(pooled_output)
41 pooled_output = self.dropout(pooled_output)
42
43 logits_main = self.classifier_main(pooled_output)
44 logits_sub = self.classifier_sub(pooled_output)
45 logits_interv = self.classifier_interv(pooled_output)
46 logits_priority = self.classifier_priority(pooled_output)
47
48 return logits_main, logits_sub, logits_interv, logits_priority
49
50# Downloading the class labels for mapping
51from huggingface_hub import hf_hub_download
52
53main_categories = json.load(open(hf_hub_download(model_name, "main_categories.json")))
54sub_categories = json.load(open(hf_hub_download(model_name, "sub_categories.json")))
55interventions = json.load(open(hf_hub_download(model_name, "interventions.json")))
56priorities = json.load(open(hf_hub_download(model_name, "priorities.json")))
57
58model = MultiTaskDistilBert.from_pretrained(
59 model_name,
60 num_main=len(main_categories),
61 num_sub=len(sub_categories),
62 num_interv=len(interventions),
63 num_priority=len(priorities)
64)
65
66device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
67model = model.to(device)
68
69
70# inference
71def classify_multitask_case(narrative: str):
72 """
73 Classifies a given narrative text into multiple categories using a multitask model.
74
75 Args:
76 narrative (str): The input text to be classified.
77
78 Returns:
79 dict: A dictionary containing the predicted labels for each task:
80 - "main_category": Predicted main category label.
81 - "sub_category": Predicted sub-category label.
82 - "intervention": Predicted intervention label.
83 - "priority": Predicted priority label.
84
85 Notes:
86 - The function preprocesses the input text by lowercasing and removing non-alphanumeric characters.
87 - It uses a tokenizer and a multitask classification model to generate predictions.
88 - Probabilities for each class are computed using softmax, and the label with the highest probability is selected for each task.
89 """
90 text = narrative.lower().strip()
91 text = re.sub(r'[^a-z0-9\s]', '', text)
92
93 inputs = tokenizer(
94 text,
95 truncation=True,
96 padding="max_length",
97 max_length=256,
98 return_tensors="pt"
99 ).to(device)
100
101 with torch.no_grad():
102 logits_main, logits_sub, logits_interv, logits_priority = model(**inputs)
103
104 # Convert to probabilities (softmax per task)
105 probs_main = torch.softmax(logits_main, dim=1).cpu().numpy()[0]
106 probs_sub = torch.softmax(logits_sub, dim=1).cpu().numpy()[0]
107 probs_interv = torch.softmax(logits_interv, dim=1).cpu().numpy()[0]
108 probs_priority = torch.softmax(logits_priority, dim=1).cpu().numpy()[0]
109
110 # Get predicted labels (argmax)
111 pred_main = int(np.argmax(probs_main))
112 pred_sub = int(np.argmax(probs_sub))
113 pred_interv = int(np.argmax(probs_interv))
114 pred_priority = int(np.argmax(probs_priority))
115
116 return {
117 "main_category": {
118 main_categories[pred_main],
119 # "probabilities": dict(zip(main_categories, probs_main.tolist()))
120 },
121 "sub_category": {
122 sub_categories[pred_sub],
123 # "probabilities": dict(zip(sub_categories, probs_sub.tolist()))
124 },
125 "intervention": {
126 interventions[pred_interv],
127 # "probabilities": dict(zip(interventions, probs_interv.tolist()))
128 },
129 "priority": {
130 priorities[pred_priority],
131 # "probabilities": dict(zip(priorities, probs_priority.tolist()))
132 },
133 }
134
135# test
136narrative= " Hello, I've been trying to find help for my son Ken. He's only ten years old and... he's been going through a terrible time at school. There's this boy, James Kibet, who keeps harassing him. It started with name-calling and teasing, but it's escalated to physical violence. I don't know what to do, Sarah. I can't bear to see my child suffer like this. I'm truly sorry to hear about your situation, Mary. It's never easy when our children are in pain. Can you tell me more about the school Ken attends and its location? We might be able to reach out to them for help. The school is Unknown Landmark, Muthambi Ward, Unknown Subcounty, Tharaka-Nithi County. I... I didn't want to burden anyone with this problem. But it seems like things are only getting worse. I don't know who else to turn to. Don't worry, Mary. You've taken the right step by reaching out to us today. We can help guide you through this difficult time. Let me first assure you that your call will be kept confidential. Now, I need to gather more information about the incidents. Can you describe any specific instances where Ken has been hurt or bullied? Oh, there have been so many times... one instance stands out though. About a week ago, James punched Ken during recess. He was left with a bloody lip and a black eye. The school officials were informed but they didn't seem to take any action against James.That sounds very serious, Mary. I'm afraid we may need to escalate this matter to the authorities if the school doesn't take appropriate action. We can provide you with resources and guidance on how to report this case to the police or child welfare services. Would that be alright? Yes, please. I'm willing to do whatever it takes to protect Ken. I just want him to feel safe again. Thank you for your help, Sarah"
137
138print(classify_multitask_case(narrative))
139{'main_category': {'Advice and Counselling'}, 'sub_category': {'School Related Issues'}, 'intervention': {'Counselling'}, 'priority': {2}}| Task | Micro F1-Score | Notes |
|---|---|---|
| Main Topic | High (e.g., >0.90) | Robust performance on primary categorization. |
| Sub-Topic | Good (e.g., >0.80) | Performance varies; higher on frequent sub-topics. |
| Intervention | High (e.g., >0.85) | Accurate prediction of recommended actions. |
| Priority | High (e.g., >0.88) | Critical for effective routing and escalation. |
| Task | Labels | Purpose |
|---|---|---|
| Sub-Category | Adoption, Albinism, Balanced Diet, Birth Registration, Breast Feeding, etc. | Identifies detailed case topics |
| Priority/Urgency | Low, Medium, High | Flags urgency for escalation |
| Main Category | Advice & Counselling, Child Custody, Disability, GBV, VANE, Nutrition, Information | High-level categorization |
| Intervention | Referred, Counselling, Signposting, Awareness/Information | Suggests next action |
1@software{chs_distilbert_multilabel,
2 author = {Bitz AI Team},
3 title = {DistilBERT Multi-Label Classifier for Call Center Transcripts},
4 year = {2025},
5 publisher = {Hugging Face},
6 url = {https://huggingface.co/openchlsystem/CHS_tz_classifier_distilbert}
7}