plbart-docstring-generation

uclanlp/plbart-base fully fine-tuned to generate English docstrings for Python functions, trained on semeru/code-text-python.

Usage

from transformers import AutoTokenizer, PLBartForConditionalGeneration

tokenizer = AutoTokenizer.from_pretrained("thealper2/plbart-docstring-generation", src_lang="python", tgt_lang="en_XX")
model = PLBartForConditionalGeneration.from_pretrained("thealper2/plbart-docstring-generation")

code = "def add(a, b):\n    return a + b"
inputs = tokenizer(" ".join(code.split()), max_length=512, truncation=True, return_tensors="pt")
out = model.generate(**inputs, num_beams=4, max_length=64,
                     decoder_start_token_id=model.config.decoder_start_token_id)
print(tokenizer.decode(out[0], skip_special_tokens=True))

Evaluation

Test split (14918 examples), beam search with 4 beams, max length 64.

Split	BLEU	ROUGE-1	ROUGE-2	ROUGE-L	Loss
test	5.94	34.79	12.67	32.04	2.6885
validation	5.46	33.95	12.33	31.28	3.8836

Mean generated length: 6.35 tokens (references: 11.20).

Training

Hyperparameter	Value
max_train_samples	50000
num_epochs	2.0
learning_rate	3e-05
train_batch_size	32
gradient_accumulation_steps	1
weight_decay	0.01
warmup_ratio	0.05
lr_scheduler_type	linear
label_smoothing_factor	0.1
max_source_length	512
max_target_length	128
bf16	True
seed	42

Trained examples: 50000. Training time: 0.29 h on NVIDIA GeForce RTX 5060 Ti (15.9 GiB, sm_120).

Limitations

Generated docstrings are short, single-sentence summaries; they tend to be shorter than human-written references and may describe parameters or behaviour incorrectly. Review them before use.