toxishield / README.md
imranraad's picture
Update README.md
e2f7150 verified
|
Raw
History Blame Contribute Delete
10.9 kB
---
library_name: transformers
tags:
- text-classification
- toxicity-detection
- content-moderation
- bert
- github
- code-review
---
# Model Card for toxishield
<!-- Provide a quick summary of what the model is/does. -->
A `bert-base-uncased` model fine-tuned for binary toxicity classification (`TOXIC` / `NON-TOXIC`) of GitHub pull request review comments.
## Model Details
### Model Description
<!-- Provide a longer summary of what this model is. -->
This model classifies a single GitHub pull request comment as toxic or non-toxic. It was fine-tuned from `google-bert/bert-base-uncased` on 38,761 labelled PR comments (the "38k detection dataset") as part of the ToxiShield project, which studies and filters toxicity in software engineering communication. The model is evaluated with stratified 10-fold cross-validation and also exported to ONNX (INT8-quantized) for lightweight/in-browser inference.
- **Developed by:** [More Information Needed]
- **Funded by [optional]:** [More Information Needed]
- **Shared by [optional]:** [More Information Needed]
- **Model type:** Transformer encoder (BERT), fine-tuned for binary sequence classification
- **Language(s) (NLP):** English (software-engineering / code-review domain text)
- **License:** [More Information Needed] (base model `google-bert/bert-base-uncased` is released under Apache-2.0; license for the fine-tuned weights has not been set)
- **Finetuned from model [optional]:** [`google-bert/bert-base-uncased`](https://huggingface.co/google-bert/bert-base-uncased)
### Model Sources [optional]
<!-- Provide the basic links for the model. -->
- **Repository:** `huggingface.co/imranraad/toxishield`
- **Paper [optional]:** https://arxiv.org/abs/2604.14408 (FSE 2026)
## Uses
<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
### Direct Use
<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
Classifying individual GitHub PR/code-review comments as `TOXIC` or `NON-TOXIC`, e.g. via `transformers.pipeline("text-classification", ...)`, to flag potentially toxic comments for human review.
### Downstream Use [optional]
<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
Integration into CI bots, code-review dashboards, or moderation tooling that triages or surfaces potentially toxic PR comments before/alongside human moderators. The ONNX INT8 export is intended for low-latency or in-browser inference in such tooling.
### Out-of-Scope Use
<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
- Not intended for languages other than English.
- Not validated on text outside the GitHub PR/code-review domain (e.g. social media, forums, chat, general web text) — the training distribution is short technical comments (median ~80 characters).
- Not intended as the sole basis for moderation, disciplinary, employment, or legal decisions — outputs should be reviewed by a human.
- Not a general-purpose toxicity/hate-speech classifier.
## Bias, Risks, and Limitations
<!-- This section is meant to convey both technical and sociotechnical limitations. -->
- **Class imbalance:** the training data is 73.9% non-toxic / 26.1% toxic (28,641 vs 10,120 of 38,761 comments), which can bias the model toward the majority (non-toxic) class.
- **Recall on the toxic class is the weaker metric:** across 10-fold CV, mean recall on toxic comments is 0.954 (± 0.010) vs. mean precision of 0.975 (± 0.003) — i.e. the model is somewhat more likely to miss a toxic comment than to wrongly flag a non-toxic one.
- **Manual review of false positives** (n=23 sampled) attributes most errors to:
- Nuance/context misreads, e.g. sarcasm, dry humor, mockery (12/23, ~52%)
- Technical jargon or inline code snippets read as hostile (6/23, ~26%)
- Self-deprecating humor (2/23, ~9%)
- General negative sentiment without toxicity (2/23, ~9%)
- The model was fine-tuned and evaluated on data from GitHub PR comments only, so its calibration will not necessarily transfer to other codebase-hosting platforms or review tools.
### Recommendations
<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
Use as a triage/assistive signal rather than an automated blocking mechanism, particularly given the lower recall on the toxic class. Expect false positives on sarcastic, self-deprecating, or jargon-heavy comments, and route model decisions through human review before any moderation action.
## How to Get Started with the Model
Use the code below to get started with the model.
```python
from transformers import pipeline
classifier = pipeline("text-classification", model="<path-or-hub-id-of-saved-model>")
classifier("this the stupidest code ever")[0]["label"]
```
## Training Details
### Training Data
<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
38,761 labelled GitHub PR review comments (`38k-detection-dataset`, dataset handle `[ANONYMIZED-ORG]/38k-dataset-labelled`, not yet published): 28,641 `NON-TOXIC` (label 0) and 10,120 `TOXIC` (label 1). Comment length ranges from 5 to 998 characters (median ~80). Split into an 80% train / 10% test CSV for the single-run fine-tune; the full dataset is additionally used for stratified 10-fold cross-validation.
### Training Procedure
<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
#### Preprocessing [optional]
Comments are tokenized with the `bert-base-uncased` WordPiece tokenizer with truncation. The single-run fine-tune uses dynamic padding (`DataCollatorWithPadding`); the 10-fold cross-validation run uses fixed `max_length=128` padding.
#### Training Hyperparameters
- **Single-run fine-tune:** learning rate 2e-5, per-device batch size 16, weight decay 0.01, 1 epoch, evaluate/save every epoch, best checkpoint restored at the end.
- **10-fold cross-validation run:** learning rate 2e-5, per-device batch size 16, gradient accumulation 8 (effective batch size 128), weight decay 0.01, up to 20 epochs with early stopping (patience 3), stratified 10-fold split (each fold further split ~89/11 into train/validation).
- **Training regime:** fp16 mixed precision on GPU when available, fp32 on CPU. <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
#### Speeds, Sizes, Times [optional]
<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
`bert-base-uncased` has ~110M parameters. The best checkpoint is additionally exported to ONNX and INT8-quantized for lighter-weight/in-browser inference. [More Information Needed] on wall-clock training time.
## Evaluation
<!-- This section describes the evaluation protocols and provides the results. -->
### Testing Data, Factors & Metrics
#### Testing Data
<!-- This should link to a Dataset Card if possible. -->
Stratified 10-fold cross-validation over the full 38,761-sample dataset (per-fold results in `results/kfold-metrics/cross_validation_results.csv`); the single-run fine-tune is additionally evaluated on the held-out 10% test split.
#### Factors
<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
No subpopulation disaggregation performed; results are reported per cross-validation fold and averaged.
#### Metrics
<!-- These are the evaluation metrics being used, ideally with a description of why. -->
Accuracy, precision, recall, and F1 (binary, positive class = `TOXIC`), chosen to capture both overall correctness and toxic-class-specific performance given the class imbalance.
### Results
10-fold cross-validation (mean ± std over 10 folds):
| Metric | Mean | Std |
|---|---|---|
| Accuracy | 0.9818 | 0.0023 |
| Precision (toxic) | 0.9753 | 0.0033 |
| Recall (toxic) | 0.9543 | 0.0096 |
| F1 (toxic) | 0.9647 | 0.0047 |
Baseline comparison — GPT-4o, zero-shot prompted, on the held-out test split (`comparison/openai-detection-inference/`):
| Class | Precision | Recall | F1 |
|---|---|---|---|
| Non-toxic | 0.84 | 0.99 | 0.91 |
| Toxic | 0.96 | 0.49 | 0.65 |
| **Accuracy** | | | **0.86** |
#### Summary
The fine-tuned BERT model substantially outperforms zero-shot GPT-4o prompting on this task, most notably on toxic-class recall (0.95 vs. 0.49) — GPT-4o zero-shot misses roughly half of toxic comments, while the fine-tuned model catches the large majority at comparable or better precision.
## Model Examination [optional]
<!-- Relevant interpretability work for the model goes here -->
[More Information Needed]
## Environmental Impact
<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
- **Hardware Type:** 2× NVIDIA RTX 6000 Ada Generation (49 GB), local workstation
- **Hours used:** [More Information Needed]
- **Cloud Provider:** N/A (local/on-prem)
- **Compute Region:** [More Information Needed]
- **Carbon Emitted:** [More Information Needed]
## Technical Specifications [optional]
### Model Architecture and Objective
`bert-base-uncased` (12-layer transformer encoder, 110M parameters) with a linear classification head over 2 labels (`NON-TOXIC` = 0, `TOXIC` = 1), fine-tuned with cross-entropy loss for binary sequence classification.
### Compute Infrastructure
#### Hardware
2× NVIDIA RTX 6000 Ada Generation (49 GB each).
#### Software
Python 3.11, PyTorch 2.5.1 (CUDA 12.1 build), Transformers 4.57.6, Datasets 5.0.0, 🤗 Evaluate, scikit-learn, Optimum/ONNX Runtime (for INT8 export).
## Citation [optional]
<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
**BibTeX:**
[More Information Needed]
**APA:**
[More Information Needed]
## Glossary [optional]
<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
- **PR:** Pull request (GitHub code-review unit).
- **TOXIC / NON-TOXIC:** The two output labels (id 1 / id 0 respectively).
## More Information [optional]
[More Information Needed]
## Model Card Authors [optional]
[More Information Needed]
## Model Card Contact
[More Information Needed]