PostgreSQL + VectorChord = Гибридный поиск. Часть 2. Безоблачная

В первой части мы развернули инфраструктуру на базе PostgreSQL и VectorChord, настроили базовые CRUD-операции и запустили гибридный поиск с реранкингом (алгоритм ReciprocalRankFusion). Однако эмбеддер и чанкер на базе spacy работали скорее как “заглушки”, чем какая-то production readyистория.

В этой части мы заменим игрушечные компоненты на локальные SOTA-модели, сохраняя оффлайн-архитектуру и отказываясь от облачных API. Пошагово развернём llama.cpp-сервер для мультиязычного эмбеддинга Jina v4, поднимем нативный реранкер через transformers + PyTorch, подключим чанкер на базе chonkie и, наконец, оценим качество поиска.

Кастомный эмбеддер

Начнем с кастомного эмбеддера. Пока изучал библиотеку, заодно почитал про Jina, т.к. до этого с ней не встречался. VectorChord использует модель jina-embeddings-v4 с доступом по API. Но модель то открытая, и лежит на HuggingFace. Если верить описанию, то jina-embeddings-v4 — это универсальная модель для мультимодального и многоязычного поиска. Модель специально разработана для поиска сложных, в том числе визуально насыщенных документов с диаграммами, таблицами и иллюстрациями. То что надо для RAG системы.

Одно из преимуществ jina-embeddings-v4 — поддержка технологии Matryoshka: эмбеддинги по умолчанию имеют размерность 2048, но модель изначально обучалась так, чтобы размерность можно было уменьшать. Поддерживаются размерности 128, 256, 512 и 1024 с минимальными потерями качества, что позволяет легко балансировать между скоростью, стоимостью хранения и точностью поиска.

Также разработчики выпустили три специализированные версии модели, каждая под свою задачу: retrieval, text-matching и code. Нас интересует jina-embeddings-v4-text-retrieval — она оптимизирована для поиска документов и идеально подходит для RAG. Из этой модели удалили визуальные компоненты, что позволило сократить на четверть количество параметров. И да, модели сжаты в GGUF, что облегчает их скачивание и запуск.

Качаем модель приемлемой размерности. Кладём её в какую-то папку (у меня это nlp_models/jina_embeddings_v4) и прописываем путь к модели в .env.

Для запуска будем использовать библиотеку llama.cpp, запущенную как сервер. Использовать llama.cpp, как и всё остальное, будем через докер. Для выбора образа можно воспользоваться официальной документацией. Нам нужен серверный образ с поддержкой GPU. Т.к. у меня CUDA 12.8, то я возьму ghcr.io/ggml-org/llama.cpp:server-cuda

Создадим отладочный compose для запуска модели и проверим как оно работает.

services:  jina-embeddings:    image: ghcr.io/ggml-org/llama.cpp:server-cuda    volumes:      - ./nlp_models/jina_embeddings_v4:/models    ports:      - "${EMBEDDER__PORT}:${EMBEDDER__PORT}"    command: >      -m /models/${EMBEDDER__FILE_NAME}      --embedding      --pooling mean      --host 0.0.0.0      --port ${EMBEDDER__PORT}      -ngl ${EMBEDDER__N_GPU_LAYERS}      -c ${EMBEDDER__CONTEXT_SIZE}      -b ${EMBEDDER__BATCH_SIZE}      -ub ${EMBEDDER__UBATCH_SIZE}      --flash-attn on      --verbose    deploy:      resources:        reservations:          devices:            - driver: nvidia              count: 1              capabilities: [gpu, compute, utility]

Выбор модели уже обсудили. Через volumes прокидываем папку со скачанной моделью в формате GGUF. Дальше идут параметры для настройки сервера llama.cpp:

• -m — путь к файлу модели. Задаётся от внутренней папки контейнера

• --embedding — ключ нужен, чтобы сервер возвращал эмбеддинги, а не генерировал текст

• --pooling задаёт, как не удивительно, метод пулинга эмбеддингов. В даном случае mean — усреднение по всем токенам

• --host, --port — сетевые настройки. Задаём на каком порту слушать и какие IP

• -ngl — сколько слоёв модели разместить на GPU. Можно настраивать, если нехватает памяти. Моделька маленькая, поэтому влезет целиком. Для этого задаём -1 или 99.

• -c — размер контекстного окна в токенах. Опять же, зависит от ограничений видеопамяти

