Semantic Chunking: разбиение по смысловым границам

import re
from dataclasses import dataclass, field


@dataclass
class Document:
    page_content: str
    metadata: dict = field(default_factory=dict)


# Аббревиатуры, после которых точка не означает конец предложения
_ABBREV_RU = re.compile(
    r"\b(г|ул|пр|д|кв|т|тел|факс|руб|коп|млн|млрд|тыс|чел|стр|с|п|пп|ст|ч|мин|сек"
    r"|т\.е|т\.д|т\.п|т\.к|напр|др|проф|д-р|канд|акад|рис|табл|рим)\.",
    re.IGNORECASE,
)

# Граница предложения: [.!?] + пробел + заглавная буква (RU/EN) или кавычка
_SENT_BOUNDARY = re.compile(r"(?<=[.!?])\s+(?=[А-ЯЁA-Z\"«])")


def split_sentences(text: str) -> list[str]:
    """
    Быстрая regex-разбивка на предложения для RU/EN текста.
    Защита от ложных разрывов по аббревиатурам.
    """
    # Заменяем точки в аббревиатурах на placeholder
    protected = _ABBREV_RU.sub(lambda m: m.group(0).replace(".", ""), text)
    # Делим по границам предложений
    raw_sents = _SENT_BOUNDARY.split(protected)
    # Возвращаем placeholder обратно
    return [s.replace("", ".").strip() for s in raw_sents if s.strip()]


def split_sentences_spacy(text: str, model: str = "ru_core_news_sm") -> list[str]:
    """
    Точная разбивка через spaCy (требует: pip install spacy && python -m spacy download ru_core_news_sm).
    Медленнее в 10–30×, но правильно обрабатывает все крайние случаи.
    """
    import spacy
    nlp = spacy.load(model)
    doc = nlp(text)
    return [sent.text.strip() for sent in doc.sents if sent.text.strip()]


# Пример
text = """Нейронные сети обучаются с помощью backprop. Градиент функции потерь
вычисляется относительно каждого веса. Оптимизатор Adam адаптирует скорость обучения.

Трансформеры используют механизм внимания. Self-attention учитывает зависимости."""

sentences = split_sentences(text)
for i, s in enumerate(sentences):
    print(f"[{i+1}] {s}")

import asyncio
import numpy as np
from openai import AsyncOpenAI

client = AsyncOpenAI()


async def embed_batch(texts: list[str], model: str = "text-embedding-3-small") -> np.ndarray:
    """
    Эмбеддирует список текстов батчем.
    Возвращает матрицу (N, dim).
    """
    response = await client.embeddings.create(input=texts, model=model)
    vectors = [item.embedding for item in response.data]
    return np.array(vectors, dtype=np.float32)


def build_buffered_sentences(
    sentences: list[str],
    buffer_size: int = 1,
) -> list[str]:
    """
    Строит «буферизованные» предложения: каждое предложение объединяется
    с buffer_size соседями с обеих сторон.

    buffer_size=0 → эмбеддируем отдельные предложения
    buffer_size=1 → предложение + 1 сосед слева + 1 сосед справа
    buffer_size=2 → ещё шире окно
    """
    buffered = []
    n = len(sentences)
    for i in range(n):
        lo = max(0, i - buffer_size)
        hi = min(n, i + buffer_size + 1)
        buffered.append(" ".join(sentences[lo:hi]))
    return buffered


async def embed_sentences(
    sentences: list[str],
    buffer_size: int = 1,
    model: str = "text-embedding-3-small",
    batch_size: int = 100,       # API limit для одного запроса
) -> np.ndarray:
    """
    Эмбеддирует предложения с буфером, обрабатывая большие тексты батчами.
    Возвращает матрицу (N, dim) — по одному вектору на предложение.
    """
    buffered = build_buffered_sentences(sentences, buffer_size)

    # Батчами по batch_size
    all_vectors = []
    for start in range(0, len(buffered), batch_size):
        batch = buffered[start:start + batch_size]
        vectors = await embed_batch(batch, model=model)
        all_vectors.append(vectors)

    return np.vstack(all_vectors)

import numpy as np


def cosine_similarity_matrix(vectors: np.ndarray) -> np.ndarray:
    """
    Вычисляет cosine similarity между каждой парой соседних векторов.
    Возвращает массив длины N-1.
    """
    # Нормализуем каждый вектор до единичной длины
    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
    norms = np.maximum(norms, 1e-9)          # защита от деления на ноль
    normalized = vectors / norms

    # Скалярное произведение соседних пар = cosine similarity
    similarities = np.einsum("id,id->i", normalized[:-1], normalized[1:])
    return similarities.astype(np.float64)


