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.