Docstring AI

Automated Python docstring generator using CodeT5 with LoRA/QLoRA fine-tuning.

Setup

conda create -n docstring_ai python=3.10
conda activate docstring_ai
pip install -r requirements.txt

Training

python scripts/train.py

Inference

python scripts/inference.py --model-path ./checkpoints/final_model --code "def hello(name): print(f'Hello {name}')"

Configuration

Edit configs/default_config.yaml to customize model, data, and training parameters.

References

• Transformers Library
• CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation
• LoRA: Low-Rank Adaptation of Large Language Models
• QLoRA: Efficient Finetuning of Quantized LLMs
• CodeSearchNet Dataset
• Hugging Face Datasets
• PEFT (Parameter-Efficient Fine-Tuning)

Project Structure

src/
  ├── data_loader.py      # Dataset loading and preprocessing
  └── trainer.py          # Model training class
scripts/
  ├── train.py            # Training script
  └── inference.py        # Inference script
configs/
  └── default_config.yaml # Configuration
notebooks/               # Jupyter notebooks for exploration