CodeBERT flaky-test classifier, FlakeBench fold 2 (59.12% macro-F1, single-fold reproduction)

README.md

@@ -0,0 +1,126 @@
+---
+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, fold 2)
+
+Classifies a Java/Kotlin test method into one of six categories: five kinds of flaky
+test plus non-flaky.
+
+**This is a partial, single-fold reproduction of published work — not the authors' model
+and not a match for their reported results.** Please read the limitations before using it.
+
+## 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.
+
+Trained on **one** of the paper's four project-disjoint folds (project group 2), so it is
+**not** comparable to the paper's four-fold headline number and should not be described
+as reproducing it.
+
+## Results
+
+Held-out test set: 2,181 tests from 25 projects unseen during training.
+
+| Category | This model | Paper (4-fold) |
+|---|---|---|
+| Async Wait | 55.38% | 58.37% |
+| Concurrency | 10.53% | 35.92% |
+| Time | 57.14% | 72.73% |
+| Unordered Collections | 54.55% | 73.63% |
+| Test Order Dependency | **77.11%** | 64.35% |
+| Non-flaky | 100.00% | 100.00% |
+| **Macro F1** | **59.12%** | **65.79%** |
+
+Macro-F1 is **6.67 points below** the paper. Order-Dependency is the one category that
+exceeds it. Concurrency is weak (10.53%): 12 of 17 Concurrency tests are misclassified as
+Async Wait — the two categories share `Thread`/`await` vocabulary and the model does not
+separate them.
+
+Against a faithful baseline trained under the same compute budget without rebalancing
+(45.59%), this model is **+13.53 points** better — bootstrap 95% CI [+4.48, +21.23],
+p = 0.999, 2,000 resamples.
+
+## Training
+
+| | |
+|---|---|
+| Base | `microsoft/codebert-base` |
+| Head | Linear(768→512) → ReLU → Dropout(0.3) → Linear(512→6) |
+| Loss | Focal loss, γ=2.0, balanced class weights |
+| Optimiser | AdamW, lr 1e-5, weight decay 0.01 |
+| Batch / max length | 8 / 512 |
+| Precision | fp16 |
+| Epochs | 18 (best at 6, selected on validation macro-F1) |
+| Rebalancing | non-flaky undersampled 4,972→800; each flaky class duplicated to 160 |
+| 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). It is what produced the gain over the baseline.
+
+## 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])
+```
+
+`trust_remote_code=True` is required — the MLP head is a custom architecture defined in
+`modeling_flakylens.py`.
+
+Labels: `0` async wait, `1` concurrency, `2` time, `3` unordered collections,
+`4` test order dependency, `5` non-flaky.
+
+## Limitations
+
+- **One fold, not four.** No evidence it generalises to folds 1, 3 or 4.
+- **Below the published result.** 59.12% vs 65.79%.
+- **Concurrency is unreliable** (10.53% F1) — it mostly predicts Async Wait instead.
+- **Noise floor is wide.** The test fold has 2,181 tests but only 103 flaky ones, so
+  per-category F1 carries roughly ±6 points of uncertainty. Treat small differences as
+  meaningless.
+- **Non-flaky dominates.** 95.3% of the test set is non-flaky and scores 100%, so overall
+  accuracy (97%) is not informative — macro-F1 is the metric that matters.
+- Java/Kotlin test methods only; 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).

config.json

@@ -0,0 +1,49 @@
+{
+  "architectures": [
+    "FlakyLensForTestClassification"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoConfig": "modeling_flakylens.FlakyLensConfig",
+    "AutoModelForSequenceClassification": "modeling_flakylens.FlakyLensForTestClassification"
+  },
+  "bos_token_id": 0,
+  "classifier_dropout": null,
+  "eos_token_id": 2,
+  "head_dropout": 0.3,
+  "head_hidden": 512,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "id2label": {
+    "0": "async wait",
+    "1": "concurrency",
+    "2": "time",
+    "3": "unordered collections",
+    "4": "test order dependency",
+    "5": "non-flaky"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "label2id": {
+    "async wait": 0,
+    "concurrency": 1,
+    "non-flaky": 5,
+    "test order dependency": 4,
+    "time": 2,
+    "unordered collections": 3
+  },
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 514,
+  "model_type": "flakylens",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "output_past": true,
+  "pad_token_id": 1,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.1",
+  "type_vocab_size": 1,
+  "use_cache": true,
+  "vocab_size": 50265
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd7e599447b3fe15ac3ed2c09cf6c7b51343f0cb2bcfcf591cbbb8320f9c2c33
+size 500194032

modeling_flakylens.py

@@ -0,0 +1,41 @@
+"""CodeBERT + MLP head for flaky-test category classification.
+
+Mirrors BERT_Arch from the FlakyLens artifact: RoBERTa pooled output -> Linear(768,512)
+-> ReLU -> Dropout -> Linear(512,6). The original applies LogSoftmax to the final layer;
+this returns raw logits instead, following the HF convention. argmax is identical either
+way, so predictions match exactly -- apply log_softmax if you need the original's values.
+"""
+import torch.nn as nn
+from transformers import RobertaConfig, RobertaModel, RobertaPreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+
+class FlakyLensConfig(RobertaConfig):
+    model_type = "flakylens"
+
+    def __init__(self, head_hidden=512, head_dropout=0.3, **kwargs):
+        super().__init__(**kwargs)
+        self.head_hidden = head_hidden
+        self.head_dropout = head_dropout
+
+
+class FlakyLensForTestClassification(RobertaPreTrainedModel):
+    config_class = FlakyLensConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.roberta = RobertaModel(config, add_pooling_layer=True)
+        self.fc1 = nn.Linear(config.hidden_size, config.head_hidden)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(config.head_dropout)
+        self.fc2 = nn.Linear(config.head_hidden, config.num_labels)
+        self.post_init()
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
+        pooled = outputs[1]
+        logits = self.fc2(self.dropout(self.relu(self.fc1(pooled))))
+        loss = None
+        if labels is not None:
+            loss = nn.functional.cross_entropy(logits, labels)
+        return SequenceClassifierOutput(loss=loss, logits=logits)

special_tokens_map.json

@@ -0,0 +1,51 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<mask>",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,57 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50264": {
+      "content": "<mask>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "errors": "replace",
+  "mask_token": "<mask>",
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "tokenizer_class": "RobertaTokenizer",
+  "trim_offsets": true,
+  "unk_token": "<unk>"
+}