chromadb使用千问嵌入模型在笔记本上运行效果

import numpy as np
import matplotlib.pyplot as plt
import chromadb
from chromadb.config import Settings
from chromadb.utils import embedding_functions
persist_directory= "./chroma_data"
# 使用本地Qwen3-Embedding-0.6B嵌入模型
custom_embedding = embedding_functions.SentenceTransformerEmbeddingFunction(
model_name="d:\\小说创作\\Qwen3-Embedding-0___6B",
device="cuda" # 若无有GPU则改为"cpu"
)
client = chromadb.PersistentClient(
path=persist_directory,
settings=Settings(anonymized_telemetry=False)
)
# 创建集合时指定自定义嵌入函数
collection = client.get_or_create_collection(
name="tests",
embedding_function=custom_embedding
)

cqdoc="从江雾里长出的立体迷宫：重庆的烟火与野性\n\
当飞机穿越厚厚的云层，机身下方突然跳出一片被江水缠绕、被高楼撑起来的城市时，你就知道，重庆到了。这座被称为“8D魔幻之都”的城市，\
从来不是平面的——它是一座在山水间野蛮生长的立体迷宫，每一条爬坡上坎的石阶，每一栋依山而建的高楼，每一碗翻滚着红油的火锅，都藏着属于自己的倔强与温柔。\n\
一、两江环抱的天生骨架：山水里的城市基因\n\
重庆的故事，是从两条江开始的。……"

from sentence_transformers import SentenceTransformer
cpu_embedding = SentenceTransformer(
model_name_or_path="d:\\小说创作\\Qwen3-Embedding-0___6B",
device="cpu" # 若无有GPU则改为"cpu"
)

cpu_embedding.encode(['中国']).shape
输出为：(1, 1024)

cq1doc=cqdoc.splitlines()
id1=list(range(len(cq1doc)))
id2=[]
id3=[]
for i in id1:
id2.append(-1)
id3.append(-1)
print(id1[i],id2[i],id3[i],cq1doc[i])
print('cqdoc嵌入向量维度是：',cpu_embedding.encode(np.array(cq1doc))[0].shape)
print('cpu的embed_query嵌入向量化函数耗时:',time.time()-start_time)
print('cpu的cqdoc嵌入向量是：',custom_embedding.embed_query(np.array(cq1doc))[0])
print('gpu与cpu嵌入向量各项相减的平方之和为:',
sum((np.array(cpu_embedding.encode(np.array(cq1doc))[0])-np.array(custom_embedding.embed_query(np.array(cq1doc))[0]))**2))
del(cpu_embedding)
输出为：
0 -1 -1 从江雾里长出的立体迷宫：重庆的烟火与野性
1 -1 -1 当飞机穿越厚厚的云层，机身下方突然跳出一片被江水缠绕、被高楼撑起来的城市时，你就知道，重庆到了。这座被称为“8D魔幻之都”的城市，从来不是平面的——它是一座在山水间野蛮生长的立体迷宫，每一条爬坡上坎的石阶，每一栋依山而建的高楼，每一碗翻滚着红油的火锅，都藏着属于自己的倔强与温柔。
2 -1 -1 一、两江环抱的天生骨架：山水里的城市基因
3 -1 -1 重庆的故事，是从两条江开始的。长江与嘉陵江在这……

Ids=[str(id1[j])+'-'+str(id2[j])+'-'+str(id3[j]) for j in range(len(id1))]
Ids
输出为：
['0--1--1',
'1--1--1',
'2--1--1',
'3--1--1',
……]

metas=[{"1级id":id1[j],"2级id":id2[j],"3级id":id3[j]} for j in range(len(id3))]
metas
输出为：
[{'1级id': 0, '2级id': -1, '3级id': -1},
{'1级id': 1, '2级id': -1, '3级id': -1},
{'1级id': 2, '2级id': -1, '3级id': -1},
……]

import time,torch
print('gpu:',torch.cuda.get_device_properties())
start_time=time.time()
print('cqdoc嵌入向量维度是：',custom_embedding.embed_query(np.array(cq1doc))[0].shape)
print('gpu的embed_query嵌入向量化函数耗时:',time.time()-start_time)
print('gpu的cqdoc嵌入向量是：',custom_embedding.embed_query(np.array(cq1doc))[0])
start_time=time.time()
print('cqdoc嵌入向量维度是：',cpu_embedding.encode(np.array(cq1doc))[0].shape)
print('cpu的embed_query嵌入向量化函数耗时:',time.time()-start_time)
print('cpu的cqdoc嵌入向量是：',custom_embedding.embed_query(np.array(cq1doc))[0])
print('gpu与cpu嵌入向量各项相减的平方之和为:',
sum((np.array(cpu_embedding.encode(np.array(cq1doc))[0])-np.array(custom_embedding.embed_query(np.array(cq1doc))[0]))**2))
del(cpu_embedding)
输出为：
gpu: _CudaDeviceProperties(name='NVIDIA GeForce RTX 3050 Ti Laptop GPU', major=8, minor=6, total_memory=4095MB, multi_processor_count=20, uuid=a08e96e6-83e8-080b-9af2-f991b4df0324, L2_cache_size=2MB)
cqdoc嵌入向量维度是： (1024,)
gpu的embed_query嵌入向量化函数耗时: 0.8678944110870361
gpu的cqdoc嵌入向量是： [ 0.03149414 0.02001953 -0.01391602 ... 0.04394531 0.03051758
-0.03051758]
cqdoc嵌入向量维度是： (1024,)
cpu的embed_query嵌入向量化函数耗时: 132.04911351203918
cpu的cqdoc嵌入向量是： [ 0.03149414 0.02001953 -0.01391602 ... 0.04394531 0.03051758
-0.03051758]
gpu与cpu嵌入向量各项相减的平方之和为: 0.00047348032

