赋予 AI 记忆：在 RTX 3090 上搭建本地 RAG 知识库问答系统

赋予 AI 记忆：在 RTX 3090 上搭建本地 RAG 知识库问答系统

摘要 ：微调让模型学会了"内功"（专业术语、说话风格），但如果想让它回答具体的、不断更新的知识（比如公司内部文档、最新的教程），重新微调成本太高。这时我们需要 RAG (检索增强生成) 。本文将教你如何把 llm_practical_series 文件夹里的所有 Markdown 教程变成 AI 的"外挂大脑"，让它能回答关于本系列教程的任何问题。

关键词 ：RAG LangChain ChromaDB Llama-3 知识库

---

1. 原理简介

RAG 的流程就像"开卷考试"：

1. 用户提问：例如"如何导出 GGUF？"
2. 检索 (Retrieve) ：系统去文档库里翻书，找到了 llm_series_advanced_export.md 里的相关段落。
3. 增强 (Augment)：系统把找到的段落贴在问题后面，告诉 AI："根据这些资料回答用户"。
4. 生成 (Generate)：AI 整理语言，输出答案。

---

2. 环境安装

我们需要安装 LangChain 全家桶和向量数据库。

conda activate llm_learn
pip install langchain langchain-community langchain-huggingface chromadb sentence-transformers

---

3. 编写 RAG 脚本

在工作目录下新建 rag_demo.py。

我们将使用 Unsloth 加载 Llama-3 作为生成器，使用 HuggingFace Embeddings 作为检索器。

from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from unsloth import FastLanguageModel
import torch

# ===========================
# 1. 建立知识库 (索引阶段)
# ===========================
print("📚 正在加载本地文档...")
# 加载当前目录下所有的 .md 文件
loader = DirectoryLoader("./", glob="**/*.md", loader_cls=TextLoader)
docs = loader.load()

# 切分文档 (Chunking)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
splits = text_splitter.split_documents(docs)
print(f"✅ 加载完成：共 {len(docs)} 个文件，切分为 {len(splits)} 个片段。")

# 向量化 (Embeddings)
print("🧠 正在生成向量索引 (ChromaDB)...")
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = Chroma.from_documents(documents=splits, embedding=embedding_model)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3}) # 每次只找最相关的 3 段

# ===========================
# 2. 加载生成模型 (Llama-3)
# ===========================
print("🚀 正在加载 Llama-3 模型...")
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "lora_model_llama3", # 使用我们微调过的模型
    max_seq_length = 2048,
    dtype = None,
    load_in_4bit = True,
)
FastLanguageModel.for_inference(model)

# ===========================
# 3. 定义 RAG 问答函数
# ===========================
def query_rag(question):
    # A. 检索
    print(f"\n🔍 正在搜索知识库：{question}")
    relevant_docs = retriever.invoke(question)
    
    context_text = "\n\n".join([doc.page_content for doc in relevant_docs])
    print(f"📄 找到相关资料：\n{'-'*20}\n{context_text[:200]}...\n{'-'*20}")
    
    # B. 构造 Prompt
    prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>

你是一个基于本地知识库的 AI 助手。请根据下面的【参考资料】回答用户的问题。如果资料里没有答案，请诚实说不知道。

【参考资料】
{context_text}<|eot_id|><|start_header_id|>user<|end_header_id|>

{question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
"""
    
    # C. 生成
    inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
    outputs = model.generate(
        **inputs, 
        max_new_tokens=256, 
        temperature=0.1,
        eos_token_id=[tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|eot_id|>")]
    )
    
    response = tokenizer.batch_decode(outputs)[0]
    answer = response.split("assistant<|end_header_id|>")[-1].strip().replace("<|eot_id||", "")
    return answer

# ===========================
# 4. 测试
# ===========================
if __name__ == "__main__":
    q1 = "如何把模型导出为 Ollama？"
    print(f"🤖 回答：{query_rag(q1)}")
    
    print("\n" + "="*50 + "\n")
    
    q2 = "RTX 3090 微调 Llama-3 需要多少显存？"
    print(f"🤖 回答：{query_rag(q2)}")

---

4. 运行结果

python rag_demo.py

预期输出 ：

当系统回答"如何导出"时，它会准确地引用 llm_series_advanced_export.md 中的步骤（如使用 model.save_pretrained_gguf）。

当系统回答"显存需求"时，它会引用 llm_finetune_guide.md 中的"显存占用 < 6GB"。

意义 ：

你现在拥有的不仅仅是一个聊天机器人，而是一个懂你所有技术文档的超级助手。你可以把你的 PDF 论文、代码库、笔记全部丢进去，让 AI 帮你读！

---