Визуализация RAG в 3D: как EmbeddingGemma и UMAP показывают процесс поиска в векторной базе

Почему RAG остаётся чёрным ящиком (и как это исправить)

Вы загружаете документы в векторную базу. Запускаете поиск. Получаете ответ. Всё работает. Или нет. Когда RAG начинает врать, вы смотрите в лог, видите кучу чисел с плавающей точкой и понимаете ровно ничего. Эмбеддинги — это 768 измерений у EmbeddingGemma, 1024 у BGE-M3. Человеческий мозг не может представить себе 768-мерное пространство. Мы видим максимум три измерения. Отсюда и проблема: мы не видим, как наши документы расположены в этом пространстве, не понимаем, почему поиск находит одни чанки и игнорирует другие.

UMAP: волшебный пресс для многомерных данных

Uniform Manifold Approximation and Projection. Звучит сложно, работает просто: берёт ваши 768-мерные вектора и сжимает их до 2D или 3D, сохраняя структуру соседства. Близкие в оригинальном пространстве точки остаются близкими после проекции. Далёкие — остаются далёкими. Не идеально, но достаточно, чтобы увидеть кластеры, выбросы и понять, что творится в вашей базе.

Почему UMAP, а не t-SNE? t-SNE медленнее на больших датасетах и хуже сохраняет глобальную структуру. UMAP быстрее, масштабируется лучше и даёт более стабильные результаты. Для визуализации поиска в RAG — идеальный инструмент.

Собираем установку: от текстов до 3D-графика

Вам понадобится: Python 3.9+, около 4 ГБ оперативки (EmbeddingGemma:300m не самая тяжёлая, но и не легковесная), и понимание, что вы делаете. Если хотите просто запустить код — пожалуйста. Но я объясню каждый шаг, чтобы вы понимали, зачем он нужен.

1 Готовим данные: что будем визуализировать?

Возьмём разнородный набор текстов: параграфы из технической документации, новости, отрывки из художественной литературы, вопросы пользователей. Чем разнообразнее — тем интереснее визуализация. В реальном проекте это будут ваши чанки документов.

documents = [
    "Векторные базы данных хранят эмбеддинги в многомерном пространстве.",
    "RAG системы комбинируют поиск по векторной базе с генерацией ответов.",
    "Сегодня на фондовом рынке наблюдался рост технологических компаний.",
    "Погода в Москве: облачно, возможен дождь к вечеру.",
    "Он медленно открыл дверь и вошёл в тёмную комнату.",
    "Как настроить гибридный поиск для RAG с BM25 и FAISS?",
    "Рецепт пасты карбонара: бекон, яйца, пармезан, пекорино.",
    "Криптовалюты показали волатильность после заявления регулятора.",
    "Луна светила ярко, отражаясь в тихой поверхности озера.",
    "Что такое мультимодальный RAG и как его использовать для видео?",
]

2 EmbeddingGemma:300m — локальный эмбеддер, который не стыдно показать

Почему EmbeddingGemma, а не OpenAI или Cohere? Потому что она бесплатная, локальная и даёт 768-мерные эмбеддинги хорошего качества. В статье "BGE M3 vs EmbeddingGemma vs Qwen3" я подробно сравнивал модели. Для визуализации EmbeddingGemma подходит идеально: баланс между качеством и требовательностью.

from transformers import AutoTokenizer, AutoModel
import torch
import numpy as np

model_name = "google/embedding-gemma-300m"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, torch_dtype=torch.float16)

# Переводим модель в eval режим и на GPU если есть
model.eval()
if torch.cuda.is_available():
    model = model.cuda()

def get_embeddings(texts):
    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
    
    if torch.cuda.is_available():
        inputs = {k: v.cuda() for k, v in inputs.items()}
    
    with torch.no_grad():
        outputs = model(**inputs)
        # Берем embedding последнего токена [CLS] как репрезентацию текста
        embeddings = outputs.last_hidden_state[:, 0, :].cpu().numpy()
    
    return embeddings

# Получаем эмбеддинги для всех документов
embeddings_768d = get_embeddings(documents)
print(f"Размерность эмбеддингов: {embeddings_768d.shape}")  # (10, 768)

3 UMAP: сжимаем 768 измерений до 3

Теперь самое интересное: проецируем наши 768-мерные вектора в 3D пространство. UMAP нужно настроить — два ключевых параметра: n_neighbors и min_dist.

