BraidenKirkland
commented
2 years ago Abstract
- Pythia – Python assisted code completion
- Currently used in Intellicode extension in VS code
- Designed to work on high throughput
- Predicting best matching completions on the order of 100 ms
- Compares two methods
- Frequency-based approach
- Invocation-based Markov Chain Model
- Offline evaluation results obtained on 2700 open-source Python projects on GitHub show
- A top-5 accuracy of 92%
- Surpassing baseline models by 20% average over classes
- For both intra and cross-project settings
1. Introduction
- For both intra and cross-project settings
- Code completion – one of the most widely used features of ide’s
- Reduces errors and improves productivity
- Previously, trigger characters (“.” or “=”) would cause all resulting methods to be displayed, but not ordered by likelihood
- Pythia is trained to find the most likely method given the code snippet
- Pythia uses a model capable of storing responses and then determining the most likely method
- Main contributions of this paper
- Introduce and implement several baselines of code completions systems
- Propose and deploy a novel end-to-end code completion system based on LSTM (long short-term memory) trained on source code snippets
- Evaluate the model on a dataset of 15.8 million method calls extracted from real source code (showing the model achieves 92% accuracy)
- Discusses challenges of:
- training deep neural networks and hyperparameter tuning
- model deployment on lightweight client devices at high throughput
2. Baseline Code Completion Systems
- Most basic method of code completion for an IDE is to list out all possible methods alphabetically
2.1. Frequency Models
- Make use of popularity in recommending methods, as opposed to alphabetically recommending
- Simplest frequency-based model would be to order class methods based on the frequency seen in the training set
- The model can be made stronger by incorporating context info about if-conditions (frequency-if)
- Invocations divided to two groups:
- inside if-statements – typical calls: file i/o wrappers, context managers, specifying tensor properties
- outside if-statements – typical calls: variable/placeholder definitions, setting up a NN or training routine
2.2. Invocation-based Markov Chain Model
- Invocations divided to two groups:
- Some sequences of invocations are always followed by the same methods
- Using previous invocation occurrences to predict the next invocation used is the core of their model of Markov Chain model of code completion
- P(mi |hi) = P(mi |mi−n, ...,mi−1), hi = m1,m2, ...,mi−1 is the history of previous method invocations in the same document scope for a given class
- The probability of the next invocation only depends of the previous n-1 invocations
3. Dataset
- The probability of the next invocation only depends of the previous n-1 invocations
- The training and testing set were collected from open-source repositories on Github
- A total of 2700 top-starred repositories were collected from domains including
- Scientific computing
- Machine learning
- Dataflow programming
- Web development
- Over 15.8 million method calls total
- Data was split 70/30 at the repository level for training and testing
- The training set was further split 80/20 for training and validation
- To serve predictions online, the model is retrained using the entire dataset???
4. Representing Code Snippets
- Source code is parsed to an abstract syntax tree (AST) using an explicit context-free grammar (CFG)
- Pythia model consumes partial file-level AST that correspond to code snippets containing member expressions and module function invocations as input for training
- Before feeding to a long short-term memory(LSTM), ASTs are placed in a sequence according to an in-order depth-first traversal
- When extracting sequences for training, up to T lookback tokens, preceding each method call, are retained. T is a tuning parameter of the model
- The nodes and tokens an AST must be mapped to a numeric vector before it can be consumed by the LSTM – the Word2Vec approach is used here
- The AST for each input file is extracted using the PTVS parser
- Training Samples:
- Feature Vector (BAK) - An AST serialized into a sequence with the “.” serving as the end-of-sequence character
- Label – the name of the method call token (one-hot encoded)
4.1. Leveraging Type Information
- They inferred the types using static analysis and included this info into the training sequences
5. Neural Code Completion Model
- Used a recurrent neural network (RNN) – a family of neural networks capable of processing sequential data of arbitrary length
- Long short-term memory networks are a particular type of RNN and obtain remarkable performance when trained on large datasets
6. Model Training
- Trained iteratively, making multiple passes over the entire dataset before converging to a minimum
- Back propagation through time (BPTT) is used as a gradient-based neural network training algorithm
- Pythia uses data-parallel distributed training algorithm with Adam optimizer, keeping an a copy of the entire neural model on each worker
- Offline training module of Pythia is implemented as a Python library making use of:
- CUDA 9
- GPU
- TensorFlow
6.1. Batching and Learning Rate Schedule
- # of nodes in a file-level AST ranges from O(n2) to O(n4)
- Sequence lengths in the range of 100 – 1000 are considered
- Training sequence length vary, but the majority have the maximum length of T
- The learning rate applied during gradient optimization is lowered upon completing each epoch according to exponential decay
- In the case of a distributed system, the learning rate is sometimes scaled up during the first several epochs
6.2. Hyperparameters
- Examples of some of the hyperparameters that need to be tuned
- Numerical values: learning rate, # of LSTM layers, dimension of embedding space, learning rate decay per epoch
- Categorical variables: precise model architecture, normalization algorithm
6.3. Tuning Neural Network Architecture
- The architecture of the neural network is a hyperparameter
- Choosing the best model for serving online involves making a trade-off between accuracy and model size
7. Evaluation
7.1. Evaluation Metrics
- Quality of the recommendations was measured using
- Top-k accuracy
- Top-1 means how often the top recommendation is correct
- Top-5 means how often the top 5 in the list contain the recommendation the user is looking for
- Mean reciprocal rank (MRR)
7.2. Evaluation Results
- Top-k accuracy
8. Model Deployment
- Pythia is currently part of the Intellicode extension in VS Code
- Neural network quantization was applied to reduce the memory footprint and model size on disk
8.1. Neural Network Quantization
- A process of reducing the number of bits to store weights
- Results of quantization
- Model size reduction was reduced to ¼ of the original
- Top-5 accuracy reduced from 92% to 89% (acceptable)
9. Conclusions
- Pythia deployed as part of Intellicode extension in VS Code IDE
- Uses long short-term memory networks trained on long-range code contexts extracted from ASTs
- Evaluation results on a dataset of 15.8 million method calls extracted from real-world source code showed that the best model got 92% top-5 accuracy
- One particular challenge was model deployment on lightweight client devices
qiaowangli
READ paper: Pythia: AI4CODE