Створюємо «пошуковий рушій» для інтернет-магазину за допомогою Embeddings

Всім привіт! Для тих, хто колись працював над розробкою блогів, інтернет-магазинів і будь-яких інших сайтів, де може бути текстовий пошук по контенту, не є секретом, що це не сама проста задача. Якщо у вашому магазині неможливо знайти товари, то в вас нічого не куплять :) Я теж зіткнувся з цією проблемою при розробці свого магазину.

Спочатку я вирішив спробувати повнотекстовий пошук в PostgreSQL. Завантажив словник української мови, створив відповідний індекс в таблиці з товарами. Воно чудово працює з різними словоформами, не враховує прийменники та інші неважливі частини мови. Проблеми починаються, коли в пошуковому запиті присутні слова, яких немає в назві товарів, або назви брендів написані з помилками чи транслітом. В таких випадках постгрес не знаходить нічого, що не є добре.

Я знаю про Elasticsearch, але ніколи з ним не працював і, чесно кажучи, мені це не здається простим рішенням. В цьому випадку треба створювати ще одне окреме сховище, яке треба якось синхронізувати з основою базою даних. До того ж, наскільки мені відомо, Elasticsearch сам по собі не вирішує проблему транслітерації та не буде просто так шукати «Xiaomi» по запиту «сяомі», для цього теж потрібно створювати відповідні словники. До того ж, деякі користувачі все ще використовують для пошуку російську мову, що теж погано відображається на результатах пошуку.

В той час, коли існують LLM, хотілося б знайти якесь «розумне» рішення, яке вміло б робити все це само по собі. І ігнорувати словоформи, і бути багатомовним, і враховувати транслітерацію та помилки.

На той момент я як раз займався вивченням питанням створення RAG-застосунку. Основною частиною цієї концепції є семантичний пошук в документах по запиту від користувача за допомогою embedding-моделей. І тоді я подумав — якщо ми можемо зробити такий пошук по текстовим документам, то теж саме можна зробити і для товарів в інтернет-магазині. В даній статті хочу поділитися своїм досвідом і показати, що в мене вийшло.

Отже, для реалізації «пошукового рушія» нам знадобляться дві речі — embedding model і векторне сховище.

Embedding model

Спершу розберемося, що таке embeddings в ML. Якщо казати дуже спрощено, то embeddings — це представлення різних обʼєктів, таких як зображення, текст або аудіо у вигляді масиву чисел. Це дозволяє порівнювати обєʼкти між собою і знаходити схожість або «відстань» між ними. А Embedding model, своєю чергою, це модель яка перетворює ці обʼєкти на векторне представлення.

Де нам взяти таку модель? Найпростіший варіант — це використати, наприклад, моделі від OpenAI. Також можна обрати опенсорсну модель з Hugging Face. При пошуку важливо враховувати підтримувані мови. Також треба визначити розмірність векторів, яку повертає модель. Я зупинився на моделі від Alibaba.

Використання опенсорсних моделей дозволяє розгортати модель «локально», в середині вашої інфраструктури, економити час на API-запити до сторонніх сервісів типу OpenAI. Також моделей просто багато і ви можете експериментувати та підібрати ту, яка краще справляється саме з вашими юзкейсами. З мінусів — вам потрібно додати окремий сервіс, який буде обслуговувати запити до локальної моделі.

Vector store

Для того, щоб не просто зберігати векторне представлення тексту, але і мати змогу робити семантичний пошук, потрібна спеціалізована база даних, яка це підтримує. Найбільш вичерпний список можливих варіантів мені трапився в документації до фреймворку LangChain. Оскільки в моєму проєкті вже використовується PostgreSQL, то найбільш очевидним вибором став pgvector — розширення для PostgreSQL. Це позбавило мене необхідності розгортати додаткове сховище і постійно синхронізувати дані між ними.

Реалізація

Перейдемо до реалізації. Створимо простий API-сервер для доступу до моделі. Встановимо потрібні залежності:

pip install fastapi sentence-transformers uvicorn

Створимо файл main.py з кодом серверу:

import os
from fastapi import FastAPI
from sentence_transformers import SentenceTransformer
from pydantic import BaseModel
# Load model from ./model directory or download it if it does not exist
model_name = "Alibaba-NLP/gte-multilingual-base"
model_path = os.path.join("model", model_name)

if os.path.exists(model_path):
   model = SentenceTransformer(model_path, trust_remote_code=True)
   print(f"Model loaded from {model_path}")
else:
   model = SentenceTransformer(model_name, trust_remote_code=True)
   model.save(model_path)
   print(f"Model downloaded and saved to {model_path}")

app = FastAPI()

class EmbeddingRequest(BaseModel):
   texts: list[str]

@app.post("/get-embeddings")
def get_embedding(payload: EmbeddingRequest):
   texts = payload.texts
   embeddings = model.encode(texts, normalize_embeddings=True)
   return {"embeddings": embeddings.tolist()}

if __name__ == "__main__":
   import uvicorn
   service_port = int(os.getenv("EMBEDDING_PORT", 9010))
   uvicorn.run(app, host="0.0.0.0", port=service_port)

Застосунок має один ендпоінт — «/get-embeddings», який приймає масив рядків з текстом для формування векторів і повертає масив векторів у відповідному порядку. Для того, щоб не завантажувати модель після кожного перезапуску застосунку, ми зберігаємо її в ./model.

