Векторный поиск на Python: эмбеддинги и pgvector с нуля

К концу этого материала у вас будет рабочий семантический поиск: пользователь вводит запрос вроде «как вернуть деньги за бракованный товар», а система находит документ с заголовком «возврат средств при получении дефектного заказа» — хотя ни одного общего слова в них нет. Мы посчитаем эмбеддинги через модель sentence-transformers, сохраним их в PostgreSQL с расширением pgvector и напишем запрос на поиск ближайших векторов. На всё уйдёт 30–40 минут, если Docker и Python уже стоят.

Что нужно иметь до начала

Проверьте это до того, как начнёте копировать команды, иначе застрянете в середине.

• Python 3.10+ — проверьте командой python --version.
• Docker — поднимем в нём PostgreSQL с уже установленным pgvector, чтобы не собирать расширение вручную. Проверьте docker --version.
• Примерно 500 МБ трафика — модель эмбеддингов скачается при первом запуске.
• Права ставить пакеты в виртуальное окружение — venv или аналог.

Версии, на которых всё проверялось: PostgreSQL 16 с pgvector 0.7.x, sentence-transformers 3.x, psycopg 3.x. С другими версиями шаги те же, но имена функций сверяйте по документации.

Что такое эмбеддинг

Эмбеддинг^* — это перевод текста в список чисел (вектор) фиксированной длины. Модель обучена так, что близкие по смыслу тексты дают близкие векторы. Поиск по смыслу сводится к поиску векторов, находящихся рядом в этом многомерном пространстве.

^* Эмбеддинг (embedding) — числовое представление объекта, здесь текста, в виде вектора. Модель all-MiniLM-L6-v2 выдаёт вектор из 384 чисел.

Шаг 1. Поднимаем PostgreSQL с pgvector

Готовый образ уже содержит расширение — собирать ничего не нужно.

docker run -d \
  --name vsearch-pg \
  -e POSTGRES_PASSWORD=devpass \
  -e POSTGRES_DB=vsearch \
  -p 5432:5432 \
  pgvector/pgvector:pg16

Проверка: команда docker ps должна показать контейнер vsearch-pg со статусом Up. Если порт 5432 занят локальным PostgreSQL, замените в команде -p 5432:5432 на -p 5433:5432 и дальше используйте порт 5433.

Шаг 2. Ставим Python-зависимости

python -m venv .venv
source .venv/bin/activate   # Windows: .venv\Scripts\activate
pip install sentence-transformers psycopg[binary] pgvector

Проверка: команда python -c "import sentence_transformers, psycopg, pgvector; print('ok')" должна вывести ok без трассировки ошибок.

Шаг 3. Создаём таблицу с колонкой-вектором

Расширение включается один раз на базу. Размерность вектора (384) должна совпадать с тем, что выдаёт модель, — иначе вставка упадёт.

import psycopg

# Подставьте свой порт, если меняли на шаге 1
DSN = "postgresql://postgres:devpass@localhost:5432/vsearch"

with psycopg.connect(DSN) as conn:
    with conn.cursor() as cur:
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
        cur.execute("""
            CREATE TABLE IF NOT EXISTS documents (
                id      bigserial PRIMARY KEY,
                title   text NOT NULL,
                content text NOT NULL,
                embedding vector(384)
            )
        """)
    conn.commit()
print("Таблица готова")

Проверка: в консоли появится Таблица готова. Заглянуть внутрь можно так: docker exec -it vsearch-pg psql -U postgres -d vsearch -c "\d documents" — должна показаться колонка embedding типа vector.

Шаг 4. Считаем эмбеддинги и заполняем базу

Модель загрузится при первом вызове SentenceTransformer и закешируется локально. Регистрация типа vector через register_vector позволяет передавать списки чисел прямо в запрос.

import psycopg
from pgvector.psycopg import register_vector
from sentence_transformers import SentenceTransformer

DSN = "postgresql://postgres:devpass@localhost:5432/vsearch"
model = SentenceTransformer("all-MiniLM-L6-v2")

docs = [
    ("Возврат бракованного товара",
     "Как вернуть деньги за товар с дефектом и в какой срок продавец обязан ответить."),
    ("Настройка резервного копирования",
     "Расписание бэкапов базы данных и проверка восстановления из копии."),
    ("Оформление самозанятости",
     "Регистрация плательщика налога на профессиональный доход через приложение."),
]

