[大模型实战 04] 从玩具到生产:基于 ChromaDB 打造工程级 RAG 系统
核心摘要 (TL;DR) 各位友人们,大家好,这里是阿尔。在上一节中,我们大概知道了大模型的构成,safetensor格式的大模型的文件组成,transformers库的基本使用。我们已经能够使用大模型去做一些简单对话应用了,它可以是上知天文,下知地理,中间还能知道人情冷暖。但是,我们需要加一个限定词,在训练数据截止日期前的。因为训练一次需要耗费很多的计算资源,时间和人力,当我们想让它知道一些新知识的时候,比如让它知道现在美国的总统是拜登还是特朗普,我们可以在对话中告诉他,这没问题,但是如果我们想让它知道更多,比如我的私人日记?比如我刚写的那篇博客?比如公司的员工手册, 比如自己产品的使用说明书? 这类私有数据,是大模型企业应用的痛点,毕竟大模型是基于在互联网上公开数据训练的。重新把这部分资料加进去,再训练一下模型?也不是不行,但是有点没有性价比,这时候就引出了大模型落地应用的核心技术-> RAG (Retrieval-Augmented Generation,检索增强生成)。 考试的时候,如果考到不会的知识,不知道各位友人们会不会头疼,如果这时候,允许我们翻书,现去书里找,我们也很有可能找得到对应的答案,哪怕我们可能完全没学过。这就是RAG的大致思路:不让模型凭空回忆,而是先给它找资料。 RAG,检索增强生成,字面上讲,就是 拿到考题->然后去翻书,通过目录之类的索引,快速翻到(检索)相关的内容->再根据这些内容(增强了的内容),回答出问题(生成回答)。 对比简单地把东西一股脑全部跟大模型说一遍,我们能清楚得发现,我们只用了检索到的那一部分内容,并没有让整本书大模型的脑子将占用, 这就是RAG的效率体现。 RAG技术的思路很简单,但是实现并非只是一个单一的技术能实现的,它有一套流水线(流水线)。 把这头"大象"放进冰箱,总共需要两步:准备好数据和让模型拿到数据。 在大模型能够翻书之前,咱们得先把我们想给它看的书整理好,放进书包里。 向量化 (Embedding):这是最关键的一步! 拿到书了之后,我们想要翻书,就得找到和问题有关系的内容,然后再将这些内容和我们自己的常识结合起来,对提出的问题进行答题。 检索 (Retrieval):拿着这个“问题向量”,去向量数据库里搜, 去找到关系性高的内容。 增强 (Augmentation):把这 3-5 个片段拼在一起,和用户的问题组合成一个超级长的 Prompt。 Prompt 模板示例: 好了,理论我们已经懂了,现在我们撸起袖子,准备来实操一下子吧。我们打算从零开始快速搭建一个工程级的RAG系统: 私有API助手, 在我直接告诉各位友人们我们要用到的工具前,我觉得也有必要大概让各位友人们知道还有哪些别的选择,我们为什么选择了这几个。 项目背景:假设我们是一家名叫 "DeepStar" 的初创公司,我们有一套内部绝密的 API 文档,新来的实习生总是问重复的问题。我们要用 RAG 让他自己查。 启动 Kaggle Notebook,确保 Internet: On,Accelerator: GPU T4 x2。 我们创建两份文档:一份是核心接口定义,一份是错误码说明。 提前根据自己的情况来配置待会儿用的词嵌入模型和推理模型。 利用 这是本篇最关键的代码。我们要实现:如果本地已经有数据库,就直接读;如果没有,才去解析文档。 现在,我们模拟实习生提问。注意,这个问题需要结合两个文档(接口定义 + 错误码)以及逻辑推理才能回答。 答复如下 既然用了 ChromaDB,我们就可以像查 SQL 一样查它。这在 Debug 时非常有用。 完整代码可以点击kaggle笔记获取 Q: 为什么不直接把所有文档都塞进 Prompt 里 (Long Context)? Q: LlamaIndex 和 LangChain 我该学哪个? Q: ChromaDB 的数据存在哪里了? 本文作者: Algieba[大模型实战 04] 从玩具到生产:基于 ChromaDB 打造工程级 RAG 系统
前言

1. RAG(检索增强生成)
1.1 什么是RAG?

1.2 RAG的步骤

第一个阶段:数据准备(Indexing) -> 把书装进书包
[0.1, -0.5, 0.8, ...]),是不是有点耳熟,对这和大模型训练的Embedding是一个思路,但是我们一般会使用特制的嵌入模型来做这个专业的事情。
第二个阶段:应用数据给大模型生成(Retrieval & Generation)-> 开卷答题
你是一个助手。请根据以下参考资料回答问题。
参考资料:[片段1]... [片段2]...
用户问题:火星基地吃什么?
2. RAG技术选型
2.1 框架: LlamaIndex vs. LangChain

2.2 嵌入模型 (Embedding):BGE vs. OpenAI
BAAI/bge-small-en-v1.5 或者 OpenAI 的 text-embedding-3-small
2.3 向量数据库:Chroma vs. Milvus vs. Pinecone

3. 上手实操

