DistilBERT Multi-Label Classification Model

This repository hosts a fine-tuned DistilBERT-base-uncased model for multi-label classification of call center transcripts.
It is designed for real-time case categorization, urgency detection, and intervention recommendations in helpline and child protection contexts.

Model Details

Developed by: BITZ IT CONSULTING
Funded by [optional]:
Shared by: Internal ML team
Model type: Multi-label text classifier (4 tasks, multi-task heads)
Language(s): English
License:
Finetuned from: distilbert-base-uncased

Sources

Repository: This Hugging Face Repo
Paper [optional]: N/A
Demo [optional]: Coming soon

Uses

Direct Use

Real-time classification of call transcripts
Case categorization for reporting and dashboards

Downstream Use

Fine-tuning on other multi-label customer service or support datasets
Integration in larger NLP pipelines (chatbots, QA systems, ASR + NLP pipelines)

Out-of-Scope Use

Not intended for medical, legal, or financial advice without human oversight
Not reliable for domains outside call center/customer service transcripts

Bias, Risks, and Limitations

The dataset includes anonymized transcripts but may reflect biases in annotation.
Part of the Dataset was synthetically generated.
Performance may degrade on languages other than English/Swahili.
Urgency detection relies on limited data → risk of false negatives for critical cases.

Recommendations

Use confidence thresholds wisely (see section below).
Keep a human-in-the-loop for critical interventions.
Retrain periodically with fresh data to reduce drift.

Label Files

label mappings are alongside models, and are used to map the prediction and classification logits against class labels defined in the json files.

How to Get Started with the Model

import torch
import json
import re
import numpy as np
from transformers import AutoTokenizer, AutoConfig, AutoModel

# Repo name on Hugging Face
model_name = "openchlsystem/CHS_tz_classifier_distilbert"
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Load config
config = AutoConfig.from_pretrained(model_name)

#  We re-define the distilbert Custom model class used in training
import torch.nn as nn
from transformers import DistilBertModel, DistilBertPreTrainedModel

class MultiTaskDistilBert(DistilBertPreTrainedModel):
    def __init__(self, config, num_main, num_sub, num_interv, num_priority):
        super().__init__(config)
        self.distilbert = DistilBertModel(config)
        self.pre_classifier = nn.Linear(config.dim, config.dim)
        self.classifier_main = nn.Linear(config.dim, num_main)
        self.classifier_sub = nn.Linear(config.dim, num_sub)
        self.classifier_interv = nn.Linear(config.dim, num_interv)
        self.classifier_priority = nn.Linear(config.dim, num_priority)
        self.dropout = nn.Dropout(config.dropout)
        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, **kwargs):
        distilbert_output = self.distilbert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True
        )
        hidden_state = distilbert_output.last_hidden_state
        pooled_output = hidden_state[:, 0]
        pooled_output = self.pre_classifier(pooled_output)
        pooled_output = nn.ReLU()(pooled_output)
        pooled_output = self.dropout(pooled_output)

        logits_main = self.classifier_main(pooled_output)
        logits_sub = self.classifier_sub(pooled_output)
        logits_interv = self.classifier_interv(pooled_output)
        logits_priority = self.classifier_priority(pooled_output)

        return logits_main, logits_sub, logits_interv, logits_priority
    
# Downloading the class labels  for mapping 
from huggingface_hub import hf_hub_download

main_categories = json.load(open(hf_hub_download(model_name, "main_categories.json")))
sub_categories = json.load(open(hf_hub_download(model_name, "sub_categories.json")))
interventions = json.load(open(hf_hub_download(model_name, "interventions.json")))
priorities = json.load(open(hf_hub_download(model_name, "priorities.json")))

model = MultiTaskDistilBert.from_pretrained(
    model_name,
    num_main=len(main_categories),
    num_sub=len(sub_categories),
    num_interv=len(interventions),
    num_priority=len(priorities)
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)


# inference
def classify_multitask_case(narrative: str):
    """
      Classifies a given narrative text into multiple categories using a multitask model.

      Args:
        narrative (str): The input text to be classified.

      Returns:
        dict: A dictionary containing the predicted labels for each task:
          - "main_category": Predicted main category label.
          - "sub_category": Predicted sub-category label.
          - "intervention": Predicted intervention label.
          - "priority": Predicted priority label.

      Notes:
        - The function preprocesses the input text by lowercasing and removing non-alphanumeric characters.
        - It uses a tokenizer and a multitask classification model to generate predictions.
        - Probabilities for each class are computed using softmax, and the label with the highest probability is selected for each task.
      """
    text = narrative.lower().strip()
    text = re.sub(r'[^a-z0-9\s]', '', text)

    inputs = tokenizer(
        text,
        truncation=True,
        padding="max_length",
        max_length=256,
        return_tensors="pt"
    ).to(device)

    with torch.no_grad():
        logits_main, logits_sub, logits_interv, logits_priority = model(**inputs)

    # Convert to probabilities (softmax per task)
    probs_main = torch.softmax(logits_main, dim=1).cpu().numpy()[0]
    probs_sub = torch.softmax(logits_sub, dim=1).cpu().numpy()[0]
    probs_interv = torch.softmax(logits_interv, dim=1).cpu().numpy()[0]
    probs_priority = torch.softmax(logits_priority, dim=1).cpu().numpy()[0]

    # Get predicted labels (argmax)
    pred_main = int(np.argmax(probs_main))
    pred_sub = int(np.argmax(probs_sub))
    pred_interv = int(np.argmax(probs_interv))
    pred_priority = int(np.argmax(probs_priority))

    return {
        "main_category": {
            main_categories[pred_main],
            # "probabilities": dict(zip(main_categories, probs_main.tolist()))
        },
        "sub_category": {
            sub_categories[pred_sub],
            # "probabilities": dict(zip(sub_categories, probs_sub.tolist()))
        },
        "intervention": {
            interventions[pred_interv],
            # "probabilities": dict(zip(interventions, probs_interv.tolist()))
        },
        "priority": {
            priorities[pred_priority],
            # "probabilities": dict(zip(priorities, probs_priority.tolist()))
        },
    }

