使用Infinispan实现高效的向量搜索

在本篇文章中，我们将深入探讨如何利用Infinispan作为向量存储，实现高效的文本相似度搜索。Infinispan是一个开源的键值数据网格，既可以作为单节点运行，也可以分布式运行。从15.x版本开始，Infinispan支持向量搜索功能，这为构建强大的搜索系统提供了更广阔的应用前景。

技术背景介绍

Infinispan是一种分布式的数据存储解决方案，能够有效地处理大规模的键值数据存储需求。其灵活的架构不仅支持简单的键值存储，还可以用于更复杂的数据类型。结合向量搜索功能，Infinispan可以支持如推荐系统、文档搜索等应用。

核心原理解析

向量搜索的核心在于将文本或其他类型的数据转换为向量表示，然后基于距离度量（如欧氏距离、余弦相似度等）进行相似度计算。Infinispan的向量存储模块通过支持嵌入模型（如Hugging Face的sentence-transformers），实现了这一功能。

代码实现演示

接下来，我们将通过一个完整的代码示例展示如何配置和使用Infinispan实现向量搜索。

前置准备

首先，我们安装所需的Python库：

%pip install sentence-transformers langchain langchain_core langchain_community

配置Infinispan

我们需要一个运行中的Infinispan实例，可以在Docker中启动：

%%bash
wget https://raw.githubusercontent.com/rigazilla/infinispan-vector/main/bbc_news.csv.gz

echo 'infinispan:
  cache-container: 
    name: default
    transport: 
      cluster: cluster 
      stack: tcp 
  server:
    interfaces:
      interface:
        name: public
        inet-address:
          value: 0.0.0.0 
    socket-bindings:
      default-interface: public
      port-offset: 0        
      socket-binding:
        name: default
        port: 11222
    endpoints:
      endpoint:
        socket-binding: default
        rest-connector:
' > infinispan-noauth.yaml

!docker rm --force infinispanvs-demo
!docker run -d --name infinispanvs-demo -v $(pwd):/user-config  -p 11222:11222 infinispan/server:15.0 -c /user-config/infinispan-noauth.yaml

选择嵌入模型

我们使用Hugging Face的嵌入模型来生成文本向量：

from langchain.embeddings import HuggingFaceEmbeddings

model_name = "sentence-transformers/all-MiniLM-L12-v2"
hf = HuggingFaceEmbeddings(model_name=model_name)

准备数据

接下来，我们处理新闻数据文件并准备文本和元数据：

import csv
import gzip

with gzip.open("bbc_news.csv.gz", "rt", newline="") as csvfile:
    reader = csv.reader(csvfile, delimiter=",", quotechar='"')
    texts, metas = [], []
    for i, row in enumerate(reader):
        text = row[0] + "." + row[4]
        texts.append(text)
        meta = {"text": row[4], "title": row[0]}
        metas.append(meta)
        if i >= 5000:  # 限制加载的记录数量
            break

填充向量存储

将文本和元数据添加到Infinispan向量存储中：

from langchain_community.vectorstores import InfinispanVS

ispnvs = InfinispanVS.from_texts(texts, hf, metas)

查询示例

我们可以通过简单的查询来检索相似的文档：

def print_docs(docs):
    for res, i in zip(docs, range(len(docs))):
        print("----" + str(i + 1) + "----")
        print("TITLE: " + res.metadata["title"])
        print(res.page_content)

docs = ispnvs.similarity_search("European nations", 5)
print_docs(docs)

清理环境

在完成实验后，可以删除Docker实例：

!docker rm --force infinispanvs-demo

应用场景分析

此技术方案适用于需要高效向量化处理的应用场景，如推荐系统、文本相似度搜索和个性化内容推送。在这些场景中，Infinispan的分布式特性和高效的向量搜索能力可以显著提升系统性能。

实践建议

• 在数据量较大时，建议采用分布式部署以提高系统的吞吐能力。
• 对于精度要求较高的应用，可结合不同的嵌入模型进行实验，以选择最优的模型。
• 在实际应用中，可以探索数据预处理和特征优化技术，以进一步提高搜索精度和效率。

如果遇到问题欢迎在评论区交流。

—END—