Как вы оптимизируете embedding генерацию для большого количества документов?

Краткий тезис

Embedding генерация — это процесс, при котором каждый документ (или чанк) превращается в векторное представление с помощью нейросетевой модели. Для миллионов документов наивный подход — генерировать эмбеддинги по одному — может занять дни. Оптимизация сводится к четырём ключевым приёмам: батчевая обработка (batch inference), GPU-ускорение, квантизация эмбеддингов и асинхронная отправка запросов к модели. Дополнительно можно кэшировать эмбеддинги для дублирующихся чанков. При правильной настройке время генерации для 1 млн документов сокращается с 20+ часов до 1–2 часов, а иногда и до 20 минут с использованием распределённых систем.

1. Проблема: генерация эмбеддингов для миллионов документов

Что такое embedding генерация?

Embedding (эмбеддинг) — это вектор чисел фиксированной размерности (например, 768 или 1024), который кодирует семантический смысл текста. В RAG-системах мы генерируем эмбеддинги для всех чанков документов и сохраняем их в векторной базе данных.

Почему это узкое место?

• Время одного forward pass модели на CPU ~10–100 мс для чанка длиной 512 токенов.
• Для 1 млн чанков на CPU потребуется ~14 часов (при 20 мс на чанк).
• На GPU с батчевым режимом время снижается до 1–3 часов.
• Если не оптимизировать, индексация данных становится непомерно долгой, что блокирует обновление базы знаний.

Какие ресурсы потребляются?

• Вычислительное время (CPU/GPU)
• Память (VRAM при квантизации, RAM при загрузке модели)
• Сеть (при вызове внешнего API)

2. Batch embeddings — пакетная обработка

Термин: Batch inference — это обработка не одного, а сразу нескольких документов в одном вызове модели. Современные эмбеддинг-модели (например, intfloat/e5-large-v2, sentence-transformers/all-MiniLM-L6-v2) оптимизированы для батчей.

Как работает?

Модель принимает тензор формы (batch_size, sequence_length) и параллельно вычисляет эмбеддинги для всех элементов.

Размер батча (batch size)

• Оптимальный размер зависит от объёма VRAM и длины чанков.
• Для большинства GPU (T4, V100, A10G) подходит batch size 32–256.
• Отношение: чем больше batch size, тем выше пропускная способность (throughput), но до предела насыщения.

Пример кода (sentence-transformers)

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("intfloat/e5-large-v2", device="cuda")
chunks = ["текст чанка 1", "текст чанка 2", ...]  # 1 млн чанков

embeddings = model.encode(chunks, batch_size=128, show_progress_bar=True)
# Всего ~10-20 минут на 1 млн чанков на V100.

Таблица: влияние batch size на throughput (на A100)

Замечание: при batch size > 256 часто возникает падение efficiency из-за ограничений памяти, если не используется gradient checkpointing.

3. GPU ускорение — перенос на видеокарту

Термин: CUDA — платформа параллельных вычислений NVIDIA, позволяющая запускать нейросети на GPU.

Шаги оптимизации

1. Переключение модели на GPU: model = SentenceTransformer("...").to("cuda").
2. Использование mixed precision (FP16): уменьшает использование памяти и ускоряет вычисления.

   model.half()  # модель весит в 2 раза меньше, forward pass быстрее
   embeddings = model.encode(chunks, batch_size=128, precision="float16")
   

3. Многопроцессное ускорение: если у вас несколько GPU, можно использовать DataParallel или model_devices в sentence-transformers.

Сравнение CPU vs GPU (V100)

4. Quantization — снижение точности эмбеддингов

Термин: Quantization (квантизация) — это уменьшение битности чисел, представляющих эмбеддинги. Обычно эмбеддинги хранятся в FP32 (4 байта на число). После квантизации можно использовать FP16 (2 байта) или INT8 (1 байт).

Зачем нужно?

• Уменьшает размер хранилища: FP32 → FP16 → уменьшение вдвое, INT8 → в 4 раза.
• Ускоряет поиск в векторной базе (меньше данных для чтения).
• При инференсе: если модель запускается в INT8, то forward pass тоже быстрее.

Риски

• Потеря точности: recall@k может упасть на 1–5% в зависимости от модели.
• Для многих задач это приемлемый компромисс.

Инструменты

• sentence-transformers поддерживает quantization_config (через bm25? нет, через torch.quantization).
• FAISS поддерживает индекс IVF с квантизованными векторами (через IndexIVFFlat с metric_type=faiss.METRIC_INNER_PRODUCT).
• PyTorch — torch.quantization.quantize_dynamic для модели, to(torch.float16) для эмбеддингов.

Пример: сохранение эмбеддингов в INT8

import numpy as np
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")
embeddings_fp32 = model.encode(chunks, batch_size=128)

# Квантизация до INT8: делим на макс абсолютное значение и умножаем на 127
max_val = np.max(np.abs(embeddings_fp32), axis=1, keepdims=True)
embeddings_int8 = (embeddings_fp32 / max_val * 127).astype(np.int8)
# При поиске нужно восстановить: float_emb = int_emb * (max_val / 127)

