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

Split	Loss	Accuracy	Macro F1
Validation	~0.389	~0.884	~0.882
Test	~0.422	~0.880	~0.875

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):

Model	Approx. Accuracy	Approx. Macro F1	Notes
BERT teacher (FP32)	~0.937	~0.936	Larger, ~110M parameters
DistilBERT student (FP32)	~0.93X	~0.93X	~67M parameters
DistilBERT student (4-bit)	~0.88	~0.88	Quantized, much smaller

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

🚀 How to Use

Simple inference (auto-quantization with bitsandbytes)

python

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)