50. 向量数据库实战（ChromaDB / FAISS）¶

原始数据 → Embedding 模型 → 数字向量 → 存入向量数据库
                                            ↓
查询文本 → Embedding 模型 → 查询向量 → 在数据库中搜索最近邻 → 返回结果

                  A · B           # 两个向量的点积
cos(θ) = ─────────────────        # 除以两个向量的长度之积
              ||A|| × ||B||

d(A, B) = √( Σ (Aᵢ - Bᵢ)² )    # 每个维度差值的平方和，再开根号

A · B = Σ (Aᵢ × Bᵢ)             # 对应位置相乘再求和

import numpy as np  # 导入数值计算库

# 两个示例向量
a = np.array([1.0, 2.0, 3.0])  # 向量 A
b = np.array([4.0, 5.0, 6.0])  # 向量 B

# --- 余弦相似度 ---
cosine_sim = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))  # 点积 / (模长之积)
print(f"余弦相似度: {cosine_sim:.4f}")  # 输出: 0.9746

# --- 欧氏距离（L2） ---
l2_dist = np.linalg.norm(a - b)  # 两向量差的模长
print(f"欧氏距离: {l2_dist:.4f}")  # 输出: 5.1962

# --- 点积 ---
dot_product = np.dot(a, b)  # 对应元素相乘再求和
print(f"点积: {dot_product:.4f}")  # 输出: 32.0000

# 用 pip 安装（推荐）
pip install chromadb  # 安装 chromadb 包

# 用 conda 安装
conda install -c conda-forge chromadb  # conda 方式安装

# 验证安装
python -c "import chromadb; print(chromadb.__version__)"  # 打印版本号确认安装成功

import chromadb  # 导入 chromadb

# ====== 第1步：创建客户端 ======
# Client() 是内存模式，数据在程序结束后消失
client = chromadb.Client()  # 创建内存客户端

# ====== 第2步：创建集合（Collection） ======
# 集合 = 一张"表"，用于存放相关的向量和文档
# get_or_create_collection：存在就获取，不存在就新建（幂等操作）
collection = client.get_or_create_collection(
    name="bioinfo_papers"  # 集合名称，只能包含字母数字下划线和连字符
)

# ====== 第3步：添加文档 ======
# ChromaDB 会自动用内置的 embedding 模型（默认 all-MiniLM-L6-v2）把文档转成向量
collection.add(
    documents=[  # 原始文本列表
        "宏基因组学通过对环境样本中的所有微生物DNA进行测序分析微生物群落",
        "16S rRNA 基因是细菌分类的金标准分子标记",
        "随机森林是一种集成学习方法，通过构建多棵决策树进行分类或回归",
        "肠道菌群失调与2型糖尿病的发生发展密切相关",
        "MetaPhlAn 使用 clade-specific marker genes 进行物种定量"
    ],
    metadatas=[  # 每条文档的元数据（可用于过滤）
        {"field": "metagenomics", "year": 2024},
        {"field": "amplicon", "year": 2023},
        {"field": "machine_learning", "year": 2024},
        {"field": "microbiome", "year": 2025},
        {"field": "metagenomics", "year": 2025}
    ],
    ids=["doc1", "doc2", "doc3", "doc4", "doc5"]  # 每条文档的唯一ID（必须唯一）
)

print(f"集合中文档数: {collection.count()}")  # 输出: 5

# ====== 第4步：语义查询 ======
results = collection.query(
    query_texts=["糖尿病相关的肠道微生物研究"],  # 查询文本（会自动 embedding）
    n_results=3  # 返回最相似的 3 条
)

# 打印结果
for i, (doc, dist) in enumerate(zip(results["documents"][0], results["distances"][0])):
    print(f"  Top-{i+1} (距离={dist:.4f}): {doc}")

# ====== 第5步：带元数据过滤的查询 ======
filtered = collection.query(
    query_texts=["宏基因组分析工具"],
    n_results=2,
    where={"field": "metagenomics"}  # 只搜 metagenomics 领域的文档
)
print("\n只搜 metagenomics 领域:")
for doc in filtered["documents"][0]:
    print(f"  {doc}")

import chromadb  # 导入 chromadb

# ====== 持久化客户端：数据存到磁盘 ======
# 指定一个目录路径，ChromaDB 会在这个目录下创建数据库文件
client = chromadb.PersistentClient(
    path="./chroma_db"  # 数据保存目录（不存在会自动创建）
)

# 创建或获取集合
collection = client.get_or_create_collection("my_collection")

# 添加数据
collection.add(
    documents=["这是一条持久化的数据"],
    ids=["persistent_doc1"]
)

# 数据会自动保存到 ./chroma_db/ 目录
# 下次用同样的 path 创建 PersistentClient 就能读到之前的数据
print(f"文档数: {collection.count()}")  # 输出: 1

