Fine-tuning a small LLM for domain-specific classification
When fine-tuning is the right choice
The standard advice for using LLMs in classification tasks is to start with prompting. Zero-shot and few-shot prompting against a capable frontier model is fast to implement, requires no training infrastructure, and often achieves acceptable accuracy. For many production use cases, that is enough.
Fine-tuning makes sense when you have hit specific ceilings that prompting cannot overcome. If your domain uses terminology, abbreviations, or labelling conventions that do not appear in general pre-training data, even a very well-crafted prompt will struggle. If you are classifying thousands of documents per day and cost per token is a meaningful budget line, a small fine-tuned model running on your own hardware will be dramatically cheaper. If you need deterministic, auditable classification where a legal or compliance team needs to inspect exactly what the model learned, fine-tuning gives you that — you own the weights.
The target architecture for most domain classification use cases is a model in the 1B to 8B parameter range. Qwen2.5 1.5B, Llama 3.2 3B, and Mistral 7B are common choices. They are small enough to fine-tune on a single A100 GPU in hours and fast enough to run inference at scale on modest hardware. You do not need a 70B model for a text classification task with 20 labels.
Dataset preparation
Good training data is more important than any hyperparameter. Classification fine-tuning needs labelled examples where each example includes an input text and a correct label.
The training format used by modern instruction-tuned models follows a chat template. Each example is a conversation with a system prompt explaining the task, a user message containing the document to classify, and an assistant message containing the correct label.
# Example training record structure
{
"messages": [
{
"role": "system",
"content": "You are a support ticket classifier. Classify the ticket into exactly one of these categories: billing, technical, account, feature_request, other. Respond with the category name only."
},
{
"role": "user",
"content": "I was charged twice for my subscription this month and I need a refund."
},
{
"role": "assistant",
"content": "billing"
}
]
}
The dataset preparation script:
import json
import random
from pathlib import Path
def prepare_dataset(raw_data: list[dict], output_dir: str, val_split: float = 0.15):
"""
raw_data: list of {"text": str, "label": str} dicts
"""
system_prompt = (
"You are a support ticket classifier. "
"Classify the ticket into exactly one of these categories: "
"billing, technical, account, feature_request, other. "
"Respond with the category name only."
)
records = []
for item in raw_data:
records.append({
"messages": [
{"role": "system", "content": system_prompt},
{"role": "user", "content": item["text"]},
{"role": "assistant", "content": item["label"]},
]
})
random.shuffle(records)
split_idx = int(len(records) * (1 - val_split))
train_records = records[:split_idx]
val_records = records[split_idx:]
output = Path(output_dir)
output.mkdir(parents=True, exist_ok=True)
with open(output / "train.jsonl", "w") as f:
for record in train_records:
f.write(json.dumps(record) + "\n")
with open(output / "val.jsonl", "w") as f:
for record in val_records:
f.write(json.dumps(record) + "\n")
print(f"Train: {len(train_records)} records")
print(f"Val: {len(val_records)} records")
# Check class distribution before training
from collections import Counter
def check_distribution(jsonl_path: str):
labels = []
with open(jsonl_path) as f:
for line in f:
record = json.loads(line)
label = record["messages"][-1]["content"]
labels.append(label)
counts = Counter(labels)
total = len(labels)
print("\nClass distribution:")
for label, count in sorted(counts.items(), key=lambda x: -x[1]):
print(f" {label}: {count} ({count/total*100:.1f}%)")
Class imbalance is the most common dataset problem. If your rarest class has fewer than 50 examples, oversample it or augment with paraphrases generated by a larger model. A model trained on severely imbalanced data will learn to predict the majority class nearly always and appear to have high accuracy while being useless for minority classes.
Aim for at least 200 examples per class, with 500 being a comfortable minimum for reliable results. For classification tasks, 2000-5000 total examples often saturates accuracy — more data past that point yields diminishing returns unless the task is genuinely ambiguous.
LoRA fine-tuning with the Hugging Face ecosystem
Parameter-efficient fine-tuning (PEFT) with LoRA (Low-Rank Adaptation) is the standard approach for fine-tuning on consumer or mid-range hardware. Instead of updating all model weights, LoRA inserts small trainable matrices at specific layers and freezes everything else. This reduces trainable parameters by 100x or more, fitting the training process into a fraction of the memory that full fine-tuning would require.
Install the required packages:
pip install transformers datasets peft trl accelerate bitsandbytes
The training script using TRL's SFTTrainer:
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, TaskType
from trl import SFTTrainer, SFTConfig
import torch
MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
DATA_DIR = "./data"
OUTPUT_DIR = "./checkpoints"
# 4-bit quantisation reduces VRAM from ~3 GB to ~1.2 GB for the 1.5B model
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
# LoRA configuration
# target_modules: which weight matrices to adapt — attention projections are standard
lora_config = LoraConfig(
task_type=TaskType.CAUSAL_LM,
r=16, # Rank — higher = more capacity, more VRAM
lora_alpha=32, # Scaling factor, typically 2x rank
lora_dropout=0.05,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
bias="none",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# Output: trainable params: 3,932,160 || all params: 1,548,033,024 || trainable%: 0.25
# Load dataset
dataset = load_dataset("json", data_files={
"train": f"{DATA_DIR}/train.jsonl",
"validation": f"{DATA_DIR}/val.jsonl",
})
# Training configuration
training_args = SFTConfig(
output_dir=OUTPUT_DIR,
num_train_epochs=3,
per_device_train_batch_size=4,
per_device_eval_batch_size=4,
gradient_accumulation_steps=4, # Effective batch size = 16
learning_rate=2e-4,
lr_scheduler_type="cosine",
warmup_ratio=0.1,
evaluation_strategy="steps",
eval_steps=100,
save_strategy="steps",
save_steps=100,
load_best_model_at_end=True,
metric_for_best_model="eval_loss",
logging_steps=10,
bf16=True,
max_seq_length=512,
)
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=dataset["train"],
eval_dataset=dataset["validation"],
args=training_args,
)
trainer.train()
trainer.save_model(f"{OUTPUT_DIR}/final")
Training 3000 examples for 3 epochs on a single A100 40GB GPU takes roughly 15-20 minutes for a 1.5B model. On a consumer RTX 4090 with 24GB VRAM, expect 30-45 minutes. On a MacBook Pro M3 Max using MPS acceleration, roughly 2-3 hours.
Evaluating the fine-tuned model
After training, systematic evaluation on a held-out test set is essential. Validation loss during training tells you whether the model is improving, but it does not tell you per-class accuracy or where the model still fails.
from transformers import pipeline
import json
from sklearn.metrics import classification_report, confusion_matrix
import numpy as np
def load_test_data(path: str) -> tuple[list[str], list[str]]:
texts, labels = [], []
with open(path) as f:
for line in f:
record = json.loads(line)
messages = record["messages"]
texts.append(messages[1]["content"]) # user message
labels.append(messages[2]["content"]) # assistant label
return texts, labels
def evaluate_model(model_path: str, test_path: str):
pipe = pipeline(
"text-generation",
model=model_path,
tokenizer=model_path,
device_map="auto",
max_new_tokens=10, # Labels are short — no need for more
do_sample=False, # Greedy decoding for deterministic results
)
system_prompt = (
"You are a support ticket classifier. "
"Classify the ticket into exactly one of these categories: "
"billing, technical, account, feature_request, other. "
"Respond with the category name only."
)
texts, true_labels = load_test_data(test_path)
predicted_labels = []
for text in texts:
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": text},
]
result = pipe(messages)
raw_output = result[0]["generated_text"][-1]["content"].strip().lower()
# Normalise to valid label or "unknown"
valid_labels = {"billing", "technical", "account", "feature_request", "other"}
label = raw_output if raw_output in valid_labels else "unknown"
predicted_labels.append(label)
print(classification_report(true_labels, predicted_labels))
# Confusion matrix
all_labels = sorted(set(true_labels + predicted_labels))
cm = confusion_matrix(true_labels, predicted_labels, labels=all_labels)
print("\nConfusion Matrix:")
print("Labels:", all_labels)
print(cm)
# Accuracy
correct = sum(t == p for t, p in zip(true_labels, predicted_labels))
print(f"\nOverall accuracy: {correct / len(true_labels) * 100:.1f}%")
evaluate_model("./checkpoints/final", "./data/test.jsonl")
Pay attention to per-class F1 scores rather than overall accuracy. A model that achieves 92% accuracy on a dataset where 90% of examples are one class has learned nothing useful about the minority classes. Weighted F1 is a better summary metric for imbalanced classification.
Compare your fine-tuned model against a baseline: the same base model without fine-tuning, prompted with the same system prompt. This tells you how much fine-tuning actually helped versus what prompting alone could achieve.
Merging LoRA adapters and deploying with Ollama
Training produces a set of LoRA adapter weights separate from the base model. For deployment, you have two options: load the adapter dynamically on top of the base model, or merge the adapter into the base model to produce a single set of weights.
Merging is cleaner for production — you get a single model file that can be deployed with Ollama, vLLM, or any other inference server without special adapter handling.
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch
BASE_MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
ADAPTER_PATH = "./checkpoints/final"
MERGED_OUTPUT = "./merged-model"
# Load base model in full precision for merging
base_model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL_ID,
torch_dtype=torch.bfloat16,
device_map="cpu", # Merge on CPU to avoid VRAM issues
trust_remote_code=True,
)
# Load adapter and merge
model_with_adapter = PeftModel.from_pretrained(base_model, ADAPTER_PATH)
merged_model = model_with_adapter.merge_and_unload()
# Save merged model
merged_model.save_pretrained(MERGED_OUTPUT)
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)
tokenizer.save_pretrained(MERGED_OUTPUT)
print(f"Merged model saved to {MERGED_OUTPUT}")
To deploy with Ollama, create a Modelfile that points at the merged weights. Ollama accepts models in the GGUF format, so you first convert with llama.cpp:
# Clone llama.cpp
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp && pip install -r requirements.txt
# Convert HuggingFace format to GGUF
python convert_hf_to_gguf.py ../merged-model --outfile ../ticket-classifier.gguf --outtype q8_0
Create the Modelfile:
FROM ./ticket-classifier.gguf
SYSTEM "You are a support ticket classifier. Classify the ticket into exactly one of these categories: billing, technical, account, feature_request, other. Respond with the category name only."
PARAMETER temperature 0
PARAMETER num_predict 10
Register and run:
ollama create ticket-classifier -f Modelfile
ollama run ticket-classifier "My payment failed and I was charged anyway"
Continuous improvement patterns
A fine-tuned model is not a one-time artifact. As your product evolves and new categories emerge, or as the model starts failing on new types of inputs, you need processes to improve it over time.
Collect model mistakes in production. Log every prediction along with a confidence score. For classification, a simple heuristic for confidence is token probability of the predicted label token. Route low-confidence predictions to a human reviewer. The disagreements between human review and model prediction are your most valuable training examples.
Keep a versioned dataset registry. Every model version should correspond to a specific, frozen dataset version. This makes regression testing possible — if a new model version performs worse on an older test set, you know something regressed.
Evaluate on subsets. Track accuracy separately for each class and for each data source or customer segment. Overall accuracy averages hide failure modes that matter enormously for specific users.
Periodically re-run the baseline comparison. As frontier models improve, the bar for "better than prompting" moves. What required fine-tuning in 2024 might be achievable with zero-shot prompting against a newer model in 2026. Re-evaluate whether fine-tuning is still the right architectural choice every six to twelve months.
Fine-tuning a small model is a legitimate engineering investment when the conditions are right — enough labelled data, a stable task definition, and clear constraints on latency or cost. The Hugging Face ecosystem with LoRA and TRL makes the mechanics accessible to any engineer who can write Python, without requiring deep ML research expertise.