第 12 章 企业 RAG 知识库与 DeepSeek 源码级集成
·
第 12 章 企业 RAG 知识库与 DeepSeek 源码级集成
12.1 向量库(Milvus/FAISS)推理层深度融合改造
12.1.1 RAG 架构概述
RAG(Retrieval-Augmented Generation)是一种将检索与生成相结合的技术,通过从知识库中检索相关信息来增强模型的生成能力。
RAG 流程:
- 文档索引:将文档转换为向量并存储到向量库
- 查询检索:将用户查询转换为向量,在向量库中检索相似文档
- 上下文构建:将检索到的文档构建为上下文
- 生成回答:基于上下文和查询生成回答
12.1.2 Milvus 集成
from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility
import numpy as np
class MilvusRetriever:
def __init__(self, collection_name="deepseek_rag", dim=768):
self.collection_name = collection_name
self.dim = dim
connections.connect("default", host="localhost", port="19530")
if not utility.has_collection(self.collection_name):
self._create_collection()
self.collection = Collection(self.collection_name)
self.collection.load()
def _create_collection(self):
fields = [
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.dim),
FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
FieldSchema(name="metadata", dtype=DataType.JSON)
]
schema = CollectionSchema(fields=fields, description="DeepSeek RAG Knowledge Base")
Collection(name=self.collection_name, schema=schema)
def add_documents(self, documents, embeddings):
entities = [
embeddings,
[doc["text"] for doc in documents],
[doc.get("metadata", {}) for doc in documents]
]
self.collection.insert(entities)
self.collection.flush()
def retrieve(self, query_embedding, top_k=5):
search_params = {
"metric_type": "COSINE",
"params": {"ef": 128}
}
results = self.collection.search(
data=[query_embedding],
anns_field="embedding",
param=search_params,
limit=top_k,
output_fields=["text", "metadata"]
)
retrieved_docs = []
for hit in results[0]:
retrieved_docs.append({
"text": hit.entity.get("text"),
"metadata": hit.entity.get("metadata"),
"score": hit.score
})
return retrieved_docs
12.1.3 FAISS 集成
import faiss
import numpy as np
class FAISSRetriever:
def __init__(self, index_path=None, dim=768):
self.dim = dim
self.index = faiss.IndexFlatIP(dim)
if index_path and os.path.exists(index_path):
self.index = faiss.read_index(index_path)
self.documents = []
def add_documents(self, documents, embeddings):
self.documents.extend(documents)
embeddings_np = np.array(embeddings, dtype=np.float32)
self.index.add(embeddings_np)
def retrieve(self, query_embedding, top_k=5):
query_np = np.array([query_embedding], dtype=np.float32)
distances, indices = self.index.search(query_np, top_k)
retrieved_docs = []
for i, idx in enumerate(indices[0]):
if idx < len(self.documents):
retrieved_docs.append({
"text": self.documents[idx]["text"],
"metadata": self.documents[idx].get("metadata", {}),
"score": float(distances[0][i])
})
return retrieved_docs
def save_index(self, index_path):
faiss.write_index(self.index, index_path)
12.1.4 向量库对比
| 特性 | Milvus | FAISS |
|---|---|---|
| 分布式支持 | 支持 | 不支持 |
| 持久化存储 | 支持 | 需手动管理 |
| 查询性能 | 高 | 极高 |
| 部署复杂度 | 中 | 低 |
| 适用场景 | 大规模生产 | 小规模实验 |
12.2 检索增强生成缓存优化源码
12.2.1 RAG 缓存设计
from functools import lru_cache
from datetime import datetime, timedelta
class RAGCache:
def __init__(self, max_size=10000, ttl_hours=24):
self.cache = {}
self.max_size = max_size
self.ttl = timedelta(hours=ttl_hours)
def get(self, query):
if query in self.cache:
entry = self.cache[query]
if datetime.now() - entry["timestamp"] < self.ttl:
return entry["result"]
del self.cache[query]
return None
def set(self, query, result):
if len(self.cache) >= self.max_size:
oldest_key = min(self.cache.keys(), key=lambda k: self.cache[k]["timestamp"])
del self.cache[oldest_key]
self.cache[query] = {
"result": result,
"timestamp": datetime.now()
}
def clear(self):
self.cache.clear()
def size(self):
return len(self.cache)
12.2.2 RAG 流程优化
class OptimizedRAG:
def __init__(self, retriever, model, tokenizer, embedding_model):
self.retriever = retriever
self.model = model
self.tokenizer = tokenizer
self.embedding_model = embedding_model
self.cache = RAGCache()
def generate(self, query, max_tokens=512, top_k=5):
cached_result = self.cache.get(query)
if cached_result:
return cached_result
query_embedding = self.embedding_model.encode(query)
retrieved_docs = self.retriever.retrieve(query_embedding, top_k=top_k)
context = "
".join([doc["text"] for doc in retrieved_docs])
prompt = f"""基于以下上下文回答问题:
上下文:
{context}
问题:{query}
请根据上下文内容回答问题,如果上下文没有相关信息,请说明。"""
input_ids = self.tokenizer.encode(prompt, return_tensors="pt").cuda()
with torch.no_grad():
output = self.model.generate(input_ids, max_new_tokens=max_tokens)
response = self.tokenizer.decode(output[0], skip_special_tokens=True)
self.cache.set(query, response)
return response
12.2.3 缓存命中率统计
class CacheStats:
def __init__(self):
self.hits = 0
self.misses = 0
self.total_queries = 0
def record_hit(self):
self.hits += 1
self.total_queries += 1
def record_miss(self):
self.misses += 1
self.total_queries += 1
def get_hit_rate(self):
if self.total_queries == 0:
return 0.0
return self.hits / self.total_queries
def reset(self):
self.hits = 0
self.misses = 0
self.total_queries = 0
12.3 私有文档解析、分段重排定制代码
12.3.1 文档解析
import PyPDF2
import docx
from pptx import Presentation
import markdown
class DocumentParser:
def __init__(self):
self.parsers = {
".pdf": self._parse_pdf,
".docx": self._parse_docx,
".pptx": self._parse_pptx,
".md": self._parse_md,
".txt": self._parse_txt
}
def parse(self, file_path):
ext = os.path.splitext(file_path)[1].lower()
if ext not in self.parsers:
raise ValueError(f"不支持的文件格式: {ext}")
return self.parsers[ext](file_path)
def _parse_pdf(self, file_path):
text = ""
with open(file_path, "rb") as f:
reader = PyPDF2.PdfReader(f)
for page in reader.pages:
text += page.extract_text() + "
"
return text
def _parse_docx(self, file_path):
doc = docx.Document(file_path)
text = ""
for paragraph in doc.paragraphs:
text += paragraph.text + "
"
for table in doc.tables:
for row in table.rows:
for cell in row.cells:
text += cell.text + " "
text += "
"
return text
def _parse_pptx(self, file_path):
prs = Presentation(file_path)
text = ""
for slide in prs.slides:
for shape in slide.shapes:
if hasattr(shape, "text"):
text += shape.text + "
"
return text
def _parse_md(self, file_path):
with open(file_path, "r", encoding="utf-8") as f:
md_text = f.read()
html = markdown.markdown(md_text)
text = "".join(BeautifulSoup(html, "html.parser").findAll(text=True))
return text
def _parse_txt(self, file_path):
with open(file_path, "r", encoding="utf-8") as f:
return f.read()
12.3.2 文档分段
class DocumentChunker:
def __init__(self, chunk_size=512, chunk_overlap=50):
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
def chunk(self, text):
chunks = []
words = text.split()
total_words = len(words)
start = 0
while start < total_words:
end = start + self.chunk_size
if end < total_words:
chunk_words = words[start:end]
else:
chunk_words = words[start:]
chunk_text = " ".join(chunk_words)
chunks.append(chunk_text)
start += self.chunk_size - self.chunk_overlap
return chunks
def chunk_with_overlap(self, text):
chunks = []
sentences = text.split("。")
current_chunk = ""
current_length = 0
for sentence in sentences:
sentence_length = len(sentence)
if current_length + sentence_length <= self.chunk_size:
current_chunk += sentence + "。"
current_length += sentence_length
else:
if current_chunk:
chunks.append(current_chunk)
current_chunk = sentence + "。"
current_length = sentence_length
if current_chunk:
chunks.append(current_chunk)
return chunks
12.3.3 文档重排
class DocumentReranker:
def __init__(self, rerank_model):
self.rerank_model = rerank_model
def rerank(self, query, documents):
texts = [doc["text"] for doc in documents]
scores = self.rerank_model.compute_score(query, texts)
ranked_docs = sorted(
zip(documents, scores),
key=lambda x: x[1],
reverse=True
)
return [doc for doc, _ in ranked_docs]
12.4 幻觉抑制底层逻辑修改
12.4.1 幻觉检测
class HallucinationDetector:
def __init__(self, embedding_model, threshold=0.7):
self.embedding_model = embedding_model
self.threshold = threshold
def detect(self, generated_text, context_documents):
context_embeddings = [
self.embedding_model.encode(doc["text"])
for doc in context_documents
]
generated_embedding = self.embedding_model.encode(generated_text)
similarities = []
for ctx_emb in context_embeddings:
similarity = self._cosine_similarity(generated_embedding, ctx_emb)
similarities.append(similarity)
avg_similarity = sum(similarities) / len(similarities)
is_hallucination = avg_similarity < self.threshold
return {
"is_hallucination": is_hallucination,
"similarity_score": avg_similarity,
"threshold": self.threshold
}
def _cosine_similarity(self, vec1, vec2):
dot_product = np.dot(vec1, vec2)
norm1 = np.linalg.norm(vec1)
norm2 = np.linalg.norm(vec2)
return dot_product / (norm1 * norm2) if norm1 > 0 and norm2 > 0 else 0.0
12.4.2 幻觉抑制
class HallucinationSuppressor:
def __init__(self, detector, max_retries=3):
self.detector = detector
self.max_retries = max_retries
def suppress(self, query, context_documents, generate_fn):
for attempt in range(self.max_retries):
response = generate_fn(query, context_documents)
detection = self.detector.detect(response, context_documents)
if not detection["is_hallucination"]:
return response, detection
print(f"检测到幻觉,尝试重新生成 (第 {attempt + 1} 次)")
return response, detection
12.4.3 防幻觉提示词
class AntiHallucinationPrompt:
def __init__(self):
self.system_prompt = """你是一个专业的问答助手,必须严格基于提供的上下文回答问题。
规则:
1. 只使用上下文提供的信息回答问题
2. 如果上下文没有相关信息,明确说明"根据提供的上下文,无法回答该问题"
3. 不要编造任何信息
4. 如果不确定答案,说明"无法确定"
请严格遵守以上规则。"""
def build_prompt(self, query, context_documents):
context = "
".join([doc["text"] for doc in context_documents])
prompt = f"""{self.system_prompt}
上下文:
{context}
问题:{query}
回答:"""
return prompt
本章小结:
本章详细介绍了企业 RAG 知识库与 DeepSeek 的源码级集成,包括向量库集成、缓存优化、文档解析分段和幻觉抑制。这些技术为企业构建智能问答系统提供了完整的解决方案。
更多资讯:lxb20110121
更多推荐



所有评论(0)