File size: 3,735 Bytes
0337b7f cd6eef4 0337b7f 669a784 0337b7f 86fe9a6 d2dfc53 cd6eef4 d2dfc53 0337b7f 535d5c7 4c0e4c5 0337b7f d2dfc53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | ---
license: mit
language:
- code
base_model: microsoft/codebert-base
pipeline_tag: text-classification
tags:
- flaky-tests
- software-testing
- code
- codebert
- reproduction
---
# CodeBERT for Flaky Test Categorisation (FlakeBench)
Classifies a Java/Kotlin test method into one of six categories: five kinds of flaky
test plus non-flaky.
## What this is
A fine-tune of `microsoft/codebert-base` on the FlakeBench dataset from
[*Understanding and Improving Flaky Test Classification*](https://utexas.app.box.com/v/august-shi-OOPSLA2025)
(OOPSLA 2025), trained as a reproduction exercise on a single 8 GB consumer GPU.
The uploaded weights are the "Balanced" configuration below.
## Training configurations
| Parameter | Baseline | lr 2e-5 | Balanced | Augmented | Paper |
| --- | --- | --- | --- | --- | --- |
| Encoder | codebert-base | codebert-base | codebert-base | codebert-base | codebert-base |
| Learning rate | 1e-5 | 2e-5 | 1e-5 | 1e-5 | 1e-5 |
| Batch size | 8 | 8 | 8 | 8 | 8 |
| Max length | 512 | 512 | 512 | 512 | 512 |
| Loss | focal γ=2.0 | focal γ=2.0 | focal γ=2.0 | focal γ=2.0 | focal γ=2.0 |
| Class weights | balanced | balanced | balanced | balanced | balanced |
| Optimizer | AdamW wd 0.01 | AdamW wd 0.01 | AdamW wd 0.01 | AdamW wd 0.01 | AdamW wd 0.01 |
| Precision | fp16 | fp16 | fp16 | fp16 | fp32 |
| Non-flaky rows | 4,972 | 4,972 | 800 | 800 | full |
| Minority handling | none | none | ×160 copies | ×200 variants | none |
| Train rows | 5,114 | 5,114 | 1,600 | 1,800 | 5,114 |
| Epochs run | 8 | 8 | 18 | 13 | 40 |
| Dynamic padding | no | no | no | no | no |
Hardware: 1× RTX 4060 Laptop (8 GB). Class rebalancing is the one deviation from the
paper's method, which trains on the raw distribution (97% non-flaky).
## Results (per-category F1)
| Category | Baseline | lr 2e-5 | Balanced | Augmented | Paper |
| ---------------- | ---------: | ---------: | ---------: | ---------: | ---------: |
| Async Wait | 76.92% | 78.26% | 74.07% | 64.52% | 58.37% |
| Concurrency | 0.00% | 0.00% | 0.00% | 0.00% | 35.92% |
| Time | 57.14% | 66.67% | 66.67% | 40.00% | 72.73% |
| Unordered Coll. | 75.00% | 83.33% | 83.33% | 72.73% | 73.63% |
| Order Dep. | 82.35% | 86.96% | 95.24% | 73.68% | 64.35% |
| Non-flaky | 100.00% | 99.92% | 99.51% | 100.00% | 100.00% |
| **Macro F1** | **65.24%** | **69.19%** | **69.89%** | **58.49%** | **65.79%** |
The **Balanced** configuration (uploaded weights) achieves the best macro-F1 of
**69.89%** .
## Usage
```python
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
name = "Ariful1904129/codebert-flakytest-fold2"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForSequenceClassification.from_pretrained(name, trust_remote_code=True).eval()
code = """@Test
public void testConnect() throws Exception {
Thread.sleep(1000);
assertTrue(client.isConnected());
}"""
x = tok(code, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
pred = model(**x).logits.argmax(-1).item()
print(model.config.id2label[pred])
```
Scope: Java/Kotlin test methods; inputs longer than 512 tokens are truncated.
## Citation
Please cite the original paper. This model is a third-party reproduction and is not
endorsed by its authors.
```bibtex
@inproceedings{flakylens2025,
title = {Understanding and Improving Flaky Test Classification},
booktitle = {OOPSLA},
year = {2025}
}
```
Dataset and method: [UT-SE-Research/FlakyLens](https://github.com/UT-SE-Research/FlakyLens). |