Модели embeddings: text-embedding-3, nomic-embed, BGE, E5

import numpy as np
from openai import AsyncOpenAI

client = AsyncOpenAI()


async def embed_openai(
    texts: list[str],
    model: str = "text-embedding-3-small",
    dimensions: int | None = None,
) -> np.ndarray:
    """
    Эмбеддинг через OpenAI API.

    Параметр dimensions (MRL) — сокращает вектор без переобучения:
      dimensions=1536  → полное качество (default для small)
      dimensions=512   → MTEB −0.7, память −3x
      dimensions=256   → MTEB −1.2, память −6x
      dimensions=64    → MTEB −4.4, память −24x  ← только для прототипов

    Модели:
      "text-embedding-3-small"  — $0.02/1M, dim=1536, MTEB=62.3
      "text-embedding-3-large"  — $0.13/1M, dim=3072, MTEB=64.6
    """
    kwargs: dict = {"input": texts, "model": model}
    if dimensions is not None:
        kwargs["dimensions"] = dimensions

    response = await client.embeddings.create(**kwargs)
    # API возвращает в порядке index, сортируем на всякий случай
    items = sorted(response.data, key=lambda x: x.index)
    matrix = np.array([item.embedding for item in items], dtype=np.float32)
    # Нормализация (API уже нормализует, но явная гарантия)
    norms = np.linalg.norm(matrix, axis=1, keepdims=True)
    return matrix / np.where(norms == 0, 1, norms)


# Сравнение storage при MRL
import asyncio

async def demo_mrl():
    texts = ["Как настроить FastAPI с PostgreSQL?"] * 10  # 10 одинаковых

    v_full = await embed_openai(texts, dimensions=1536)
    v_half = await embed_openai(texts, dimensions=512)
    v_small = await embed_openai(texts, dimensions=64)

    print(f"Full  (1536d): {v_full.nbytes / 1024:.1f} КБ — shape {v_full.shape}")
    print(f"Half  ( 512d): {v_half.nbytes / 1024:.1f} КБ — shape {v_half.shape}")
    print(f"Small (  64d): {v_small.nbytes / 1024:.1f} КБ — shape {v_small.shape}")

    # Проверяем: насколько full и half схожи на одном запросе?
    sim = float(np.dot(v_full[0], v_full[0]))  # всегда 1.0 (нормализован)
    # cross-dim similarity нельзя считать — разные пространства!
    # MRL гарантирует, что в рамках одной dim всё работает корректно

asyncio.run(demo_mrl())

"""
pip install sentence-transformers torch
Размер модели: ~550 МБ
"""

import numpy as np
from sentence_transformers import SentenceTransformer

# trust_remote_code=True обязателен — nomic использует кастомный код модели
_model_nomic: SentenceTransformer | None = None

def get_nomic_model() -> SentenceTransformer:
    global _model_nomic
    if _model_nomic is None:
        _model_nomic = SentenceTransformer(
            "nomic-ai/nomic-embed-text-v1.5",
            trust_remote_code=True,
        )
    return _model_nomic


def embed_nomic(
    texts: list[str],
    task: str = "search_document",
    dimensions: int | None = None,
) -> np.ndarray:
    """
    task варианты:
      "search_document"  — для индексируемых чанков
      "search_query"     — для пользовательских запросов
      "clustering"       — для кластеризации документов
      "classification"   — для задач классификации

    dimensions: MRL-усечение (768 → 512 → 256 → 128 → 64)
    """
    model = get_nomic_model()
    prefixed = [f"{task}: {t}" for t in texts]

    vectors = model.encode(
        prefixed,
        normalize_embeddings=True,
        batch_size=32,
        show_progress_bar=len(texts) > 100,
    )

    if dimensions is not None and dimensions < vectors.shape[1]:
        vectors = vectors[:, :dimensions]
        # После усечения нужна повторная нормализация
        norms = np.linalg.norm(vectors, axis=1, keepdims=True)
        vectors = vectors / np.where(norms == 0, 1, norms)

    return vectors.astype(np.float32)


def embed_nomic_query(query: str) -> np.ndarray:
    return embed_nomic([query], task="search_query")[0]

