1
# RAG-Sequence
2
def rag_sequence(question, retriever, generator, top_k=5):
3
    docs = retriever.retrieve(question, top_k=top_k)
4

5
    answers = []
6
    for doc in docs:
7
        # 每个文档独立生成完整答案
8
        answer = generator.generate(question, context=doc)
9
        score = retriever.score(question, doc) * generator.score(answer)
10
        answers.append((answer, score))
11

12
    # 边际化：选择得分最高的答案
13
    best_answer = max(answers, key=lambda x: x[1])
14
    return best_answer[0]

1
# RAG-Token（更常用）
2
def rag_token(question, retriever, generator, top_k=5):
3
    docs = retriever.retrieve(question, top_k=top_k)
4

5
    answer_tokens = []
6
    for step in range(max_length):
7
        token_probs = []
8
        for doc in docs:
9
            # 每个文档对下一个 token 的概率
10
            prob = generator.next_token_prob(question, doc, answer_tokens)
11
            doc_weight = retriever.score(question, doc)
12
            token_probs.append(prob * doc_weight)
13

14
        # 边际化后选择 token
15
        next_token = merge_and_select(token_probs)
16
        answer_tokens.append(next_token)
17

18
    return tokens_to_text(answer_tokens)

1
# 查询重写方法
2
def rewrite_query(original_query, llm):
3
    # 方法1：LLM 改写
4
    prompt = f"请将以下问题改写为更适合搜索引擎的查询：{original_query}"
5
    rewritten = llm.generate(prompt)
6

7
    # 方法2：HyDE（Hypothetical Document Embedding）
8
    # 先让模型生成假设性答案，用答案去检索
9
    hypothetical_answer = llm.generate(f"请回答：{original_query}")
10
    hyde_query = embed(hypothetical_answer)
11

12
    # 方法3：多查询扩展
13
    sub_queries = llm.generate(f"将以下问题分解为3个子问题：{original_query}")
14

15
    return rewritten, hyde_query, sub_queries

1
# 混合检索
2
def hybrid_retrieve(query, dense_index, sparse_index, alpha=0.7):
3
    # 稠密检索：语义相似度
4
    dense_results = dense_index.search(embed(query), top_k=20)
5

6
    # 稀疏检索：关键词匹配（BM25）
7
    sparse_results = sparse_index.search(query, top_k=20)
8

9
    # Reciprocal Rank Fusion (RRF) 融合
10
    fused = reciprocal_rank_fusion(
11
        dense_results, sparse_results,
12
        k=60  # RRF 参数
13
    )
14

15
    return fused[:10]
16

17
def reciprocal_rank_fusion(*rankings, k=60):
18
    """RRF 融合多个排序列表"""
19
    scores = {}
20
    for ranking in rankings:
21
        for rank, doc in enumerate(ranking):
22
            scores[doc] = scores.get(doc, 0) + 1.0 / (k + rank + 1)
23
    return sorted(scores.keys(), key=lambda x: scores[x], reverse=True)

1
# 重排序
2
def rerank(query, documents, cross_encoder):
3
    """使用交叉编码器重排序"""
4
    pairs = [(query, doc.text) for doc in documents]
5
    scores = cross_encoder.predict(pairs)
6

7
    # 按交叉编码器分数重新排序
8
    ranked = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)
9
    return [doc for doc, score in ranked]

1
# 分块策略对比
2
chunking_strategies = {
3
    "固定长度": {
4
        "方法": "每 512 token 切一块",
5
        "优点": "简单",
6
        "缺点": "可能切断语义完整性",
7
    },
8
    "递归字符分割": {
9
        "方法": "按段落 → 句子 → 字符递归分割",
10
        "优点": "保持段落完整性",
11
        "缺点": "块大小不均匀",
12
    },
13
    "语义分块": {
14
        "方法": "用 embedding 相似度确定边界",
15
        "优点": "语义完整",
16
        "缺点": "计算开销大",
17
    },
18
    "父-子文档": {
19
        "方法": "小块检索，大块生成",
20
        "优点": "检索精准 + 上下文丰富",
21
        "缺点": "索引复杂",
22
    },
23
}

1
# NTK-aware 位置缩放
2
def ntk_aware_rope_extension(base_theta, scale_factor, dim):
3
    """通过调整 RoPE 频率实现外推"""
4
    # 原始频率
5
    freqs = 1.0 / (base_theta ** (torch.arange(0, dim, 2) / dim))
6

7
    # NTK 缩放：低频分量缩放更多，高频分量保持不变
8
    # 这保证了局部位置关系不被破坏
9
    scale = scale_factor ** (dim / (dim - 2))
10
    freqs_scaled = freqs / scale
11

12
    return freqs_scaled
13

14
# YaRN (Yet another RoPE extensioN) 方法
15
def yarn_rope(base_theta, scale_factor, dim, gamma=0.1, beta=0.5):
16
    """YaRN：结合 NTK 和注意力温度调整"""
17
    freqs = 1.0 / (base_theta ** (torch.arange(0, dim, 2) / dim))
18

19
    # 三段式频率缩放
20
    # 高频段：保持不变
21
    # 中频段：插值
22
    # 低频段：完全缩放
23
    critical_freq = gamma / scale_factor
24

25
    freqs_scaled = torch.where(
26
        freqs < critical_freq,
27
        freqs / scale_factor,  # 低频：缩放
28
        freqs                   # 高频：不变
29
    )
30

31
    # 注意力温度调整
32
    attention_scale = 1.0 + 0.1 * math.log(scale_factor)
33

34
    return freqs_scaled, attention_scale

1
# Sliding Window Attention (Mistral 使用)
2
def sliding_window_attention(Q, K, V, window_size=4096):
3
    """每个 token 只关注局部窗口"""
