基于检索的代码生成：用相似代码片段辅助生成

text // 类似片段: // def load_json(path): ... // 当前任务: 编写读取JSON文件的函数

生成器可直接“模仿”片段中的模式。

2. **交叉注意力融合（Cross-Attention）**  
生成器的解码器对检索片段执行额外的交叉注意力，动态提取有用信息。常见于Encoder-Decoder架构（如RAG模型）。  
- 检索片段被编码成键值对，解码时每个生成步可自适应关注不同片段。

3. **基于提示的融合（Prefix-Tuning）**  
将检索片段编码成连续的提示向量，插入到生成模型的嵌入层。  
- 优点：保持模型结构不变，仅学习一个映射层将检索结果转为虚拟token。

### 代表性工作  
- **REDCODER**：早期检索增强代码生成模型，使用双塔检索+CodeGPT生成，拼接片段到输入。  
- **ReCode**：迭代检索：用模型初步生成的结果作为新查询进行二次检索，逐步精炼。  
- **RepoCoder**：多级检索，先从本地文件检索，再从全局代码库检索，适合仓库级代码补全。  
- **RAG for Code**：将RAG架构迁移至代码，端到端训练检索与生成模块。

## 实践示例：搭建一个简单的检索增强代码生成系统  

我们将用Python实现一个最小系统：使用BM25检索相似函数，并将它们作为上下文喂给CodeGen模型。  
**准备**：安装 `pyserini`（或直接用 `rank_bm25`）、`transformers`。

### 步骤1：构建代码库索引  
假设我们有一个函数列表作为语料库。  
```python
from rank_bm25 import BM25Okapi
import re

# 示例语料——若干函数定义
corpus = [
 "def read_json(file_path: str) -> dict:",
 "    import json",
 "    with open(file_path, 'r') as f:",
 "        return json.load(f)",

 "def write_json(data: dict, file_path: str):",
 "    import json",
 "    with open(file_path, 'w') as f:",
 "        json.dump(data, f, indent=4)",

 "def fetch_url(url: str) -> str:",
 "    import requests",
 "    response = requests.get(url)",
 "    return response.text",
 # ... 更多函数
]

# 简单分词：按空格和标点分割，保留标识符
def tokenize(code: str) -> list:
 return re.findall(r"[A-Za-z_]\w*|[^\s\w]+", code)

tokenized_corpus = [tokenize(doc) for doc in corpus]
bm25 = BM25Okapi(tokenized_corpus)

步骤2：定义检索函数

输入查询描述，返回Top-3最相似函数。

def retrieve(query: str, top_k=3):
    query_tokens = tokenize(query)
    scores = bm25.get_scores(query_tokens)
    # 按分数降序取索引
    top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_k]
    return [corpus[i] for i in top_indices]

步骤3：将检索结果注入生成模型

使用预训练的CodeGen-350M（或任何文本-代码生成模型）。

from transformers import pipeline

generator = pipeline('text-generation', model='Salesforce/codegen-350M-mono')

def generate_with_retrieval(prompt: str, top_k=3):
    retrieved_snippets = retrieve(prompt, top_k)
    # 构造增强后的prompt
    context = "// Similar code examples:\n"
    for idx, snippet in enumerate(retrieved_snippets):
        context += f"// Example {idx+1}:\n{snippet}\n"
    full_prompt = context + "\n// Task: " + prompt + "\n"
    result = generator(full_prompt, max_new_tokens=100, do_sample=False)[0]['generated_text']
    return result

步骤4：测试

query = "load json from a file"
output = generate_with_retrieval(query)
print(output)