Казалось бы, в посгресе и так есть неплохой полнотекстовый поиск (tsvector/tsquery), и вы из коробки можете проиндексировать ваши тексты, а потом поискать по ним. Но на самом деле это не совсем то, что нужно — такой поиск работает лишь по чётким совпадениям слов. Т.е. postgres не догадается, что "кошка гонится за мышью" — это довольно близко к "котёнок охотится на грызуна". Как же победить такую проблему?

TLDR:

1. Преобразовываем наши тексты в наборы чисел (векторы) при помощи API openAI.
2. Сохраняем векторы в базе с помощью pgvector.
3. Легко ищем близкие друг к другу векторы или ищем их по вектору-запросу.
4. Ускоряем индексами.

Делаем вектор из текста

Итак, разберёмся с преобразованием текста в вектор (это называется embedding). Необязательно закапываться глубоко в теорию, на практике всё делается очень просто — можно тупо использовать API OpenAI. На текущий момент это стоит от 2 до 13 центов за миллион токенов. Почти даром.

На входе берём набор слов (одно слово или много сразу), а на выходе получаем вектор.

Что вообще это за вектор такой? Это просто набор чисел, по сути — координаты в 1500-мерном пространстве. Причём преобразование из текста в embedding происходит на заранее обученной модели так, чтобы близкие по смыслу тексты имели схожее направление векторов.

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

Немного практики:

Получить embeddings можно, просто послав запрос на API OpenAI (токен можно получить на их сайте, закинув 10 баксов)

curl https://api.openai.com/v1/embeddings \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $OPENAI_API_KEY" \
  -d '{
    "input": "Your text string goes here",
    "model": "text-embedding-3-large"
  }'

Но вообще, конечно же, есть библиотеки-обёртки, например, github.com/sashabaranov/go-openai

strings := []string{"котик гонится за мышкой", "кошка ловит грызуна"}

client := openai.NewClient("your-token")
ctx := context.Background()
response, err := client.CreateEmbeddings(ctx, openai.EmbeddingRequestStrings{
    Input: strings,
    Model: "text-embedding-3-large",
    Dimensions: 2000,
})

По умолчанию модель text-embedding-3-large выдаёт вектор с размерностью 3072, а text-embeddng-3-small — 1536. Но в запросе можно указать, если хочешь поменьше (параметр dimensions)

Пихаем это в посгрес

Несколько месяцев назад Каруна послала меня на конференцию по базам данных, где рассказали про pgvector.
В общем, в посгресе есть расширение, которое предоставляет тип данных vectorи ряд операций: сложение, вычитание, умножение (поэлементное), косинусное расстояние, евклидово расстояние, скалярное произведение. А также функции sum(), avg(), и два типа индексов.

Способы установки расширения смотрите на их страничке, а для того, чтобы пощупать, можно просто использовать докер образ pgvector/pgvector:pg16 — это постргес 16 с pgvector.

Перед использованием расширение надо активировать

CREATE EXTENSION IF NOT EXISTS vector

Создадим таблицу для хранения текстов и их embeddings

CREATE TABLE embeddings (
    id SERIAL PRIMARY KEY, 
    text TEXT UNIQUE, 
    embedding vector(2000) NOT NULL
);

и запишем туда эмбединги для описаний таких товаров:

"котик гонится за мышкой"
"котик гонится за мышками"
"собака лает, караван идёт"
"однажды в студёную зимнюю пору"
"кошка охотится на грызуна"
"котёнок ловит крысу"
"на марсе не растут грибы"
"умирает конь"

Вот, накидал на коленке программу на Go, которая получает эмбеддинги для фраз и записывает в базу. Кстати, есть библиотека https://github.com/pgvector/pgvector-go, упрощающая взаимодействие с типом vector, но в таком простом примере можно обойтись и без неё:

package main

import (
    "context"
    "log"

    "github.com/jackc/pgx/v4/pgxpool"
    "github.com/sashabaranov/go-openai"
)

