向量检索（RAG）之向量数据库研究

from annoy import AnnoyIndex
import random

f = 40  # Length of item vector that will be indexed

# 构建过程
t = AnnoyIndex(f, 'angular') # 支持5种距离计算方式："angular", "euclidean", "manhattan", "hamming", or "dot"
for i in range(1000):
    v = [random.gauss(0, 1) for z in range(f)]
    t.add_item(i, v) # 添加随机向量信息
t.build(10) # 10 trees, 构建由n_trees树组成的森林。更多的树在查询时有更高的精度。调用build后，不能再添加任何项目。n_jobs指定用于构建树的线程数。n_jobs=-1使用所有可用的CPU核。
t.save('test.ann') # 将索引保存到磁盘

# 加载
u = AnnoyIndex(f, 'angular')
u.load('test.ann') # super fast, will just mmap the file 从磁盘加载（mmaps）一个索引
print(u.get_nns_by_item(0, 1000)) # 用于根据 item id=0 查找最相似的 k=1000 个 item

from nomic import atlas
import numpy as np

# 1、构建向量索引
num_embeddings = 10000
embeddings = np.random.rand(num_embeddings, 256)
project = atlas.map_embeddings(embeddings=embeddings)
# 检索向量
neighbors, distances = map.vector_search(ids=['42'])
print(project.get_data(ids=neighbors[0]))

# 2、构建文本索引，由 atlas 内部模型生成向量，也可以使用自己的模型（例如：HuggingFace model、Cohere model）
dataset = load_dataset('ag_news')['train'] # https://huggingface.co/datasets/ag_news
max_documents = 10000
subset_idxs = np.random.randint(len(dataset), size=max_documents).tolist()
documents = [dataset[i] for i in subset_idxs]
project = atlas.map_text(data=documents,
                          indexed_field='text',
                          name='News 10k Example',
                          colorable_fields=['label'],
                          description='News 10k Example.'
                          )

sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")

openai_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key="YOUR_API_KEY",
                model_name="text-embedding-ada-002"
            )

cohere_ef  = embedding_functions.CohereEmbeddingFunction(api_key="YOUR_API_KEY",  model_name="large")
cohere_ef(texts=["document1","document2"])

ef = embedding_functions.InstructorEmbeddingFunction(model_name="hkunlp/instructor-xl", device="cuda")

palm_embedding = embedding_functions.GooglePalmEmbeddingFunction(
    api_key=api_key, model=model_name)

from chromadb.api.types import Documents, EmbeddingFunction, Embeddings

class MyEmbeddingFunction(EmbeddingFunction):
    def __call__(self, texts: Documents) -> Embeddings:
        # embed the documents somehow
        return embeddings

# pip install chromadb

import chromadb

# step 1 Get the Chroma Client
chroma_client = chromadb.Client()

# step 2 Create a collection, store your embeddings, documents, and any additional metadata
collection = chroma_client.create_collection(name="my_collection")

# step 3 Add some text documents to the collection
# 第一种：采用默认方式生成 embeddings
collection.add(
    documents=["This is a document", "This is another document"],
    metadatas=[{"source": "my_source"}, {"source": "my_source"}],
    ids=["id1", "id2"]
)
# 第二种：自己生成 embeddings
collection.add(
    embeddings=[[1.2, 2.3, 4.5], [6.7, 8.2, 9.2]],
    documents=["This is a document", "This is another document"],
    metadatas=[{"source": "my_source"}, {"source": "my_source"}],
    ids=["id1", "id2"]
)

# step 4 Query the collection
results = collection.query(
    query_texts=["This is a query document"],
    n_results=2
)
results = collection.query(
    query_embeddings=[[11.1, 12.1, 13.1],[1.1, 2.3, 3.2] ...]
    n_results=10,
    where={"metadata_field": "is_equal_to_this"}, # 条件过滤
    where_document={"$contains":"search_string"} # 搜索文本
)

# 同时，可以进行 update 和 delete 操作

