chromadb emmbedding 向量检索

前言

最近做项目,其实就是poc,做一个知识向量检索,为rag(检索增强生成)做准备,其实就是LLM不具备的知识,先检索出来,然后提示词喂给大模型,让LLM开卷作答,检索的内容准确性决定LLM答题的准确性。其实选项无非是开源的向量数据库选择,一般而言示例poc使用chromadb就够了,如果需要大量数据检索,业界一般使用Milvus,当然其他ck等也可以使用,甚至可以把检索和存储分离实现。

准备

以chromadb为例,根据官方文档:Chroma Clients - Chroma Docs

为了方便,使用clients模式,实际生产需要服务器模式。

其中核心参数为:向量检索算法

默认使用l2 欧式距离,还有内积和余弦算法,先默认l2吧

使用的embedding function其实默认也有内嵌模型,这里使用

Qwen3-Embedding-4B-4bit-DWQ(其实有8b的,耗资源相对较多)

其实除了向量检索,还可以分词,词检索,类似ES,这里的Bm25(best match 25个版本)即词检索,其实也可以混用,取决于实际场景。

使用openai的标准api,Python代码实现poc。

demo

以mac ipad 等涨价的新闻为例

文本准备

bash 复制代码
涨价新闻
6 月 25 日消息,苹果公司宣布上调 Mac、iPad 等产品价格,以应对 AI 数据中心扩张引发的内存芯片及存储器空前短缺所带来的成本压力。

苹果在声明中表示:"消费电子行业正面临前所未有的挑战。AI 数据中心的迅猛扩张导致存储需求激增,我们从未见过零部件价格以如此幅度和速度上涨。此前我们一直在内部消化成本压力,但目前已不得不开始上调多款产品售价,其中就包括今天公布的 iPad 和 Mac。"

附起售价变化如下:

Mac系列
MacBook Neo 的起售价从 4599 元上调至 5499 元,涨价 900 元

MacBook Air 的起售价从 8499 元上调至 9999 元,涨价 1500 元

MacBook Pro 的起售价从 13499 元上调至 15999 元,涨价 2500 元

iMac 的起售价从 10999 元上调至 12499 元,涨价 1500 元

Mac mini 的起售价从 4499 元上调至 5999 元,涨价 1500 元

Mac Studio 的起售价从 16499 元上调至 19999 元,涨价 3500 元

苹果提高 Mac、iPad、Vision Pro、HomePod 等产品价格,以应对内存短缺

iPad系列
iPad Pro 的起售价从 8999 元上调至 10799 元,涨价 1800 元

iPad Air 的起售价从 4799 元上调至 5999 元,涨价 1200 元

iPad mini 的起售价从 3999 元上调至 4799 元,涨价 800 元

iPad(A16)的起售价从 2999 元上调至 3799 元,涨价 800 元

苹果提高 Mac、iPad、Vision Pro、HomePod 等产品价格,以应对内存短缺

其他系列
Vision Pro
Vision Pro 的起售价从 29999 元上调至 31999 元,涨价 2000 元

HomePod
HomePod 的起售价从 2299 元上调至 2699 元,涨价 400 元

HomePod mini 的起售价从 749 元上调至 999 元,涨价 250 元

苹果提高 Mac、iPad、Vision Pro、HomePod 等产品价格,以应对内存短缺

Apple TV
Apple TV 的起售价从 129 美元上调至 199 美元(IT之家注:现汇率约合 1357 元人民币),涨价 70 美元(现汇率约合 477.2 元人民币)

来向量搜索,其实就是用了模型的向量化能力,需要模型,笔者使用开源的qwen3

模型准备

做向量检索需要使用embedding模型,开启openai标准

chromadb操作

使用openai和chromadb,执行embedding-function的嵌入和集合的新增和查找

python 复制代码
from typing import Dict, Any
from chromadb import Documents, EmbeddingFunction, Embeddings
from chromadb.utils.embedding_functions import register_embedding_function
from openai import OpenAI


@register_embedding_function
class CustomEmbeddingFunction(EmbeddingFunction[Documents]):

    def __init__(self, model: str = "Qwen3-Embedding-4B-4bit-DWQ", *args: Any, **kwargs: Any):
        self.model = model
        self.client = OpenAI(api_key="xxx", base_url="http://127.0.0.1:8000/v1")

    def __call__(self, input: Documents) -> Embeddings:
        if not input:
            return Embeddings.empty()
        # embed the documents somehow
        embeddings = []
        for i in range(0, len(input), 10):
            text = input[i:i + 10]
            data = self.client.embeddings.create(input=text, model=self.model).data
            embeddings.extend(y.embedding for y in data)
        return embeddings

    @staticmethod
    def name() -> str:
        return "Qwen3-Embedding-4B-4bit-DWQ"

    def get_config(self) -> Dict[str, Any]:
        return dict(model=self.model)

    @staticmethod
    def build_from_config(config: Dict[str, Any]) -> "EmbeddingFunction":
        return CustomEmbeddingFunction(config["Qwen3-Embedding-4B-4bit-DWQ"])

