Embeddings向量模型使用示例

from sentence_transformers import SentenceTransformer

model_path = r"D:\workspace-ai\models\BAAI-bge-m3"
model = SentenceTransformer(model_path)

# 示例文本数据
documents = [
    "沙漠仙人掌通过肥厚茎干储存水分，表面蜡质层减少蒸发，刺由叶片退化形成。夜间开花吸引蝙蝠传粉，耐60℃高温与零下低温，是荒漠生态系统的关键物种。",
    "蜂鸟是唯一能悬停飞行的鸟类，每秒振翅50次以上，以花蜜为主食。其羽毛具有虹彩效应，阳光下呈现金属光泽。主要分布于美洲热带地区，体型最小的吸蜜蜂鸟仅5厘米长，筑巢材料使用蜘蛛丝增强弹性。",
    "月球运行至地球与太阳之间时发生的天文奇观，分日全食、环食和偏食三种类型。最长全食阶段约7分钟，观测需专用滤光片保护眼睛，古代文明常将其视为重大预兆记录在甲骨文中。"
]

# 生成密集向量
embeddings = model.encode(documents, normalize_embeddings=True).tolist()
# [[1024...],[1024...],[1024...]]

import chromadb
from chromadb.config import Settings

# 初始化 Chroma 客户端
chroma_client = chromadb.Client(
    # is_persistent 是否持久化到磁盘持
    Settings(persist_directory="./chroma_db", allow_reset=True, is_persistent=False)
)

# 创建或加载集合
collection = chroma_client.get_or_create_collection(
    name="bge-m3-demo",
    metadata={"hnsw:space": "cosine"}  # 使用余弦相似度
)

# 添加数据到 Chroma
ids = [f"doc_{i}" for i in range(len(documents))]
results = collection.get(ids)
if len(results.get("ids")) > 0:
    collection.delete(results.get("ids"))

collection.add(
    documents=documents,  # 原始文本
    embeddings=embeddings,  # 文本生成的向量
    ids=ids  # 数据ID
)

# 查询条件
query = "可以悬停飞行的动物？"
query_embedding = model.encode(query, normalize_embeddings=True).tolist()

# 执行查询
from chromadb.api import types

results = collection.query(
    query_embeddings=[query_embedding],
    n_results=3,
    include=[types.IncludeEnum.metadatas, types.IncludeEnum.documents, types.IncludeEnum.distances, ]
)

# 打印结果
print("0.000 :", query)
for doc, distance in zip(results["documents"][0], results["distances"][0]):
    print(f"{1 - distance:.3f} : {doc}")