4
    seq_len = Q.size(1)
5
    output = torch.zeros_like(Q)
6

7
    for i in range(seq_len):
8
        start = max(0, i - window_size + 1)
9
        # 只计算窗口内的注意力
10
        local_Q = Q[:, i:i+1, :]
11
        local_K = K[:, start:i+1, :]
12
        local_V = V[:, start:i+1, :]
13

14
        attn = torch.matmul(local_Q, local_K.transpose(-2, -1))
15
        attn = F.softmax(attn / math.sqrt(Q.size(-1)), dim=-1)
16
        output[:, i:i+1, :] = torch.matmul(attn, local_V)
17

18
    return output
19

20
# 通过多层堆叠，窗口外的信息可以逐层传递
21
# 窗口 4096 × 层数 32 = 理论感受野 131K

1
# Mixture of Depths (MoD)
2
def mixture_of_depths(layer, x, capacity_ratio=0.5):
3
    """只让部分 token 参与完整注意力计算"""
4
    # 路由器判断哪些 token 需要"深度思考"
5
    importance_scores = router(x)  # [batch, seq_len, 1]
6

7
    # 选择前 50% 重要的 token
8
    topk_mask = importance_scores.topk(
9
        int(x.size(1) * capacity_ratio), dim=1
10
    )
11

12
    # 重要 token 做完整注意力
13
    important_tokens = x[topk_mask.indices]
14
    processed = layer.attention(important_tokens)
15

16
    # 不重要的 token 用残差连接跳过
17
    output = x.clone()
18
    output[topk_mask.indices] = processed
19

20
    return output

1
# Needle in a Haystack 测试
2
def needle_in_haystack_test(model, context_length, needle_position):
3
    """在长文本中的特定位置插入关键信息，测试模型能否检索到"""
4

5
    # 1. 构建长上下文（填充无关文本）
6
    haystack = generate_filler_text(context_length - 200)
7

8
    # 2. 在指定位置插入"针"（关键信息）
9
    needle = "独角兽的喜欢的食物是草莓味的棉花糖"
10

11
    # 3. 将针插入到指定位置
12
    full_context = insert_at_position(haystack, needle, needle_position)
13

14
    # 4. 提问
15
    question = "独角兽喜欢什么食物？"
16

17
    # 5. 测试模型能否回答
18
    answer = model.generate(f"{full_context}\n\n{question}")
19

20
    return needle in answer  # 是否正确检索到针

1
# 混合方案：RAG 为主 + Long Context 辅助
2
def hybrid_rag_long_context(query, knowledge_base, model):
3
    # 1. RAG：从大规模知识库中检索相关文档
4
    retrieved_docs = retrieve(query, knowledge_base, top_k=10)
5

6
    # 2. 上下文评估：检索结果是否能放进上下文窗口
7
    total_tokens = count_tokens(retrieved_docs)
8

9
    if total_tokens <= model.context_limit * 0.7:
10
        # 3a. Long Context 路径：直接放入全部检索结果
11
        context = format_docs(retrieved_docs)
12
        answer = model.generate(query, context=context)
13
    else:
14
        # 3b. RAG 优化路径：压缩后放入
15
        # 先用小模型做摘要压缩
16
        summaries = [compress(doc, ratio=0.3) for doc in retrieved_docs]
17
        context = format_docs(summaries)
18
        answer = model.generate(query, context=context)
19

20
    # 4. 引用验证
21
    answer_with_citations = add_citations(answer, retrieved_docs)
22

23
    return answer_with_citations

1
# 向量数据库性能参考（百万级向量）
2
performance = {
3
    "Milvus": {"QPS": "~10,000", "延迟": "<10ms", "召回率": ">95%"},
4
    "Qdrant": {"QPS": "~8,000", "延迟": "<5ms", "召回率": ">95%"},
5
    "Weaviate": {"QPS": "~5,000", "延迟": "<15ms", "召回率": ">93%"},
6
    "ChromaDB": {"QPS": "~2,000", "延迟": "<20ms", "召回率": ">90%"},
7
}

1
# Embedding 模型选择建议
2
def choose_embedding_model(scenario):
3
    if scenario == "生产环境、质量优先":
4
        return "text-embedding-3-large"  # OpenAI，成本较高
5
    elif scenario == "中文为主、开源部署":
6
        return "bge-m3"  # BAAI，多语言
7
    elif scenario == "英文、开源、高质量":
8
        return "gte-Qwen2-7B-instruct"  # 阿里，开源最强
9
    elif scenario == "低成本原型":
10
        return "text-embedding-3-small"  # OpenAI，便宜

1
# 索引构建最佳实践
2
def build_quality_index(documents):
3
    index_entries = []
4

5
    for doc in documents:
6
        # 1. 语义分块
7
        chunks = semantic_chunk(doc, max_tokens=512, overlap=50)
8

9
        for chunk in chunks:
10
            # 2. 丰富元数据
11
            entry = {
12
                "text": chunk.text,
13
                "embedding": embed(chunk.text),
14
                "metadata": {
15
                    "source": doc.source,
16
                    "title": doc.title,
17
                    "page": chunk.page_number,
18
                    "section": chunk.section_header,
19
                    "timestamp": doc.last_updated,
20
                }
21
            }
22
            index_entries.append(entry)
23

24
    # 3. 构建多级索引
25
    # 一级：文档级摘要索引（用于粗筛）
26
    doc_summaries = [summarize(doc) for doc in documents]
27
    summary_index = build_index(doc_summaries)
28

29
    # 二级：段落级详细索引（用于精检）
30
    chunk_index = build_index(index_entries)
31

32
    return {"summary_index": summary_index, "chunk_index": chunk_index}