# test
narrative= " 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"

print(classify_multitask_case(narrative))

Expected output `{'main_category': {'Advice and Counselling'}, 'sub_category': {'School Related Issues'}, 'intervention': {'Counselling'}, 'priority': {2}}`

Training Details

Training Data

The model was fine-tuned on a proprietary dataset of over 11,000 call transcripts (1,000+ anonymized real calls and 10,000+ synthetic calls). The data was annotated for four multi-label tasks:

Main Topic: 7 classes (e.g., GBV, Nutrition) Sub-Topic: 50+ classes (e.g., Albinism, Birth Registration) Intervention: 4 classes (e.g., Referred, Counselling) Priority: 3 classes (Low, Medium, High)

The dataset was stratified and balanced to minimize bias and maximize generalization.

Training Procedure

Preprocessing

Text was tokenized using the DistilBERT tokenizer (distilbert-base-uncased) with a maximum sequence length of 512 tokens. Standard normalization (lowercasing) was applied.

Training Hyperparameters

Training regime: fp16 mixed precision Learning Rate: 2e-5 Train Batch Size: 16 Epochs: 12 Optimizer: AdamW Weight Decay: 0.01 Loss Function: Combined Cross-Entropy Loss (Multi-Task Learning) Evaluation

Testing Data, Factors & Metrics

Testing Data

Model evaluation was performed on a held-out validation set (10% of the total data), stratified by sub-category to ensure representative distribution.

Metrics

Primary Metric: Micro F1-Score (aggregated across all labels)

Secondary Metrics: Precision, Recall, and F1-score for each individual class and task.

Overall Metric: eval_avg_acc (average accuracy across all four tasks) was used for model selection.

Results

The model achieves strong performance across all tasks, with a high Micro F1-Score. Performance is consistently high on majority classes. Continuous learning cycles are used to improve performance on minority classes.

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.

Training Procedure

Loss Function: Binary Cross-Entropy with logits (multi-label)
Optimizer: AdamW
Scheduler: Warm-up + linear decay
Epochs: 12
Batch Size: 16
Learning Rate: 2e-5

Classification Tasks

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

Integration Guide: NLP Pipeline

ASR (Whisper) → transcribes call audio
DistilBERT Classifier → assigns categories, urgency, interventions
Case Management → routes to appropriate queues & dashboards

Technical Specifications

Model Architecture

Base: DistilBERT-base-uncased
Custom heads: 4 classification heads (main category, sub-category, priority, intervention)
Loss Function: Multi-task weighted Cross-Entropy

Compute Infrastructure

Hardware: NVIDIA GPUs (training), CPU/GPU (inference)
Software: PyTorch, Hugging Face Transformers, MLflow for tracking, Label Studio for data annotation and labeling and DVC for versioning the training data.

Citation

If you use this model, please cite:

@software{chs_distilbert_multilabel,
  author = {Bitz AI Team},
  title = {DistilBERT Multi-Label Classifier for Call Center Transcripts},
  year = {2025},
  publisher = {Hugging Face},
  url = {https://huggingface.co/openchlsystem/CHS_tz_classifier_distilbert}
}

Contact

Maintainer: Bitz AI Team
Email: [[email protected]]

Model Sources [optional]

Repository: [https://huggingface.co/openchlsystem/CHS_tz_classifier_distilbert]
Paper [optional]: [More Information Needed]
Demo [optional]: [More Information Needed]

Environmental Impact

Carbon emissions can be estimated using the Machine Learning Impact calculator presented in Lacoste et al. (2019).

Hardware Type: NVIDIA GeForce RTX 4060
Hours used: [More Information Needed]
Cloud Provider: N/A
Compute Region: N/A
Carbon Emitted: N/As

DistilBERT Multi-Label Classification Model

Model Details

Sources

Uses

Direct Use

Downstream Use

Out-of-Scope Use

Bias, Risks, and Limitations

Recommendations

Label Files

How to Get Started with the Model

Expected output {'main_category': {'Advice and Counselling'}, 'sub_category': {'School Related Issues'}, 'intervention': {'Counselling'}, 'priority': {2}}

Training Details

Training Data

Training Procedure

Preprocessing

Testing Data, Factors & Metrics

Testing Data

Metrics

Results

Training Procedure

Classification Tasks

Integration Guide: NLP Pipeline

Technical Specifications

Model Architecture

Compute Infrastructure

Citation

Contact

Model Sources [optional]

Environmental Impact

Expected output `{'main_category': {'Advice and Counselling'}, 'sub_category': {'School Related Issues'}, 'intervention': {'Counselling'}, 'priority': {2}}`