The full length Python autocompletion Videoand aTwitter thread describing how it works
This is a learning/demo project to show how deep learning can be used to auto complete Python code. You can experiment with LSTM and Transformer models. We also have built a simple VSCode extension to try out the trained models.
It gives quite decent results by saving above 30% key strokes in most files, and close to 50% in some. We calculated key strokes saved by making a single (best) prediction and selecting it with a single key.
The dataset we use is the Python code found in repos linked in Awesome-pytorch-list.We download all the repositories as zip files, extract them, remove non Python files and split them randomly to build training and validation datasets.
We train a character level model without any tokenization of the source code, since it's the simplest.
- Clone this repo
- Install requirements from
requirements.txt
- Run
Python _autocomplete/create_dataset.py
.- It collects repos mentioned in PyTorch awesome list
- Downloads the zip files of the repos
- Extract the zips
- Remove non Python files
- Collect all Python code to
data/train.py
and,data/eval.py
- Run
Python _autocomplete/train.py
to train the model. Try changing hyper-parameters like model dimensions and number of layers. - Run
evaluate.py
to evaluate the model.
You can also run the training notebook on Google Colab.
-
Clone this repo
-
Install requirements from
requirements.txt
-
Install npm packages
You need to haveNode.JSinstalled
cdvscode_extension
npm install#This will install the NPM packages
-
Start the server
Python _autocomplete/serve.py
-
Open the extension project (folder) inVSCode
cdvscode_extension
code.#This will open vscode_extension in VSCode
If you don't haveVSCode command line launcher
start VSCode and open the project withFile > Open
- Run the extension from VSCode
Run > Start Debugging
This will open another VSCode editor window, with the extension
- Create or open a Python file and start editing!
Here's a sample evaluation of a trained transformer model.
Colors:
- yellow:the token predicted is wrong and the user needs to type that character.
- blue:the token predicted is correct and the user selects it with a special key press, such as TAB or ENTER.
- green:autocompleted characters based on the prediction