vector_store_path = 'vector_store'
CHUNK_SIZE = 1000
# text
chunked_text = []
for i in range(0, CHUNK_SIZE):
    chunked_text.append("text" + str(i + 1))
# embedding
embedding = []
for i in range(0, CHUNK_SIZE):
    embedding.append([random.random(), random.random(), random.random()])
# meta
meta = []
for i in range(0, CHUNK_SIZE):
    d = {'key': i + 1}
    meta.append(d)

vector_store = DeepLakeVectorStore(
    path=vector_store_path,
)
 
vector_store.add(text=chunked_text,
                 embedding=embedding,
                 metadata=meta)
# vector_store.summary()

# prompt = "What are the first programs he tried writing?"
search_results = vector_store.search(embedding=[0.1, 0.2, 0.3])

# 也可以使用 openai 的模型
def embedding_function(texts, model="text-embedding-ada-002"):
   if isinstance(texts, str):
       texts = [texts]
   texts = [t.replace("\n", " ") for t in texts]
   return [data['embedding']for data in openai.Embedding.create(input = texts, model=model)['data']]

PUT /hotels-index
{
  "settings": {
    "index": {
      "knn": true,
      "knn.algo_param.ef_search": 100,
      "number_of_shards": 1,
      "number_of_replicas": 0
    }
  },
  "mappings": {
    "properties": {
      "location": {
        "type": "knn_vector",
        "dimension": 2,
        "method": {
          "name": "hnsw",
          "space_type": "l2",
          "engine": "lucene", # 这里指定 lucene
          "parameters": {
            "ef_construction": 100,
            "m": 16
          }
        }
      }
    }
  }
}

POST /_bulk
{ "index": { "_index": "hotels-index", "_id": "1" } }
{ "location": [5.2, 4.4], "parking" : "true", "rating" : 5 }
{ "index": { "_index": "hotels-index", "_id": "2" } }
{ "location": [5.2, 3.9], "parking" : "false", "rating" : 4 }
{ "index": { "_index": "hotels-index", "_id": "3" } }
{ "location": [4.9, 3.4], "parking" : "true", "rating" : 9 }

POST /hotels-index/_search
{
  "size": 3,
  "query": {
    "knn": {
      "location": {
        "vector": [
          5,
          4
        ],
        "k": 3,
        "filter": {
          "bool": {
            "must": [
              {
                "range": {
                  "rating": {
                    "gte": 8,
                    "lte": 10
                  }
                }
              },
              {
                "term": {
                  "parking": "true"
                }
              }
            ]
          }
        }
      }
    }
  }
}

cd /tmp
git clone --branch v0.4.2 https://github.com/pgvector/pgvector.git
cd pgvector
make
make install # may need sudo

# 启用扩展功能
CREATE EXTENSION vector;

# 创建一个具有3个维度的向量列
CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3));

# 插入向量
INSERT INTO items (embedding) VALUES ('[1,2,3]'), ('[4,5,6]');

# 通过L2距离获得最近的结果，还支持内积（<#>）和余弦距离（<=>）。
SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5;

# 获取一定距离内的记录
SELECT * FROM items WHERE embedding <-> '[3,1,2]' < 5;

# 创建索引：在表有一些数据后创建索引，可以对多达2000个维度的向量进行索引
# L2 distance
CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);
# Inner product
CREATE INDEX ON items USING ivfflat (embedding vector_ip_ops) WITH (lists = 100);
# Cosine distance
CREATE INDEX ON items USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

# 查询
SELECT * FROM items WHERE category_id = 123 ORDER BY embedding <-> '[3,1,2]' LIMIT 5;

pip install pinecone-client # 安装 pinecone python 客户端

import pinecone

# 初始化，验证 API key
pinecone.init(api_key="YOUR_API_KEY", environment="YOUR_ENVIRONMENT")
# 创建一个名为"quickstart"的索引，它使用8维向量的欧氏距离度量来执行近似的近邻搜索。创建索引大约需要一分钟时间。
pinecone.create_index("quickstart", dimension=4, metric="euclidean") # euclidean、cosine、dotproduct
# 返回现有的索引列表
pinecone.list_indexes()

