迈向真实世界的软件智能体:如何将语义高亮功能融入智能体编程?
引言
我们训练了一个代码语义高亮模型SWE-Pruner并开源了配套的Agentic接入框架,在真实的Coding Agent多轮任务(SWE-bench, SWE-QA)上,在保持性能不变甚至提升的情况下提供多轮场景下30%甚至更多的token开销减少。该模型基于语义理解,自动识别并高亮检索到的文档中语义相关的句子。而框架作为agentic的接入层,可以轻松接入claude code,openhands等SOTA agent系统
模型发布
- HuggingFace: https://huggingface.co/ayanami-kitasan/code-pruner
- License: MIT (commercial-friendly)
- Architecture: 0.6B Dual-Head model based on Qwen3-Reranker-0.6B
- Supported Programming Languages: Python
- Github: https://github.com/Ayanami1314/swe-pruner
- arxiv:https://arxiv.org/abs/2601.16746
在本文中,我们将分享我们的技术方案。
问题:编码代理的上下文膨胀
在先前的RAG工作中,上 下文的粗粒度膨胀已经是一个很严重的问题 https://huggingface.co/blog/zilliz/zilliz-semantic-highlight-model
在生产环境中的 RAG 系统中,一个典型的查询会检索 10 个文档,每个文档包含数千个词元,每次查询消耗数万个词元。问题在于: 只有几十个句子真正包含相关信息 ,其余的都是噪声,这会增加成本并降低答案质量。
然而,对于真实环境下的软件工程任务,问题甚至更加严重:一个典型的查询可能涉及上百个代码文件,而一些大型项目中,巨型文件的代码token包含的远不止数千个,现代的编程agent却仍然在广泛使用关键词匹配的方法,即grep,来对大型代码仓库进行探索,匹配项的过分庞大使得现代agent中都配备了相当多的上下文硬截断逻辑


在人工智能代理场景中,这个问题会变得更加严重,因为查询是经过推理和分解后的复杂指令。传统的高亮显示方式只是机械地标记匹配的词语,却会遗漏真正有价值的分析结论。
这就迫切需要一种有针对性的高亮模型,该模型仅保留上下文相关的代码,并高亮显示它们,同时剪掉所有无关的噪声内容——这种技术也被广泛称为上下文剪枝。类似的纯文本任务也被Provence为代表的一系列工作较为完善地解决,但对于究竟如何在代码任务上引入语义高亮模型,社区尚未给出一个答案。
现有模型的困境
我们调研了现有的解决方案,但发现它们并不完全符合我们的需求。
Provence/XProvence,OpenProvence,zilliz/semantic-highlight-bilingual-v1沿用Provence的思路继续优化Provence的上下文长度和多语言性能。然而,它们并没有指出在代码等非纯文本场景下,如何构造合理的粗细粒度,如何满足额外的AST约束等。

