| | --- |
| | language: |
| | - code |
| | library_name: transformers |
| | pipeline_tag: text-classification |
| | tags: |
| | - code-review |
| | - bug-detection |
| | - codebert |
| | - python |
| | - security |
| | - static-analysis |
| | datasets: |
| | - code_search_net |
| | base_model: microsoft/codebert-base |
| | metrics: |
| | - f1 |
| | - accuracy |
| | --- |
| | |
| | # 🔍 CodeSheriff Bug Classifier |
| |
|
| | A fine-tuned **CodeBERT** model that classifies Python code snippets into five bug categories. Built as the classification engine inside [CodeSheriff](https://github.com/jayansh21/CodeSheriff) — an AI system that automatically reviews GitHub pull requests. |
| |
|
| | **Base model:** `microsoft/codebert-base` · **Task:** 5-class sequence classification · **Language:** Python |
| |
|
| | --- |
| |
|
| | ## Labels |
| |
|
| | | ID | Label | Example | |
| | |----|-------|---------| |
| | | 0 | Clean | Well-formed code, no issues | |
| | | 1 | Null Reference Risk | `result.fetchone().name` without a None check | |
| | | 2 | Type Mismatch | `"Error: " + error_code` where `error_code` is an int | |
| | | 3 | Security Vulnerability | `"SELECT * FROM users WHERE id = " + user_id` | |
| | | 4 | Logic Flaw | `for i in range(len(items) + 1)` | |
| |
|
| | --- |
| |
|
| | ## Usage |
| |
|
| | ```python |
| | from transformers import AutoTokenizer, AutoModelForSequenceClassification |
| | import torch |
| | |
| | tokenizer = AutoTokenizer.from_pretrained("jayansh21/codesheriff-bug-classifier") |
| | model = AutoModelForSequenceClassification.from_pretrained("jayansh21/codesheriff-bug-classifier") |
| | |
| | LABELS = { |
| | 0: "Clean", |
| | 1: "Null Reference Risk", |
| | 2: "Type Mismatch", |
| | 3: "Security Vulnerability", |
| | 4: "Logic Flaw" |
| | } |
| | |
| | code = """ |
| | def get_user(uid): |
| | query = "SELECT * FROM users WHERE id=" + uid |
| | return db.execute(query) |
| | """ |
| | |
| | inputs = tokenizer(code, return_tensors="pt", truncation=True, max_length=512) |
| | with torch.no_grad(): |
| | logits = model(**inputs).logits |
| | |
| | probs = torch.softmax(logits, dim=-1) |
| | pred = logits.argmax(dim=-1).item() |
| | confidence = probs[0][pred].item() |
| | |
| | print(f"{LABELS[pred]} ({confidence:.1%})") |
| | # Security Vulnerability (99.3%) |
| | ```` |
| |
|
| | --- |
| |
|
| | ## Training |
| |
|
| | **Dataset:** [CodeSearchNet](https://huggingface.co/datasets/code_search_net) 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 |
| |
|
| | --- |
| |
|
| | ## Links |
| |
|
| | - GitHub: [jayansh21/CodeSheriff](https://github.com/jayansh21/CodeSheriff) |
| | - Live demo: [huggingface.co/spaces/jayansh21/CodeSheriff](https://huggingface.co/spaces/jayansh21/CodeSheriff) |
| |
|
| | ``` |
| | ```` |