• -b и -ub — на мой взгляд самые странные параметры. Расшифровываются как batch-size и ubatch-size, но к привычному по PyTorch батчу не имеют никакого отношения — здесь это именно число токеном, которое обрабатывает модель за раз. -ub должен позполить посчитать количество частей, на которые делится -b входных токенов при вычислении.

• --flash-attn on включает технологию Flash Attention для ускорения вычислений

Секция deploy даёт доступ контейнеру к видеокарте хоста. Теперь дополним .env файл.

EMBEDDER__FILE_NAME=jina-embeddings-v4-text-retrieval-Q8_0.ggufEMBEDDER__PORT=8080EMBEDDER__N_GPU_LAYERS=99EMBEDDER__CONTEXT_SIZE=8192EMBEDDER__BATCH_SIZE=8192EMBEDDER__UBATCH_SIZE=8192

У меня сервер выбрасывал ошибку, если длина чанка была больше -ub токенов. Сделал вывод, что деление не работало, поэтому задал их одинаковыми. Число 8192 взял с карточки модели.

Запустим контейнер и проверим работоспособность модели:

curl -f http://localhost:8080/healthcurl -X POST http://localhost:8080/v1/embeddings -H "Content-Type: application/json" -d '{"input": ["Query: Hello World"]}'

Должны получить "status":"ok" на первый запрос и вот такую структуру на второй:

{  "model": "jina-embeddings-v4-text-retrieval-Q8_0.gguf",  "object": "list",  "usage": {    "prompt_tokens": 7,    "total_tokens": 7  },  "data": [    {      "embedding": [2048 чисел],      "index": 0,      "object": "embedding"    }  ]}

Теперь напишем класс LocalJinaEmbedding. Начнем с конструктора

from vechord.embedding import BaseEmbeddingfrom core import settingsclass LocalJinaEmbedding(BaseEmbedding):    SUPPORTED_DIMS = [128, 256, 512, 1024, 2048]    def __init__(self, dim: int = 1024, timeout: float = 30.0):        if dim not in self.SUPPORTED_DIMS:            raise ValueError(                f'Dimension {dim} not supported. Choose from {self.SUPPORTED_DIMS}')        self.base_url = settings.embedder.url        self.dim = dim        self.timeout = timeout        self._client = None

Класс унаследуем от vechord.embedding.BaseEmbedding. В конструктор передадим необходимую размерность эмбеддингов, а также таймаут сетевого соединения. URL нашего llama.cpp сервера передадим через настройки. Реализуем клиент для установки соединения с сервером:

import httpxclass LocalJinaEmbedding(BaseEmbedding):    # предыдущий код    async def _get_client(self) -> httpx.AsyncClient:        if self._client is None:            self._client = httpx.AsyncClient(                timeout=httpx.Timeout(self.timeout),                limits=httpx.Limits(max_keepalive_connections=10, max_connections=100))        return self._client    async def close(self):        if self._client:            await self._client.aclose()            self._client = None

Для клиента реализуем ленивую загрузку — соединение будет создано при первом обращении к методу getclient. Также реализуем метод close для корректного закрытия пула соединений. Следующий на очереди — метод получения эмбеддингов. Метод будет универсальным, как для одной стройки, так и для батча:

from loguru import loggerimport numpy as npclass LocalJinaEmbedding(BaseEmbedding):    # предыдущий код    async def _get_embeddings(self, texts: str | list[str]) -> list[np.ndarray]:        if isinstance(texts, str):            texts = [texts]        payload = {            'input': texts,            'model': settings.embedder.model_name,            'encoding_format': 'float'        }        if self.dim != 2048:            payload['dimensions'] = self.dim            try:                client = await self._get_client()                response = await client.post(                  url=f'{self.base_url}/v1/embeddings',                  json=payload                )                response.raise_for_status()                data = response.json()                  embeddings = [np.array(item['embedding'], dtype=np.float32) for item in data['data']]                return embeddings                        except httpx.HTTPStatusError as e:                logger.error(f'Failed to get embeddings: {e}')                raise

В методе ничего сверхъестественного — формируем словарь нашего запроса, делаем POST-запрос к llama серверу, преобразуем ответ. Формат np.array выбран с учетом того, что методы класса BaseEmbedding должны возвращать именно np.ndarray. Осталось переопределить базовые методы vectorize_chunk и vectorize_query. Тут есть одна особенность выбранной модели — она по-разному считает эмбеддинги для запросов и чанков документов (которые называют странным словом passage). Соответственно, модель ожидает разный формат входного запроса. Вот выдержка из документации:

Этим и будут отличатся наши методы:

class LocalJinaEmbedding(BaseEmbedding):    # предыдущий код    async def vectorize_chunk(            self,            text: str | None = None,            image: bytes | None = None,            image_url: str | None = None,    ) -> np.ndarray:        self.verify(text=text, image=image, image_url=image_url)        prefixed_text = f'Passage: {text}'        return (await self._get_embeddings(prefixed_text))[0]    async def vectorize_query(            self,            text: str | None = None,            image: bytes | None = None,            image_url: str | None = None,    ) -> np.ndarray:        self.verify(text=text, image=image, image_url=image_url)        prefixed_text = f'Query: {text}'        return (await self._get_embeddings(prefixed_text))[0]

Для улучшения производительности обработки дополним наш эмбеддер методом для пакетной обработки

class LocalJinaEmbedding(BaseEmbedding):    # предыдущий код    async def vectorize_batch(        self, texts: list[str], is_query: bool = False    ) -> list[np.ndarray]:        if not texts:            return []        prefix = 'Query: ' if is_query else 'Passage: '        prefixed_texts = [f'{prefix}{text}' for text in texts]        return await self._get_embeddings(prefixed_texts)

В методе формируем префикс, исходя из переданного параметра is_query, далее вызываем метод получения эмбеддингов для всех полученных на вход текстов, дополненных префиксом.

Внесём новые поля в настройки:

class EmbedderConfig(BaseModel):    model_name: str    file_name: str    host: str    port: int    @property    def url(self) -> str:        return f'http://{self.host}:{self.port}'

Теперь доработаем метод get_embedder, сделав из него фабрику эмбеддеров

@lru_cachedef get_embedder() -> BaseEmbedding:    if 'core' in settings.embedder.model_name.lower():        emb = SpacyDenseEmbedding(model='ru_core_news_md')    elif 'jina' in settings.embedder.model_name.lower():        emb = LocalJinaEmbedding(dim=settings.db.embedding_dim)    else:        raise ValueError('Embedding model not supported')    return emb

Можно попробовать запустить сервис и посоздавать документы через http://127.0.0.1:8000/docs. Главное — не забыть поменять размерность эмбеддингов и удалить старую базу данных.

Кастомный реранкер

Как писал в предыдущей статье, реранкер позволяет повысить качество гибридного поиска. Сейчас в коде используется RRF. Он без API ключей, но сам метод прям слишком просто. Поэтому, по аналогии с эмбеддером, сделаем локальный реранкер на базе доступной на HuggingFace модели jina-reranker-v3.

Для начала модель необходимо скачать. Так как это полноценная модель, а не GGUF файл, её просто так скачать не получится. Необходимо либо ставить huggingface-cli, либо писать модуль загрузки. Но, если использовать uv, то всё гораздо проще — можно просто запустить утилиту от HuggingFace с помощью uvx, без всяких дополнительных установок:

uvx hf download jinaai/jina-reranker-v3 --local-dir ./папка/для/скачивания

Модель есть. Теперь установим необходимые зависимости. А необходимы нам torch, transformers и accelerate. transformers нужна для запуска нашей модели, accelerate обеспечит оптимизацию ресурсов, ну а torch необходим для работы transformers. Причём для использования ресурсов видеокарты нужна не простая версия torch, а собранная под конкретную версию CUDA. У меня 12.8, поэтому

uv add torch --index https://download.pytorch.org/whl/cu128uv add transformers accelerate

Теперь код для запуска модели. Так как библиотека transformers — это синхронная библиотека, а в vechord.rerank.BaseReranker метод rerank (да и FastAPI тоже) определён как асинхронный, придётся реализовывать дополнительные методы — обёртки над синхронными вызовами.

Для начала напишем конструктор класса:

import asyncioimport torchfrom vechord.rerank import BaseRerankerfrom core import settingsclass LocalJinaReranker(BaseReranker):    def __init__(self, use_fp16: bool = True,):        self.model_path = settings.reranker.model_path        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'        self.use_fp16 = use_fp16        self._model = None        self._lock = asyncio.Lock()

В конструкторе мы тернарным оператором назначаем вычислительное устройство на основе доступности CUDA. Это нужно, чтобы в дальнейшем перенести модель на GPU. Параметр use_fp16 служит для оптимизации размера модели ра GPU. Последние два поля класса нужны для организации ленивой загрузки модели реранкера. Далее код загрузки модели:

from transformers import AutoModelclass LocalJinaReranker(BaseReranker):    # предыдущий код    def _load_model_sync(self):        if self._model is not None:            return        dtype = torch.float16 if (self.use_fp16 and self.device == 'cuda') else None        self._model = AutoModel.from_pretrained(            str(self.model_path), dtype=dtype, trust_remote_code=True,        ).to(self.device)        self._model.eval()