Запускаємо сервер:

python main.py

І перевіряємо, що все працює:

Тепер можна перейти до підготовки нашої БД. В документації до pgvector детально описані всі способи встановлення розширення, а також в них є Docker Image. Я просто замінив образ postgres:14 на pgvector/pgvector:pg14 в своєму docker-compose файлі. Все запустилося без проблем. Далі треба увімкнути розширення для потрібної бази даних (це робиться один раз для кожної потрібної бази):

CREATE EXTENSION vector;

Після цього додаємо колонку в потрібну таблицю:

ALTER TABLE public.product ADD COLUMN embedding VECTOR(768);

Тут 768 — це розмірність векторів, яку генерує наша модель. Ви не зможете записати в цю колонку вектор іншої розмірності.

Подальша реалізація бізнес-логіки буде дуже залежати від вашого конкретного проєкту і мої фрагменти коду будуть нерелевантні, тому далі я буду використовувати псевдокод і спрощені SQL запити, щоб передати суть.

Проєкт написаний на TypeScript/Node.js, тому я скористався бібліотекою pgvector-node для формування запитів в БД повʼязаних з векторами. Повний список бібліотек для різних мов програмування тут.

Далі, нам треба згенерувати емебеднігі для всіх товарів і заповнити нову колонку даними. Для цього робимо наступне:

1. Дістаємо з бази всі товари, де колонка embedding порожня.
2. Збираємо назви товарів.
3. Відправляємо їх нашій моделі.
4. Отримуємо векторні представлення.
5. Зберігаємо значення в БД.

Приблизний код виглядає так:

let productChunk: Array<{ id: number, title: string }>;
const limit = 50;
do {
  productChunk = await ProductRepository.getProductInfoForEmbeddingGeneration(limit);
  if (!productChunk.length) {
    break;
  }
  
  const dataToUpdate: Array<{ id: number, embedding: Array<number> }> = [];
  // Send product titles to the model to get embeddings
  const embeddings = await EmbeddingModelAPI.getEmbeddingsByText(
    productChunk.map(product => product.title)
  );
  // Combine product id with embeddings
  for (let i = 0; i < productChunk.length; i++) {
    const product = productChunk[i];
    const embedding = embeddings.embeddings[i];
    dataToUpdate.push({ id: product.id, embedding });
  }
  
  // Update product embeddings
  await ProductRepository.updateProductEmbeddings(dataToUpdate);
} while (productChunk.length === limit);

Запит на оновлення товару з використанням TypeORM і pgvector-node:

dbConnection.query(
 `UPDATE public.product SET embedding = $1 WHERE id = $2`,
 [pgvector.toSql(product.embedding), product.id]
);

Після цього можемо побачити результат в БД:

Тепер ми можемо робити семантичний пошук товарів. Але спочатку треба зазначити що, коли ми шукаємо записи за схожістю, то насправді ми не фільтруємо результати, а просто упорядковуємо від більш схожих до менш схожих. Тобто ми робимо запити не через WHERE, а додаємо ORDER BY. Це означає, що наш SELECT буде завжди повертати всі товари, просто вони будуть ранжовані за релевантністю. Це важливо для побудови UI вашого застосунку. По-перше, можливо, ви захочете обмежити кількість результатів пошуку. Я зробив просто 5 сторінок. По-друге, щоб не написали користувачі в пошуку, вони ніколи не побачать порожні результати, просто вони будуть не дуже релевантні.

Базова логіка пошуку виглядає так:

1. Відправляємо пошуковий запит в нашу модель для формування ембедінгу.
2. Використовуємо цей ембедінг для пошуку схожих товарів в БД.
3. Для ранжування я використовую оператор <-> (L2 Distance). Детальніше про всі оператори тут.

Приблизний код:

const embeddings = await EmbeddingModelAPI.getEmbeddingsByText([searchQuery]);
const products = await ProductRepository.getClientProductListByEmbeddingSearch(
  embeddings.embeddings[0],
  limit,
  offset
);

Запит на пошук товарів з використанням TypeORM і pgvector-node:

dbConncection.query(
  `
    SELECT id, title, price,
      (embedding <-> $1) AS "embeddingDistance"
    FROM public.product
    WHERE embedding IS NOT NULL
    ORDER BY "embeddingDistance"
    LIMIT $2 OFFSET $3
  `,
  [pgvector.toSql(embedding), limit, offset]
);

Після всього цього мені вдалося значно покращити пошук в магазині. Ось приклади результатів:

Як можна побачити, модель добре справляється з помилками, транслітом, а також підтримує різні мови.

Наприкінці ще хочу додати можливі шляхи покращення перфомансу пошуку:

1. Кешувати ембедінги для пошукового запиту на короткий час, щоб не звертатися до моделі при переході між сторінками результатів пошуку.
2. Формувати ембедінги для товарів на основі не тільки назви, але й, наприклад, опису і атрибутів. В моєму випадку це складно через відсутність цього у 90% товарів.
3. Спробувати додати індекс на колонку з вектором. Але він впливає не тільки на швидкість, а також і на результати пошуку.

Підписуйтеся на Telegram-канал «DOU #tech», щоб не пропустити нові технічні статті