Large language models (LLMs) have shown impressive in-context learning (ICL) ability in code generation. LLMs take a prompt consisting of requirement-code examples and a new requirement as input, and output new programs. Existing studies have found that ICL is highly dominated by the examples and thus arises research on example selection. However, existing approaches randomly select examples or only consider the textual similarity of requirements to retrieve, leading to sub-optimal performance. In this paper, we propose a novel learning-based selection approach named LAIL (LLM-Aware In-context Learning) for code generation. Given a candidate example, we exploit LLMs themselves to estimate it by considering the generation probabilities of ground-truth programs given a requirement and the example. We then label candidate examples as positive or negative through the probability feedback. Based on the labeled data, we import a contrastive learning objective to train an effective retriever that acquires the preference of LLMs in code generation. We apply LAIL to three LLMs and evaluate it on three representative datasets (e.g., MBJP, MBPP, and MBCPP). LATA outperforms the state-of-the-art baselines by 11.58%, 6.89%, and 5.07% on CodeGen, and 4.38%, 2.85%, and 2.74% on GPT-3.5 in terms of Pass@1, respectively.

翻译：大语言模型（LLMs）在代码生成中展现出令人印象深刻的上下文学习（ICL）能力。LLMs接收由需求-代码示例和新需求组成的提示作为输入，并输出新的程序。现有研究发现，ICL在很大程度上由示例主导，由此催生了关于示例选择的研究。然而，现有方法随机选择示例或仅考虑需求的文本相似性进行检索，导致性能次优。本文提出一种名为LAIL（LLM感知上下文学习）的新型基于学习的选择方法，用于代码生成。针对候选示例，我们利用LLMs自身通过考虑给定需求和示例时的真实程序生成概率来评估它。然后，我们根据概率反馈将候选示例标记为正例或负例。基于标记数据，我们引入对比学习目标来训练一个有效检索器，使其获取LLMs在代码生成中的偏好。我们将LAIL应用于三个LLMs，并在三个代表性数据集（如MBJP、MBPP和MBCPP）上进行评估。在Pass@1指标上，LAIL在CodeGen上分别超过最先进基线11.58%、6.89%和5.07%，在GPT-3.5上分别超过4.38%、2.85%和2.74%。