def embed_nomic_docs(docs: list[str], **kwargs) -> np.ndarray:
    return embed_nomic(docs, task="search_document", **kwargs)


# Пример
docs = [
    "FastAPI — высокопроизводительный веб-фреймворк",
    "Установка: pip install fastapi uvicorn",
    "История Второй мировой войны",
]

doc_vecs = embed_nomic_docs(docs)
query_vec = embed_nomic_query("как установить fastapi")

scores = doc_vecs @ query_vec  # cosine sim (нормализованные)
print("Результаты nomic-embed:")
for score, doc in sorted(zip(scores, docs), reverse=True):
    print(f"  {score:.3f}  {doc}")

"""
pip install FlagEmbedding torch
Размер модели: ~1.2 ГБ
FlagEmbedding даёт доступ к dense, sparse и colbert-векторам.
"""

import numpy as np
from FlagEmbedding import BGEM3FlagModel
from dataclasses import dataclass


@dataclass
class BGEOutput:
    dense: np.ndarray                       # (N, 1024) — dense vectors
    sparse: list[dict[str, float]]          # [{token: weight}, ...]


_bge_model: BGEM3FlagModel | None = None

def get_bge_model() -> BGEM3FlagModel:
    global _bge_model
    if _bge_model is None:
        _bge_model = BGEM3FlagModel(
            "BAAI/bge-m3",
            use_fp16=True,      # FP16: ×2 быстрее на GPU, качество не теряется
        )
    return _bge_model


def embed_bge(
    texts: list[str],
    return_sparse: bool = False,
    batch_size: int = 12,
    max_length: int = 8192,
) -> BGEOutput:
    """
    Dense-only (return_sparse=False): быстрее, достаточно для большинства задач.
    Dense+Sparse (return_sparse=True): hybrid retrieval без отдельного BM25.

    Sparse-вектор — словарь {token_id: вес}, аналог SPLADE.
    Qdrant и Weaviate поддерживают sparse-векторы нативно.
    """
    model = get_bge_model()
    output = model.encode(
        texts,
        batch_size=batch_size,
        max_length=max_length,
        return_dense=True,
        return_sparse=return_sparse,
        return_colbert_vecs=False,  # ColBERT требует много памяти, пропускаем
    )
    dense = output["dense_vecs"].astype(np.float32)

    if return_sparse:
        sparse = output["lexical_weights"]  # list[dict] — token → weight
    else:
        sparse = [{} for _ in texts]

    return BGEOutput(dense=dense, sparse=sparse)


def hybrid_score(
    query_dense: np.ndarray,
    query_sparse: dict[str, float],
    doc_dense: np.ndarray,
    doc_sparse: dict[str, float],
    alpha: float = 0.5,
) -> float:
    """
    Гибридный скор: alpha * dense_score + (1-alpha) * sparse_score.
    alpha=1.0 → только dense, alpha=0.0 → только sparse.
    """
    dense_score = float(np.dot(query_dense, doc_dense))

    # Sparse score: dot product по общим токенам
    common_tokens = set(query_sparse) & set(doc_sparse)
    sparse_score = sum(
        query_sparse[t] * doc_sparse[t] for t in common_tokens
    )
    # Нормализуем sparse в [0,1] — грубая оценка
    sparse_norm = max(
        sum(v**2 for v in query_sparse.values()) ** 0.5 *
        sum(v**2 for v in doc_sparse.values()) ** 0.5,
        1e-9,
    )
    sparse_score /= sparse_norm

    return alpha * dense_score + (1 - alpha) * sparse_score


# Пример hybrid retrieval
docs = [
    "BGE-M3 supports dense retrieval",
    "Установка FastAPI: pip install fastapi",
    "История Python: создан Гвидо ван Россумом",
]

query = "как установить fastapi"

# Эмбеддируем с sparse
doc_out   = embed_bge(docs,    return_sparse=True)
query_out = embed_bge([query], return_sparse=True)

print("BGE-M3 Hybrid scores:")
for i, doc in enumerate(docs):
    score = hybrid_score(
        query_out.dense[0], query_out.sparse[0],
        doc_out.dense[i],   doc_out.sparse[i],
        alpha=0.6,
    )
    print(f"  {score:.3f}  {doc}")