import time
start_time=time.time()
collection.add(ids=Ids,documents=cq1doc,metadatas=metas)
print(time.time()-start_time)
输出为：0.8988301753997803

collection.get(where={'$or':[{'1级id':{'$gt':24}},{'1级id':{'$eq':8}},{'1级id':{'$lt':4}}]})
输出为：
{'ids': ['0--1--1',
'1--1--1',
'2--1--1',
'3--1--1',
'8--1--1',
'25--1--1',
'26--1--1',
'27--1--1'],
'embeddings': None,
'documents': ['从江雾里长出的立体迷宫：重庆的烟火与野性',
'当飞机穿越厚厚的云层，机身下方突然跳出一片被江水缠绕、被高楼撑起来的城市时，你就知道，重庆到了。这座被称为“8D魔幻之都”的城市，从来不是平面的——它是一座在山水间野蛮生长的立体迷宫，每一条爬坡上坎的石阶，每一栋依山而建的高楼，每一碗翻滚着红油的火锅，都藏着属于自己的倔强与温柔。',
'一、两江环抱的天生骨架：山水里的城市基因',
'重庆的故事，是从两条江开始的。长江与嘉陵江在这里交汇，像两条蓝色的丝带，把城市分割成渝中、江北、南岸三个半岛，也赋予了重庆与生俱来的江湖气质。',
'如果说山水是重庆的骨架，那美食就是重庆的灵魂。这座城市的味道，是麻辣鲜香的，是热情奔放的，是刻在每一个重庆人骨子里的。',
'重庆的文化事业也在蓬勃发展。近年来，重庆举办了一系列具有影响力的文化活动，比如重庆国际马拉松赛、重庆国际音乐节、重庆国际艺术双年展等，吸引了来自世界各地的游客和艺术家。同时，重庆还注重保护和传承传统文化，比如川剧、铜梁龙舞、秀山花灯等，这些传统文化在新时代焕发出了新的活力。',
'结语：永远让人着迷的重庆重庆是一座复杂而多元的城市，它既有山水的诗意，又有江湖的野性；既有市井的烟火，又有历史的厚重；既有传统的韵味，又有现代的活力。它像一本永远读不完的书，每一页都有新的惊喜。',
'如果你问我，重庆到底是什么样子的？我会告诉你，重庆是清晨小面摊的热气，是深夜火锅店的喧闹；是洪崖洞的灯光，是李子坝的轻轨；是十八梯的石板路，是黄桷垭的老街；是红岩村的红色记忆，是解放碑的繁华热闹。它是一座让人来了就不想走的城市，一座让人永远着迷的城市。'],……

collection.get(include=["documents", "embeddings"],where={'1级id': 8})
输出为：
{'ids': ['8--1--1'],
'embeddings': array([[-0.06445313, 0.00405884, -0.0090332 , ..., 0.02685547,
0.02404785, -0.03564453]]),
'documents': ['如果说山水是重庆的骨架，那美食就是重庆的灵魂。这座城市的味道，是麻辣鲜香的，是热情奔放的，是刻在每一个重庆人骨子里的。'],
'uris': None,
'included': ['documents', 'embeddings'],
'data': None,
'metadatas': None}

可以用collection.query进行全集合模糊搜索，库中没有但语义有一定相似的内容前若干个都会被列出来，这个检索过程简称为召回，召回结果中有些就肯定是错的，但也可能有近似对的。简单理解，chromadb这样的向量库，把长文本应分层拆开成段落、子段落，召回结果再用AI处理，就可得到想要结果了，但AI处理提示词是精华。

AI处理召回的大堆结果，提示词例：{召回结果}是{…问题}检索到的多个可能有关信息，其中部分可能完全无关甚至有错，请参考后作出回答

交互界面中可使用召回后AI效果评分或收录快捷键，对检索问答优秀的特别是命名实体可添加进附加第三层子段落，以提高专有问题响应质量与速度，这也算人类反馈强化学习以改进知识库了。