3.1 环境配置 (Kaggle)
# 1. 更新transformers及其相关库
!pip install -U transformers peft accelerate bitsandbytes sentence-transformers
# 2. 安装 LlamaIndex 核心及相关插件
!pip install llama-index-core llama-index-llms-huggingface llama-index-embeddings-huggingface
# 3. 安装 ChromaDB 向量库支持
!pip install llama-index-vector-stores-chroma chromadb
下载依赖库可能会需要一点时间,之后我看看能不能在kaggle上用uv去做包管理。3.2 造数据:模拟企业内部文档
import os
data_path = "/kaggle/working/data"
# 创建数据目录
os.makedirs(data_path, exist_ok=True)
# 文档 1: 核心 API 定义
api_doc = """
[机密] DeepStar 核心交易接口 v2.0
1. 创建订单 API: POST /api/v2/order/create
- 必填参数: 'user_id' (String), 'amount' (Decimal), 'token' (X-Auth-Token)
- 特殊逻辑: 如果 amount > 10000, 必须额外传递 'audit_code' (审计码)。
- 频率限制: 单用户每秒最多 5 次调用。
2. 查询余额 API: GET /api/v2/balance
- 缓存策略: 默认缓存 5 秒。传递 'no-cache=true' 可强制刷新。
"""
# 文档 2: 错误码字典
error_doc = """
[机密] DeepStar 全局错误码字典
- E1001: 签名验证失败。请检查 X-Auth-Token 是否过期。
- E2009: 余额不足。注意:冻结金额不计入可用余额。
- E5003: 审计风控拦截。大额交易未通过自动审计,请联系人工客服。
"""
with open(f"{/data_path}/api_specs.txt", "w") as f:
f.write(api_doc)
with open(f"/{data_path}/error_codes.txt", "w") as f:
f.write(error_doc)
print("[Success] 企业文档库已就绪!")
3.3 初始化大脑与眼睛 (Settings)
embedding_model ="BAAI/bge-small-zh-v1.5"
llm = "Qwen/Qwen2.5-7B-Instruct"
# 在本地服务器,可以用modelscope下载下来, 把路径配置在这儿Settings 全局配置,将默认的 OpenAI 替换为本地模型。import torch
from llama_index.core import Settings
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
# 1. 设置 Embedding (眼睛)
# 使用 BGE-Small,显存占用极低,检索中文效果极佳
print("正在加载 Embedding 模型...")
Settings.embed_model = HuggingFaceEmbedding(
model_name=embedding_model
)
# 2. 设置 LLM (大脑)
# 使用 Qwen2.5-7B-Instruct
print("正在加载 LLM 模型...")
Settings.llm = HuggingFaceLLM(
model_name=llm,
tokenizer_name=llm,
context_window=30000,
max_new_tokens=512,
generate_kwargs={"temperature": 0.1, "do_sample": True}, # 技术文档要求严谨,温度调低
device_map="auto",
model_kwargs={"dtype": torch.float16, "trust_remote_code": True}
)
print("[Success] 模型加载完毕!")
3.4 核心组件:ChromaDB 持久化流水线
import chromadb
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore
# 定义持久化路径
CHROMA_DB_PATH = "/kaggle/working/chroma_db"
COLLECTION_NAME = "deepstar_docs"
# 1. 初始化 Chroma 客户端 (PersistentClient 实现了写硬盘功能)
db_client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
# 2. 创建或获取集合 (Collection)
chroma_collection = db_client.get_or_create_collection(COLLECTION_NAME)
# 3. 将 Chroma 对接给 LlamaIndex
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
# 4. 智能加载逻辑 (幂等性设计)
if chroma_collection.count() == 0:
print("[Info] 数据库为空,开始初始化...")
# 读取 data 目录下的所有文件
documents = SimpleDirectoryReader("./data").load_data()
# 建立索引并自动存入 Chroma (Ingestion)
index = VectorStoreIndex.from_documents(
documents, storage_context=storage_context
)
print("[Success] 数据写入完成!")
else:
print(f"[Info] 发现 {chroma_collection.count()} 条存量数据,直接加载...")
# 直接从 Vector Store 加载,无需重新计算 Embedding
index = VectorStoreIndex.from_vector_store(
vector_store, storage_context=storage_context
)
print("[Success] 索引加载完成!")
3.5 验收测试:复杂逻辑问答
# 创建查询引擎
query_engine = index.as_query_engine(similarity_top_k=3)
# 实习生的提问
questions = [
"创建订单时,如果你只有 100 块钱,能传 amount=20000 吗?为什么?",
"我收到了 E5003 错误,这是什么意思?该怎么办?"
]
print("======== 开始 RAG 问答测试 ========")
for q in questions:
print(f"\n[Question] {q}")
response = query_engine.query(q)
print(f"[Answer]\n{str(response)}")
# 打印引用源 (Debug 必备,看看它参考了哪个文件)
source_file = response.source_nodes[0].metadata.get('file_name')
print(f"[Source]: {source_file}")

4. 进阶技巧:如何管理你的数据库?
# 偷看数据库里的前 2 条记录
data = chroma_collection.peek(limit=2)
print("\n[Debug] 数据库抽查:")
for i, doc in enumerate(data['documents']):
print(f"--- 片段 {i} ---")
print(f"内容: {doc[:50]}...") # 只打印前50个字
print(f"来源: {data['metadatas'][i]}")5. 完整代码
5. 常见问题 (Q&A)
A: 虽然现在很多模型支持长文本(比如 128k),但直接塞文档有三个问题:
A:
A: 在上面代码中,我们通过 PersistentClient 指定了路径 /kaggle/working/chroma_db。
它就像 SQLite 一样,数据就存在这个文件夹里的 .sqlite3 和 .bin 文件中。咱们可以把这个文件夹拷贝到任何电脑上,无需重新向量化就能直接使用。
本文链接: https://blog.algieba12.cn/llm04-rag-llamaindex-chromadb/
版权声明: 本博客所有文章除特别声明外,均采用 BY-NC-SA 许可协议。转载请注明出处!