"""
pip install sentence-transformers
Варианты: multilingual-e5-small (120 МБ), -base (470 МБ), -large (560 МБ)
"""

import numpy as np
from sentence_transformers import SentenceTransformer
from functools import lru_cache


@lru_cache(maxsize=3)
def get_e5_model(size: str = "large") -> SentenceTransformer:
    """Кешируем модель: один экземпляр на процесс."""
    return SentenceTransformer(f"intfloat/multilingual-e5-{size}")


def embed_e5_documents(
    texts: list[str],
    size: str = "large",
    batch_size: int = 32,
) -> np.ndarray:
    """
    Префикс "passage: " для документов — ОБЯЗАТЕЛЕН.
    Без него MTEB Retrieval падает на ~10 пунктов.
    """
    model = get_e5_model(size)
    prefixed = [f"passage: {t}" for t in texts]
    vecs = model.encode(prefixed, batch_size=batch_size, normalize_embeddings=True)
    return vecs.astype(np.float32)


def embed_e5_query(query: str, size: str = "large") -> np.ndarray:
    """Префикс "query: " для запросов."""
    model = get_e5_model(size)
    vec = model.encode([f"query: {query}"], normalize_embeddings=True)
    return vec[0].astype(np.float32)


# Выбор варианта по задаче
# small:  CPU-friendly, 120 МБ, ~2x быстрее large, MTEB 58.7
# base:   баланс скорость/качество, 470 МБ
# large:  максимальное качество, 560 МБ, MTEB 60.8

# Пример: сравнение small vs large на русских текстах
texts_ru = [
    "Установка FastAPI в виртуальное окружение Python",
    "История создания Эйфелевой башни",
    "Нейронные сети и машинное обучение",
]
query_ru = "установка веб-фреймворка"

for size in ["small", "large"]:
    docs_v = embed_e5_documents(texts_ru, size=size)
    query_v = embed_e5_query(query_ru, size=size)
    scores = docs_v @ query_v
    top = texts_ru[scores.argmax()]
    print(f"E5-{size}: top-1 = «{top[:50]}» ({scores.max():.3f})")

"""
Бенчмарк: сравниваем модели по скорости и MRR@10 на своих данных.
Запускайте на репрезентативной выборке своего корпуса.
"""

import asyncio
import time
import numpy as np
from dataclasses import dataclass, field


@dataclass
class BenchResult:
    model: str
    mrr_at_10: float
    embed_ms_per_doc: float
    dim: int
    total_docs: int


def compute_mrr(query_vecs: np.ndarray, doc_vecs: np.ndarray, qrels: list[list[int]]) -> float:
    """
    qrels[i] = список индексов релевантных документов для запроса i.
    """
    mrrs = []
    for q_idx, (qvec, relevant) in enumerate(zip(query_vecs, qrels)):
        sims = doc_vecs @ qvec
        ranked = np.argsort(-sims)[:10]
        for rank, doc_idx in enumerate(ranked, 1):
            if doc_idx in relevant:
                mrrs.append(1.0 / rank)
                break
        else:
            mrrs.append(0.0)
    return float(np.mean(mrrs))


async def benchmark_all(
    docs: list[str],
    queries: list[str],
    qrels: list[list[int]],
) -> list[BenchResult]:
    results = []

    # ── OpenAI text-embedding-3-small ──
    t0 = time.time()
    doc_vecs   = await embed_openai(docs,    model="text-embedding-3-small")
    query_vecs = await embed_openai(queries, model="text-embedding-3-small")
    ms = (time.time() - t0) / len(docs) * 1000
    mrr = compute_mrr(query_vecs, doc_vecs, qrels)
    results.append(BenchResult("text-embedding-3-small", mrr, ms, 1536, len(docs)))

    # ── nomic-embed-text-v1.5 ──
    t0 = time.time()
    doc_vecs   = embed_nomic_docs(docs)
    query_vecs = np.array([embed_nomic_query(q) for q in queries])
    ms = (time.time() - t0) / len(docs) * 1000
    mrr = compute_mrr(query_vecs, doc_vecs, qrels)
    results.append(BenchResult("nomic-embed-v1.5", mrr, ms, 768, len(docs)))

    # ── BGE-M3 dense ──
    t0 = time.time()
    doc_out   = embed_bge(docs)
    query_out = embed_bge(queries)
    ms = (time.time() - t0) / len(docs) * 1000
    mrr = compute_mrr(query_out.dense, doc_out.dense, qrels)
    results.append(BenchResult("bge-m3-dense", mrr, ms, 1024, len(docs)))

    # ── multilingual-e5-large ──
    t0 = time.time()
    doc_vecs   = embed_e5_documents(docs)
    query_vecs = np.array([embed_e5_query(q) for q in queries])
    ms = (time.time() - t0) / len(docs) * 1000
    mrr = compute_mrr(query_vecs, doc_vecs, qrels)
    results.append(BenchResult("multilingual-e5-large", mrr, ms, 1024, len(docs)))

    return results