# 先连接索引
index = pinecone.Index("quickstart")
# 新增数据 Upsert sample data (5 8-dimensional vectors)
upsert_response = index.upsert(
    vectors=[
        (
         "vec1",                # Vector ID 
         [0.1, 0.2, 0.3, 0.4],  # Dense vector values
         {"genre": "drama"}     # Vector metadata
        ),
        (
         "vec2", 
         [0.2, 0.3, 0.4, 0.5], 
         {"genre": "action"}
        )
    ],
    namespace="example-namespace"
)
# 查询数据，使用条件过滤+欧氏距离得到前三个相似向量
query_response = index.query(
    namespace="example-namespace",
    top_k=3,
    include_values=True,
    include_metadata=True,
    vector=[0.1, 0.2, 0.3, 0.4],
    filter={
        "genre": {"$in": ["comedy", "documentary", "drama"]}
    }
)

from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams

# 创建 collection
client = QdrantClient("localhost", port=6333)
client.recreate_collection(
    collection_name="test_collection",
    vectors_config=VectorParams(size=4, distance=Distance.DOT),
)
# 查看 collection
collection_info = client.get_collection(collection_name="test_collection")

# 添加 points
operation_info = client.upsert(
    collection_name="test_collection",
    wait=True,
    points=[
        PointStruct(id=1, vector=[0.05, 0.61, 0.76, 0.74], payload={"city": "Berlin"}),
        PointStruct(id=2, vector=[0.19, 0.81, 0.75, 0.11], payload={"city": ["Berlin", "London"]}),
        PointStruct(id=3, vector=[0.36, 0.55, 0.47, 0.94], payload={"city": ["Berlin", "Moscow"]}),
        PointStruct(id=4, vector=[0.18, 0.01, 0.85, 0.80], payload={"city": ["London", "Moscow"]}),
        PointStruct(id=5, vector=[0.24, 0.18, 0.22, 0.44], payload={"count": [0]}),
        PointStruct(id=6, vector=[0.35, 0.08, 0.11, 0.44]),
    ]
)
# 查询
search_result = client.search(
    collection_name="test_collection",
    query_vector=[0.2, 0.1, 0.9, 0.7], 
    query_filter=Filter(
        must=[
            FieldCondition(
                key="city",
                match=MatchValue(value="London")
            )
        ]
    ),
    search_params=models.SearchParams(
        hnsw_ef=128,
        exact=False # exact - 不使用近似搜索 (ANN) 的选项。如果设置为 true，搜索可能会运行很长时间，因为它会执行完整扫描以检索准确的结果。（暴力检索）
    ),
    limit=3
)

 hnsw_index:
    # 索引图中每个节点的边数。值越大 - 搜索越准确，需要的空间就越大。
    m: 16
    # 建立索引时要考虑的邻居数量。值越大——搜索越准确，建立索引所需的时间越长。
    ef_construct: 100
    # Minimal size (in KiloBytes) of vectors for additional payload-based indexing.
    # If payload chunk is smaller than `full_scan_threshold_kb` additional indexing won't be used -
    # in this case full-scan search should be preferred by query planner and additional indexing is not required.
    # Note: 1Kb = 1 vector of size 256
    full_scan_threshold_kb: 10000
    # Number of parallel threads used for background index building. If 0 - auto selection.
    max_indexing_threads: 0
    # Store HNSW index on disk. If set to false, index will be stored in RAM. Default: false
    on_disk: false
    # Custom M param for hnsw graph built for payload index. If not set, default M will be used.
    payload_m: null

# Load data
import weaviate
import json

