t5-base-code-summarization
google-t5/t5-base (223M parameters) fine-tuned to
generate a one-sentence natural-language summary (docstring) for a Python function.
- Input:
"summarize code: " + <python source code>(the prefix is required) - Output: a short English summary of what the function does
- Language: Python only
Usage
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
repo = "thealper2/t5-base-code-summarization"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForSeq2SeqLM.from_pretrained(repo)
code = """def calculate_average(numbers):
return sum(numbers) / len(numbers)"""
inputs = tokenizer("summarize code: " + code, return_tensors="pt",
truncation=True, max_length=512)
# Decoding settings (beam search etc.) are loaded from generation_config.json.
output = model.generate(**inputs)
print(tokenizer.decode(output[0], skip_special_tokens=True))
Evaluation
Scores on 5,000 held-out test functions, never seen during training or model selection. Validation scores are from the in-training evaluation subset.
| Metric | Test | Validation |
|---|---|---|
| BLEU (sacreBLEU, corpus) | 3.92 | 4.69 |
| Smoothed BLEU-4 (sentence avg.) | 6.65 | 7.21 |
| ROUGE-1 | 36.03 | 37.42 |
| ROUGE-2 | 12.61 | 14.14 |
| ROUGE-L | 32.97 | 34.02 |
| Semantic similarity (MiniLM cosine) | 54.02 | – |
| Avg. generated length (words) | 6.17 | 5.92 |
| Avg. reference length (words) | 10.02 | 9.90 |
BLEU/ROUGE reward lexical overlap with a single reference docstring, so a correct summary phrased differently scores low. Read them alongside the examples below. CodeBLEU is not reported: it scores generated code, while this model generates English.
Examples from the test split
def validate_flavor_data(self, expected, actual):
self.log.debug('Validating flavor data...')
self.log.debug('actual: {}'.format(repr(actual)))
act = [a.name for a in actual]
return self._validate_list_data(expected, act)
- Reference: Validate flavor data.
- Generated: Validate flavor data.
def check(text):
err = "hedging.misc"
msg = "Hedging. Just say it."
narcissism = [
"I would argue that",
", so to speak",
"to a certain degree",
]
return existence_check(text, narcissism, err, msg)
- Reference: Suggest the preferred forms.
- Generated: Check if hedging is valid.
def on_source_directory_chooser_clicked(self):
title = self.tr('Set the source directory for script and scenario')
self.choose_directory(self.source_directory, title)
- Reference: Autoconnect slot activated when tbSourceDir is clicked.
- Generated: Sets the source directory for script and scenario.
Training data
sentence-transformers/codesearchnet (pair config),
code → comment pairs. The dataset mixes about six languages without a label, so
Python functions were detected by parsing with ast. Leading docstrings were stripped
from the code (otherwise the target leaks into the input), summaries were cut to their
leading prose, and broken, non-English and boilerplate rows were dropped.
| Split | Examples |
|---|---|
| train | 20,000 |
| validation | 5,000 |
| test | 5,000 |
Training procedure
| Hyper-parameter | Value |
|---|---|
| Learning rate | 0.0003 |
| Scheduler / warmup | linear / 0.03 |
| Optimizer | adamw_torch |
| Effective batch size | 32 (per-device 8 × accumulation 4) |
| Epochs | 3.0 |
| Weight decay | 0.01 |
| Max source / target length | 512 / 64 tokens |
| Precision | bf16 |
| Gradient checkpointing | True |
| Seed | 42 |
| Training time | 50.49 min |
| Peak GPU memory | 4.17 GB |
Best checkpoint selected on validation ROUGE-L with early stopping.
Generation
num_beams=4, max_length=64, min_length=4, length_penalty=1.0, no_repeat_ngram_size=3, early_stopping=True, do_sample=False
Limitations
- Trained on Python only; other languages are out of distribution.
- Inputs longer than 512 tokens are truncated, so the end of long functions is not seen.
- Summaries tend to be shorter and more generic than human-written docstrings.
- Docstrings in CodeSearchNet are noisy; the model inherits their style and errors.