Full Line Code Completion: Bringing AI to Desktop

Code completion, a prevalent feature across various IDEs [1, 2], is aimed at optimizing the work of a developer, not only by speeding up code typing but also by facilitating API exploration. Common in modern IDEs, single-token, or standard, code completion suggests individual tokens via a lookup window (see Figure 1a). This feature, relying on static code analysis, generates suggestions from multiple candidate providers, which include both built-in IDE modules and external plugins.

In addition to recent advancements in single-token code completion, there is a niche of multi-token completion that is being actively studied in academia [25, 26, 27, 28] and employed in products [3, 29, 4, 5]. Similar to single-token completion, multi-token completion can be obtained from a corresponding candidate provider. However, in contrast to standard code completion, it is usually shown as inline completion (also referred to as gray text completion, see Figure 1b) and thus usually provides only one suggestion.

Most multi-token code completion solutions in both academia and industry tend to use language modelling approaches to generate suggestions. Specifically, they utilize Transformer-based neural networks [13] like GPT-2 [14] and LLaMA [24]. Such models usually have billions of parameters, implying gigabytes of memory footprint. Moreover, they are computationally intensive, so cloud-based multi-GPU servers are required to provide candidate suggestions.

Over the past few years, Transformer-based neural networks started dominating in language modelling [14, 15, 24]. Specifically, models based solely on the decoder part of a Transformer (also known as autoregressive Transformers) are utilized for multi-token generation. A typical multi-token generation pipeline includes tokenization and sequence generation algorithms. Below, we briefly describe these parts for clearer understanding of the challenges we encountered when implementing FLCC.

Tokenization. There exist many techniques for breaking down source sequences into tokens, including character-level, word-level, and subword-level tokenizations. The latter is typically used [13, 30] with a Transformer-based model of any kind. Specifically, byte-pair encoding (BPE) [14, 24] is used for decoder-based autoregressive Transformers like GPTs and LLaMAs. This algorithm iteratively combines the most frequent pairs of adjacent tokens (bytes, characters, or character sequences) to form new, longer tokens. This process is repeated a predefined number of times or until a desired vocabulary size is reached. Then, during tokenization, the source sequence is transformed into a list of integers that correspond to an ordinal number of each token is the combined vocabulary.

Sequence generation algorithm. Transformer decoders compute pair-wise relations between all tokens in the given context to produce a probability distribution over the next token. To obtain multiple tokens using such a model, one needs to apply it several times in a row, with each next distribution depending on the token generated on the previous step.

This raises the problem of the search for the most probable token sequence. Since BPE vocabulary usually consists of thousands of tokens, the generation of the next step distribution for every possible next token is not feasible. One of the popular solutions to this problem is the beam search algorithm [31].

The overall scheme of the algorithm is presented in Figure 2. Beam search is a heuristic algorithm that focuses on traversing the graph of all possible sequences by expanding the node with the highest potential. This approach maintains only a fixed number of potential solutions at each step, called beams, or hypotheses. By focusing on the most promising hypotheses and discarding the less likely ones, beam search efficiently navigates through the graph. Usually, top $k$ most probable beams are supported at each step, and the next hypothesis score is obtained from the previous one, multiplying it by the probability of the next token. For Transformers, this procedure has an efficient implementation: the relations between context tokens can be computed once and reused afterward.

At JetBrains, an established IDE release cycle comprises 3 major releases per year, enumerated by year and year’s version in the following format: $YYYY.X$. Each release is preceded by a free early access program (EAP), allowing users to download experimental versions of the IDEs. Participation in the EAP includes user consent for the anonymized submission of usage logs to JetBrains. This data is crucial for our analysis, enabling us to determine the users’ experiences with new features. EAP releases are also used as a platform for A/B testing. The EAPs of versions 2023.3, 2024.1, and 2024.2 were used to conduct the online evaluation for Full Line Code Completion.

Many existing datasets for training code completion models, particularly for Python [32, 33, 34], have limitations due to licensing issues and quality concerns. Open-source code, often collected from platforms like GitHub, may have restrictive licenses or none at all [35], which poses legal challenges [36], as seen in GitHub Copilot’s [3] usage. Furthermore, direct extraction of GitHub’s Python files yields many autogenerated and duplicated files [37], leading to potential biases and data leakage issues in model training. A notable development to address these issues is the introduction of The Stack dataset [38, 8, 37] by BigCode [39], a collaboration of over 350 researchers backed by HuggingFace and ServiceNow. This dataset, encompassing over 6TB of code in 358 languages, is extensively used for training large language models due to its permissive licensing and wide coverage.