# Запуск и вывод
async def main():
    # Ваши данные: список документов, запросы и список релевантных индексов
    docs    = ["..."] * 200   # ← вставьте свои документы
    queries = ["..."] * 20    # ← ваши запросы
    qrels   = [[0, 1], [5]]   # ← qrels[i] = индексы релевантных docs для query[i]

    results = await benchmark_all(docs, queries, qrels)

    print(f"{'Модель':<25} {'MRR@10':>7} {'ms/doc':>8} {'Dim':>6}")
    print("─" * 52)
    for r in sorted(results, key=lambda x: -x.mrr_at_10):
        print(f"{r.model:<25} {r.mrr_at_10:>7.3f} {r.embed_ms_per_doc:>8.1f} {r.dim:>6}")

asyncio.run(main())

"""
Blue-Green переиндексация при смене embedding-модели.

Стратегия:
  1. Создаём новую коллекцию (new_collection)
  2. Переиндексируем все документы новой моделью
  3. Переключаем трафик на новую коллекцию
  4. Удаляем старую коллекцию

Это обеспечивает zero-downtime: пока идёт переиндексация,
старая коллекция продолжает отвечать на запросы.
"""

import asyncio
from dataclasses import dataclass
import chromadb


@dataclass
class IndexStats:
    total: int = 0
    processed: int = 0
    failed: int = 0


async def reindex_to_new_model(
    old_collection_name: str,
    new_collection_name: str,
    new_embed_fn,          # async callable(texts) -> np.ndarray
    chroma_client: chromadb.Client,
    batch_size: int = 100,
    progress_fn=None,      # callback(processed, total)
) -> IndexStats:
    """
    Переносит все документы из старой коллекции в новую,
    эмбеддируя через new_embed_fn.
    """
    old_col = chroma_client.get_collection(old_collection_name)
    new_col = chroma_client.get_or_create_collection(new_collection_name)
    stats = IndexStats()

    # Получаем все документы из старой коллекции (без эмбеддингов)
    all_data = old_col.get(include=["documents", "metadatas"])
    stats.total = len(all_data["ids"])

    for i in range(0, stats.total, batch_size):
        batch_ids   = all_data["ids"][i : i + batch_size]
        batch_texts = all_data["documents"][i : i + batch_size]
        batch_metas = all_data["metadatas"][i : i + batch_size]

        try:
            new_vecs = await new_embed_fn(batch_texts)
            new_col.add(
                ids=batch_ids,
                embeddings=new_vecs.tolist(),
                documents=batch_texts,
                metadatas=batch_metas,
            )
            stats.processed += len(batch_ids)
        except Exception as e:
            print(f"Ошибка в батче {i}: {e}")
            stats.failed += len(batch_ids)

        if progress_fn:
            progress_fn(stats.processed, stats.total)

    return stats


# Использование
async def migrate():
    client = chromadb.Client()

    # Новая модель — BGE-M3 вместо text-embedding-3-small
    async def new_embed(texts: list[str]):
        return embed_bge(texts).dense

    stats = await reindex_to_new_model(
        old_collection_name="docs_v1",
        new_collection_name="docs_v2",
        new_embed_fn=new_embed,
        chroma_client=client,
        progress_fn=lambda n, t: print(f"\r{n}/{t}", end=""),
    )
    print(f"\nПереиндексировано: {stats.processed}, ошибок: {stats.failed}")
    # После успеха: переключаем приложение на "docs_v2", удаляем "docs_v1"

asyncio.run(migrate())