Chroma 对单次调用的记录数有上限,用 client.get_max_batch_size() 能查到:
import chromadb
client = chromadb.EphemeralClient()
print(client.get_max_batch_size()) # 5461这个数不是给你当目标用的。实际经验是每批 100~1000 条比较舒服,再大内存和超时都会出问题(尤其是 HttpClient 模式,payload 太大会被服务端拒绝,默认上限 40 MiB)。
一个通用的分批工具:
from typing import Iterator, Sequence, TypeVar
T = TypeVar("T")
def chunked(seq: Sequence[T], size: int) -> Iterator[Sequence[T]]:
"""把序列按 size 切成一段一段,最后一段可能不够 size
例如:list(chunked([1, 2, 3, 4, 5], 2))
[[1, 2], [3, 4], [5]]
"""
# 步长取 size,起点 0 / size / 2*size ...
# 末尾不足 size 时自动只切到结尾
for i in range(0, len(seq), size):
yield seq[i:i + size]
def add_in_batches(col, ids, embeddings, documents=None, metadatas=None,
batch_size=500, verbose=False):
"""分批写入,返回实际写入条数"""
written = 0
# ids / embeddings / documents / metadatas 是按下标一一对应的,
# 所以先切下标、再按同一批下标去取各自的分片,保证一批里各字段条数一致
for idx_batch in chunked(range(len(ids)), batch_size):
sl = slice(idx_batch[0], idx_batch[-1] + 1)
col.add(
ids=ids[sl],
embeddings=embeddings[sl] if embeddings else None,
documents=documents[sl] if documents else None,
metadatas=metadatas[sl] if metadatas else None,
)
written += len(idx_batch)
if verbose:
print(f"已写入 {written}/{len(ids)}")
return written
if __name__ == "__main__":
import chromadb
# 20 条数据按 6 条一批写入
ids = [f"id{i}" for i in range(20)]
embeddings = [[float(i), float(i % 3), 1.0] for i in range(20)]
documents = [f"第 {i} 条文档" for i in range(20)]
metadatas = [{"idx": i, "even": i % 2 == 0} for i in range(20)]
# 自己给向量,就不用下载默认的嵌入模型了
col = chromadb.EphemeralClient().get_or_create_collection(
"batch_demo", embedding_function=None
)
print("写入条数:", add_in_batches(
col, ids, embeddings, documents, metadatas, batch_size=6, verbose=True
))
print("集合里现在有:", col.count(), "条")运行脚本,输出如下:
已写入 6/20
已写入 12/20
已写入 18/20
已写入 20/20
写入条数: 20
集合里现在有: 20 条注意:分批的时候千万别用 embeddings[i:i+n] 这种写法去索引 numpy 的 2D 数组再 tolist(),容易把维度搞丢。建议统一转成 List[List[float]] 再切片。例如:
import numpy as np
# 2维数组:3条向量,每条向量长度4
embeddings = np.array([
[0.1,0.2,0.3,0.4],
[0.5,0.6,0.7,0.8],
[0.9,1.0,1.1,1.2]
])
print(embeddings.shape) # (3,4)使用 embeddings[i:i+n] 进行 numpy 原生切片:
batch = embeddings[0:1]
print(batch.shape) # (1,4) 二维,没问题
# 如果用单索引 embeddings[0],不是切片
single = embeddings[0]
print(single.shape) # (4,) 变成一维数组很多人代码写着写着混淆区间切片 i:i+n 和单点索引 [i],为了避开这个坑,建议统一使用分片处理。