使用 ChromaDB 创建向量数据库

使用 OpenAI API 的 Embeddings 入门

Emmanuel Pire

Senior Software Engineer, DataCamp

安装 ChromaDB

  • ChromaDB 是一个简单而强大的向量数据库
  • 两种模式：
    • 本地：适合开发与原型
    • 客户端/服务器：适合生产

本地模式与客户端/服务器模式下的 ChromaDB。客户端/服务器模式中客户端与服务器分离显示。

使用 OpenAI API 的 Embeddings 入门

连接数据库

import chromadb


client = chromadb.PersistentClient(path="/path/to/save/to")
  • 数据将持久化到磁盘
使用 OpenAI API 的 Embeddings 入门

创建集合

  • 集合相当于表
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction

collection = client.create_collection(
    name="my_collection",

    embedding_function=OpenAIEmbeddingFunction(
        model_name="text-embedding-3-small",
        api_key="..."
    )

)
  • 集合可自动创建嵌入
使用 OpenAI API 的 Embeddings 入门

查看集合

client.list_collections()

[Collection(name=my_collection)]
使用 OpenAI API 的 Embeddings 入门

插入嵌入

单个文档

collection.add(ids=["my-doc"], documents=["This is the source text"])
  • 必须提供 ID
  • 嵌入将由集合创建！

多个文档

collection.add(
  ids=["my-doc-1", "my-doc-2"], 
  documents=["This is document 1", "This is document 2"]
)
使用 OpenAI API 的 Embeddings 入门

检查集合

统计集合中的文档数

collection.count()

3
使用 OpenAI API 的 Embeddings 入门

检查集合

查看前 10 项

collection.peek()

{'ids': ['my-doc', 'my-doc-1', 'my-doc-2'],
 'embeddings': [[...], [...], [...]],
 'documents': ['This is the source text',
  'This is document 1',
  'This is document 2'],
 'metadatas': [None, None, None]}
使用 OpenAI API 的 Embeddings 入门

检索条目

collection.get(ids=["s59"])

{'ids': ['s59'],
 'embeddings': None,
 'metadatas': [None],
 'documents': ['Title: Naruto Shippûden the Movie: The Will of Fire (Movie)\nDescription: When ...'],
 'uris': None,
 'data': None}
使用 OpenAI API 的 Embeddings 入门

Netflix 数据集

 

Title: Kota Factory (TV Show)
Description: In a city of coaching centers known to train India's finest...
Categories: International TV Shows, Romantic TV Shows, TV Comedies

Title: The Last Letter From Your Lover (Movie)
Description: After finding a trove of love letters from 1965, a reporter sets...
Categories: Dramas, Romantic Movies
使用 OpenAI API 的 Embeddings 入门

估算嵌入成本

  • 嵌入模型（text-embedding-3-small）费用为 $0.00002/1k tokens
cost = 0.00002 * len(tokens)/1000
  • 使用 tiktoken 统计 token 数
    • pip install tiktoken
1 https://openai.com/pricing
使用 OpenAI API 的 Embeddings 入门

估算嵌入成本

import tiktoken

enc = tiktoken.encoding_for_model("text-embedding-3-small")


total_tokens = sum(len(enc.encode(text)) for text in documents)


cost_per_1k_tokens = 0.00002

print('Total tokens:', total_tokens)
print('Cost:', cost_per_1k_tokens * total_tokens/1000)

Total tokens: 444463
Cost: 0.00888926
使用 OpenAI API 的 Embeddings 入门

让我们来练习！

使用 OpenAI API 的 Embeddings 入门

Preparing Video For Download...