如何在LangChain+ChromaDB中基于Base Embeddings实现自定义本地嵌入模型

核心逻辑 #

LangChain的BaseEmbeddings是抽象基类，本身不能直接实例化，必须通过子类继承并实现它规定的两个核心方法：

• embed_documents(texts: List[str]) -> List[List[float]]：批量处理文档生成嵌入向量
• embed_query(text: str) -> List[float]：处理单个查询生成嵌入向量

如果需要异步调用，还可以选择性重写aembed_documents和aembed_query方法。

具体实现步骤 #

1. 导入依赖 #

from langchain.embeddings.base import BaseEmbeddings
from langchain.vectorstores import Chroma
from typing import List

2. 子类化BaseEmbeddings实现自定义嵌入逻辑 #

这里用一个简单示例演示，实际场景替换成你的模型推理代码即可：

class CustomEmbeddings(BaseEmbeddings):
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        # 替换成你的批量文档嵌入逻辑
        # 示例：生成128维随机向量（仅作演示，实际用真实模型输出）
        import random
        return [[random.random() for _ in range(128)] for _ in texts]

    def embed_query(self, text: str) -> List[float]:
        # 替换成你的单个查询嵌入逻辑
        import random
        return [random.random() for _ in range(128)]
    
    # 可选：实现异步方法（需要异步调用时用）
    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
        return self.embed_documents(texts)
    
    async def aembed_query(self, text: str) -> List[float]:
        return self.embed_query(text)

3. 和ChromaDB结合使用 #

实例化自定义嵌入类，直接传入Chroma初始化即可：

# 初始化自定义嵌入模型
custom_embeddings = CustomEmbeddings()

# 基于文本创建Chroma向量库
texts = ["自定义嵌入测试文档1", "自定义嵌入测试文档2"]
db = Chroma.from_texts(texts, custom_embeddings)

# 执行相似性查询
query = "测试查询"
docs = db.similarity_search(query)
print(docs)

注意事项 #

• 确保自定义嵌入的向量维度和ChromaDB配置一致（示例用128维，实际要和你的模型输出匹配）
• 如果用本地部署的模型做嵌入源，直接在embed_documents和embed_query里调用模型推理接口就行
• 批量处理大量文档时，可以在方法里加并发逻辑提升效率

内容的提问来源于stack exchange，提问作者Vardhan Gupta