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