Векторные базы данных

[0.12, -0.56, 0.44, ..., -0.03]

from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection
from pymilvus.model.dense import SentenceTransformerEmbeddingFunction

# 1. Подключение к Milvus (по умолчанию localhost:19530)
connections.connect("default", host="localhost", port="19530")

# 2. Инициализация функции эмбеддинга с моделью e5-base-v2
# Эта модель требует:
# - Префикс "passage: " для документов
# - Префикс "query: " для поисковых запросов
ef = SentenceTransformerEmbeddingFunction("intfloat/e5-base-v2")

# 3. Определение схемы коллекции:
# - "id" — целочисленный идентификатор (первичный ключ)
# - "text" — исходный текст документа (строка)
# - "emb" — эмбеддинг-вектор размерности 768
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=False),
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=512),
    FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=768)
]
schema = CollectionSchema(fields, description="Коллекция с эмбеддингами от e5-base-v2")

# 4. Создание коллекции в Milvus с заданной схемой
collection = Collection("e5_collection", schema)

# 5. Подготовка и вставка документов
# Важно: перед подачей в модель нужно добавить префикс "passage: "
raw_docs = ["Hello world", "Milvus vector database", "Semantic search with e5 model"]
docs = [f"passage: {d}" for d in raw_docs]  # добавляем префикс
ids = [1, 2, 3]

# Вычисляем эмбеддинги для документов используя e5‑base
embs = ef.encode_documents(docs)

# Вставляем в коллекцию:
# - идентификаторы
# - исходные (чистые) тексты без префиксов
# - эмбеддинги
collection.insert([ids, raw_docs, embs])

# 6. Создание индекса по полю "emb" для ускорения поиска
collection.create_index(
    field_name="emb",
    index_params={
        # тип индекса
        "index_type": "IVF_FLAT",   
      
        # параметр разбиения на кластеры
        "params": {"nlist": 128}, 
      
        # метрика расстояния (евклидово расстояние)
        "metric_type": "L2"            
    }
)

# 7. Загрузка коллекции в оперативную память
# Без этого поиск работать не будет
collection.load()

# 8. Поисковый запрос
# Аналогично — используем префикс "query: " перед текстом запроса
query_docs = ["query: vector database"]
q_emb = ef.encode_queries(query_docs)

# Выполняем семантический поиск по эмбеддингам
results = collection.search(
    # эмбеддинг поискового запроса
    data=q_emb,      
  
    # поле, по которому осуществляется поиск
    anns_field="emb",    
  
    # параметры поиска
    param={"metric_type": "L2", "params": {"nprobe": 10}},  
  
    # количество ближайших соседей
    limit=2,      
  
    # дополнительные поля, которые нужно вернуть
    output_fields=["text"]           
)

# 9. Выводим результаты поиска
for i, hits in enumerate(results):
    print(f"Результаты для запроса: '{query_docs[i]}'")
    if not hits:
        print("Ничего не найдено")
        continue
    for rank, hit in enumerate(hits, start=1):
        print(f"  {rank}:")
        print(f"    ID: {hit.id}")
        print(f"    Текст: {hit.entity.get('text')}")
        print(f"    Расстояние: {hit.distance:.4f}")
      
# Результаты для запроса: 'query: vector database'
#  1:
#    ID: 2
#    Текст: Milvus vector database
#    Расстояние: 2.8374
#  2:
#    ID: 3
#    Текст: Semantic search with e5 model
#    Расстояние: 5.4931

# 10. Удаление документа по ID
# В данном случае удаляется документ с id = 1
collection.delete(expr="id in [1]")

# 11. Удаление всей коллекции (если больше не нужна)
collection.drop()