向量数据库-lancedb

简介

lancedb是由Rust编写的向量数据库，主要用于人工智能的矢量数据存储与查询。用于检索大规模多模式的数据嵌入，经常用在人工智能的检索增强生成(RAG)中作为知识库从存储与检索。LanceDB是一种高效的向量数据库，支持文本和嵌入向量的存储与查询。

LanceDB是一款新型无服务器向量数据库，专为AI应用而设计
存储、查询和过滤向量、元数据和多模态数据：支持文本、图像、视频、点云等。
支持向量相似搜索、全文搜索和 SQL。
原生 Python 和 Javascript/Typescript 支持。
零拷贝自动版本管理：无需额外基础设施即可管理数据版本。
GPU 支持：在构建向量索引时。
生态系统集成：与 LangChain、LlamaIndex、Apache-Arrow、Pandas、Polars、DuckDB 等。
LanceDB 已在Github上开源

安装

pip install lancedb

简单用法

  
import lancedb
import ollama # 用于转化向量
import pyarrow as pa

# 创建数据库连接
uri = "data/sample-lancedb"
db = lancedb.connect(uri)   

# 用于转化向量
cli = ollama.Client(host=host)
res = cli.embeddings(model="qwen2",prompt="你好，这里是要转会的文本。")
vector = res["embedding"]

# 创建一张空表
schema = pa.schema([
    pa.field("vector", pa.list_(pa.float32(), list_size=3584)),
    pa.field("content", pa.string()),
    pa.field("id", pa.int32()),
])
tbl = db.create_table("empty_table", schema=schema)

# 创建表,带数据
data = [
    { "vector": vector, "content": "你好，这里是要转会的文本。", "id": 100},
]
tbl = db.create_table("my_table", data=data, exist_ok=True)

# 打开表，添加数据
tbl = db.open_table("my_table")
tbl.add(data)

# 修改数据
table.update(where='id = 100', values={"vector": vector})

# 查询数据：metric="L2"表示使用欧式距离，cosine表示余弦距离，目前好像只支持这两个
tbl.search(vector).metric("L2").limit(2).to_pandas()

# 删除数据
tbl.delete('id = 100')

# 删除表
db.drop_table("my_table")

高级用法

多种数据类型

  
import pyarrow as pa
schema = pa.schema(
    [
        pa.field("vector", pa.list_(pa.float16(), 2)),
        pa.field("text", pa.string())
    ]
)

自定义数据类型

  
from lancedb.pydantic import Vector, LanceModel

class Content(LanceModel):
    movie_id: int
    vector: Vector(128)
    genres: str
    title: str
    imdb_id: int

    @property
    def imdb_url(self) -> str:
        return f"https://www.imdb.com/title/tt{self.imdb_id}"

复合数据类型

  
class Document(BaseModel):
    content: str
    source: str

class NestedSchema(LanceModel):
    id: str
    vector: Vector(1536)
    document: Document

tbl = db.create_table("nested_table", schema=NestedSchema, mode="overwrite")

索引

lancedb支持创建倒排索引的乘积量化。num_partitions是索引中的分区数，默认值是行数的平方根。num_sub_vectors是子向量的数量，默认值是向量的维度除以16。

  
# 创建IVF_PQ索引
tbl.create_index(num_partitions=256, num_sub_vectors=96)

支持CUDA的GPU或者Apple的MPS加速

  
# 使用GPU创建
accelerator="cuda"
# accelerator="mps"

使用索引加速近似查找

  
tbl.search(np.random.random((1536))).limit(2).nprobes(20).refine_factor(10).to_pandas()

nprobes是探针数量，默认为20，增加探针数量则会提高查找的精度并相应增加计算耗时。refine_factor是一个粗召的数量，用于读取额外元素并重新排列，以此来提高召回。

向量化模型

内置向量模型

  
import lancedb
from lancedb.pydantic import LanceModel, Vector
from lancedb.embeddings import get_registry

model = get_registry().get("sentence-transformers").create(name="BAAI/bge-small-en-v1.5", device="cpu")

class Words(LanceModel):
    text: str = model.SourceField() # 指定这个字段为需要模型进行向量化的字段
    vector: Vector(model.ndims()) = model.VectorField() # 指定这个字段为模型向量化的结果

table = db.create_table("words", schema=Words)
table.add(
    [
        {"text": "hello world"},
        {"text": "goodbye world"}
    ]
)

query = "greetings"
actual = table.search(query).limit(1).to_pydantic(Words)[0]
print(actual.text)

官方支持了多种sentence-transformers的向量化模型。用上述方法调用内置模型需要指定模型的SourceField和VectorField。

自定义向量模型

  
from lancedb.embeddings import register
from lancedb.util import attempt_import_or_raise

@register("sentence-transformers")
class SentenceTransformerEmbeddings(TextEmbeddingFunction):
    name: str = "all-MiniLM-L6-v2"
    # set more default instance vars like device, etc.

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._ndims = None

    def generate_embeddings(self, texts):
        return self._embedding_model().encode(list(texts), ...).tolist()

    def ndims(self):
        if self._ndims is None:
            self._ndims = len(self.generate_embeddings("foo")[0])
        return self._ndims

    @cached(cache={}) 
    def _embedding_model(self):
        return sentence_transformers.SentenceTransformer(name)
from lancedb.pydantic import LanceModel, Vector

registry = EmbeddingFunctionRegistry.get_instance()
stransformer = registry.get("sentence-transformers").create()

class TextModelSchema(LanceModel):
    vector: Vector(stransformer.ndims) = stransformer.VectorField()
    text: str = stransformer.SourceField()

tbl = db.create_table("table", schema=TextModelSchema)

tbl.add(pd.DataFrame({"text": ["halo", "world"]}))
result = tbl.search("world").limit(5)

官方提供了模板用于自定义模型，但是我觉得直接调用模型进行向量化表示更直接吧，这样感觉有点追求格式化的统一了。