LLM应用开发——向量数据库

# 以 Qdrant 为例的混合搜索逻辑
from qdrant_client import QdrantClient
from qdrant_client.http import models

client = QdrantClient(":memory:")

# 混合搜索查询
search_result = client.query_points(
    collection_name="products",
    prefetch=[
        # 1. 向量搜索 (Dense)
        models.Prefetch(
            query=[0.1, 0.2, 0.3, ...], 
            limit=20
        ),
        # 2. 关键词搜索 (Sparse/BM25)
        models.Prefetch(
            query=models.SparseVector(indices=[1, 5], values=[0.5, 0.3]),
            using="sparse",
            limit=20
        ),
    ],
    # 3. 使用 RRF (Reciprocal Rank Fusion) 合并结果
    query=models.FusionQuery(fusion=models.Fusion.RRF),
)

import chromadb
from chromadb.utils import embedding_functions

# 创建客户端
client = chromadb.Client()  # 内存模式
# client = chromadb.PersistentClient(path="./chroma_db")  # 持久化

# 创建集合（使用默认Embedding）
collection = client.create_collection(
    name="my_documents",
    metadata={"description": "文档集合"}
)

# 添加文档
collection.add(
    documents=[
        "机器学习是人工智能的一个分支",
        "深度学习使用神经网络进行学习",
        "向量数据库用于存储和检索向量",
        "RAG结合检索和生成提升回答质量"
    ],
    ids=["doc1", "doc2", "doc3", "doc4"],
    metadatas=[
        {"source": "wiki", "topic": "ML"},
        {"source": "wiki", "topic": "DL"},
        {"source": "blog", "topic": "DB"},
        {"source": "paper", "topic": "LLM"}
    ]
)

print(f"集合中有 {collection.count()} 个文档")

# 语义搜索
results = collection.query(
    query_texts=["什么是神经网络"],
    n_results=3,
    include=["documents", "distances", "metadatas"]
)

print("搜索结果:")
for i, (doc, distance, metadata) in enumerate(zip(
    results['documents'][0],
    results['distances'][0],
    results['metadatas'][0]
)):
    print(f"{i+1}. [{distance:.4f}] {doc}")
    print(f"   元数据: {metadata}")

# 带过滤条件的查询
results = collection.query(
    query_texts=["人工智能技术"],
    n_results=2,
    where={"topic": "ML"}  # 元数据过滤
)

from sentence_transformers import SentenceTransformer

# 自定义Embedding函数
class CustomEmbeddingFunction:
    def __init__(self, model_name="BAAI/bge-large-zh-v1.5"):
        self.model = SentenceTransformer(model_name)
    
    def __call__(self, input):
        return self.model.encode(input).tolist()

# 创建使用自定义Embedding的集合
custom_ef = CustomEmbeddingFunction()
collection = client.create_collection(
    name="custom_embedding_collection",
    embedding_function=custom_ef
)

from pymilvus import (
    connections, Collection, FieldSchema, 
    CollectionSchema, DataType, utility
)

# 连接Milvus
connections.connect(host="localhost", port="19530")

# 定义Schema
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=768),
    FieldSchema(name="source", dtype=DataType.VARCHAR, max_length=256),
]
schema = CollectionSchema(fields=fields, description="文档集合")

# 创建集合
collection_name = "documents"
if utility.has_collection(collection_name):
    utility.drop_collection(collection_name)

collection = Collection(name=collection_name, schema=schema)

# 创建索引
index_params = {
    "metric_type": "COSINE",
    "index_type": "IVF_FLAT",
    "params": {"nlist": 128}
}
collection.create_index(field_name="embedding", index_params=index_params)

print(f"集合 {collection_name} 创建成功")

from sentence_transformers import SentenceTransformer
import numpy as np

# 加载Embedding模型
model = SentenceTransformer("BAAI/bge-large-zh-v1.5")

