下载 gte-large-zh → 将文档转换成向量 → 保存到 Chroma → 根据问题检索相似文档 → 按元数据过滤。

下载模型：gte-large-zh

python 复制代码

from pathlib import Path
from huggingface_hub import snapshot_download

# 当前路径是 notebook/rag，因此 parents[1] 是项目根目录
project_root = Path.cwd().parents[1]

model_path = (
    project_root
    / "llm_app"
    / "embedding_models"
    / "gte-large-zh"
)

snapshot_download(
    repo_id="thenlper/gte-large-zh",
    local_dir=str(model_path)
)

print("模型保存位置：", model_path)

下载 gte-large-zh → 将文档转换成向量 → 保存到 Chroma → 根据问题检索相似文档 → 按元数据过滤。

python 复制代码

from pathlib import Path

import numpy as np
import chromadb

from huggingface_hub import snapshot_download
from chromadb.utils.embedding_functions import (
    SentenceTransformerEmbeddingFunction
)


# =========================================================
# 1. 设置项目、模型和数据库路径
# =========================================================

project_root = Path(
    r"F:\宝信论文\MK-RAG\rag_full_stack_course_notebooks"
)

model_path = (
    project_root
    / "llm_app"
    / "embedding_models"
    / "gte-large-zh"
)

chroma_data_path = (
    project_root
    / "data"
    / "chroma_db"
)

model_path.parent.mkdir(parents=True, exist_ok=True)
chroma_data_path.mkdir(parents=True, exist_ok=True)

print("项目路径：", project_root)
print("模型路径：", model_path)
print("数据库路径：", chroma_data_path)


# =========================================================
# 2. 检查并下载 gte-large-zh
# =========================================================

weight_exists = (
    (model_path / "model.safetensors").is_file()
    or
    (model_path / "pytorch_model.bin").is_file()
)

model_complete = (
    (model_path / "config.json").is_file()
    and (model_path / "modules.json").is_file()
    and weight_exists
)

if not model_complete:
    print("本地模型不完整，开始下载 gte-large-zh......")

    snapshot_download(
        repo_id="thenlper/gte-large-zh",
        local_dir=str(model_path)
    )

    print("模型下载完成。")
else:
    print("本地模型已经存在，不需要重新下载。")


# 再次检查模型
weight_exists = (
    (model_path / "model.safetensors").is_file()
    or
    (model_path / "pytorch_model.bin").is_file()
)

if not (model_path / "config.json").is_file():
    raise FileNotFoundError(
        f"模型缺少 config.json：{model_path}"
    )

if not weight_exists:
    raise FileNotFoundError(
        f"模型缺少权重文件：{model_path}"
    )


# =========================================================
# 3. 创建本地 Chroma 客户端
# =========================================================

# PersistentClient 不需要启动 Chroma 服务
chroma_client = chromadb.PersistentClient(
    path=str(chroma_data_path)
)

print("Chroma 状态：", chroma_client.heartbeat())


# =========================================================
# 4. 创建 Embedding 函数
# =========================================================

embedding_function = SentenceTransformerEmbeddingFunction(
    model_name=str(model_path),
    device="cpu",                # 有合适的 CUDA 环境可改为 "cuda"
    normalize_embeddings=True
)

print("Embedding 模型加载完成。")


# =========================================================
# 5. 创建或获取集合
# =========================================================

collection = chroma_client.get_or_create_collection(
    name="rag_db_gte_large_zh",
    embedding_function=embedding_function,
    metadata={
        "hnsw:space": "cosine"
    }
)

print("集合名称：", collection.name)


# =========================================================
# 6. 准备测试文档
# =========================================================

documents = [
    (
        "在向量搜索领域，我们拥有多种索引方法和向量处理技术，"
        "它们使我们能够在召回率、响应时间和内存使用之间做出权衡。"
    ),
    (
        "虽然单独使用特定技术如倒排文件（IVF）、乘积量化（PQ）"
        "或分层导航小世界（HNSW）通常能够带来满意的结果。"
    ),
    (
        "GraphRAG 本质上就是 RAG，只不过与一般 RAG 相比，"
        "其检索路径上多了一个知识图谱。"
    )
]

ids = [
    "id1",
    "id2",
    "id3"
]

metadatas = [
    {
        "chapter": 3,
        "verse": 16
    },
    {
        "chapter": 4,
        "verse": 5
    },
    {
        "chapter": 12,
        "verse": 5
    }
]


# =========================================================
# 7. 添加或更新文档
# =========================================================

# upsert 可以重复运行，相同 ID 会更新而不是报错
collection.upsert(
    ids=ids,
    documents=documents,
    metadatas=metadatas
)

print("文档添加完成。")
print("文档总数：", collection.count())


# =========================================================
# 8. 查看集合中的数据
# =========================================================

peek_result = collection.peek(limit=1)

print("\n集合中的第一条数据：")
print("ID：", peek_result["ids"])
print("文档：", peek_result["documents"])
print("元数据：", peek_result["metadatas"])


# =========================================================
# 9. 根据 ID 获取文档和向量
# =========================================================

id_result = collection.get(
    ids=["id2"],
    include=[
        "documents",
        "embeddings",
        "metadatas"
    ]
)

print("\n根据 ID 查询：")
print("文档：", id_result["documents"])
print("元数据：", id_result["metadatas"])

if id_result["embeddings"] is not None:
    embedding_array = np.asarray(
        id_result["embeddings"]
    )

    print("向量形状：", embedding_array.shape)


# =========================================================
# 10. 相似度检索
# =========================================================

query = "索引技术有哪些？"

query_result = collection.query(
    query_texts=[query],
    n_results=2,
    include=[
        "documents",
        "metadatas",
        "distances"
    ]
)

print("\n相似度检索结果：")

for index, (
    document,
    metadata,
    distance
) in enumerate(
    zip(
        query_result["documents"][0],
        query_result["metadatas"][0],
        query_result["distances"][0]
    ),
    start=1
):
    print(f"\n第 {index} 条：")
    print("文档：", document)
    print("元数据：", metadata)
    print("余弦距离：", distance)


# =========================================================
# 11. 根据元数据过滤后检索
# =========================================================

filter_result = collection.query(
    query_texts=[query],
    n_results=2,
    include=[
        "documents",
        "metadatas",
        "distances"
    ],
    where={
        "verse": 5
    }
)

print("\nverse=5 的过滤检索结果：")

for index, (
    document,
    metadata,
    distance
) in enumerate(
    zip(
        filter_result["documents"][0],
        filter_result["metadatas"][0],
        filter_result["distances"][0]
    ),
    start=1
):
    print(f"\n第 {index} 条：")
    print("文档：", document)
    print("元数据：", metadata)
    print("余弦距离：", distance)