Средний

13 мин

6.4K

В первой части мы развернули инфраструктуру на базе 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.gguf
EMBEDDER__PORT=8080
EMBEDDER__N_GPU_LAYERS=99
EMBEDDER__CONTEXT_SIZE=8192
EMBEDDER__BATCH_SIZE=8192
EMBEDDER__UBATCH_SIZE=8192

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

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

curl -f http://localhost:8080/health

curl -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 BaseEmbedding
from core import settings

class 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 httpx

class 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 logger
import numpy as np

class 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_cache
def 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/cu128

uv add transformers accelerate

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

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

import asyncio
import torch
from vechord.rerank import BaseReranker
from core import settings

class 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 AutoModel

class 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 = 64
class 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 Pipeline
from vechord.chunk import BaseChunker
from core import settings

class 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 позволяет работать с мультимодальными эмбеддингами, что открывает ещё больше вариантов реализации поиска.

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

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

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