LLM 本地知识库构建：文档分割与向量化方案探讨

伪代码示例

from transformers import BartForConditionalGeneration, BartTokenizer
from transformers import LongformerModel, LongformerTokenizer
import torch

# 假设已按段落拆分
document_paragraphs = [
    "First long paragraph ...",
    "Second long paragraph ...",
    # ...
]

# 初始化 BART 模型用于摘要
summary_model_name = 'facebook/bart-large-cnn'
summary_tokenizer = BartTokenizer.from_pretrained(summary_model_name)
summary_model = BartForConditionalGeneration.from_pretrained(summary_model_name)

# 初始化 Longformer 模型用于编码
longformer_tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')
longformer_model = LongformerModel.from_pretrained('allenai/longformer-base-4096')

def summarize_paragraphs(paragraphs):
    summarized_paragraphs = []
    for paragraph in paragraphs:
        inputs = summary_tokenizer(paragraph, return_tensors="pt", max_length=1024, truncation=True)
        summary_ids = summary_model.generate(inputs['input_ids'], num_beams=4, max_length=200, early_stopping=True)
        summary = summary_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
        summarized_paragraphs.append(summary)
    return summarized_paragraphs

def encode_paragraphs_with_longformer(paragraphs):
    inputs = longformer_tokenizer(paragraphs, return_tensors="pt", padding=True, truncation=True)
    outputs = longformer_model(**inputs)
    return outputs.last_hidden_state

summarized_paragraphs = summarize_paragraphs(document_paragraphs)
document_encoding = encode_paragraphs_with_longformer(summarized_paragraphs)

优缺点分析：

• 优点：高效处理长文档，保持关键信息，利用稀疏注意力机制维持上下文联系。
• 缺点：摘要准确性依赖模型能力，资源消耗较大，自动摘要可能丢失细节。

2. 滑动窗口

滑动窗口是最简单且易于实现的方式。对长段落使用滑动窗口分块，确保窗口间有重叠部分（Overlap），使每个窗口包含相邻窗口的上下文信息。

核心逻辑

1. 分块：设置窗口大小（Window Size）和步长（Step）。步长小于窗口大小可产生重叠。
2. 向量化：对每个窗口进行向量化。
3. 聚合：将窗口向量合成为文档向量，常用方法包括平均池化、最大池化或加权平均。

代码示例

import torch
from transformers import BertTokenizer, BertModel
import numpy as np
import os

def sliding_window_split(text, window_size, step):
    chunks = []
    start = 0
    text = text.strip()
    while start < len(text):
        end = start + window_size
        chunks.append(text[start:end])
        start += step
        if end >= len(text) and start < len(text):
            chunks.append(text[start:])
            break
    return chunks

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

def vectorize(chunk):
    inputs = tokenizer(chunk, return_tensors='pt', max_length=512, truncation=True, padding='max_length')
    outputs = model(**inputs)
    # 取最后一个隐藏层的均值作为句子向量
    chunk_vector = outputs.last_hidden_state.mean(dim=1)
    return chunk_vector

def aggregate_vectors(vectors):
    return torch.mean(torch.stack(vectors), dim=0)

# 模拟文件读取与处理
dir_path = './data'
file_vectors_aggr = {}

for file in os.listdir(dir_path):
    file_path = os.path.join(dir_path, file)
    with open(file_path, 'r') as f:
        text = f.read()
    
    chunks = sliding_window_split(text, 1024, 555)
    vectors = [vectorize(c) for c in chunks]
    aggregated_vector = aggregate_vectors(vectors)
    file_vectors_aggr[file] = aggregated_vector

# 检索示例
question = '什么是 MVCC'
q_vec = vectorize(question).detach()

for file, vector in file_vectors_aggr.items():
    vector_np = vector.detach().numpy().flatten()
    q_vec_np = q_vec.numpy().flatten()
    similarity = np.dot(vector_np, q_vec_np) / (np.linalg.norm(vector_np) * np.linalg.norm(q_vec_np))
    if similarity > 0.48:
        print(f"Matched: {file}, Score: {similarity}")

