AI学习：第3天

太棒了！API 调用成功，你已经拥有了一个本地的、可编程的 LLM 服务。接下来，我们进入第三天：实现 RAG（检索增强生成）。

🎯 第三天目标：做一个能基于你自己的文档回答问题的系统

比如，你可以上传一段关于你后端项目的描述，然后问 "我的项目用了哪些技术栈？"，模型会根据你提供的文档来回答，而不是凭空想象。

你需要做什么（总耗时约 50 分钟）

1️⃣ 安装必要的库（5 分钟）

在 ai_env 环境中执行：

pip install sentence-transformers chromadb

• sentence-transformers：提供嵌入模型（将文本转成向量）
• chromadb：轻量级向量数据库（用于存储和检索文档片段）

2️⃣ 准备一个示例文档（5 分钟）

创建一个 my_knowledge.txt 文件，写入一些你自己的项目介绍。例如：

我的后端项目使用了 Spring Boot 3.0，数据库是 MySQL 8.0，缓存用 Redis，消息队列用 RabbitMQ。项目部署在阿里云 Kubernetes 集群上。

3️⃣ 编写 RAG API 代码（30 分钟）

在 app.py 的基础上，增加一个 /rag 端点。你可以复制以下完整代码到 app_rag.py：

# app_rag.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from modelscope import AutoModelForCausalLM, AutoTokenizer
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.utils import embedding_functions
import torch
import time
import os

# ---------- 1. 加载大模型 ----------
print("Loading LLM...")
model_id = 'qwen/Qwen2-0.5B-Instruct'
llm = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
print("LLM loaded.")

# ---------- 2. 加载嵌入模型 ----------
print("Loading embedding model...")
embedder = SentenceTransformer('BAAI/bge-small-zh')  # 中文嵌入模型，约 33MB
print("Embedding model loaded.")

# ---------- 3. 初始化向量数据库 ----------
chroma_client = chromadb.Client()
# 使用嵌入函数
embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name='BAAI/bge-small-zh')
collection = chroma_client.create_collection(name="docs", embedding_function=embedding_fn)

# ---------- 4. 加载并切分你的文档 ----------
def load_and_chunk(file_path, chunk_size=200):
    with open(file_path, 'r', encoding='utf-8') as f:
        text = f.read()
    # 简单按字符数切分，也可以按句号分割
    chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
    return chunks

# 请将你的文档放在同目录下的 my_knowledge.txt
if os.path.exists("my_knowledge.txt"):
    chunks = load_and_chunk("my_knowledge.txt", chunk_size=200)
    # 存入 ChromaDB
    for i, chunk in enumerate(chunks):
        collection.add(
            documents=[chunk],
            ids=[f"doc_{i}"]
        )
    print(f"Loaded {len(chunks)} document chunks.")
else:
    print("Warning: my_knowledge.txt not found. RAG will work but no context.")

app = FastAPI()

class RAGRequest(BaseModel):
    question: str
    top_k: int = 3
    max_tokens: int = 200
    temperature: float = 0.7

@app.post("/rag")
async def rag_endpoint(request: RAGRequest):
    # 1. 检索相关文档片段
    results = collection.query(query_texts=[request.question], n_results=request.top_k)
    retrieved_docs = results['documents'][0] if results['documents'] else []
    
    # 2. 构造增强的 prompt
    context = "\n".join(retrieved_docs) if retrieved_docs else "无相关文档。"
    prompt = f"""基于以下已知信息，简洁地回答用户的问题。如果无法从已知信息中找到答案，请说"根据已有信息无法回答"。
已知信息：
{context}
问题：{request.question}
回答："""
    
    # 3. 调用大模型生成
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = llm.generate(
            inputs.input_ids,
            max_new_tokens=request.max_tokens,
            temperature=request.temperature,
            do_sample=True
        )
    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    assistant_response = full_response.split("assistant")[-1].strip()
    
    return {
        "answer": assistant_response,
        "retrieved_docs": retrieved_docs,
        "inference_time": time.time()
    }

@app.get("/health")
def health():
    return {"status": "ok"}

由于我的本地下载不下来这个模型，所以更换了其他模型

# app_rag.py
from fastapi import FastAPI
from pydantic import BaseModel
from modelscope import AutoModelForCausalLM, AutoTokenizer, snapshot_download
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.utils import embedding_functions
import torch
import time
import os

