thealper2/t5-base-code-summarization
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.
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.
Training procedure
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.