import umap

# Создаем UMAP редуктор для 3D
reducer = umap.UMAP(
    n_components=3,           # Проецируем в 3D
    n_neighbors=5,            # Сколько соседей учитывать
    min_dist=0.1,             # Минимальное расстояние между точками
    metric='cosine',          # Косинусное расстояние (стандарт для эмбеддингов)
    random_state=42           # Для воспроизводимости
)

# Применяем преобразование
embeddings_3d = reducer.fit_transform(embeddings_768d)
print(f"3D эмбеддинги: {embeddings_3d.shape}")  # (10, 3)

Что делают параметры? n_neighbors=5 означает, что UMAP будет пытаться сохранить локальную структуру — отношения между каждым текстом и его 5 ближайшими соседями. min_dist=0.1 контролирует, насколько плотно будут сгруппированы точки. Чем меньше значение, тем плотнее кластеры. Для визуализации поиска я рекомендую начать с этих значений и поэкспериментировать.

4 Plotly: интерактивная 3D визуализация, где можно покрутить

Matplotlib для 3D — это боль. Неинтерактивно, неудобно, выглядит уныло. Plotly даёт интерактивный график, который можно крутить, приближать, нажимать на точки и видеть текст.

import plotly.graph_objects as go

# Создаем 3D scatter plot
fig = go.Figure(data=[go.Scatter3d(
    x=embeddings_3d[:, 0],
    y=embeddings_3d[:, 1],
    z=embeddings_3d[:, 2],
    mode='markers+text',
    marker=dict(
        size=8,
        color=range(len(documents)),  # Разный цвет для каждой точки
        colorscale='Viridis',
        opacity=0.8
    ),
    text=documents,
    textposition="top center",
    hoverinfo="text",
    hovertext=[f"Документ {i}: {doc[:50]}..." for i, doc in enumerate(documents)]
)])

# Настраиваем layout
fig.update_layout(
    title="3D визуализация эмбеддингов документов",
    scene=dict(
        xaxis_title="UMAP измерение 1",
        yaxis_title="UMAP измерение 2",
        zaxis_title="UMAP измерение 3"
    ),
    width=1000,
    height=800
)

# Сохраняем в HTML для интерактивного просмотра
fig.write_html("rag_3d_visualization.html")

# Или показываем в Jupyter
# fig.show()

Что вы увидите на графике (и почему это важно)

Запустите HTML файл в браузере. Покрутите график. Вы увидите не случайные точки, а структуру:

• Технические тексты про RAG и векторные базы сгруппируются в одном районе
• Новости (фондовый рынок, криптовалюты, погода) образуют другой кластер
• Художественные описания уйдут в третью область
• Вопросы пользователей окажутся ближе к техническим текстам (потому что семантически ближе)

Это и есть семантическое пространство EmbeddingGemma. Тексты с похожим значением — близко. Разные — далеко. Теперь представьте, что у вас не 10 документов, а 10 000. И вы видите, как они разбиваются на кластеры по темам, доменам, стилям.

Добавляем поиск: как выглядит запрос в этом пространстве

Теперь самое главное: визуализируем сам процесс поиска. Берём пользовательский запрос, получаем его эмбеддинг, проецируем в 3D и смотрим, к каким документам он ближе всего.

# Пользовательский запрос
query = "Как работает поиск в векторных базах данных?"

# Получаем эмбеддинг запроса
query_embedding_768d = get_embeddings([query])

# Проецируем в 3D используя уже обученный UMAP редуктор
query_embedding_3d = reducer.transform(query_embedding_768d)

# Вычисляем косинусные расстояния до всех документов
from sklearn.metrics.pairwise import cosine_similarity

similarities = cosine_similarity(query_embedding_768d, embeddings_768d)[0]
top_k_indices = np.argsort(similarities)[::-1][:3]  # Топ-3 ближайших документа

print(f"Запрос: {query}")
print("Топ-3 похожих документа:")
for idx in top_k_indices:
    print(f"  [{similarities[idx]:.3f}] {documents[idx]}")

Теперь добавим запрос на график:

# Добавляем точку запроса на существующий график
fig.add_trace(go.Scatter3d(
    x=[query_embedding_3d[0, 0]],
    y=[query_embedding_3d[0, 1]],
    z=[query_embedding_3d[0, 2]],
    mode='markers',
    marker=dict(
        size=12,
        color='red',
        symbol='diamond'
    ),
    name="Запрос",
    hoverinfo="text",
    hovertext=[f"Запрос: {query}"]
))