# 准备数据
texts = [
    "Python是一种流行的编程语言",
    "机器学习需要大量数据训练",
    "向量数据库支持相似性搜索",
    "RAG可以提升LLM的回答质量"
]
embeddings = model.encode(texts)
sources = ["doc", "doc", "blog", "paper"]

# 插入数据
data = [
    texts,
    embeddings.tolist(),
    sources
]
collection.insert([texts, embeddings.tolist(), sources])
collection.flush()

print(f"插入 {collection.num_entities} 条数据")

# 加载集合到内存
collection.load()

# 搜索
query = "如何进行语义搜索"
query_embedding = model.encode([query])

search_params = {"metric_type": "COSINE", "params": {"nprobe": 10}}
results = collection.search(
    data=query_embedding.tolist(),
    anns_field="embedding",
    param=search_params,
    limit=3,
    output_fields=["text", "source"]
)

print(f"\n查询: {query}")
print("搜索结果:")
for hits in results:
    for hit in hits:
        print(f"  [{hit.score:.4f}] {hit.entity.get('text')}")

import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

# 准备数据
model = SentenceTransformer("BAAI/bge-large-zh-v1.5")
texts = [
    "机器学习基础概念",
    "深度学习神经网络",
    "自然语言处理技术",
    "计算机视觉应用",
    "强化学习算法"
]
embeddings = model.encode(texts).astype('float32')

# 创建索引
dimension = embeddings.shape[1]
index = faiss.IndexFlatIP(dimension)  # 内积索引

# 归一化向量（用于余弦相似度）
faiss.normalize_L2(embeddings)
index.add(embeddings)

print(f"索引中有 {index.ntotal} 个向量")

# 搜索
query = "什么是神经网络"
query_vec = model.encode([query]).astype('float32')
faiss.normalize_L2(query_vec)

k = 3
distances, indices = index.search(query_vec, k)

print(f"\n查询: {query}")
print("搜索结果:")
for i, (idx, dist) in enumerate(zip(indices[0], distances[0])):
    print(f"  {i+1}. [{dist:.4f}] {texts[idx]}")

# IVF索引（适合大规模数据）
nlist = 100  # 聚类中心数量
quantizer = faiss.IndexFlatL2(dimension)
index_ivf = faiss.IndexIVFFlat(quantizer, dimension, nlist)
index_ivf.train(embeddings)  # 需要训练
index_ivf.add(embeddings)
index_ivf.nprobe = 10  # 搜索时探测的聚类数

# HNSW索引（高召回率）
index_hnsw = faiss.IndexHNSWFlat(dimension, 32)  # 32是图的连接数
index_hnsw.add(embeddings)

# PQ索引（压缩存储）
m = 8  # 子向量数量
index_pq = faiss.IndexPQ(dimension, m, 8)
index_pq.train(embeddings)
index_pq.add(embeddings)

# 保存索引
faiss.write_index(index, "my_index.faiss")

# 加载索引
loaded_index = faiss.read_index("my_index.faiss")

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
from sentence_transformers import SentenceTransformer

# 连接Qdrant
client = QdrantClient(host="localhost", port=6333)

# 创建集合
collection_name = "documents"
client.recreate_collection(
    collection_name=collection_name,
    vectors_config=VectorParams(size=768, distance=Distance.COSINE)
)

# 准备数据
model = SentenceTransformer("BAAI/bge-large-zh-v1.5")
texts = [
    "Python编程语言入门",
    "机器学习算法详解",
    "深度学习实战项目"
]
embeddings = model.encode(texts)

# 插入数据
points = [
    PointStruct(
        id=i,
        vector=embedding.tolist(),
        payload={"text": text, "index": i}
    )
    for i, (text, embedding) in enumerate(zip(texts, embeddings))
]
client.upsert(collection_name=collection_name, points=points)

# 搜索
query = "如何学习编程"
query_vec = model.encode([query])[0]

results = client.search(
    collection_name=collection_name,
    query_vector=query_vec.tolist(),
    limit=3
)

print(f"查询: {query}")
for result in results:
    print(f"  [{result.score:.4f}] {result.payload['text']}")