YAML Metadata Warning:The pipeline tag "text-to-text-generation" is not in the official list: text-classification, token-classification, table-question-answering, question-answering, zero-shot-classification, translation, summarization, feature-extraction, text-generation, fill-mask, sentence-similarity, text-to-speech, text-to-audio, automatic-speech-recognition, audio-to-audio, audio-classification, audio-text-to-text, voice-activity-detection, depth-estimation, image-classification, object-detection, image-segmentation, text-to-image, image-to-text, image-to-image, image-to-video, unconditional-image-generation, video-classification, reinforcement-learning, robotics, tabular-classification, tabular-regression, tabular-to-text, table-to-text, multiple-choice, text-ranking, text-retrieval, time-series-forecasting, text-to-video, image-text-to-text, image-text-to-image, image-text-to-video, visual-question-answering, document-question-answering, zero-shot-image-classification, graph-ml, mask-generation, zero-shot-object-detection, text-to-3d, image-to-3d, image-feature-extraction, video-text-to-text, keypoint-detection, visual-document-retrieval, any-to-any, video-to-video, other
T5-small (52M) for automated code review: CSCI 680 Assignment 1
Three T5 models fine-tuned for two code-review tasks on Java methods, from a course assignment comparing the effect of pre-training and of multi-task vs single-task fine-tuning. Each model is in its own folder:
| folder | condition | starts from | fine-tuned on |
|---|---|---|---|
finetune-A |
A (nopt) | random weights (no pre-training) | both tasks (multi-task) |
finetune-B |
B (pt) | our pre-trained T5 | both tasks (multi-task) |
finetune-C |
C (pt-st) | our pre-trained T5 | code2comment only (single-task) |
Tasks and input format
Inputs carry a task prefix:
- code&comment2code: given a method and a reviewer's comment, produce the revised method.
code&comment2code: <code>…method, with the commented lines between <START> and <END>…</code><technical_language>…comment…</technical_language> - code2comment: given a method, produce a review comment.
code2comment: …method…
Usage
from transformers import AutoTokenizer, T5ForConditionalGeneration
repo, folder = "nswhay/ai4se_t5", "finetune-B"
tokenizer = AutoTokenizer.from_pretrained(repo, subfolder=folder)
model = T5ForConditionalGeneration.from_pretrained(repo, subfolder=folder)
source = "code2comment: public int size() { return list.size(); }"
ids = tokenizer(source, return_tensors="pt", truncation=True, max_length=512).input_ids
out = model.generate(ids, max_new_tokens=512, num_beams=1)
print(tokenizer.decode(out[0], skip_special_tokens=True))
The results below use greedy decoding (beam 1), up to 512 new tokens, inputs truncated at 512 tokens.
Model
T5 architecture trained from scratch: 6 encoder and 6 decoder layers, d_model 512, d_ff 2048, 8 heads, ReLU feed-forward, dropout 0.1, 52.2M parameters. Our own 16,000-token tokenizer, trained on the pre-training corpus.
Training
Pre-training (B and C): span-corruption denoising (15% of tokens masked, mean span 3) on Java methods mined from GitHub, 600,000 steps at an effective batch of 128 (about 7.9B tokens), peak learning rate 1e-4, on one RTX 3080.
Fine-tuning: AdamW, learning rate 1e-4, weight decay 0.01, 1,000 warm-up steps, effective batch 32, up to 100,000 steps. Every 2,000 steps the model was validated on all 5,000 validation rows of each task it was trained on, and the checkpoint published here is the one with the highest validation exact match: the mean of the two tasks' exact match for A and B, code2comment exact match for C. Training stopped early when neither exact match nor validation loss had improved for 8 validations.
| condition | best step | stopped at | why | training hours |
|---|---|---|---|---|
| A | 88,000 | 100,000 | step limit | 12.5 |
| B | 72,000 | 88,000 | early stopping | 10.3 |
| C | 98,000 | 100,000 | step limit | 5.0 |
Results (test set, 5,000 rows per task)
| task | condition | exact match | CodeBLEU / BLEU-4 |
|---|---|---|---|
| code&comment2code | A (nopt) | 4.04% | 79.96 |
| code&comment2code | B (pt) | 9.26% | 80.93 |
| code2comment | A (nopt) | 0.06% | 1.07 |
| code2comment | B (pt) | 0.14% | 1.45 |
| code2comment | C (pt-st) | 0.64% | 1.68 |
Exact match is after whitespace normalisation. CodeBLEU: codebleu 0.7.0, Java, default weights, 0–100. BLEU-4:
sacrebleu 2.6.0, nrefs:1|case:mixed|eff:no|tok:13a|smooth:exp.
Limitations
- CodeBLEU barely measures the edit. Returning the input method unchanged scores CodeBLEU 81.87 on test, above both models, so exact match is the meaningful measure for code&comment2code.
- Code2comment exact matches are mostly memorised comments. 29 of C's 32 test exact matches are comments that appear word for word in the training set (reviewers repeating the same note across a project); the models rarely write a new comment exactly as the reviewer did.
- Runaway generations. Some outputs never produce the end-of-sequence token within 512 tokens (on code2comment test: 79 for A, 182 for B, 12 for C).
- Small models trained on one GPU for a course assignment; not intended for production code review.