# ====== 更新文档 ======
collection.update(
    ids=["doc1"],  # 要更新的文档 ID
    documents=["更新后的宏基因组文档内容"],  # 新内容
    metadatas=[{"field": "metagenomics", "year": 2026}]  # 新元数据
)

# ====== 用 upsert 代替 add（存在就更新，不存在就插入）======
collection.upsert(
    ids=["doc1", "doc6"],  # doc1 已存在会更新，doc6 不存在会插入
    documents=["更新的文档1", "全新的文档6"],
    metadatas=[{"field": "metagenomics", "year": 2026}, {"field": "new", "year": 2026}]
)

# ====== 删除文档 ======
collection.delete(ids=["doc3"])  # 按 ID 删除
collection.delete(where={"field": "amplicon"})  # 按元数据条件删除

# CPU 版本（推荐初学者）
pip install faiss-cpu  # 安装 FAISS CPU 版

# GPU 版本（需要 CUDA 环境）
pip install faiss-gpu  # 安装 FAISS GPU 版

# conda 安装
conda install -c conda-forge faiss-cpu  # conda 方式安装 CPU 版

# 验证安装
python -c "import faiss; print(faiss.__version__)"  # 打印版本号

import numpy as np  # 数值计算库
import faiss  # 导入 FAISS

# ====== 准备数据 ======
d = 128           # 向量维度
nb = 10000        # 数据库中的向量数
nq = 5            # 查询向量数

np.random.seed(42)  # 固定随机种子，保证可复现
xb = np.random.random((nb, d)).astype('float32')  # 数据库向量（必须是 float32）
xq = np.random.random((nq, d)).astype('float32')  # 查询向量

# ====== 第1步：创建索引 ======
# IndexFlatL2 = 暴力搜索（精确），使用 L2（欧氏距离）
index = faiss.IndexFlatL2(d)  # 创建一个 128 维的 L2 索引

print(f"索引是否已训练: {index.is_trained}")  # True（FlatL2 不需要训练）
print(f"索引中的向量数: {index.ntotal}")  # 0（还没添加数据）

# ====== 第2步：添加向量 ======
index.add(xb)  # 把 10000 个向量加入索引
print(f"添加后的向量数: {index.ntotal}")  # 10000

# ====== 第3步：搜索 ======
k = 4  # 返回最近的 4 个邻居
D, I = index.search(xq, k)  # D=距离矩阵, I=索引矩阵

# D 的形状: (nq, k) = (5, 4)，每行是一个查询的 k 个最近邻的距离
# I 的形状: (nq, k) = (5, 4)，每行是一个查询的 k 个最近邻的 ID

print("前2个查询的最近邻 ID:")
print(I[:2])  # 例: [[7468 3412 8905 1234], [5621 9087 ...]]

print("前2个查询的最近邻距离:")
print(D[:2])  # 距离值，越小越相似

import numpy as np
import faiss

d = 128           # 向量维度
nb = 100000       # 10万条数据
nq = 10           # 10个查询

np.random.seed(42)
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')

# ====== 创建 IVF 索引 ======
nlist = 100  # 聚类中心数（把向量空间分成 100 个区域）

# 先创建一个量化器（用于确定向量属于哪个簇）
quantizer = faiss.IndexFlatL2(d)  # 量化器使用暴力搜索
# 然后创建 IVF 索引
index = faiss.IndexIVFFlat(
    quantizer,  # 量化器
    d,          # 维度
    nlist       # 聚类中心数
)

# ====== IVF 需要训练 ======
print(f"训练前 is_trained: {index.is_trained}")  # False
index.train(xb)  # 用数据库向量训练（学习聚类中心）
print(f"训练后 is_trained: {index.is_trained}")  # True

# ====== 添加数据并搜索 ======
index.add(xb)  # 添加向量
print(f"向量数: {index.ntotal}")  # 100000

# nprobe 控制搜索时访问多少个簇（默认 1）
index.nprobe = 10  # 搜索最近的 10 个簇（精度更高，但更慢）

k = 4
D, I = index.search(xq, k)
print("IVF 搜索结果（前3个查询的最近邻ID）:")
print(I[:3])

import faiss

# ====== 保存索引到文件 ======
faiss.write_index(index, "my_index.faiss")  # 保存到磁盘
print("索引已保存")

# ====== 从文件加载索引 ======
loaded_index = faiss.read_index("my_index.faiss")  # 从磁盘加载
print(f"加载后向量数: {loaded_index.ntotal}")  # 应该和保存前一样

# 加载后可以直接搜索
D, I = loaded_index.search(xq, k)

import faiss

d = 128  # 向量维度

# index_factory 用字符串描述索引结构，一行代码创建复杂索引
# 格式: "预处理,粗量化,精量化"

# 暴力搜索
index1 = faiss.index_factory(d, "Flat")  # 等价于 IndexFlatL2(d)

