Что такое embedding: векторное представление текста

import numpy as np
from openai import AsyncOpenAI

client = AsyncOpenAI()  # OPENAI_API_KEY из переменной окружения


async def embed_one(text: str, model: str = "text-embedding-3-small") -> np.ndarray:
    """Эмбеддинг одного текста. Возвращает нормализованный вектор."""
    response = await client.embeddings.create(input=text, model=model)
    vector = np.array(response.data[0].embedding, dtype=np.float32)
    # API уже возвращает нормализованные векторы, но явная нормализация не помешает
    return vector / np.linalg.norm(vector)


async def embed_batch(
    texts: list[str],
    model: str = "text-embedding-3-small",
    batch_size: int = 100,
) -> np.ndarray:
    """
    Batch-эмбеддинг списка текстов.
    Разбиваем на батчи: API принимает до 2048 входов, но 100 — надёжный лимит.
    Возвращает матрицу (N, dim).
    """
    all_vectors: list[list[float]] = []

    for i in range(0, len(texts), batch_size):
        batch = texts[i : i + batch_size]
        response = await client.embeddings.create(input=batch, model=model)
        # Ответы возвращаются в том же порядке, что входы
        batch_vectors = [item.embedding for item in sorted(response.data, key=lambda x: x.index)]
        all_vectors.extend(batch_vectors)

    matrix = np.array(all_vectors, dtype=np.float32)
    # Нормализуем все строки сразу
    norms = np.linalg.norm(matrix, axis=1, keepdims=True)
    return matrix / norms


# Пример
import asyncio

async def main():
    texts = [
        "Как установить FastAPI?",
        "FastAPI installation guide",   # другой язык, похожий смысл
        "Рецепт торта Наполеон",        # другая тема
        "Как НЕ нужно устанавливать FastAPI",  # с отрицанием!
    ]

    vectors = await embed_batch(texts)
    print(f"Размерность: {vectors.shape}")  # (4, 1536)

    # Сравниваем с первым запросом
    query = vectors[0]
    for i, (text, vec) in enumerate(zip(texts[1:], vectors[1:]), 1):
        sim = float(np.dot(query, vec))  # dot = cosine sim для нормализованных
        print(f"sim(query, text[{i}]) = {sim:.3f}  ← «{text[:40]}»")

    # Ожидаемый вывод:
    # sim(query, text[1]) = 0.921  ← «FastAPI installation guide»
    # sim(query, text[2]) = 0.124  ← «Рецепт торта Наполеон»
    # sim(query, text[3]) = 0.873  ← «Как НЕ нужно устанавливать FastAPI»  ← !

asyncio.run(main())

"""
sentence-transformers — библиотека для локальных embedding-моделей.
Поддерживает CPU и GPU, возвращает numpy-массивы.

pip install sentence-transformers
"""

import numpy as np
from sentence_transformers import SentenceTransformer


# Модель загружается один раз и переиспользуется
# multilingual-e5-large: 1.2 ГБ, хорошее качество для RU/EN
model = SentenceTransformer("intfloat/multilingual-e5-large")


def embed_local(
    texts: list[str],
    batch_size: int = 32,
    prefix: str = "passage: ",
) -> np.ndarray:
    """
    Эмбеддинг текстов локальной моделью.

    E5-модели требуют префиксы:
      "query: "   — для поисковых запросов
      "passage: " — для индексируемых документов

    Это важно: при неправильном префиксе качество снижается на 5-15%.
    """
    prefixed = [prefix + t for t in texts]
    return model.encode(
        prefixed,
        batch_size=batch_size,
        normalize_embeddings=True,  # вернёт векторы длиной 1
        show_progress_bar=len(texts) > 100,
    )


def embed_query(query: str) -> np.ndarray:
    return embed_local([query], prefix="query: ")[0]


def embed_documents(docs: list[str]) -> np.ndarray:
    return embed_local(docs, prefix="passage: ")


# Пример: поиск ближайшего документа
documents = [
    "FastAPI — современный веб-фреймворк для Python",
    "Установка FastAPI: pip install fastapi uvicorn",
    "Django — полнофункциональный веб-фреймворк",
    "PostgreSQL — реляционная база данных",
    "Машинное обучение с PyTorch",
]

doc_vectors = embed_documents(documents)
query_vector = embed_query("как установить FastAPI?")

# Cosine similarity = dot product для нормализованных
similarities = doc_vectors @ query_vector  # матричное умножение (N,)
ranked = sorted(enumerate(similarities), key=lambda x: -x[1])

print("Результаты поиска:")
for idx, score in ranked:
    print(f"  {score:.3f}  {documents[idx]}")

# Ожидаемый вывод:
#   0.912  Установка FastAPI: pip install fastapi uvicorn  ← top-1
#   0.874  FastAPI — современный веб-фреймворк для Python
#   0.731  Django — полнофункциональный веб-фреймворк
#   0.312  PostgreSQL — реляционная база данных
#   0.219  Машинное обучение с PyTorch

"""
Визуализация embedding-пространства через UMAP.
pip install umap-learn matplotlib
"""

import asyncio
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import umap


async def visualize_embeddings():
    # Тестовые тексты с метками тематики
    data = [
        # Python
        ("Python — язык программирования", "python"),
        ("pip install пакет", "python"),
        ("def функция(аргументы):", "python"),
        ("FastAPI веб-фреймворк", "python"),
        # Базы данных
        ("SQL запрос SELECT * FROM", "database"),
        ("PostgreSQL база данных", "database"),
        ("индекс в базе данных", "database"),
        ("JOIN таблиц в SQL", "database"),
        # ML
        ("обучение нейронной сети", "ml"),
        ("функция потерь loss", "ml"),
        ("градиентный спуск оптимизация", "ml"),
        ("точность модели accuracy", "ml"),
        # История
        ("Пётр I реформы России", "history"),
        ("Французская революция 1789", "history"),
        ("Второй мировой войны", "history"),
        ("история древнего Рима", "history"),
    ]

    texts, labels = zip(*data)
    vectors = await embed_batch(list(texts))

    # Снижаем размерность до 2D
    reducer = umap.UMAP(n_components=2, random_state=42, n_neighbors=5)
    points_2d = reducer.fit_transform(vectors)

    # Рисуем
    colors = {"python": "#7c6fff", "database": "#60a5fa", "ml": "#4ade80", "history": "#fbbf24"}
    fig, ax = plt.subplots(figsize=(10, 8), facecolor="#0d1117")
    ax.set_facecolor("#0d1117")

    for (x, y), text, label in zip(points_2d, texts, labels):
        color = colors[label]
        ax.scatter(x, y, color=color, s=80, alpha=0.8, zorder=3)
        ax.annotate(text[:25], (x, y), fontsize=7, color=color,
                    xytext=(4, 4), textcoords="offset points", alpha=0.9)

    # Легенда
    patches = [mpatches.Patch(color=c, label=l) for l, c in colors.items()]
    ax.legend(handles=patches, facecolor="#1a1a2e", edgecolor="#374151", labelcolor="white")
    ax.set_title("Embedding space (UMAP 2D)", color="white", pad=15)
    ax.tick_params(colors="#374151")
    for spine in ax.spines.values():
        spine.set_edgecolor("#374151")

    plt.tight_layout()
    plt.savefig("embedding_space.png", dpi=150, bbox_inches="tight", facecolor="#0d1117")
    print("Сохранено: embedding_space.png")


asyncio.run(visualize_embeddings())