def find_breakpoints(
    similarities: np.ndarray,
    method: str = "percentile",
    percentile_threshold: float = 5.0,
    absolute_threshold: float = 0.5,
) -> list[int]:
    """
    Находит индексы разрывов (границ чанков).

    method="percentile":  граница где sim < percentile(similarities, percentile_threshold)
    method="absolute":    граница где sim < absolute_threshold
    method="gradient":    граница где резкое падение sim (нижний percentile_threshold% градиентов)

    Возвращает список индексов i таких, что граница стоит ПОСЛЕ предложения i.
    """
    if method == "percentile":
        threshold = float(np.percentile(similarities, percentile_threshold))
        return [i for i, s in enumerate(similarities) if s < threshold]

    if method == "absolute":
        return [i for i, s in enumerate(similarities) if s < absolute_threshold]

    if method == "gradient":
        if len(similarities) < 2:
            return []
        # Падение: sim[i] - sim[i+1] > 0 → тема переключилась
        gradients = np.diff(similarities)         # длина N-2
        neg_grads = -gradients                    # инвертируем: большой = резкое падение
        threshold  = float(np.percentile(neg_grads, 100 - percentile_threshold))
        return [i + 1 for i, g in enumerate(neg_grads) if g > threshold]

    raise ValueError(f"Неизвестный метод: {method!r}")


def sentences_to_chunks(
    sentences: list[str],
    breakpoints: list[int],
    min_chunk_size: int = 50,
) -> list[str]:
    """
    Собирает предложения в чанки по найденным границам.
    Слишком маленькие чанки (< min_chunk_size символов) объединяются с соседом.
    """
    if not sentences:
        return []

    breakpoint_set = set(breakpoints)
    chunks: list[str] = []
    current: list[str] = []

    for i, sent in enumerate(sentences):
        current.append(sent)
        # Граница стоит ПОСЛЕ предложения i
        if i in breakpoint_set and i < len(sentences) - 1:
            chunk_text = " ".join(current).strip()
            if len(chunk_text) >= min_chunk_size:
                chunks.append(chunk_text)
                current = []
            # Иначе продолжаем накапливать (чанк слишком маленький)

    # Последний чанк
    if current:
        last = " ".join(current).strip()
        if chunks and len(last) < min_chunk_size:
            chunks[-1] = chunks[-1] + " " + last  # прикрепляем к предыдущему
        elif last:
            chunks.append(last)

    return chunks

import asyncio
from dataclasses import dataclass, field

import numpy as np


@dataclass
class SemanticSplitterConfig:
    """Конфигурация семантического сплиттера."""
    # Параметры эмбеддинга
    embedding_model:  str   = "text-embedding-3-small"
    buffer_size:      int   = 1        # окно усреднения вокруг каждого предложения
    batch_size:       int   = 100      # предложений за один API-запрос

    # Параметры порога
    breakpoint_method:        str   = "percentile"  # percentile | absolute | gradient
    percentile_threshold:     float = 5.0           # нижние 5% переходов → граница
    absolute_threshold:       float = 0.5           # только для method="absolute"

    # Постобработка
    min_chunk_size:   int   = 50       # символов, маленькие чанки склеиваются
    max_chunk_size:   int   = 3000     # символов, большие чанки дополнительно режутся


class SemanticSplitter:
    """
    Async семантический сплиттер.
    Разбивает текст на чанки по семантическим границам через embedding-модель.
    """

    def __init__(self, config: SemanticSplitterConfig | None = None):
        self.cfg = config or SemanticSplitterConfig()

    async def split_text(self, text: str) -> list[str]:
        """Разбивает один текст на семантические чанки."""
        # 1. Предложения
        sentences = split_sentences(text)
        if len(sentences) <= 1:
            return [text.strip()] if text.strip() else []

        # 2. Embeddings с буфером
        vectors = await embed_sentences(
            sentences,
            buffer_size=self.cfg.buffer_size,
            model=self.cfg.embedding_model,
            batch_size=self.cfg.batch_size,
        )

        # 3. Cosine similarity между соседними
        similarities = cosine_similarity_matrix(vectors)

        # 4. Поиск границ
        breakpoints = find_breakpoints(
            similarities,
            method=self.cfg.breakpoint_method,
            percentile_threshold=self.cfg.percentile_threshold,
            absolute_threshold=self.cfg.absolute_threshold,
        )

        # 5. Сборка чанков
        chunks = sentences_to_chunks(
            sentences, breakpoints,
            min_chunk_size=self.cfg.min_chunk_size,
        )

        # 6. Дополнительное разбиение слишком больших чанков
        final: list[str] = []
        for chunk in chunks:
            if len(chunk) > self.cfg.max_chunk_size:
                # Fallback: recursive splitting для больших чанков
                from recursive_splitting import RecursiveCharacterSplitter
                sub = RecursiveCharacterSplitter(
                    chunk_size=self.cfg.max_chunk_size,
                    chunk_overlap=self.cfg.max_chunk_size // 10,
                ).split_text(chunk)
                final.extend(sub)
            else:
                final.append(chunk)

        return final

    async def split_documents(self, documents: list[Document]) -> list[Document]:
        """Разбивает список документов, параллельно обрабатывая их."""
        # Параллельная обработка всех документов
        results = await asyncio.gather(*[
            self._split_document(doc) for doc in documents
        ])
        return [chunk for doc_chunks in results for chunk in doc_chunks]

    async def _split_document(self, doc: Document) -> list[Document]:
        chunks = await self.split_text(doc.page_content)
        return [
            Document(
                page_content=chunk,
                metadata={
                    **doc.metadata,
                    "chunk_index": i,
                    "chunk_total": len(chunks),
                    "chunk_size":  len(chunk),
                    "chunking_method": "semantic",
                },
            )
            for i, chunk in enumerate(chunks)
        ]