# IVF + Flat（100个聚类中心）
index2 = faiss.index_factory(d, "IVF100,Flat")  # 等价于 IndexIVFFlat

# IVF + PQ（100个聚类中心 + 乘积量化压缩到 16 字节）
index3 = faiss.index_factory(d, "IVF100,PQ16")  # 内存友好的方案

# HNSW
index4 = faiss.index_factory(d, "HNSW32")  # 32 个邻居的 HNSW 图

# PCA降维 + IVF + PQ
index5 = faiss.index_factory(d, "PCA64,IVF100,PQ8")  # 先降到64维再索引

# pip install pymilvus  # 安装 Milvus Python SDK

from pymilvus import MilvusClient  # 导入轻量客户端

# Milvus Lite 模式：单文件嵌入式运行（类似 SQLite）
client = MilvusClient("milvus_demo.db")  # 数据存到本地文件

# 创建集合
client.create_collection(
    collection_name="demo",
    dimension=128  # 向量维度
)

# 插入数据（需要自己提供向量）
import numpy as np
vectors = np.random.random((100, 128)).tolist()  # 100条 128维向量
data = [{"id": i, "vector": vectors[i]} for i in range(100)]
client.insert(collection_name="demo", data=data)

# 搜索
query_vector = np.random.random((1, 128)).tolist()
results = client.search(
    collection_name="demo",
    data=query_vector,
    limit=5  # 返回 Top-5
)
print(results)

你的向量数量？
├── < 10万 → ChromaDB（最简单，开箱即用）
├── 10万~1000万 → ChromaDB 或 Qdrant
├── 1000万~1亿 → FAISS（纯搜索）或 Qdrant/Milvus（需要数据库特性）
└── > 1亿 → Milvus（自建）或 Pinecone（托管）

你需要内置 embedding 吗？
├── 是 → ChromaDB（自动处理）
└── 否 → FAISS / Milvus / Qdrant（自己生成 embedding）

你需要 GPU 加速吗？
├── 是 → FAISS（faiss-gpu）
└── 否 → 都可以

# 概念示例：用向量数据库加速序列相似性搜索
# 实际中需要用 ProtTrans、ESM-2 等模型生成序列 embedding

import chromadb
import numpy as np

client = chromadb.Client()

# 假设我们已经用 ESM-2 模型将蛋白质序列转成了 embedding
collection = client.get_or_create_collection(
    name="protein_sequences",
    metadata={"hnsw:space": "cosine"}  # 使用余弦距离
)

# 添加蛋白质序列（实际中 embeddings 由 ESM-2 生成）
collection.add(
    documents=[  # 原始序列描述
        "Insulin receptor substrate 1 - Homo sapiens",
        "Glucokinase - Mus musculus",
        "GLUT4 glucose transporter - Rattus norvegicus"
    ],
    embeddings=[  # 这里用随机向量示意，实际应为模型输出
        np.random.random(320).tolist(),
        np.random.random(320).tolist(),
        np.random.random(320).tolist()
    ],
    ids=["prot1", "prot2", "prot3"]
)

# 查询：找与目标蛋白最相似的序列
query_embedding = np.random.random(320).tolist()  # 实际为目标序列的 embedding
results = collection.query(
    query_embeddings=[query_embedding],
    n_results=2
)
print(results["documents"])

原因：添加的文档 ID 已经存在于集合中
解决：用 upsert() 代替 add()，或先 delete() 再 add()

# 错误写法
collection.add(ids=["doc1"], documents=["新内容"])  # doc1 已存在，报错！

# 正确写法
collection.upsert(ids=["doc1"], documents=["新内容"])  # 存在就更新，不存在就插入

原因：FAISS 只接受 float32 类型的 numpy 数组
解决：确保数据类型为 np.float32

# 错误写法
data = np.random.random((100, 128))  # 默认 float64
index.add(data)  # 报错！

# 正确写法
data = np.random.random((100, 128)).astype('float32')  # 转成 float32
index.add(data)  # 正常

原因：IVF/PQ 等索引需要先 train() 再 add()
解决：在添加数据前先用训练数据调用 index.train()

index = faiss.IndexIVFFlat(quantizer, d, nlist)
# index.add(xb)  # 直接 add 会报错！

index.train(xb)  # 必须先训练
index.add(xb)    # 训练后才能添加数据

原因：documents 列表中包含 None 值
解决：确保所有文档都是非空字符串，或传入 embeddings 代替 documents

原因：FAISS 未正确安装，或安装了 faiss-gpu 但没有 CUDA 环境
解决：
  pip install faiss-cpu   # CPU 版本（不需要 GPU）
  # 或
  conda install -c conda-forge faiss-cpu  # conda 方式

原因：使用了 Client() 而不是 PersistentClient()
解决：
  # 内存模式（数据会丢失）
  client = chromadb.Client()

  # 持久化模式（数据保存到磁盘）
  client = chromadb.PersistentClient(path="./my_chroma_db")