python - 如何将 torch.tensor 或 np.array 保存到 redis 并搜索向量相似度？

Question

我在使用 python 代码将数据保存到 redis 时遇到了麻烦。 只需使用 redis 和 r.ft()

上传数据会是这样的。我还想刷新相同 id 中不同值的嵌入。

id 是数据索引，嵌入将在所有数据之间以相同的形状展平。 (例如 1024) id 嵌入 0 [3.1515, 4.5562, ..., ] 1 [3, 8.62, ..., ]

上传redis后，我想用redis搜索某个批处理的嵌入。

如果输入批处理形状为 [3, 1024]，则搜索应该迭代到批处理并返回 [3, top-k] 与 redis 中的嵌入相似的 id。

现在我真的很难做到这一点。等待帮助。

Answer 1

首先是一些有用的链接:This notebook有一些有用的例子，这里是 RediSearch docs for using vector similarity ，最后，这是一个 example app所有这些都汇集在一起​​。

要将 numpy 数组存储为 Redis 中的向量字段，您需要首先在架构中使用 VectorField 创建搜索索引:

import numpy as np
import redis

from redis.commands.search.indexDefinition import (
    IndexDefinition,
    IndexType
)
from redis.commands.search.query import Query
from redis.commands.search.field import (
    TextField,
    VectorField
)

# connect
r = redis.Redis(...)

# define vector field
fields = [VectorField("vector",
    "FLAT", {
        "TYPE": "FLOAT32",
        "DIM": 1024,  # 1024 dimensions
        "DISTANCE_METRIC": "COSINE",
        "INITIAL_CAP": 10000, # approx initial count of docs in the index
    }
)]

# create search index
r.ft(INDEX_NAME).create_index(
    fields = fields,
    definition = IndexDefinition(prefix=["doc:"], index_type=IndexType.HASH)
)

有了索引后，您可以使用 hset 和 pipeline 将数据写入 Redis。 Redis 中的向量存储为字节字符串(请参阅下面的 tobytes()):

# random vectors
vectors = np.random.rand(10000, 1024).astype(np.float32)

pipe = r.pipeline(transaction=False)
for id_, vector in enumerate(vectors):
    pipe.hset(key=f"doc:{id_}", mapping={"id": id_, "vector": vector.tobytes()})
    if id_ % 100 == 0:
        pipe.execute() # write batch
pipe.execute() # cleanup

开箱即用，您可以使用pipeline调用通过一次API调用来多次查询Redis:

base_query = f'*=>[KNN 5 @vector $vector AS vector_score]'
query = (
    Query(base_query)
    .sort_by("vector_score")
    .paging(0, 5)
    .dialect(2)
)
query_vectors = np.random.rand(3, 1024).astype(np.float32)

# pipeline calls to redis
pipe = r.pipeline(transaction=False)
for vector in query_vectors:
    pipe.ft(INDEX_NAME).search(query, {"vector": query_vector.tobytes()})
res = pipe.execute()

然后，您需要解压 res 对象，其中包含来自 Redis 的所有三个查询的原始响应。希望这会有所帮助。