client = weaviate.Client(
    url = "https://some-endpoint.weaviate.network",  # Replace with your endpoint
    auth_client_secret=weaviate.AuthApiKey(api_key="YOUR-WEAVIATE-API-KEY"),  # Replace w/ your Weaviate instance API key
    additional_headers = {
        "X-HuggingFace-Api-Key": "YOUR-HUGGINGFACE-API-KEY"  # Replace with your inference API key
    }
)

# ===== add schema ===== # 如果您使用不同的矢量化器，我们还提供建议的 vectorizer 模块配置：Cohere、Hugging Face、OpenAI、PaLM
class_obj = {
    "class": "Question",
    "vectorizer": "text2vec-huggingface",  # If set to "none" you must always provide vectors yourself. Could be any other "text2vec-*" also.
    "moduleConfig": {
        "text2vec-huggingface": {
            "model": "sentence-transformers/all-MiniLM-L6-v2",  # Can be any public or private Hugging Face model.
            "options": {
                "waitForModel": True
            }
        }
    }
}
client.schema.create_class(class_obj)

# ===== import data =====
# Load data
import requests
url = 'https://raw.githubusercontent.com/weaviate-tutorials/quickstart/main/data/jeopardy_tiny.json'
resp = requests.get(url)
data = json.loads(resp.text)

# Configure a batch process
with client.batch(
    batch_size=100
) as batch:
    # Batch import all Questions
    for i, d in enumerate(data):
        print(f"importing question: {i+1}")

        properties = {
            "answer": d["Answer"],
            "question": d["Question"],
            "category": d["Category"],
        }

        client.batch.add_data_object(
            properties,
            "Question",
            # vector=custom_vector  # 如果没有指定模型，需要自己设置向量
        )

import weaviate
import json

client = weaviate.Client(
    url = "https://some-endpoint.weaviate.network",  # Replace with your endpoint
    auth_client_secret=weaviate.AuthApiKey(api_key="YOUR-WEAVIATE-API-KEY"),  # Replace w/ your Weaviate instance API key
    additional_headers = {
        "X-HuggingFace-Api-Key": "YOUR-HUGGINGFACE-API-KEY"  # Replace with your inference API key
    }
)

nearText = {"concepts": ["biology"]}

response = (
    client.query
    .get("Question", ["question", "answer", "category"])
    .with_near_text(nearText)   # 指定向量 .with_near_vector({"vector": [-0.0125526935]})
    .with_limit(2)
    .with_where({ # 条件过滤器
        "path": ["round"],
        "operator": "Equal",
        "valueText": "Double Jeopardy!"
    })
    .do()
)

print(json.dumps(response, indent=4))

client.schema.update_config("DeepImage", {
    "vectorIndexConfig": {
        "pq": {
            "enabled": True,
        }
    }
})

from pymilvus import MilvusClient

# Initialize a MilvusClient instance
# Replace uri and API key with your own
client = MilvusClient(
    uri="<CLUSTER-ENDPOINT>", # Cluster endpoint obtained from the console
    token="<API-KEY>"
)

# Create a collection
client.create_collection(
    collection_name="medium_articles_2020",
    dimension=768
)

# Insert a single entity
res = client.insert(
    collection_name="medium_articles_2020",
    data={
        'id': 0, 
        'title': 'The Reported Mortality Rate of Coronavirus Is Not Important', 
        'link': 'https://medium.com/swlh/the-reported-mortality-rate-of-coronavirus-is-not-important-369989c8d912', 
        'reading_time': 13, 
        'publication': 'The Startup', 
        'claps': 1100, 
        'responses': 18, 
        'vector': [0.041732933, 0.030096486]
    }
)

with open("https://s3.us-west-2.amazonaws.com/publicdataset.zillizcloud.com/medium_articles_2020_dpr/medium_articles_2020_dpr.json") as f:
  data = json.load(f)
res = client.search(
    collection_name="medium_articles_2020",
    data=[data["rows"][0]["title_vector"]],
    filter='claps > 100 and publication in ["The Startup", "Towards Data Science"]',
    output_fields=["title", "claps", "publication"]
)