代码库检索的准确性上限是什么?

为了自动确定代码库中最重要的上下文,我们刚刚发布了 Continue 的代码库检索功能的 v1 版本。您可以提出诸如“我应该在哪里为服务器添加一个新端点?”、“构建此项目使用了哪些脚本?”或“我在哪里编辑用户个人资料页面?”之类的问题。要回答这些问题,必须解决以下问题:给定用户查询(如上所述),从整个代码库中找到 N 个最有用的标记。
这是一个极其开放的搜索问题,提供了多种潜在解决方案。一开始,一个简单的测试查询会指出明显的问题:整个文件太长,无法放入上下文;有些文件是空的;而且文件顶部的几十个导入常常是不相关的。但很快,明显的解决方案就会饱和,需要自动且明确地比较不同策略。是时候引入一个指标了。
\[ F_1 = \frac{P\cdot R}{P + R} \]
F1 分数是衡量 \(P\)(精确率,即所选文件中相关文件的百分比)和 \(R\)(召回率,即相关文件中所选文件的百分比)的指标。总的来说,它让我们了解检索系统的准确性。在优化检索管道时,我们可以分别测试每个阶段的 F1 分数,以了解哪个阶段是瓶颈。
虽然进展尚早(我们尚不知道这篇博客文章标题的答案!),而且我们还没有足够大的基准来分享有用的数字,但定性进展已经很大,并且还有数不清的改进空间!在这里,我将对我们看到的可能性以及我们迄今为止所做的工作进行高层次的概述(请注意加粗的文字, 表示我们目前所处的状态)。将来,我将跟进我们实验的路线,并分享它们对 F1 分数的具体影响。
评估
如上所述,构建一个好的检索管道最重要的一点是能够衡量其质量。以下是我们现在以及将来会用来测试系统的一些方法。
朴素:尝试少量输入并检查包含的内容
基础:编码少量示例问题以及期望包含的文件。基于此计算 F1 分数。
更好:做同样的事情,但确保从不同的代码库和问题中选择示例
未来:生成或找到大型数据集并进行昂贵的评估
分块
构建检索系统时遇到的第一个细节是,有些文件太长,无法放入 LLM 的上下文长度限制,因此必须将其缩短为可管理的“分块”。此外,OpenAI 的 ada-002 嵌入模型通常在输入长度为 512 个或更少标记时表现更好。为了将潜在的大型代码文件强制分割成更小的部分,我们需要一个分块策略。
朴素:截断文件
基础:将文件分割成等长的片段
更好:我们的分块策略使用 tree-sitter
来解析任何编程语言的 AST,从而能够理解类和函数的结构。其工作原理如下:
- 检查文件是否已经适合,如果适合则使用整个文件。
- 否则,提取所有顶级函数和类。
- 对于每个提取的内容,检查其是否适合。如果适合,则添加。
- 如果不适合,则截断子方法的内容以构建如下所示的分块:
class CodebaseIndex:
directory: str
def __init__(self, directory: str):
...
async def exists(self) -> bool:
"""Returns whether the index exists (has been built)"""
...
async def build(
self,
ignore_files: List[str] = [],
) -> AsyncGenerator[float, None]:
"""Builds the index, yielding progress as a float, 0-1"""
...
async def update(self) -> AsyncGenerator[float, None]:
"""Updates the index, yielding progress as a float, 0-1"""
...
async def query(self, query: str, n: int = 4) -> List[Chunk]:
"""Queries the index, returning the top n results"""
...
- 递归地处理所有被隐藏的子方法,为它们创建自己的分块。
未来:在每个分块的顶部包含某些元数据可能会有帮助,例如ctags,用于检索、重新排序或两者。我们还可以包含由语言模型生成的代码片段摘要,或可能被问到的关于该文件的问题示例。
检索
为了选择最重要的分块,我们遵循一个两阶段过程。第一阶段是检索,在此阶段我们收集大量结果。从根本上说,这个过程力求实现高召回率。
嵌入模型
虽然句子转换器已经存在一段时间,并且仍然能够达到顶尖性能,但有多种方法可以针对特定用例改进它们。
朴素:使用 OpenAI 的 ada-002 嵌入模型或现成的句子转换器模型(我们使用 all-MiniLM-L6-v2 进行完全本地的嵌入)
基础:学习一个矩阵,用其乘以向量嵌入,以强调对您的用例特别有用的维度
更好:训练一个完整的自定义嵌入模型
未来:为不同情况学习多个矩阵,可能每个代码库对应一个
假设文档嵌入 (HyDE)
在比较用户的问题和候选代码片段时,应用相似性搜索来查找最相关的分块。但存在一个重要的不对称性:用户的问题是一个问题,而必需的代码片段是代码——这是两类不同的文本,可能在嵌入模型的潜在空间中是分离的。为了解决这个问题,最初的 HyDE 论文意识到,您可以要求语言模型生成一个假想的响应,并将其与嵌入的文档进行比较。
朴素:使用原始查询执行相似性搜索
基础:请求一个虚构的代码片段并将其嵌入以进行比较
更好:请求一个特定语言的代码片段
未来:使用代码库结构等信息,在编写假设文档时为模型提供额外知识。例如,如果看到 main.tsx
和一个 components
文件夹,模型可以估计它应该编写一个 React 代码片段。
其他检索源
使用向量嵌入称为密集检索。虽然它是一种很好的捕获不完全匹配词语的相关概念的方法,但通过与其他搜索方法结合,我们可以做得更好。以下是我们目前正在实验的一些方法:
精确搜索
使用 ripgrep 搜索单词或短语的精确匹配。可以选择要求 LLM 生成要搜索的关键词。
Meilisearch
Meilisearch 是一个稍微模糊的搜索引擎,可以处理拼写错误和其他微小差异,同时仍然提供极快的搜索来增强密集检索。
工具使用
让 LLM 自行编写查询。例如,可以提示它输出要搜索的关键词、要匹配的 AST 模式,或者给定代码库的 ctags 后要调查的文件夹。
标注
投入预处理以描述每个代码片段的特征。例如,它是函数、类还是变量?它是用什么编程语言编写的?有注释吗?函数是异步的吗?甚至可以让代码库的作者亲自参与这个标注过程。然后可以使用标注来过滤或排序结果。
集成检索
我们如何结合所有这些来源?
朴素:从每个索引检索相同数量的分块,使其总数达到所需
基础:从每个来源选择略多一些结果,并更重地加权从多个来源检索到的结果
更好:确保每个来源返回定量、标准化的结果,以便可以将每个分块的总和用作总分
未来:使用回归学习模型来结合分数
重新排序
在检索到最初的一组分块后,我们希望执行第二步将其精简到顶部几个结果,力求实现高精确率。在我们的初步测试中,通过从约 50 个初始结果减少到约 10 个最终结果,我们取得了不错的效果。随着我们找到更可靠的检索和重新排序方法,以及模型上下文长度的增加,这个数字可能会改变。
朴素:将所有结果展示给 gpt-3.5-turbo,并要求其选出最佳 10 个(可能需要分组并行运行以将所有结果放入上下文窗口)
基础:并行请求每个代码片段,询问一个表示其是否相关的单标记回答:是或否
更好:使用标记“是”返回的 logprobs 来获得连续的排名
未来:训练自定义重新排序模型。这很可能比功能齐全的 LLM 需要更少的参数。例如,Cohere 提供这样的模型。
使用本地模型重新排序
我们目前使用的重新排序过程并行地向模型发送许多请求。对于本地模型,我们没有大量批处理的便利,因此需要采用其他方法。一种可能性是使用句子转换器作为交叉编码器。
提示
管道的最后一步是实际回答问题!如果您只是将所有上下文和原始输入传递给提示,您可能会得到一个不错的响应。但是,由于使用多个文档来回答高层次问题是您可以提供给 LLM 的提示的一个非平凡子集,因此一些特定的指令可能会有所帮助。例如,“请使用提供给您的信息回复查询,引用具体的文件和代码片段,并且不要讨论未向您展示的任何内容。”
朴素:仅使用代码片段和问题进行提示
基础:包含上述基本指令
更好:添加在没有足够信息回答问题时如何操作的指令。例如,“如果没有足够信息回答问题,请建议用户可以在哪里查找以了解更多。”
未来:包含一些额外信息,例如仅显示包含文件最近共同祖先目录的文件树。
未来工作
以下是我们很乐意尝试的一些更具推测性的想法:
处理
收到响应后,我们目前处理任何文件引用,使其成为指向文件本身的单击链接。我们在这里的下一个任务是将代码中的对象(如函数和类)也链接起来。更进一步,我们可以鼓励模型输出可链接的文本。
图关系
代码具有固有的图结构,可以使用语言服务器协议的“跳转到定义”功能或分析提交中的文件关系来理解。如果检索过程选择了 5 个不同的文件,而所有这些文件都导入了某个共同函数,那么是否也应该包含该函数?如果选定的 README 部分链接到某个文件,我们是否应该包含该文件?如果两个文件在提交中经常一起编辑,它们是否应该在检索中经常一起返回?
使用用户反馈
如果选择了不相关的文件,并且用户点击了“踩”按钮表示,我们如何使用这些信息进行改进?目前,我们想到三种方法:
- 微调重新排序模型。这只是构建一个数据集的问题,其中提示与模型回复“是”的提示相同,但正确响应是“否”。
- 更新计算的嵌入。这里的想法是朝着用户请求的方向“惩罚”向量嵌入,以便将来将其视为相似性较低。这可以通过从文档向量中减去用户输入向量的某个倍数来实现。嵌入会随着时间在特定代码库中针对特定用户或团队而改进,但不会整体改进。
- 更新嵌入函数。一个已知技巧是通过使用额外的矩阵乘法来强调潜在空间中对您的用例更重要的方面,从而改进嵌入。我们不是更新向量嵌入本身,而是通过更新这个矩阵来更新其计算过程(并重新计算它们)。
预先定制
开发者改进其检索体验的另一种方法是投入时间编写“LLM 的文档”。这可以采取多种形式,在大多数情况下,与现有文档或其他常见实践没有太大区别。一些想法:
- 除了
.gitignore
之外,选择某些文件或文件夹进行忽略。我们目前支持使用名为.continueignore
的文件。 - 编写一个“FAQ”文件,其中包含问题列表以及应该用于回答问题的关联文件,以及问题的答案。然后可以将每个问题进行嵌入,如果通过相似性搜索选择了它,我们可以将其替换为答案和链接的文件。
- 编写每个重要目录的摘要,以便 LLM 可以阅读这些摘要来决定通过工具使用探索哪些文件夹,或者以便可以检索摘要并将其用作初始检索过程的一部分。
- 使用特定关键词预先标注文件和文件夹。例如,我可以将一个目录标记为与“服务器”相关,而另一个标记为“webapp”。
从其他上下文来源检索
没有理由只限制在代码上——Continue 预装了一些上下文提供程序,可以轻松引用 GitHub 问题、当前提交的差异、文档网站或通过我们的 SDK 配置的任何其他来源。所有这些都可以进行嵌入并包含在检索过程中。
结论
我们在检索方面的工作尚处于早期阶段,但这将是未来重要的工作领域,我们很高兴继续分享我们的进展更新。如果您有兴趣了解以上哪些想法影响最大,请尽快回来查看,如果您有任何自己的想法,也请随时联系我们!