Модель загружаем с помощью класса библиотеки transformers. Параметр trust_remote_code нужен для обеспечения загрузки модели с кастомной архитектуры. Загруженную модель, с помощью .to(self.device) переносим на определённое в конструкторе устройство. Последним вызываем метод .eval(), который переводит модель из режима обучения в режим инференса. Теперь обёртка:

class LocalJinaReranker(BaseReranker):    # предыдущий код    async def _ensure_model(self):        if self._model is None:            async with self._lock:                if self._model is None:                    await asyncio.to_thread(self._load_model_sync)

Данный метод проверяет существование модели. Если её нет, то в потоке запускаем синхронный загрузчик. В отличие от эмбеддера, здесь будет немного другой подход — в асинхронном методе реранкинга мы будем асинхронно же убеждаться, что модель загружена, после чего в синхронном методе работы с transformers будем использовать гарантировано существующее поле класса self._model.

В синхронном методе реранкера просто вызовем API модели:

class LocalJinaReranker(BaseReranker):    # предыдущий код    def _rerank_sync(self, query: str, documents: list[str]) -> list[int]:        results = self._model.rerank(query=query, documents=documents, return_embeddings=False,)        return [result['index'] for result in results]

Метод возвращает только список индексов, как и прописано в BaseReranker. Теперь асинхронная обертка метода rerank:

MAX_DOCUMENTS = 64class LocalJinaReranker(BaseReranker):    # предыдущий код    async def rerank(self, query: str, chunks: list[str]) -> list[int]:        if not chunks:            return []        if len(chunks) > MAX_DOCUMENTS:            logger.warning(f'Received {len(chunks)} chunks, but model may only support up to 64.')            chunks = chunks[:MAX_DOCUMENTS]        await self._ensure_model()        indices = await asyncio.to_thread(self._rerank_sync, query, chunks)        return indices

Ограничение MAX_DOCUMENTS взято из карточки модели. В методе сначала, как и писал ранее, сначала проверяем модель, а за тем вызываем синхронный метод в потоке rerank. Закончим наш реранкер деструктором:

class LocalJinaReranker(BaseReranker):    # предыдущий код    async def close(self):        if self._model is not None:            await asyncio.to_thread(lambda: self._model.to('cpu'))            del self._model            if torch.cuda.is_available():                torch.cuda.empty_cache()            self._model = None            logger.info('Model resources released')

Память, если модель была на GPU, сама по себе не освободиться. Для этого мы сначала переносим модель на CPU (в отдельном потоке, т.к. это может занять длительное время), затем удаляем объект модели и чистим кеш на GPU.

Код реранкера готов. Теперь дополним код нашего поискового сервиса:

class SearchService:    # предыдущий код    @staticmethod    async def hybrid_search_rerank(request: HybridSearchRequest) -> list[ChunkResponse]:        vector = await _embedder.vectorize_query(request.query)        text_retrieves = await vr.search_by_keyword(            Chunk, keyword=request.query, topk=request.topk * request.boost)        vec_retrieves = await vr.search_by_vector(            Chunk, vec=vector, topk=request.topk * request.boost, probe=request.probe)        chunks = list(            {chunk.uid: chunk for chunk in text_retrieves + vec_retrieves}.values()        )        indices = await _reranker.rerank(request.query, [chunk.text for chunk in chunks])        return [chunks[i] for i in indices[:request.topk]]

И создадим соответствующий эндпоинт

@router.post(    path='/hybrid',    response_model=list[ChunkResponse],    summary='Гибридный поиск с нейросетевым ранжированием')async def hybrid_search(request: HybridSearchRequest):    return await SearchService.hybrid_search_rerank(request)

Кастомный чанкер

Раз уж сделали кастомные эмбеддер и реранкер, сделаем и чанкер. SpacyChunker неплохо справляется с разбиением, но слишком уж меленько крошит текст. Воспользуемся чанкером chonkie. Документацию переписывать не буду, просто отмечу что есть несколько методов чанкирования и постобработки. Мы воспользуемся методом SemanticChunker, который измеряет семантическую (смысловую) близость между предложениями, чтобы понять, где лучше сделать разрыв. Для этого ему нужна какая-нибудь моделька для получения ембеддингов. Так как сервис должен работать локально, опять загрузим модель с HuggingFace:

uvx hf download minishlab/potion-base-32M --local-dir ./nlp_models/chonkie_chunker

