aiXcoder-7B-v2: Training LLMs to Fully Utilize the Long Context in Repository-level Code Completion

Large Language Models (LLMs) have shown promising results in repository-level code completion, which completes code based on the in-file and cross-file context of a repository. The cross-file context typically contains different types of information (e.g., relevant APIs and similar code) and is lengthy. In this paper, we found that LLMs struggle to fully utilize the information in the cross-file context. We hypothesize that one of the root causes of the limitation is the misalignment between pre-training (i.e., relying on nearby context) and repo-level code completion (i.e., frequently attending to long-range cross-file context). To address the above misalignment, we propose Code Long-context Alignment - COLA, a purely data-driven approach to explicitly teach LLMs to focus on the cross-file context. Specifically, COLA constructs a large-scale repo-level code completion dataset - COLA-132K, where each sample contains the long cross-file context (up to 128K tokens) and requires generating context-aware code (i.e., cross-file API invocations and code spans similar to cross-file context). Through a two-stage training pipeline upon COLA-132K, LLMs learn the capability of finding relevant information in the cross-file context, thus aligning LLMs with repo-level code completion. We apply COLA to multiple popular LLMs (e.g., aiXcoder-7B) and extensive experiments on COLA-132K and a public benchmark - CrossCodeEval. Our experiments yield the following results. 1) Effectiveness. COLA substantially improves the performance of multiple LLMs in repo-level code completion. For example, it improves aiXcoder-7B by up to 19.7% in exact match. 2) Generalizability. The capability learned by COLA can generalize to new languages. 3) Enhanced Context Utilization Capability. We design two probing experiments, which show COLA improves the capability of LLMs in utilizing the information in cross-file context.