另一个困境是,现有的模型沿用BGE的类BERT管道,但大部分现代代码嵌入模型都是LLM架构,decoder-only的模型如何设计裁剪头来处理特征的交互和输出的合理性仍是一个待解决的问题——BERT自然双向注意,但decoder模型需要某些“代偿”因果掩码单向注意的方案。
我们的选择
既然市面上没有这样的代码剪枝模型,那我们自己训练一个
第一个问题是在代码中,如何将纯文本里面的细粒度“句子”转换成代码的对应概念?
我们提出两套方案,基于代码行和基于AST statements,鉴于数据的生产难度,我们选择了以行取代句子作为细粒度
第二个问题是,选择什么样的模型?
BGE-m3是非常好的模型,但它没有针对代码数据进行大规模预训练;BAAI的另一个优秀模型是BGE-Code-v1,但社区对它的适配相对较少
我们注意到了社区的另一个工作modernbert,然而出于数据分布或是模型表达能力等原因,我们并没有在这个模型上训练成功
我们的猜想是——代码任务对基础模型的能力提出了可能超越大部分旧BERT模型的要求
最后我们采用的是Qwen的新模型,Qwen/Qwen3-Reranker-0.6B, 它的海量预训练数据和世界知识产生了相当好的代码嵌入,0.6B的大小又和传统bert类模型开销接近,LLM模型作为基底的另一个好处是自然完成了多语言query的适配。
类似Qwen这种LLM based reranker依然采用Casual Mask,这对token打分实际造成了一些影响,添加一个Full self-attention模块能让模型在测试集上的F1上升接近0.1
我们实验了不同的架构设计和特征trick,发现几个训练技巧对这个任务比较好用:
- 冻结大部分层:我们发现只解冻最后两层和全参数微调差距极小,而冻结之后更能避免过拟合,训练速度也更快
- 特征融合:对于LLM不同层数的特征作用,在先前已经有了相当多的工作进行讨论,一般认为浅层会有更多的语法信息,而深层则更重语义,将浅层、中间层和深层做简单的concat接入下游即可显著提升token打分头性能(我们使用的是7, 14和28层)
- Sample Loss Norm: 代码片段的长短有着很大的差别,更关键的是,与文本领域大部分的信息都是稀疏不同,保留的代码行既有稀疏的片段(大部分行被删除)、又有稠密的片段(大部分行被保留),且长度上也有比较大的区别。如果简单在token level计算,模型训练相对不稳定,将每一个样本的loss使用样本的长度归一化能更好的平衡长短样本
- CRF Layer:我们特殊设计了Token打分头 的模型结构,使用CRF层进行序列建模,取代BERT风格的mlp层,实验证明对于改进模型的稳定性有一定帮助
模型架构&训练
我们设计了一个如下图的模型架构,它在思路上借鉴了Provence,但细节做了大量改动

此外,我们使用了和Provence相同的做法保留了基础模型的原始rerank头, 即图中的Path B,这种Rerank-While-Prune的方式进一步降低了在RAG管道等场景接入模型的成本。
最终,我们使用8 * H100 训练了3个epoch,用时大约是4个小时。
我们的具体训练超参数可以在论文中找到。
推理过程:
推理过程非常直接:
- 将输入拼接为
指令 + 查询 + 代码 - 对上下文中的每个 token 进行打分(范围在 0 到 1 之间)
- 对每行内所有 token 的得分取平均,得到该行的得分
- 根据阈值保留得分高的行,同时移除得分低的行
一个有趣的细节是阈值的选择,在Provence的开源模型之中,他们推荐将threshold设置为0.1来保守裁剪;而我们阈值调优实验发现,和训练保持一致的0.5阈值就能在下游任务中达到最佳性能,而0.1似乎过于保守——0.3已经会保留非常多的部分,这也可能是代码数据的分布特点和文本不同所导致的,在下文中会详细说明。
训练数据:如何构造“部分相关”的查询-代码对?
我们使用Qwen-Coder-30B-A3B-Instruct,从github 2025年的新仓库上自行构建数据(再次感谢qwen团队的优秀工作!),在来自6k个仓库、19w个文件的20w条代码片段上,使用了9个种子任务构建数据,并使用Qwen3-Next-80B-A3B-Thinking进行质量筛选,最后得到了约5万条高质量“部分关联query code对”
为什么是“部分相关”?
与文本领域不同的是,代码领域的QA数据集相对较少,缺少类似ms macro这样规模的query,code对;另一个问题是,由于缺少搜索引擎商开源的数据,现有的code query数据集的query往往分布偏窄,常是“这个函数的功能是什么?”“这段代码的运行结果是什么?”这种与整段代码相关的查询
但我们想要在代码之中引入semantic highlight功能,需要的是对“部分代码”的查询,例如,“这一条if分支会产出什么结果?” “这个类的属性A的作用是什么?”,这样我们才能划分出highlight的部分
也就是说,我们最后的数据集格式会是类似 <查询,原始代码,highlight部分> 的三元组