词向量相似检索 - RAG初探

使用BERT的词向量生成

from transformers import BertTokenizer, BertModel
import torch

# 初始化分词器和模型
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

# 编码文本
input_text = "Hello, world!"
encoded_input = tokenizer(input_text, return_tensors='pt')

# 提取嵌入向量
with torch.no_grad():
    outputs = model(**encoded_input)
    # 获取最后一层的隐藏状态，它的维度是 [batch_size, sequence_length, hidden_size]
    last_hidden_states = outputs.last_hidden_state
    # 通常取[CLS]标记的嵌入向量，它位于序列的开头
    embedding_vector = last_hidden_states[:, 0, :].numpy()

# `embedding_vector` 就是文本的嵌入表示

使用opensearch的词向量检索

Opensearch使用词向量进行数据存储和相似查找的功能是通过它的向量搜索特性来实现的。这主要是通过在索引创建时定义一个特殊的字段类型来完成的，这个字段类型可以存储词向量（例如，通过深度学习模型得到的嵌入向量），然后可以利用这些向量进行相似度搜索。这个过程大致可以分为以下几个步骤：

• 安装和配置：首先，确保你使用的Opensearch版本支持向量搜索。你可能需要安装额外的插件或确保使用的是包含这一功能的Opensearch版本。

• 创建带有向量字段的索引：在创建索引的映射(mapping)时，你需要定义一个类型为**dense_vector**的字段，这个字段用于存储向量数据。例如，你可以为文档中的每个项目定义一个嵌入向量。

• 索引数据：在向索引中添加文档时，将每个文档的向量数据存储在先前定义的**dense_vector**字段中。

• 相似度搜索：一旦数据被索引，就可以使用特定的查询进行相似度搜索了。Opensearch提供了不同的方式来执行这种类型的查询，例如通过计算向量之间的余弦相似度。

下面是一个简单的例子，展示了如何在Opensearch中创建一个带有**dense_vector**字段的索引，索引一些数据，并执行一个基于向量相似度的搜索。

1. 创建索引

jsonCopy code
PUT /my_index
{
  "mappings": {
    "properties": {
      "my_vector": {
        "type": "dense_vector",
        "dims": 128  // 假设你的向量长度为128
      },
      "my_text": {
        "type": "text"
      }
    }
  }
}

2. 索引文档

jsonCopy code
PUT /my_index/_doc/1
{
  "my_text": "这是一个示例文档。",
  "my_vector": [0.5, 0.8, ...]  // 你的向量数据
}

3. 执行向量相似度搜索

jsonCopy code
POST /my_index/_search
{
  "query": {
    "script_score": {
      "query": {"match_all": {}},
      "script": {
        "source": "cosineSimilarity(params.query_vector, 'my_vector') + 1.0",
        "params": {
          "query_vector": [0.5, 0.8, ...]  // 你要搜索的向量
        }
      }
    }
  }
}

在这个例子中，我们首先创建了一个名为**my_index的索引，其中包含一个名为my_vector的dense_vector字段和一个名为my_text的文本字段。然后，我们向这个索引中添加了一个文档，包含文本和向量数据。最后，我们执行了一个搜索查询，使用script_score**查询来计算查询向量和文档中向量之间的余弦相似度，以找到最相似的文档。