5. Асинхронность — параллельные запросы к API

Если embedding модель — внешний сервис (OpenAI, Cohere, собственный microservice), то узким местом становится сетевая задержка.

Термин: Asyncio — библиотека для асинхронного ввода-вывода в Python. Она позволяет отправлять много запросов одновременно, не блокируя поток.

Код с aiohttp

import asyncio
import aiohttp

async def embed_chunk(session, api_url, chunk):
    async with session.post(api_url, json={"text": chunk}) as resp:
        return await resp.json()

async def main(chunks):
    async with aiohttp.ClientSession() as session:
        tasks = [embed_chunk(session, URL, chunk) for chunk in chunks]
        results = await asyncio.gather(*tasks)
    return results

Эффект

• Вместо последовательных запросов (0.1–0.5 сек на запрос) можно получить ~1000 запросов в секунду при наличии большого пула соединений.
• Важно соблюдать rate limits (лимиты вызовов) API. Используйте asyncio.Semaphore для ограничения параллелизма.

Когда не нужна асинхронность?

• Если модель развёрнута локально на GPU — лучше использовать batch inference, асинхронность даст выигрыш только при распределённом деплое (несколько инстансов).

6. Кэширование эмбеддингов для дублирующихся чанков

В больших датасетах часто встречаются повторяющиеся фрагменты (например, “Контакты: …” в каждом письме, общие секции в документах). Вместо того чтобы пересчитывать эмбеддинг, его можно закэшировать.

Реализация

• Хеш-таблица на основе хэша текста (например, hashlib.sha256).
• Key-value хранилище: Redis или dictionary в памяти.
• При добавлении нового чанка проверяем хэш; если такой уже есть — берём эмбеддинг из кэша.

Пример

import hashlib
cache = {}  # или Redis

def get_embedding_with_cache(chunk, model):
    sha = hashlib.sha256(chunk.encode()).hexdigest()
    if sha in cache:
        return cache[sha]
    emb = model.encode([chunk])[0]
    cache[sha] = emb
    return emb

Когда кэш эффективен?

• Доля дубликатов > 20% — кэш может сократить время генерации на 10–30%.
• При частом переиндексировании одних и тех же документов.

7. Дополнительные техники: раннее фильтрование и отложенная векторизация

7.1 Фильтрация чанков перед эмбеддингом

• Удаление пустых, слишком коротких (< 10 символов) или шаблонных чанков.
• Это снижает количество вызовов модели на 5–15%.

7.2 Отложенная векторизация (lazy embedding)

• Сначала сохраняем сырые тексты в базу, а эмбеддинги считаем только для запросов? Нет, это не для индексации. Для поиска нужны эмбеддинги всех документов. Но если вы используете sparse retrieval (BM25) параллельно с Dense, можно для части документов использовать только BM25, а эмбеддинги генерировать по требованию. Это trade-off скорость/точность.

8. Инструментарий и пайплайн оптимизации

Полный пайплайн (1 млн документов)

1. Прочитать документы → 1 млн чанков
2. Отфильтровать дубликаты с помощью хэшей (кэш)
3. Разбить на батчи по 128
4. Загрузить модель на GPU, переключить в FP16
5. Для каждого батча: forward pass → сохранить эмбеддинги на диск (numpy, parquet)
6. Если используем внешний API: асинхронные запросы с ограничением параллелизма
7. После получения всех эмбеддингов — квантизация до INT8 (опционально)
8. Запись в векторную БД (FAISS, Qdrant, Weaviate)

Оценка времени (на V100, 16 GB VRAM)

Без оптимизации (CPU, batch=1, FP32) то же заняло бы ~22 часа — экономия 90% времени.

9. Метрики для оценки эффективности

• Throughput: количество чанков в секунду.
• Latency p50/p99: время генерации одного чанка.
• VRAM usage: сколько памяти занимает модель и батч.
• Recall@k (качество эмбеддингов после квантизации): не должно просесть более чем на 2%.

Пет-проект для закрепления

Задача: оптимизировать генерацию эмбеддингов для датасета из 500 000 новостей (используйте open-source датасет ag_news).

Инструменты: sentence-transformers, PyTorch, tqdm, asyncio (опционально), FAISS.

Шаги:

1. Загрузите датасет, разбейте на чанки (если длинные — по 512 токенов).
2. Реализуйте базовое решение: CPU, batch_size=1, замерьте время.
3. Перенесите на GPU, увеличьте batch_size до 128, замерьте.
4. Переключите модель в FP16, замерьте.
5. Реализуйте кэширование по хэшу, посчитайте процент дубликатов.
6. Сохраните эмбеддинги в INT8, сравните recall на 1000 запросов с FP32.
7. (Опционально) Если датасет большой, используйте асинхронный HTTP-клиент к локальной модели через Flask/Triton.

Ожидаемый результат: вы получите работающий конвейер, который генерирует эмбеддинги в 10–20 раз быстрее исходного, с отчётом по метрикам и потерям качества.

Связь с другими вопросами

Навигация

• Предыдущий: 76
• Следующий: 78
• Индекс: 00. Индекс разборов