1、通过Docker快速部署ChromaDB
docker run -d --name chromadb-container -p 8000:8000 chromadb/chroma
2、打开ChromaDB提供的swagger地址 http://localhost:8000/docs
3、构建代码测试向量
以下代码做了基础测试,创建集合、文字转向量、插入集合、检索数据具体的可以查看Swagger文档提供了更多的操作,运行当前代码可以根据语意搜索到对应的内容,文字转向量需要部署应用Xinference
import requests
import hashlib
# 文字转向量
def embedding(data):
url = "http://10.10.40.102:19997/v1/embeddings"
headers = {
"Content-Type": "application/json"
}
data = {
"input": data,
"model": "bge-large-zh-v1.5"
}
response = requests.post(url, headers=headers, json=data)
vectorData = [item['embedding'] for item in response.json()['data']]
return vectorData
# 创建集合
def create_collections(json):
result = requests.post("http://localhost:8000/api/v1/collections?tenant=default_tenant&database=default_database",
json=json)
return result
# 获取集合列表
def get_collections():
result = requests.get("http://localhost:8000/api/v1/collections?tenant=default_tenant&database=default_database")
return result
# 插入向量
def add_collections_vectory(collections_id, json):
result = requests.post(f"http://localhost:8000/api/v1/collections/{collections_id}/add", json=json)
print(result.json())
# 根据名称获取集合
def name_get_collections(collections_name):
result = requests.get(
f"http://localhost:8000/api/v1/collections/{collections_name}?tenant=default_tenant&database=default_database")
return result
def delete_collections(collections_name):
result = requests.delete(
f"http://localhost:8000/api/v1/collections/{collections_name}?tenant=default_tenant&database=default_database")
print(result.json())
# 数组转md5
def convert_to_md5(data):
return [hashlib.md5(item.encode('utf-8')).hexdigest() for item in data]
# 添加向量数据
def conversion_data(raw_data):
# 转唯一ID
ids = convert_to_md5(raw_data)
# 文字转向量
embeddings = embedding(raw_data)
# 转向量
return {"embeddings": embeddings, "documents": raw_data, "ids": ids}
def retrieval(collections_name, name):
collections_id = name_get_collections(collections_name).json().get("id")
fruit = embedding(name)
result = requests.post(f"http://localhost:8000/api/v1/collections/{collections_id}/query", json={
"query_embeddings": fruit,
"n_results": 5, # 返回结果数量
"include": [
"metadatas",
"distances",
"documents"
]
})
print(result.json().get("distances"))
print(result.json().get("documents"))
def generate(collections_name, raw_data):
# 删除之前的集合
delete_collections(collections_name)
# 创建新集合
create_collections({"name": collections_name})
# 获取集合ID
collections_id = name_get_collections(collections_name).json().get("id")
# 把向量插入到chromadb
add_collections_vectory(collections_id, conversion_data(raw_data))
# 生成数据插入到集合
generate("test", ["苹果", "香蕉", "汽车", "狗", "猫", "大象", "飞机", "坦克"])
# 检索数据
retrieval("test", "交通工具")