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