基于Python+LanceDB实战向量搜索

本篇实战演示向量搜索的实现和示例。

预期效果

给出一个查询的字符串，通过向量搜索，在下面三个语句中搜索出关联性最大的那句。

    "熊猫是中国的国宝，主要栖息在四川山区。","长城是古代中国建造的军事防御工事，全长超过2万公里。","量子计算利用量子力学原理进行信息处理，相比经典计算机有显著优势。"

最终的实现效果如下图：

给出查询的语句：“熊猫栖息地?“，查找到关联性最大的那句。

实现手段

嵌入模型使用 sentence-transformers/all-MiniLM-L6-v2，384维向量

关于 all-MiniLM-L6-v2 模型的介绍，可以参考：一篇吃透模型：all-MiniLM-L6-v2
向量数据库：LanceDb

实现步骤

通过嵌入模型将语句转换为向量。
将原文本和向量存入到向量数据库
将查询的语句转换为向量，通过向量搜索找到关联的那句。

实现的完整代码如下：

import lancedb
from sentence_transformers import SentenceTransformer
import pyarrow as pamodel = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
documents = ["熊猫是中国的国宝，主要栖息在四川山区。","长城是古代中国建造的军事防御工事，全长超过2万公里。","量子计算利用量子力学原理进行信息处理，相比经典计算机有显著优势。"
]
# 获取嵌入向量
embeddings = model.encode(documents)
vector_dim =  len(embeddings[0]) #获取向量的维度，这里是384db = lancedb.connect("data/mylancedb")
schema = pa.schema([pa.field("text", pa.string()),pa.field("embedding", pa.list_(pa.float32(), vector_dim))
])
datas=[{"text": doc, "embedding": embedding} for doc, embedding in zip(documents, embeddings)]
# 创建表并插入数据
table = db.create_table("my_table",datas,schema=schema, mode="overwrite")query = "熊猫栖息地?"
# 转换为向量
query_embed = model.encode(query)
# 使用向量搜索
result = (table.search(query_embed).select(["text"]).limit(1).to_pandas()
)
print("Search results:")
print(result)