# Добавляем линии от запроса к топ-3 документам
for idx in top_k_indices:
    fig.add_trace(go.Scatter3d(
        x=[query_embedding_3d[0, 0], embeddings_3d[idx, 0]],
        y=[query_embedding_3d[0, 1], embeddings_3d[idx, 1]],
        z=[query_embedding_3d[0, 2], embeddings_3d[idx, 2]],
        mode='lines',
        line=dict(color='red', width=2, dash='dash'),
        showlegend=False
    ))

fig.write_html("rag_3d_with_query.html")

Откройте новый HTML файл. Вы увидите красный ромбик — ваш запрос. И красные пунктирные линии, которые тянутся к трём ближайшим документам. Теперь вы буквально видите, как работает поиск: запрос попадает в определённую область семантического пространства, и векторная база находит ближайших соседей.

LanceDB в картинках: как выглядит реальная векторная база

До этого мы работали с массивами в памяти. В реальном проекте у вас LanceDB, Chroma, Qdrant или другая векторная БД. Давайте загрузим документы в LanceDB и визуализируем их оттуда.

import lancedb
import pyarrow as pa

# Создаем базу в памяти (или укажите путь к файлу для постоянного хранения)
db = lancedb.connect("/tmp/lancedb_rag_viz")

# Создаем таблицу
schema = pa.schema([
    pa.field("id", pa.int64()),
    pa.field("text", pa.string()),
    pa.field("vector", pa.list_(pa.float32(), 768)),
])

table = db.create_table("documents", schema=schema, exist_ok=True)

# Подготавливаем данные для вставки
data = []
for i, (text, embedding) in enumerate(zip(documents, embeddings_768d)):
    data.append({
        "id": i,
        "text": text,
        "vector": embedding.astype(np.float32).tolist(),
    })

# Вставляем документы
table.add(data)

# Выполняем поиск
results = table.search(query_embedding_768d[0].astype(np.float32)).limit(3).to_list()
print("Результаты поиска из LanceDB:")
for r in results:
    print(f"  [{r['_distance']:.3f}] {r['text']}")

Теперь у вас есть полноценная векторная база. И вы можете визуализировать не только поиск, но и структуру всей базы. Особенно полезно, когда база большая и вы хотите понять, нет ли проблем с кластеризацией.

Ошибки, которые вы увидите первыми (и как их исправить)

Зачем это всё? (Неочевидные применения)

Визуализация — не просто красивая картинка. Это инструмент отладки и анализа:

1. Отладка качества эмбеддингов. Видите, что технические документы и новости перемешались? Значит, модель плохо улавливает разницу между доменами. Может, стоит дообучить на ваших данных? В статье "Self-supervised обучение в Colab" я показывал, как это сделать.
2. Анализ покрытия базы знаний. Есть области в пространстве, где нет документов? Это слепые зоны вашего RAG. Запросы, попадающие туда, не найдут релевантных чанков.
3. Визуализация дрейфа. Делайте снимки пространства раз в неделю. Видите, как кластеры смещаются? Это дрейф данных или модели. Вовремя заметите.
4. Объяснение стейкхолдерам. Покажите продукт-менеджеру или заказчику не таблицу с метриками, а 3D-график, где видно, как запрос находит документы. Это убеждает лучше любых цифр.

Когда следующий раз ваш RAG начнёт галлюцинировать, не гадайте на кофейной гуще. Визуализируйте. Посмотрите, где находится запрос относительно документов. Увидите проблему. Поймёте, как её исправить.

Что дальше? От визуализации к интерактивному дебаггеру

Соберите дашборд на Streamlit или Gradio. Загружайте документы, смотрите их в 3D, вводите запросы и наблюдайте, как работает поиск в реальном времени. Добавьте возможность выделять проблемные зоны, помечать выбросы, пересчитывать эмбеддинги другой моделью.

Или интегрируйте с VectorDBZ — инструментом из статьи "VectorDBZ: Твой отладчик для векторных БД". Получите полный стек для отладки RAG: от низкоуровневого просмотра индексов до высокоуровневой 3D визуализации семантического пространства.

Главное — перестать воспринимать векторный поиск как чёрный ящик. Это геометрия. Её можно увидеть. А что видно — то можно понять, улучшить и объяснить.