In recent years, code completion became a major field of study for deep learning applications. Many researchers and companies trained and released open-source models, resulting in the availability of training code, configuration specifics, and models’ weights.

The employed models are typically decoders of the Transformer architecture, including GPT and LLaMA architectures representing specific types of decoders. The most famous examples of such models are InCoder [27, 40], StarCoder [28, 41], CodeLLaMA [26, 42], CodeGen and CodeGen-2 [25, 43, 44], CodeParrot [45], and PolyCoder [46, 47]. Despite the fact that the models are open-source, they are not suitable for local code completion engines like FLCC for several reasons. First of all, they are usually too big to ship to users and do not meet the memory footprint restrictions for desktop applications like IDEs. Secondly, the license for the model’s weights might not be permissive for commercial use. Last but not least, only specific architectures like GPT-2 and LLaMA have corresponding inference engines that allow to run them on an arbitrary end-user device fast enough.

Lately, several products for multi-token completion appeared on the market, including Copilot [3], TabNine [4], Codeuim [5], CodeWhisperer [29], aiXcoder [7], Sourcegraph Cody [6], and others. Some of them remain free to use (e.g., Codeium and CodeWhisperer), while others require paid subscription. More importantly, all such products, except for TabNine and aiXcoder, offer only cloud-based solutions that might have certain restrictions for users behind the firewall, users with limited internet access, and users who have privacy concerns regarding sending their code to the cloud. Cloud-based solutions are unavailable not only for huge markets like China, but also for enterprise clients that want the source code to stay inside their network. Moreover, none of these solutions have a common API for inline completion, which makes them conflict with others and requires re-implementing the same functionality, which leads to poor UX and many confirmed bugs.

For Python, we used the subset of The Stack [38] — the StarCoder dataset [48], which we additionally processed to keep only permissive licenses. This completely fulfilled our needs in terms of handling licenses, as well as the filtering of duplicated files.

For training a model, the data is divided into training, validation, and testing segments. The training set is used for the primary adjustment of the model. The validation set allows for intermittent evaluation of the model as it is being trained. After training, the testing set evaluates the model’s performance.

After collecting the dataset, we found a representative set of repositories to be separated into the test part, thus ensuring that the model is not trained on them. This data is used later in a specialized offline evaluation setup described in Section IV-A. Additionally, to avoid potential data leaks, we ensure that all the forks of the given repository are placed in the same data segment. For the training of the Python model, the corresponding proportions of the train-validation-test sets of the data were 80%-5%-15%, adding up to slightly more than 50 GB of data.

In FLCC, we use Transformer-based models for the source code, so we chose to use the byte-pair encoding (BPE) [30] as a standard approach for such models [14, 15, 24].

Character-pair encoding. The traditional BPE is supposed to be used for bytes directly, so it is applicable to all languages if we choose an encoding like UTF-8. However, we decided to use character-pair encoding with slight modifications. Compared to natural language texts, source code primarily consists of common statements or sequences of more than one token [50]. Some good examples of such recurring “idioms” in Python are the following:

• •

  for i in range(

• •

  return True

• •

  if __name__ == "__main__":

Thus, in our approach, we implement character-pair encoding on top of the YouTokenToMe (YTTM) library [51], allowing to merge pairs over tabs and spaces but forbidding merges over the newline (\n). This approach allows us to create better representations of the lines in the source code and compress it more efficiently. This becomes significantly important when generating full lines of code because a full line then requires fewer tokens, hence leading to less computational time spent on executing the model. A similar approach is used by Meta Research for the Incoder model [27].

Insight 4: To compress the code better, we can augment the default byte-pair encoding to be character-pair, as well as allow the merging of tokens over the spaces and tabs. This way, recurring idioms become single tokens and can be generated more easily.

We chose the vocabulary size of 16,384 tokens because it allows for a decent model quality while maintaining an acceptable memory footprint: for our sizes of models, a significant part of the model’s weights depends on the vocabulary size.

Cleaning vocabulary. After the initial experiments with tokenization, we observed an important peculiarity of the obtained tokens that had to be dealt with. For a vocabulary of 16,384 tokens, as many as 800 tokens consisted of characters from various languages, mostly Chinese, but also Japanese, Korean, Arabic, and others. Since there are no comments in the dataset, these characters came from string literals and are not suited for code generation. This observation led us to force only ASCII symbols to be used for the vocabulary while leaving other symbols as <UNK> tokens and merging sequences of them into a single <UNK> token.

Insight 5: Even without the comments, raw code also contains characters that can hinder the model’s ability to suggest proper tokens. To combat this, we focused on the ASCII symbols.

We construct each training example as follows. Firstly, we sample a file from the train set. Then, we concatenate file extension, <LANG_SEP_CHAR>, file path, <METAINFO_SEP_CHAR>, and the part of the code above caret, such that the total amount of tokens in the combined three parts fits the maximum context size of the model. This way, the meta information is always present in the context given to the model, like a header.

Having experimented with several model architectures, including GPT-2 and LLaMA, we ended up with a LLaMA-like model for FLCC for all languages. In our evaluations, LLaMA-like models performed best not only in terms of quality but also speed, thanks to llama.cpp [52] — a multiplatform inference engine for LLaMA-like models.

We use the HuggingFace [53] implementations to train models on the described dataset. The models are trained in an automatic mixed precision mode [54] for several days on 8 NVidia H100 GPUs. Following recent advancements in choosing the number of training steps given the model and dataset sizes, we used HuggingFace’s approximation [55] for every trained model. For our models, we used the maximum context length of 1,536 tokens, since it performed the best according to our benchmarks, given memory footprint restrictions from target IDEs. Unlike in many similar circumstances, we do not employ knowledge distillation [56], because we discovered that for our sizes of the models, this technique does not provide any quality improvement.

To control the model’s overfitting while training, we use a standard cross-entropy loss function evaluated on the validation dataset. However, we pay little attention to well-known metrics like Recall at k (R@k) or Mean Reciprocal Rank at k (MRR@k) because they evaluate only single-token prediction quality, whereas FLCC generates multi-token suggestions.

There are memory footprint and speed constraints among restrictions for FLCC. Thus, to compress the model’s size and speed up its inference, a well-known technique called quantization [57] was used. Quantization refers to changing the way the model weights are stored and changing the precision that is used for model-related computations. For LLaMA-like models, quantization from FP32 (32-bit floating point numbers) to INT4 (4-bit integer numbers) is supported via llama.cpp. This transformation decreases the size of the model several times, trading it off to a negligible quality loss.

The quantization allowed us to achieve two major goals. Firstly, the size of the model decreased from almost 400 MB to slightly more than 100 MB, which fits within the desired restrictions. Secondly, llama.cpp allows us to shorten the CPU inference time due to more efficient usage of the registers for INT4 data: e.g., the quantized model runs 2 times faster compared to full precision model on M2 CPU on MacOS and 4 times faster on Intel i9 CPU on Linux.

In summary, for Full Line Code Completion, INT4-quantized LLaMA-like models with 100M parameters were used. The models utilize a modified BPE tokenizer implemented on top of the YouTokenToMe library.

In real-world code completion, the caret position can arbitrarily divide a token, leading to tokenization being inconsistent with the model’s training. An example case is presented in Figure 4. Here, code completion is invoked after “for i”, when the programmer intended to write “for i in range(count):”, and this creates a tokenization issue. During training, the model sees this as a single “for i in range(” token. However, in this context, this might be split into “for ” and “i”, differing from the training scenario and reducing the generation quality. This problem escalates with larger average token lengths, increasing the likelihood of mid-token caret positions.

To address this issue in FLCC, we find an offset in the line, such that tokenization of the code to the left of the offset resembles training tokenization more. This involves backtracking left from the caret position for as long as the slice from the offset to the caret position is a full substring of one of the vocabulary tokens. Then, during model execution, we can filter out hypotheses that do not match the cut-out characters.

For example, in Figure 4, starting from “for i” and moving left until the start of the line, the prefix matches the vocabulary token “for i in range(”. During code suggestion, the process begins with an empty prefix, discarding hypotheses that do not align with the “for i” prefix. This method, called token healing, resolves generation issues from mid-token positions by reverting to one token prior and constraining the initial generation to the incomplete token.

Insight 6: Having longer tokens allows us to suggest full lines more conveniently, however, this also introduces a problem of generation starting from mid-token position, which can be solved by the token healing technique.

During beam search, one obtains a set of the most probable hypotheses that continue the context initially fed to the model. In our experiments, we found that collecting only those hypotheses that end with \n (even if they are not the most probable ones) is the most beneficial approach. We call such hypotheses terminated.

The rationale beyond this change is the following. The sequences that were not finished with \n in the maximum allowed number of steps in the sequence generation algorithm are probably too long and might not finish at all even if the algorithm continued working. Moreover, the user expects a code-suggesting tool to show finished constructions.

Insight 7: Collecting end-of-line-finished hypotheses allows suggesting better completion candidates.

Another modification to the standard algorithm that we implemented is that the number of iterations we perform in the beam search is not pre-defined but dynamic. More precisely, the generation algorithm has the following set of stopping criteria:

• •

  The number of iterations reached the maximum allowed number $n$.

• •

  All current hypotheses are terminated.

• •

  All current non-terminated hypotheses have $k$ times lower probability than the best terminated hypothesis.

The rationale for this is the following. FLCC is expected to generate sequences until the end of the line, however, a plain increase in the number of iterations might lead to higher computational costs. Thus, we need to continue generating or collecting hypotheses only in scenarios that look promising enough. Our experiments showed that $n=20$ and $k=3$ are the best values in our setup.

Insight 8: A dynamic approach in beam search with carefully selected stopping criteria allows generating longer suggestions while keeping the mean computation time almost the same.

The sequence generation algorithm is used both for studying preliminary metrics during training and inside the FLCC plugin.

Full Line Code Completion is a computationally intensive plugin, because it involves the execution of a neural network. To achieve the best possible efficiency and speed, we employed the latest developments of inference engines for Transformer-based models — native CPU runtime for model inference. The implementation is as follows:

• •

  Alongside FLCC in the running IDE, a separate instance of a native C++ server is launched.

• •

  The server is a wrapper for the model inference runtime llama.cpp.

• •

  The server manages all model-related inference computations.

• •

  gRPC [59] protocol is used to establish the connection between the plugin and the server.

In mid-2023, we compared (1) ONNX Runtime with GPT-2-like models and (2) llama.cpp with LLaMA-like models, both with our native server implementation on MacOS with an M2 Max processor. The benchmark showed that the mean time from the invocation of completion to the showing of the suggestion with ONNX Runtime was approximately 75 ms. For llama.cpp, the same procedure took approximately 100 ms. Meanwhile, llama.cpp continued to be developed, and the same comparison dramatically changed in the late 2023: it became possible to infer the same LLaMA model in just 50 ms instead of 100 ms, so our development focus changed accordingly.

Insight 9: It is crucial to create a generalized infrastructure and regularly benchmark quickly developing frameworks to achieve the best possible product quality.

The last unique feature of our implementation of inference is the caching of the model’s hidden state. The beam search algorithm that we use consists of two steps: context processing and iterative generation. Context processing, which computes token interrelations, is observed to be 3–10 times slower than generation (influenced by hardware and numerical precision).

To optimize the algorithm, we initialize the model at 50% maximum context capacity, enhancing the speed of generation in two cases:

• •

  When the caret is moved within the initially processed context, generation proceeds with existing hidden states, conserving computational resources despite a potentially reduced context.

• •

  Continuous code writing beyond the initial position leverages existing context computations, updating hidden states only for new code without extra processing.

To match the current context with the cached one, we maintain a list of integers corresponding to cached tokens.

This approach significantly improves efficiency by reducing redundant computations in the generation pipeline. According to our observations, this allows to significantly increase model inference speed in over 90% of the real-world scenarios, while keeping the quality at almost the same level.

Insight 10: Sophisticated caching strategies allow to speed up code completion without significant quality loss.

In general, one can formulate the following two tasks for Full Line Code Completion:

• •

  Generate code suggestions so that the developer writes code easier.

• •

  Do not disturb the developer with incorrect or unnecessary code suggestions.

Post-processing mostly targets the second goal: after the initial set of hypotheses is generated, we should decide whether it contains a useful suggestion, and if it does, which one it is.

Filtering. Our filtering mechanisms include:

• •

  Safety. Removes dangerous code (e.g., rm -rf, DROP DATABASE), secrets (emails, access keys, etc. [60]), and profanity.

• •

  Low scored suggestions. Eliminates suggestions that have a beam search score lower than a certain value.

• •

  Incorrect code. Filters out unresolved references, type mismatches, incorrect argument counts, etc.

The Incorrect code filter is the slowest one since it involves running code inspections, and thus it is launched the last and is restricted with maximum execution time. Suggestions with undefined correctness status are not shown to the user. Our research for the 2023.3 PyCharm Pro release showed that FLCC benefits from filtering out incorrect code. Only 1% of valuable suggestions are lost, and as a trade-off for such a negligible loss, we establish the distinguishing feature of Full Line Code Completion — suggesting only valid code.

Transformations. After redundant suggestions are filtered out, we are left with a set of “surviving” hypotheses, among which we pick the one with the highest beam search score. The chosen suggestion might still need to be tweaked a little bit, e.g., if not all paired symbols like brackets and quotes are correctly closed. For this reason, in a final step before showing the code suggestion, we fix such incomplete code closures.