func main() {

    strings := []string{
        "котик гонится за мышкой",
        "кошка охотится на грызуна",
        "котёнок ловит крысу",
        "котик гонится за мышками",
        "собака лает, караван идёт",
        "однажды в студёную зимнюю пору",
        "на марсе не растут грибы",
        "умирает конь",
    }

    ctx := context.Background()
    embeddings := getEmbeddingsFromStrings(ctx, strings)

    db := initDb(ctx)

    if _, err := db.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector"); err != nil {
        log.Fatalf("couldn't create extension: %v", err)
    }

    if _, err := db.Exec(ctx, `
        CREATE TABLE IF NOT EXISTS embeddings (
            id SERIAL PRIMARY KEY, 
            text TEXT UNIQUE, 
            embedding vector(2000) NOT NULL
        )`); err != nil {
        log.Fatalf("couldn't create table: %v", err)
    }

    for str, embedding := range embeddings {
        _, err := db.Exec(ctx,
            `INSERT INTO embeddings 
                (text, embedding) VALUES 
                ($1, $2::float4[]::vector)`, str, embedding)
        if err != nil {
            log.Fatalf("couldn't insert: %v", err)
        }
    }
}

func getEmbeddingsFromStrings(ctx context.Context, strings []string) map[string][]float32 {
    client := openai.NewClient("your-token")

    response, err := client.CreateEmbeddings(ctx, openai.EmbeddingRequestStrings{
        Input:      strings,
        Model:      "text-embedding-3-large",
        Dimensions: 2000,
    })
    if err != nil {
        log.Fatalf("error: %v", err)
    }

    result := make(map[string][]float32)

    for _, data := range response.Data {
        embedding := data.Embedding
        result[strings[data.Index]] = embedding
    }

    return result
}

func initDb(ctx context.Context) *pgxpool.Pool {
    db, err := pgxpool.Connect(ctx, "postgres://pgvector_test:pgvector_test@localhost:5432/pgvector_test?sslmode=disable")
    if err != nil {
        log.Fatalf("couldn't connect to database: %v", err)
    }

    return db
}

После чего возьмём вектор для фразы "котик гонится за мышкой" и поищем, какие вектора ближе всего в смысле косинусного расстояния <=>:

select
    text,
    embedding <=> (
        select embedding
        from embeddings
        where text = 'котик гонится за мышкой'
    ) distance
from embeddings
order by distance;

и получаем такой результат:

Видно, что расстояние до самого себя равно нулю, а дальше postgres нашёл вполне близкие по смыслу фразы. Результат, кстати, сильно зависит от размерности векторов. При 1500 "умирает конь" поднимается довольно высоко в списке. Не знаю, возможно, конь тоже в каком-то смысле грызун, а умирает он не хуже, чем мышь в зубах кошки. Но при 2000 всё выглядит немного логичнее.

Так или иначе, в любом случае результаты намного интереснее, чем мог бы найти на этих данных встроенный полнотекстовый поиск посгреса.

В примере я ищу, насколько один существующий текст похож на другой, но можно сделать и поиск произвольной фразы: просто взять текст поисковой строки, преобразовать его в вектор и сунуть в запрос.

Ускоряем с помощью индексов

В экстеншене есть 2 индекса:

HNSW — медленно индексирует, использует больше памяти, но запросы работают быстрее. Создаёт многослойный граф.

-- пример индекса для ускорения расчёта косинусного расстояния
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops);

IVFFlat работает по другому принципу, он делит векторы на списки, а затем ищет подмножество этих списков, наиболее близкое к вектору запроса.

CREATE INDEX ON items USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

Индекс IVFFlate можно строить только после того, как таблица наполнена данными.

Нюансы и заключение

Так как в посгресе лежат уже просто векторы определённой размерности, то, как мне кажется, никто не помешает в них добавить ещё пару измерений. Например, если мы ищем похожие товары в интернет-магазине, то, наверно, можно попробовать добавить в вектора такие числа как "длина", "ширина", "цвет", и т. д.

Понятно, что зависимость от API чужого сервиса — это плохо, поэтому можно попробовать сделать свою систему для получения embeddings. Я нагуглил пару способов, как это сделать на языке Go, но глубоко в эту сторону не копал.

При использовании индексов результат может чуть-чуть отличаться от запроса без индексов. Это компромис для скорости работы.

Подписывайтесь на мой канал Cross Join, если вам интересен подобный контент.