索引和查询向量

了解如何在 Redis 中索引和查询向量嵌入

Redis 查询引擎允许您索引哈希JSON对象中的向量字段(更多信息请参阅向量参考页面)。

向量字段可以存储文本嵌入,即 AI 生成的文本内容的向量表示。两个嵌入之间的向量距离衡量它们的语义相似度。当您比较查询嵌入与存储的嵌入的相似度时,Redis 可以检索与查询含义密切相关的文档。

在下面的示例中,我们使用@xenova/transformers库生成向量嵌入,并将其存储和索引到 Redis 查询引擎中。代码首先演示了如何用于哈希文档,然后有一个单独的部分解释与 JSON 文档的区别

初始化

安装所需的依赖项

  1. 如果尚未安装,请安装node-redis
  2. 安装 @xenova/transformers
npm install @xenova/transformers

在你的 JavaScript 源文件中,导入所需的类

import * as transformers from '@xenova/transformers';
import {VectorAlgorithms, createClient, SchemaFieldTypes} from 'redis';

@xenova/transformers 模块处理嵌入模型。本示例使用 all-distilroberta-v1 模型,该模型

  • 生成 768 维向量
  • 将输入截断为 128 个 token
  • 使用 Word piece 分词(有关详细信息,请参阅 Hugging Face 文档中的Word piece 分词

pipe 函数生成嵌入。pipeOptions 对象指定如何从 token 嵌入生成句子嵌入(详细信息请参阅all-distilroberta-v1 文档)

let pipe = await transformers.pipeline(
    'feature-extraction', 'Xenova/all-distilroberta-v1'
);

const pipeOptions = {
    pooling: 'mean',
    normalize: true,
};

创建索引

首先,连接到 Redis 并删除任何名为 vector_idx 的现有索引

const client = createClient({url: 'redis://localhost:6379'});
await client.connect();

try { 
    await client.ft.dropIndex('vector_idx'); 
} catch (e) {
    // Index doesn't exist, which is fine
}

接下来,使用以下 schema 创建索引

  • content: 用于索引内容的文本字段
  • genre: 标签字段,表示文本的流派
  • embedding: 向量字段,具有
    • HNSW 索引
    • L2 距离指标
    • Float32 值
    • 768 维 (与嵌入模型匹配)
await client.ft.create('vector_idx', {
    'content': {
        type: SchemaFieldTypes.TEXT,
    },
    'genre': {
        type: SchemaFieldTypes.TAG,
    },
    'embedding': {
        type: SchemaFieldTypes.VECTOR,
        TYPE: 'FLOAT32',
        ALGORITHM: VectorAlgorithms.HNSW,
        DISTANCE_METRIC: 'L2',
        DIM: 768,
    }
}, {
    ON: 'HASH',
    PREFIX: 'doc:'
});

添加数据

使用 hSet() 将数据对象添加到索引。索引会自动处理带有 doc: 前缀的对象。

对于每个文档

  1. 使用 pipe() 函数和 pipeOptions 生成嵌入
  2. 使用 Buffer.from() 将嵌入转换为二进制字符串
  3. 使用 hSet() 存储文档

使用 Promise.all() 批量处理命令并减少网络往返次数

const sentence1 = 'That is a very happy person';
const doc1 = {
    'content': sentence1, 
    'genre': 'persons', 
    'embedding': Buffer.from(
        (await pipe(sentence1, pipeOptions)).data.buffer
    ),
};

const sentence2 = 'That is a happy dog';
const doc2 = {
    'content': sentence2, 
    'genre': 'pets', 
    'embedding': Buffer.from(
        (await pipe(sentence2, pipeOptions)).data.buffer
    )
};

const sentence3 = 'Today is a sunny day';
const doc3 = {
    'content': sentence3, 
    'genre': 'weather', 
    'embedding': Buffer.from(
        (await pipe(sentence3, pipeOptions)).data.buffer
    )
};

await Promise.all([
    client.hSet('doc:1', doc1),
    client.hSet('doc:2', doc2),
    client.hSet('doc:3', doc3)
]);

运行查询

查询索引

  1. 为查询文本生成嵌入
  2. 将嵌入作为参数传递给搜索
  3. Redis 计算向量距离并对结果进行排序

查询返回一个文档对象数组。每个对象包含

  • id: 文档的键
  • value: 一个包含 RETURN 选项中指定字段的对象
const similar = await client.ft.search(
    'vector_idx',
    '*=>[KNN 3 @embedding $B AS score]',
    {
        'PARAMS': {
            B: Buffer.from(
                (await pipe('That is a happy person', pipeOptions)).data.buffer
            ),
        },
        'RETURN': ['score', 'content'],
        'DIALECT': '2'
    },
);

for (const doc of similar.documents) {
    console.log(`${doc.id}: '${doc.value.content}', Score: ${doc.value.score}`);
}

await client.quit();

首次运行可能需要更长时间,因为它会下载模型数据。输出显示按分数(向量距离)排序的结果,分数越低表示相似度越高

doc:1: 'That is a very happy person', Score: 0.127055495977
doc:2: 'That is a happy dog', Score: 0.836842417717
doc:3: 'Today is a sunny day', Score: 1.50889515877

与 JSON 文档的区别

JSON 文档支持更丰富的数据模型,包含嵌套字段。与哈希文档的主要区别在于:

  1. 在 schema 中使用路径标识字段
  2. 使用 AS 选项声明路径的别名
  3. 创建索引时将 ON 设置为 JSON
  4. 向量使用数组而不是二进制字符串
  5. 使用 json.set() 而不是 hSet()

使用路径别名创建索引

await client.ft.create('vector_json_idx', {
    '$.content': {
        type: SchemaFieldTypes.TEXT,
        AS: 'content',
    },
    '$.genre': {
        type: SchemaFieldTypes.TAG,
        AS: 'genre',
    },
    '$.embedding': {
        type: SchemaFieldTypes.VECTOR,
        TYPE: 'FLOAT32',
        ALGORITHM: VectorAlgorithms.HNSW,
        DISTANCE_METRIC: 'L2',
        DIM: 768,
        AS: 'embedding',
    }
}, {
    ON: 'JSON',
    PREFIX: 'jdoc:'
});

使用 json.set() 添加数据。使用扩展运算符将 Float32Array 转换为标准 JavaScript 数组

const jSentence1 = 'That is a very happy person';
const jdoc1 = {
    'content': jSentence1,
    'genre': 'persons',
    'embedding': [...(await pipe(jSentence1, pipeOptions)).data],
};

const jSentence2 = 'That is a happy dog';
const jdoc2 = {
    'content': jSentence2,
    'genre': 'pets',
    'embedding': [...(await pipe(jSentence2, pipeOptions)).data],
};

const jSentence3 = 'Today is a sunny day';
const jdoc3 = {
    'content': jSentence3,
    'genre': 'weather',
    'embedding': [...(await pipe(jSentence3, pipeOptions)).data],
};

await Promise.all([
    client.json.set('jdoc:1', '$', jdoc1),
    client.json.set('jdoc:2', '$', jdoc2),
    client.json.set('jdoc:3', '$', jdoc3)
]);

使用相同的语法查询 JSON 文档,但请注意向量参数仍必须是二进制字符串

const jsons = await client.ft.search(
    'vector_json_idx',
    '*=>[KNN 3 @embedding $B AS score]',
    {
        "PARAMS": {
            B: Buffer.from(
                (await pipe('That is a happy person', pipeOptions)).data.buffer
            ),
        },
        'RETURN': ['score', 'content'],
        'DIALECT': '2'
    },
);

结果与哈希文档查询相同,只是前缀为 jdoc:

jdoc:1: 'That is a very happy person', Score: 0.127055495977
jdoc:2: 'That is a happy dog', Score: 0.836842417717
jdoc:3: 'Today is a sunny day', Score: 1.50889515877

了解更多

有关索引选项、距离指标和查询格式的更多信息,请参阅向量搜索

评价本页
回到顶部 ↑