File size: 2,690 Bytes
2617f02 cdaeab0 2617f02 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | ---
language:
- en
- code
library_name: transformers
pipeline_tag: text-generation
base_model: uclanlp/plbart-base
datasets:
- semeru/code-text-python
tags:
- plbart
- code-summarization
- docstring-generation
- python
metrics:
- bleu
- rouge
model-index:
- name: thealper2/plbart-docstring-generation
results:
- task:
type: text2text-generation
name: Python docstring generation
dataset:
name: semeru/code-text-python
type: semeru/code-text-python
split: test
metrics:
- type: bleu
value: 5.9444
name: BLEU
- type: rouge1
value: 34.7862
name: ROUGE-1
- type: rouge2
value: 12.6668
name: ROUGE-2
- type: rougeL
value: 32.0423
name: ROUGE-L
---
# plbart-docstring-generation
[`uclanlp/plbart-base`](https://huggingface.co/uclanlp/plbart-base) fully fine-tuned to generate English docstrings for Python functions, trained on [`semeru/code-text-python`](https://huggingface.co/datasets/semeru/code-text-python).
## Usage
```python
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.
|