A fine-tuned CodeBERT model that classifies Python code snippets into five bug categories. Built as the classification engine inside CodeSheriff — an AI system that automatically reviews GitHub pull requests.
Base model:microsoft/codebert-base · Task: 5-class sequence classification · Language: Python
Dataset:CodeSearchNet Python split with heuristic labeling, augmented with seed templates for underrepresented classes. Final training set: 4,600 balanced samples across all five classes. Stratified 80/10/10 train/val/test split.
Key hyperparameters:
Parameter
Value
Epochs
4
Effective batch size
16 (8 × 2 grad accum)
Learning rate
2e-5
Optimizer
AdamW + linear warmup
Max token length
512
Class weighting
Yes — balanced
Hardware
NVIDIA RTX 3050 (4GB)
Evaluation
Test set: 840 samples (stratified).
Class
Precision
Recall
F1
Support
Clean
0.92
0.88
0.90
450
Null Reference Risk
0.63
0.78
0.70
120
Type Mismatch
0.96
0.95
0.95
75
Security Vulnerability
0.99
0.92
0.95
75
Logic Flaw
0.96
0.97
0.97
120
Macro F1
0.89
0.90
0.89
Confusion matrix:
Clean NullRef TypeMis SecVuln Logic
Actual Clean [ 394 52 1 1 2 ]
Actual NullRef [ 23 93 1 0 3 ]
Actual TypeMis [ 3 1 71 0 0 ]
Actual SecVuln [ 4 1 1 69 0 ]
Actual Logic [ 3 0 0 0 117 ]
Logic Flaw and Security Vulnerability are the strongest classes — both have clear lexical patterns. Null Reference Risk is the weakest (precision 0.63) because null-risk code closely resembles clean code structurally. Most misclassifications there are false positives rather than missed bugs.
Limitations
Python only — not trained on other languages
Function-level input — works best on 5–50 line snippets
Heuristic labels — training data was pattern-matched, not expert-annotated
Not a SAST replacement — probabilistic classifier, not a sound static analysis tool
