distilbert-IT-ticket-student-4bit

📘 DistilBERT IT Ticket Classifier – 4-bit NF4 Student

This repository contains a 4-bit–quantized version of my DistilBERT IT ticket classifier:

• Base student model: saketgarodia1/bert-it-ticket-student
• Teacher model: saketgarodia1/bert-IT-ticket-classifier-full

The goal of this repo is to provide a much smaller, GPU-friendly model that still achieves strong performance on IT ticket topic classification.

---

🧠 Task

Task: Multi-class text classification
Domain: IT service / support tickets
Labels (8 classes): internal IT “Topic_group” categories (e.g. VPN, Email, Storage, Hardware, etc.)

---

🔧 Model Details

• Architecture: DistilBERT (distilbert-base-uncased)
• Head: AutoModelForSequenceClassification with 8 labels
• Training scheme (base student):
  • Knowledge distillation from the BERT teacher
  • Combined loss:
    [ \mathcal{L} = \alpha \cdot \text{CE}(\text{student}, y) + (1 - \alpha) \cdot \text{KL}(p_T, p_S) ]
    • α = 0.3 (more weight on hard labels)
    • Temperature T = 2 for softened logits
  • Optimizer: AdamW
  • Learning rate: 2e-5
  • Weight decay: 0.01
  • Epochs: 3
  • Batch size: 16 / 32
  • Teacher is frozen (used only to produce logits)

The full-precision student lives in saketgarodia1/bert-it-ticket-student.
This repo is a post-training 4-bit quantized copy of that model.

---

🧮 Quantization (4-bit NF4 with bitsandbytes)

Quantization is done using bitsandbytes with NF4 (NormalFloat4):

• load_in_4bit=True
• bnb_4bit_quant_type="nf4"
• bnb_4bit_compute_dtype=torch.bfloat16
• bnb_4bit_use_double_quant=True

Rough behavior:

• Quantized to 4-bit (stored as torch.uint8):
  • Most large linear weights in the transformer (attention & FFN weights, pre-classifier, etc.).
• Kept in FP16 (torch.float16):
  • Token & position embeddings
  • LayerNorm weights & biases
  • All bias terms
  • Quantization metadata (scales/zero-points, etc.)

This gives a large memory reduction compared to the FP32 teacher and FP16 student, while keeping accuracy close to the original student.

---

📊 Evaluation

Evaluated on the same validation/test splits of
Dataset: saketgarodia1/IT-service-topic-classification-data

Metrics are accuracy and macro-averaged F1.

🔍 Student 4-bit NF4

Values are computed with a standard cross-entropy loss on the hard labels, and sklearn.metrics.accuracy_score / f1_score(average="macro").

🔗 Comparison (for context)

From the full-precision models (trained in earlier notebooks):

(Replace 0.93X with your final student numbers if you log them.)

---

🚀 How to Use

Simple inference (auto-quantization with bitsandbytes)

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    BitsAndBytesConfig,
)

model_id = "saketgarodia1/distilbert-IT-ticket-student-4bit"

# 4-bit NF4 config (same as used for export)
nf4_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForSequenceClassification.from_pretrained(
    model_id,
    quantization_config=nf4_config,
    device_map="auto",   # place layers on available GPU(s)
)

text = "VPN not connecting to corporate WiFi"
inputs = tokenizer(text, return_tensors="pt").to(model.device)

with torch.no_grad():
    logits = model(**inputs).logits

pred_id = logits.argmax(dim=-1).item()
print("Predicted class id:", pred_id)