根据官方文档,写了一个自定义embedding_function,其实核心就是使用embedding模型的向量化能力,然后编写chromadb的操作逻辑

python 复制代码
import chromadb
import openai
from chromadb import Settings, ReadLevel
from openai import OpenAI

import custom_embedding_func


class ChromaDBDataBase:
    def __init__(self, path, collection_name):
        chromadb_client = chromadb.PersistentClient(path=path)
        self.collection = chromadb_client.get_or_create_collection(name=collection_name,
                                                                   embedding_function=custom_embedding_func.CustomEmbeddingFunction(),
                                                                   metadata={"hnsw:space":"l2"})

    def addDoc(self, docs, ids):
        self.collection.add(documents=docs, ids=ids)

    def query(self, query):
        result = self.collection.query(query_texts=[query], n_results=2)
        for i,doc in enumerate(result['documents'][0]):
            distance = result['distances'][0][i]
            print(f'doc{i+1}: {doc}')
            print(f'distance: {distance:4f}\n')
        return result


    def search(self, query):
        return self.collection.search(searches=[query], read_level = ReadLevel.INDEX_AND_WAL)

至此核心的逻辑写完,写个main方法,然后简单的写写分割文本逻辑,这里使用langchain的api

python 复制代码
from chroma_db import ChromaDBDataBase
from langchain_text_splitters import RecursiveCharacterTextSplitter

if __name__ == '__main__':
    with open("./demo.txt", "r", encoding="utf-8") as f:
        content = f.read()

    c_db = ChromaDBDataBase("./chromadb_demo", "apple_price")
    splitter = RecursiveCharacterTextSplitter(chunk_size=200,chunk_overlap=50, separators=['涨价新闻','Mac系列','ipad系列','其他系列', '\n\n', '\n', ' ', '?', ',', ','])
    texts = splitter.split_text(content)
    ids = [f'id{i}' for i in range(len(texts))]
    c_db.addDoc(texts, ids)

    query_text = 'mac mini 涨价了吗'
    res = c_db.query(query_text)
    print(res)

    search_text = 'mac studio涨价多少?'
    search_res = c_db.search(search_text)
    print(search_res)

执行后,看看这种情况的效果怎么样,当然实际场景的文本切分可以定制场景,这里比较简单粗暴

结果分析

bash 复制代码
doc1: iMac 的起售价从 10999 元上调至 12499 元,涨价 1500 元

Mac mini 的起售价从 4499 元上调至 5999 元,涨价 1500 元

Mac Studio 的起售价从 16499 元上调至 19999 元,涨价 3500 元

苹果提高 Mac、iPad、Vision Pro、HomePod 等产品价格,以应对内存短缺
distance: 0.515388

doc2: Mac系列
MacBook Neo 的起售价从 4599 元上调至 5499 元,涨价 900 元

MacBook Air 的起售价从 8499 元上调至 9999 元,涨价 1500 元

MacBook Pro 的起售价从 13499 元上调至 15999 元,涨价 2500 元

iMac 的起售价从 10999 元上调至 12499 元,涨价 1500 元
distance: 0.586438