# ---------- 1. 加载LLM ----------
print("Loading LLM...")
model_id = 'qwen/Qwen2-0.5B-Instruct'
llm = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
print("LLM loaded.")

# ---------- 2. 加载Embedding模型 ----------
print("Loading embedding model from ModelScope...")
# 确保这里使用的model_id是正确的
embedding_model_id = 'BAAI/bge-small-zh-v1.5'
# 下载模型到本地，如果已经下载过，snapshot_download会直接使用缓存
local_model_dir = snapshot_download(embedding_model_id)
# 使用 SentenceTransformer 加载下载好的本地模型
embedder = SentenceTransformer(local_model_dir)
print("Embedding model loaded.")

# ---------- 3. 初始化向量数据库 ----------
# 注意：为了演示，这里使用了In-Memory的Client，服务重启后数据会丢失。
# 如需持久化，请使用 chromadb.PersistentClient(path="./chroma_db")
chroma_client = chromadb.Client()
collection = chroma_client.create_collection(name="docs")

# ---------- 4. 加载并切分你的文档 ----------
def load_and_chunk(file_path, chunk_size=200):
    with open(file_path, 'r', encoding='utf-8') as f:
        text = f.read()
    chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
    return chunks

# 请确保同目录下有一个 my_knowledge.txt 文件
if os.path.exists("my_knowledge.txt"):
    chunks = load_and_chunk("my_knowledge.txt", chunk_size=200)
    for i, chunk in enumerate(chunks):
        # 手动生成向量
        embedding = embedder.encode(chunk).tolist()
        collection.add(
            embeddings=[embedding],
            documents=[chunk],
            ids=[f"doc_{i}"]
        )
    print(f"Loaded {len(chunks)} document chunks into ChromaDB.")
else:
    print("Warning: my_knowledge.txt not found. RAG will work but no context.")

app = FastAPI()

class RAGRequest(BaseModel):
    question: str
    top_k: int = 3
    max_tokens: int = 200
    temperature: float = 0.7

@app.post("/rag")
async def rag_endpoint(request: RAGRequest):
    # 1. 检索：将问题转为向量，在ChromaDB中搜索
    question_embedding = embedder.encode(request.question).tolist()
    results = collection.query(query_embeddings=[question_embedding], n_results=request.top_k)
    retrieved_docs = results['documents'][0] if results['documents'] else []

    # 2. 构造增强的 prompt
    context = "\n".join(retrieved_docs) if retrieved_docs else "无相关文档。"
    prompt = f"""基于以下已知信息，简洁地回答用户的问题。如果无法从已知信息中找到答案，请说"根据已有信息无法回答"。
已知信息：
{context}
问题：{request.question}
回答："""

    # 3. 调用大模型生成
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = llm.generate(
            inputs.input_ids,
            max_new_tokens=request.max_tokens,
            temperature=request.temperature,
            do_sample=True
        )
    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    assistant_response = full_response.split("assistant")[-1].strip()

    return {
        "answer": assistant_response,
        "retrieved_docs": retrieved_docs,
        "inference_time": time.time()
    }

@app.get("/health")
def health():
    return {"status": "ok"}

4️⃣ 启动新服务并测试（10 分钟）

uvicorn app_rag:app --reload --port 8000

然后另开终端测试：

import requests
resp = requests.post("http://127.0.0.1:8000/rag", json={"question": "我的项目用了什么数据库？"})
print(resp.json()["answer"])

你会看到模型基于 my_knowledge.txt 中的内容回答 "MySQL 8.0"。

---

🧭 后续计划（第4~7天）

• 第四天：让 RAG 支持多文档（上传 PDF、Word）和更好的文本切分（按段落/语义）。
• 第五天：增加对话记忆（多轮对话）。
• 第六天：优化推理速度（使用 vLLM 或量化模型）。
• 第七天：写一个简单的 Gradio 前端界面，展示你的完整作品。

你也可以根据自己的节奏调整。目前你已完成了 环境+API+RAG基础，这在简历上已经是一个很完整的 AI 工程化项目了。

---

💎 今天的核心产出

• 一个能基于私有文档问答的 RAG API。
• 理解了"检索"+"生成"的基本流程。
• 为简历增加了一条亮点："基于 Qwen 和 ChromaDB 实现 RAG 系统，提升问答准确性"。