为什么用本地 RAG

文档不上云，全在本地处理。免费，不需要 API 额度。模型、向量库、切分策略都能自己调。

流程就两步：检索 + 生成。先从文档里找到相关段落，喂给 AI 让它基于这些内容回答。

环境

先装 Ollama，拉好模型：

ollama pull qwen2.5:7b
ollama pull nomic-embed-text

嵌入模型用 nomic-embed-text，274MB，专门做向量嵌入的，比用主模型快很多。

Python 环境：

mkdir ~/rag-demo && cd ~/rag-demo
python3 -m venv venv && source venv/bin/activate
pip install langchain langchain-community langchain-ollama chromadb pypdf

完整代码

直接一个文件搞定，保存为 rag.py：

#!/usr/bin/env python3
import os, sys, shutil, subprocess
from langchain_community.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaEmbeddings, ChatOllama
from langchain_community.vectorstores import Chroma
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate

MODEL = "qwen2.5:7b"
EMBED_MODEL = "nomic-embed-text"
DB_PATH = "./chroma_db"
DOCS_DIR = "./docs"

def ensure_model(name):
    r = subprocess.run(["ollama", "list"], capture_output=True, text=True)
    if name not in r.stdout:
        print(f"下载模型 {name}...")
        subprocess.run(["ollama", "pull", name], check=True)

def build_db():
    print(f"扫描文档目录: {DOCS_DIR}")
    loader = DirectoryLoader(DOCS_DIR, glob="**/*")
    docs = loader.load()
    print(f"加载了 {len(docs)} 个文档")

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500, chunk_overlap=50,
        separators=["nn", "n", "。", "！", "？", ".", "!", "?", " ", ""]
    )
    chunks = splitter.split_documents(docs)
    print(f"切分成 {len(chunks)} 个文本块")

    if os.path.exists(DB_PATH):
        shutil.rmtree(DB_PATH)

    Chroma.from_documents(
        documents=chunks,
        embedding=OllamaEmbeddings(model=EMBED_MODEL),
        persist_directory=DB_PATH,
    )
    print(f"向量库已保存")

def load_db():
    return Chroma(
        persist_directory=DB_PATH,
        embedding_function=OllamaEmbeddings(model=EMBED_MODEL),
    )

def main():
    ensure_model(MODEL)
    ensure_model(EMBED_MODEL)

    if not os.path.exists(DB_PATH) or "--rebuild" in sys.argv:
        build_db()

    vectorstore = load_db()
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
    llm = ChatOllama(model=MODEL, temperature=0)

    prompt = ChatPromptTemplate.from_messages([
        ("system", "基于以下参考资料回答问题。如果资料中没有相关信息，直接说"文档中没有相关内容"。nn参考资料：n{context}"),
        ("human", "{input}"),
    ])

    chain = create_retrieval_chain(
        retriever,
        create_stuff_documents_chain(llm, prompt)
    )

    print("输入问题开始查询（quit 退出 / rebuild 重建索引）")
    while True:
        q = input("n你的问题：").strip()
        if q.lower() in ('quit', 'exit', '退出'):
            break
        if q.lower() == 'rebuild':
            build_db()
            continue

        print("⏳ 回答中...")
        result = chain.invoke({"input": q})
        print(f"n💬 {result['answer']}")

        print("n📎 引用来源：")
        for i, doc in enumerate(result["context"], 1):
            print(f"  [{i}] {doc.metadata.get('source', 'unknown')}")
            print(f"      {doc.page_content[:200]}...")

if __name__ == "__main__":
    main()

用法

mkdir docs
cp /path/to/your/document.pdf docs/

cd ~/rag-demo
source venv/bin/activate
python3 rag.py

添加新文档后 python3 rag.py --rebuild 重建索引。

支持的文件格式：PDF、Word、Markdown、TXT、HTML、Excel 都能读。

调优

回答不准：chunk_size 改小（500→300），TOP_K 改大（3→5），temperature 降低。

速度慢：嵌入模型别用主模型，nomic-embed-text 快很多。向量库只在首次构建时耗时，后续直接加载。

7B 模型在 CPU 上回答一次大概 10~30 秒，GPU 上 3~8 秒。

想给别人用？接个 Gradio 界面：

pip install gradio

代码最后加上：

import gradio as gr

def answer(q):
    result = chain.invoke({"input": q})
    return result["answer"]

gr.Interface(fn=answer, inputs=gr.Textbox(label="问题"), outputs=gr.Textbox(label="回答")).launch()