{'ids': [['id3', 'id2']], 'embeddings': None, 'documents': [['iMac 的起售价从 10999 元上调至 12499 元,涨价 1500 元\n\nMac mini 的起售价从 4499 元上调至 5999 元,涨价 1500 元\n\nMac Studio 的起售价从 16499 元上调至 19999 元,涨价 3500 元\n\n苹果提高 Mac、iPad、Vision Pro、HomePod 等产品价格,以应对内存短缺', 'Mac系列\nMacBook Neo 的起售价从 4599 元上调至 5499 元,涨价 900 元\n\nMacBook Air 的起售价从 8499 元上调至 9999 元,涨价 1500 元\n\nMacBook Pro 的起售价从 13499 元上调至 15999 元,涨价 2500 元\n\niMac 的起售价从 10999 元上调至 12499 元,涨价 1500 元']], 'uris': None, 'included': ['metadatas', 'documents', 'distances'], 'data': None, 'metadatas': [[None, None]], 'distances': [[0.5153882503509521, 0.5864384770393372]]}
Traceback (most recent call last):
  File "/Users/huahua/PycharmProjects/rag-demo/main.py", line 20, in <module>
    search_res = c_db.search(search_text)
  File "/Users/huahua/PycharmProjects/rag-demo/chroma_db.py", line 29, in search
    return self.collection.search(searches=[query], read_level = ReadLevel.INDEX_AND_WAL)
           ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/huahua/PycharmProjects/rag-demo/.venv/lib/python3.14/site-packages/chromadb/api/models/Collection.py", line 433, in search
    return self._client._search(
           ~~~~~~~~~~~~~~~~~~~~^
        collection_id=self.id,
        ^^^^^^^^^^^^^^^^^^^^^^
    ...<3 lines>...
        read_level=read_level,
        ^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/Users/huahua/PycharmProjects/rag-demo/.venv/lib/python3.14/site-packages/chromadb/api/rust.py", line 387, in _search
    raise NotImplementedError("Search is not implemented for Local Chroma")
NotImplementedError: Search is not implemented for Local Chroma

从结果看是非常准确的,这里search居然不支持local模式

比如我问一下ipad air涨价了吗

python 复制代码
doc1: 苹果提高 Mac、iPad、Vision Pro、HomePod 等产品价格,以应对内存短缺

iPad系列
iPad Pro 的起售价从 8999 元上调至 10799 元,涨价 1800 元

iPad Air 的起售价从 4799 元上调至 5999 元,涨价 1200 元

iPad mini 的起售价从 3999 元上调至 4799 元,涨价 800 元
distance: 0.539649

doc2: iPad mini 的起售价从 3999 元上调至 4799 元,涨价 800 元

iPad(A16)的起售价从 2999 元上调至 3799 元,涨价 800 元

苹果提高 Mac、iPad、Vision Pro、HomePod 等产品价格,以应对内存短缺
distance: 0.579633

{'ids': [['id4', 'id5']], 'embeddings': None, 'documents': [['苹果提高 Mac、iPad、Vision Pro、HomePod 等产品价格,以应对内存短缺\n\niPad系列\niPad Pro 的起售价从 8999 元上调至 10799 元,涨价 1800 元\n\niPad Air 的起售价从 4799 元上调至 5999 元,涨价 1200 元\n\niPad mini 的起售价从 3999 元上调至 4799 元,涨价 800 元', 'iPad mini 的起售价从 3999 元上调至 4799 元,涨价 800 元\n\niPad(A16)的起售价从 2999 元上调至 3799 元,涨价 800 元\n\n苹果提高 Mac、iPad、Vision Pro、HomePod 等产品价格,以应对内存短缺']], 'uris': None, 'included': ['metadatas', 'documents', 'distances'], 'data': None, 'metadatas': [[None, None]], 'distances': [[0.5396493077278137, 0.5796329379081726]]}

效果还是不错的,这里的chromadb实际上使用的sqlite充当结构化数据库存储,实际场景需要一些类似PG等数据库来存储,向量数据库本质是一个id,一份数据存储向量,一份数据存储文本(分割后)信息。

集合信息,可以创建多个集合,通过关系型数据库关联关系

算法说明

向量化embedding:一般大学的数据向量是2维或者3维的,但是embedding模型的向量化维度可以是成百上千维,比如qwen,默认1024维,维度越高效果越好,但是计算复杂度越高,取舍问题。

上面说了向量检索算法,余弦,欧式距离,内积

余弦:其实就是向量的夹角,夹角越小越相似,余弦有范围限制-1~1,归一化后性能会提升

欧式距离:绝对距离,当然越小越好

内积:其实也是向量的夹角,其实内积和余弦的公式是有相关性的,缺少归一化处理

总结

其实向量检索以前就有了,只不过现在结合embedding模型做向量化,再结合LLM做AI处理。本质上LLM训练后的数据也是embedding(词向量),先tokenizer分词,然后编码,通过自注意力机制,softmax等算法,生产式AI,其实检索也支持多维度检索,比如词检索和向量检索同时使用,当然chromadb本身也支持bm25 embedding。

这里就需要简单说一下为什么需要检索,llm本身能力这么强,但是知识是变化的,模型训练的知识和时间决定了llm本身能力的上限,当遇到没有训练的知识,或者最新的知识,模型就识别不了,此时有3种主流的支持方式:

  1. 检索生成-开卷考试

  2. 模型微调-把知识开个旁路训练进去,可能会影响原本模型的能力

  3. 重新训练-效果最后,成本最高

一般而言大部分场景都不可能自己训练模型,一般是做微调和检索生成,检索生成成本低,开发时间最短,通过调节分块的算法优化和embedding function的处理,基本上能解决大部分场景的问题

相关推荐
cui_ruicheng1 小时前
Python从入门到实战(六):非序列容器
开发语言·python
月光刺眼1 小时前
用LangChain 打造第一个 Agent:LLM 大脑与 Tool 手脚的完整闭环
javascript·人工智能·全栈
小林ixn1 小时前
手写AI编程Agent:从0到1复刻Trae/Cursor的核心能力
人工智能·aigc·ai编程
冬哥聊AI1 小时前
腾讯三面:Claude Code 写书超过上下文长度怎么办?/compact 只是应急,五步架构才是正解
人工智能
饼干哥哥1 小时前
用Claude Code跑通跨境电商的5大场景,顶10个人的团队
人工智能·设计模式·正则表达式
ATA88881 小时前
数据库管理工具的数据安全机制对比:权限管控与审计
数据库·人工智能·数据分析·数据库管理
legendary_1631 小时前
SINK芯片:Type-C统一供电时代的小家电核心方案
c语言·开发语言·人工智能·智能手机
私人珍藏库1 小时前
[Android] 发票制作器-内置多行业专业模板
android·人工智能·app·软件·多功能
桐桐桐1 小时前
深入理解 MCP:AI 连接外部世界的标准协议
人工智能