聚合的好处：

• 保持文档完整性，捕捉综合信息。
• 提高检索效率，避免与每个片段逐一比较。
• 减少信息丢失，增强上下文理解。

3. 自定义稀疏注意力与 LoRA

对于超长文档，传统全连接注意力计算量过大。稀疏注意力模式仅计算相关 token 间的注意力。此外，可通过微调（Fine-tuning）适配特定领域。

LoRA 微调示例

LoRA（Low-Rank Adaptation）是一种高效的微调技术，通过低秩矩阵更新权重，减少显存占用。

import torch
import torch.nn as nn
from transformers import BertModel, BertConfig

class LoRA_BERT(nn.Module):
    def __init__(self, bert_model_name, rank):
        super(LoRA_BERT, self).__init__()
        self.bert = BertModel.from_pretrained(bert_model_name)
        self.rank = rank
        self.lora_adjustments = nn.ModuleDict()

        # 为 BERT 的每个注意力层添加 LoRA 调整
        for i, layer in enumerate(self.bert.encoder.layer):
            # 简化示例，实际需针对 Query, Key, Value 分别处理
            self.lora_adjustments[f"layer_{i}"] = LoRAAdjustment(layer.attention.self.query, rank)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        sequence_output = outputs.last_hidden_state
        return sequence_output

class LoRAAdjustment(nn.Module):
    def __init__(self, query_layer, rank):
        super(LoRAAdjustment, self).__init__()
        self.query_layer = query_layer
        self.rank = rank
        hidden_size = query_layer.out_features

        self.delta_W = nn.Parameter(torch.randn(hidden_size, rank))
        self.delta_b = nn.Parameter(torch.randn(hidden_size))

    def forward(self, x):
        adjusted_weight = self.query_layer.weight + self.delta_W @ self.delta_W.t()
        adjusted_bias = self.query_layer.bias + self.delta_b
        return torch.nn.functional.linear(x, adjusted_weight, adjusted_bias)

bert_model_name = 'bert-base-uncased'
model = LoRA_BERT(bert_model_name, rank=4)

实施建议：

• 需深入研究 Transformer 内部结构。
• 适合研究人员或对模型定制要求高的场景。
• 对于一般应用，LoRA 微调比从头训练更高效。

向量数据库与检索策略

构建知识库不仅涉及文档处理，还需选择合适的向量数据库进行存储与检索。

常见选型

1. FAISS：Facebook 开源的高效相似度搜索库，适合单机部署，支持多种索引类型（IVF, HNSW）。
2. Milvus：云原生向量数据库，支持分布式部署，适合大规模数据场景。
3. ChromaDB：轻量级，易于集成，适合开发测试环境。

检索优化

• 混合检索：结合关键词检索（BM25）与向量检索，提升召回率。
• 重排序（Rerank）：初步检索后使用 Cross-Encoder 模型对 Top-K 结果进行精排，提高准确率。
• 多路召回：针对不同字段或内容类型建立多个索引，合并结果。

总结

通过几种方法的尝试，除非对大模型本身能力进行优化，否则每一种方式都存在限制。如果不考虑硬件成本和技术开发成本，训练或调优出特定的领域模型才是最优解。如果基座模型足够强大，直接使用 LangChain 进行文档分割也足以应对多数场景。

不同的分割器和参数组合能找到相对较好的方案。如果有条件（数据量充足且有硬件资源），建议对开源模型进行调优，使其更好地适配本地文库知识。例如构建公司内部文档知识库时，新增公司文档后，调优后的模型会有更好的表现。

在实际落地中，应权衡精度、延迟与成本。对于初创项目，滑动窗口 + 现成 Embedding 模型是性价比最高的选择；对于高并发、高精度要求的场景，建议引入 Rerank 机制并考虑模型微调。