graykode/ai-docstring

VSCode Python AI Docstring Generator

Visual Studio Code extension to quickly generate docstrings for python functions using AI(NLP) technology. This project is forked for NilsJPWerner/autoDocstring. Previously, the description of the function had to be written by the user, but the AI would see the code and summarize.

Features

• AI Quickly generate a docstring snippet that can be tabbed through.
• Choose between several different types of docstring formats.
• Infers parameter types through pep484 type hints, default values, and var names.
• Support for args, kwargs, decorators, errors, and parameter types

Docstring Formats

• Google (default)
• docBlockr
• Numpy
• Sphinx
• PEP0257 (coming soon)

Usage

Usage is very simple. You just (1) run the container for the model inference server and (2) install extension in vscode and use.

(1) Run the container for the model inference server

1. If you have GPU machine : docker run -it -d --gpus 0 -p 5000:5000 graykode/ai-docstring:gpu, after installing nvidia-docker.
2. If you have only CPU : a. Run flask server with google colab and ngrok(Recommend!) or b. use docker cpu image : docker run -it -d -p 5000:5000 graykode/ai-docstring:cpu
   • At this time, it is very likely to cause OOM problem. We need more memory limit than roughly 2GB. So add --memory 2g --memory-swap parameter in linux and change memory limit in Preferences > Advanced more than 2GB(default) in macOS.

(2) Install extension in vscode and use

Cursor must be on the line directly below the definition to generate full auto-populated docstring

• Press enter after opening docstring with triple quotes (""" or ''')
• Keyboard shortcut: ctrl+shift+2 or cmd+shift+2 for mac
  • Can be changed in Preferences -> Keyboard Shortcuts -> extension.generateDocstring
• Command: Generate Docstring
• Right click menu: Generate Docstring

Extension Settings

Extension Settings are the same as the mother project except for autoDocstring.ServerEndpoint :

• ai-docstring.ServerEndpoint: endpoint address accessible to the server.

About training and dataset

For training data, github/CodeSearchNet was used, and as an initial model, we used Code2NL(Code to Natural Language) fine-tuning tasks in microsoft/CodeBERT. For detailed instructions, refer to the paper (CodeBERT: A Pre-Trained Model for Programming and Natural Languages) and this section.

Inference Benchmark(mean of 100 trials)

• CPU : Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
• GPU : Nvidia Tesla T4

License

This project is licensed under the Apache 2.0 License which is based on MIT License.