一架梯子,一头程序猿,仰望星空!
Chroma向量数据库教程 > 内容正文

JavaScript Chromadb 向量计算(文本嵌入模型)


Embeddings 向量是人工智能表示任何类型数据的原生方式,很多算法都使用特征向量表示数据。它们可以表示文本、图像以及音频和视频。计算特征向量有很多方法,可以使用开源的嵌入模型,也可以使用一些云服务的API,例如openai的API。

Chroma 为流行的嵌入模型提供了轻量级封装,让您可以轻松地在应用程序中使用它们。

Transformers.js

Chroma 可以使用 Transformers.js 在本地机器上运行嵌入模型,生成特征向量。Transformers 使用 “Xenova/all-MiniLM-L6-v2 “模型。在命令行中运行 npm install @xenova/transformers,确保已安装 Transformers.js 库。

const {ChromaClient} = require('chromadb');
const client = new ChromaClient({path: "http://localhost:8000"});
const {TransformersEmbeddingFunction} = require('chromadb');
const embedder = new TransformersEmbeddingFunction();

(async () => {
    // 通过embeddingFunction参数指定使用的嵌入函数
    const collection = await client.getOrCreateCollection({name: "name", embeddingFunction: embedder})

    //添加数据,不指定向量,add函数会调用embeddingFunction函数计算文本向量
    await collection.add({
        ids: ["id1", "id2", "id3"],
        metadatas: [{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}], 
        documents: ["lorem ipsum...", "doc2", "doc3"], 
    })

    // 查询数据,也会调用embeddingFunction函数计算queryTexts的文本向量,然后去查询相似的数据
    const results = await collection.query({
        nResults: 2, 
        queryTexts: ["lorem ipsum"]
    }) 
})();

OPENAI

Chroma 为 OpenAI 的嵌入模型 API 提供了方便的封装。应为需要调用openai的api,需要一个 API 密钥。你可以通过注册 OpenAI 账户获得 API 密钥。

const {OpenAIEmbeddingFunction} = require('chromadb');
const embedder = new OpenAIEmbeddingFunction({openai_api_key: "apiKey"})

// use directly 
const embeddings = embedder.generate(["document1","document2"])

// 通过embeddingFunction参数指定嵌入函数
const collection = await client.createCollection({name: "name", embeddingFunction: embedder})
const collection = await client.getCollection({name: "name", embeddingFunction: embedder})

你可以传入一个可选的 model_name 参数,让你选择要使用的 OpenAI 嵌入模型。Chroma 默认使用 text-embedding-ada-002。

自定义嵌入函数

如果你使用其他嵌入模型,可以集成EmbeddingFunction基类,调用自己的选择的模型即可

提示:自定义嵌入函数不是必须的,因为你完全可以提前使用自己的嵌入模型计算好向量,然后使用向量去读写chroma数据。

class MyEmbeddingFunction {
  private api_key: string;

  constructor(api_key: string) {
    this.api_key = api_key;
  }

  public async generate(texts: string[]): Promise<number[][]> {
    // 调用自己的模型计算向量
    return embeddings;
  }
}

关联主题