with psycopg.connect(DSN) as conn:
    register_vector(conn)
    with conn.cursor() as cur:
        for title, content in docs:
            # Векторизуем заголовок + текст вместе для более полного смысла
            vec = model.encode(f"{title}. {content}")
            cur.execute(
                "INSERT INTO documents (title, content, embedding) VALUES (%s, %s, %s)",
                (title, content, vec),
            )
    conn.commit()
print("Загружено документов:", len(docs))

Проверка: вывод Загружено документов: 3. Число строк в таблице сверьте запросом SELECT count(*) FROM documents; — должно быть 3.

Шаг 5. Ищем по смыслу

Запрос тоже превращаем в вектор той же моделью, а затем берём ближайшие строки. Оператор <=> в pgvector считает косинусное расстояние: чем меньше, тем ближе по смыслу. Есть ещё <-> (евклидово) и <#> (скалярное произведение) — начните с косинусного, оно устойчивее к длине текста.

import psycopg
from pgvector.psycopg import register_vector
from sentence_transformers import SentenceTransformer

DSN = "postgresql://postgres:devpass@localhost:5432/vsearch"
model = SentenceTransformer("all-MiniLM-L6-v2")

query = "как получить деньги обратно за сломанную вещь"
query_vec = model.encode(query)

with psycopg.connect(DSN) as conn:
    register_vector(conn)
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT title,
                   embedding <=> %s AS distance
            FROM documents
            ORDER BY distance
            LIMIT 3
            """,
            (query_vec,),
        )
        for title, distance in cur.fetchall():
            print(f"{distance:.4f}  {title}")

Проверка: первой строкой с наименьшим расстоянием окажется «Возврат бракованного товара», хотя в запросе нет ни слова из заголовка. Это и есть семантический поиск.

Полнотекстовый поиск ищет слова. Векторный ищет смысл. Запрос «сломанная вещь» находит документ про «товар с дефектом» — там нет ни одного общего слова, но векторы рядом.

Индекс: когда строк станет много

На трёх документах поиск идёт полным перебором, и это нормально. На десятках тысяч строк перебор станет медленным. pgvector предлагает два типа индексов приближённого поиска.

Пример индекса HNSW под косинусное расстояние (класс операторов должен совпадать с оператором в запросе — здесь vector_cosine_ops под <=>):

CREATE INDEX ON documents
    USING hnsw (embedding vector_cosine_ops);

Оба индекса приближённые: они могут пропустить настоящего ближайшего соседа ради скорости. Если вам нужен строго точный результат — оставляйте полный перебор и масштабируйте иначе. Параметры сборки и поиска (m, ef_construction, hnsw.ef_search) сверяйте в актуальном README pgvector — значения по умолчанию менялись между версиями.

Если не получилось

• ERROR: type "vector" does not exist — не выполнен CREATE EXTENSION vector в этой базе, либо вы подключились не к той базе. Проверьте имя базы в строке подключения.
• ERROR: expected 384 dimensions, not N — размерность колонки не совпала с моделью. У all-MiniLM-L6-v2 это 384; у другой модели другое число. Пересоздайте таблицу с нужной размерностью.
• psycopg.OperationalError: connection refused — контейнер не поднялся или порт другой. Посмотрите docker ps и логи docker logs vsearch-pg.
• Долгая пауза при первом encode — это скачивание модели, а не зависание. Один раз, дальше из кеша.
• Индекс создан, но поиск не ускорился — класс операторов индекса не совпал с оператором в ORDER BY. Под <=> нужен vector_cosine_ops.

Что делать дальше

Рабочий каркас готов. Дальше стоит смотреть в эту сторону: замерить точность на своих запросах и подобрать модель под язык (многоязычные модели sentence-transformers лучше держат русский, чем базовые английские); нормализовать длину текстов — очень длинные документы лучше резать на куски и хранить кусками; проверить, как индекс ведёт себя под вашим объёмом, замерив время запроса на реальных данных. Инвалидация и переиндексация при обновлении документов — отдельная тема, здесь мы её не трогали.