---
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.