# Использование
async def demo():
    text = """
    Нейронные сети обучаются с помощью метода обратного распространения ошибки.
    Градиент функции потерь вычисляется относительно каждого веса в сети.
    Оптимизатор Adam адаптирует скорость обучения для каждого параметра.

    Трансформеры используют механизм внимания для обработки последовательностей.
    Self-attention позволяет модели учитывать зависимости на любой дистанции.
    Позиционное кодирование добавляет информацию о порядке токенов.
    """.strip()

    splitter = SemanticSplitter(SemanticSplitterConfig(
        percentile_threshold=10.0,
        buffer_size=1,
    ))
    chunks = await splitter.split_text(text)

    for i, chunk in enumerate(chunks):
        print(f"── Chunk {i+1} ({len(chunk)} симв) ──")
        print(chunk[:150])
        print()


if __name__ == "__main__":
    asyncio.run(demo())

from sentence_transformers import SentenceTransformer
import numpy as np


class LocalEmbedder:
    """
    Локальный эмбеддер через sentence-transformers.
    Без API-ключей, без стоимости за токены.
    Для production: использовать GPU-инстанс или ONNX-runtime.
    """

    def __init__(
        self,
        model_name: str = "intfloat/multilingual-e5-large",
        device: str = "cpu",            # "cuda" если есть GPU
        batch_size: int = 32,
    ):
        self.model = SentenceTransformer(model_name, device=device)
        self.batch_size = batch_size

    def embed(self, texts: list[str]) -> np.ndarray:
        """
        Эмбеддирует список текстов.
        multilingual-e5 требует префикс "query: " для поиска и "passage: " для документов.
        При чанкинге используем "passage: " — мы индексируем, а не ищем.
        """
        prefixed = [f"passage: {t}" for t in texts]
        return self.model.encode(
            prefixed,
            batch_size=self.batch_size,
            show_progress_bar=False,
            normalize_embeddings=True,     # уже нормализованы → dot product = cosine sim
        )

    async def embed_async(self, texts: list[str]) -> np.ndarray:
        """Обёртка для использования в async-пайплайне через executor."""
        import asyncio
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self.embed, texts)


# Интеграция с SemanticSplitter
class SemanticSplitterLocal(SemanticSplitter):
    """Версия с локальной моделью вместо OpenAI API."""

    def __init__(self, config: SemanticSplitterConfig | None = None):
        super().__init__(config)
        self._embedder = LocalEmbedder()

    async def _get_vectors(self, sentences: list[str]) -> np.ndarray:
        buffered = build_buffered_sentences(sentences, self.cfg.buffer_size)
        return await self._embedder.embed_async(buffered)

import hashlib
import json
import struct
from pathlib import Path

import numpy as np


class EmbeddingCache:
    """
    Дисковый кеш для embedding-векторов.
    Ключ: SHA-256(text + model_name).
    Формат: бинарный (numpy save) для скорости.
    """

    def __init__(self, cache_dir: str = ".embedding_cache"):
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(exist_ok=True)

    def _key(self, texts: list[str], model: str) -> str:
        payload = json.dumps({"texts": texts, "model": model}, ensure_ascii=False)
        return hashlib.sha256(payload.encode()).hexdigest()[:16]

    def get(self, texts: list[str], model: str) -> np.ndarray | None:
        path = self.cache_dir / f"{self._key(texts, model)}.npy"
        if path.exists():
            return np.load(str(path))
        return None

    def set(self, texts: list[str], model: str, vectors: np.ndarray) -> None:
        path = self.cache_dir / f"{self._key(texts, model)}.npy"
        np.save(str(path), vectors)


# Патчим embed_batch для использования кеша
_cache = EmbeddingCache()

async def embed_batch_cached(
    texts: list[str],
    model: str = "text-embedding-3-small",
) -> np.ndarray:
    cached = _cache.get(texts, model)
    if cached is not None:
        return cached

    vectors = await embed_batch(texts, model=model)
    _cache.set(texts, model, vectors)
    return vectors