Данная модель используется в классе SemanticChunker по умолчанию, так что для начала сойдёт.

Библиотека chonkie позволяет проводить чанкирование в несколько этапов. Для этого можно либо определять этапы по отдельности, либо создать пайплайн обработки. Реализуем класс. И по традиции, начнём с конструктора. Класс будем наследовать от базового vechord.chunk.BaseChunker:

from chonkie import Pipelinefrom vechord.chunk import BaseChunkerfrom core import settingsclass LocalChonkie(BaseChunker):    def __init__(self):        model_path = settings.chunker.model_path.resolve()        self.pipeline = (            Pipeline()            .chunk_with(                chunker_type='semantic',                embedding_model=str(model_path),                chunk_size=settings.chunker.max_tokens,                threshold=settings.chunker.similarity_threshold,            )            .refine_with(refinery_type='overlap')        )

В конструкторе создаём пайплайн из двух этапов — чанкирование методом SemanticChunker и постобработку, которая добавит дополнительное перекрытие текстов между смежными чанками. В параметры чанкера передаём путь к модели для рассчёта эмбеддингов, максимальный размер чанка (в токенах) и порог для определения сходства текстов. Теперь метод чанкирования. В методе асинхронно запускаем пайплайн и возвращаем полученные чанки текста:

class LocalChonkie(BaseChunker):    # предыдущий код    async def segment(self, text: str) -> list[str]:        result = await self.pipeline.arun(text)        return [ch.text for ch in result.chunks]

Собственно, всё. Теперь осталось потестить как это всё вместе работает.

Тестирование

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

Для оценки качества поиска сделаем следующее: загрузим чего-нибудь в базу данных (я брал “Войну и мир” Л.Н. Толстого), создадим тестовый датасет, состоящий из пар “Вопрос” — “Развёрнутый ответ”, и оценим релевантность чанков из разных методов поиска с помощью RAGAS.

Описывать код теста я не буду (он приведён в tests/ragas/evaluate.py), т.к. в статье и так одни кодо-вставки. Остановлюсь на основных моментах.

Напрямую RAGAS не сможет работать с нашим эмбеддером. Ему нужны специальные методы, определённые в классе ragas.embeddings.base.BaseRagasEmbeddings. Для обеспечения работоспособности пришлось создать собственный класс-обёртку.

Также у меня не получилось получить нормальную оценку на маленьких локальных моделях — они все падали с ошибками в процессе оценки. Пришлось воспользоваться API GigaChat (это не реклама, просто уже был там зарегистрирован). При регистрации дают 1кк халявных токенов на младшую модель, из которых около 100к у меня сразу ушло на один прогон теста. А установив библиотеку langchain_gigachat можно получить интеграцию GigaChat с RAGAS.

В самом тесте грузим тестовый датасет, на каждый вопрос получаем POST запросом чанки текста, далее оцениваем среднее время ответа и две метрики из библиотеки RAGAS — context_precision и context_recall. Данные метрики как раз предназначены для оценки этапа поиска в RAG-сисетемах. Они не проверяют сам сгенерированный ответ, а анализируют, насколько хороши фрагменты контекста, которые были найдены.

Метрика context_precision (точность контекста) измеряет долю релевантных фрагментов среди всех извлечённых. Метрика отвечает на вопрос: “Сколько из возвращённых чанков действительно полезно для ответа на вопрос”.

Метрика context_recall (полнота контекста) измеряет насколько извлечённый контекст покрывает информацию, необходимую для формирования правильного ответа (сравнивается с эталонным ответом из ground_truth)

Тест повторяем четыре раза (по количеству вариантов поиска). Результаты в таблице:

Все тесты проводились с новым эмбеддером и чанкером. По метрикам видно, что нейросетевой реранкер даже на таких плохих данных показал себя лучше других. Результаты RRF поиска, похоже, испортили нерелевантные чанки поиска по ключевым словам. То, что время семантического поиска больше времени гибридного (который включает в себя семантический) спишем на малый размер тестового датасета.

Послесловие

В материале я постарался раскрыть основные способы запуска локальных моделей для дальнейшего использования в VectorChord: GGUF-файлы через llama.cpp-сервер, нативные PyTorch-модели напрямую из экосистемы HuggingFace, ну и просто библиотекой.

И опять затравочка на будущее — VectorChord позволяет работать с мультимодальными эмбеддингами, что открывает ещё больше вариантов реализации поиска.

Я всё. Надеюсь в этот раз метрики статьи будут получше 😁.

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

Код проекта доступен тут.