# -*- coding: utf-8 -*- """ Утилиты для работы с текстом: - Разбиение на чанки (токенизация) с помощью tiktoken для индексации документов. - Подсчёт количества токенов в тексте (для обрезки истории диалога). - Fallback-разбиение по символам, если tiktoken не установлен. Используется в: - KBService при индексации документов (split_into_chunks). - RAG-пайплайне для обрезки истории диалога по токенам (count_tokens). Особенности: - Разбиение на предложения использует библиотеку razdel для корректной обработки аббревиатур (т.е., и.т.д.) и других сложных случаев русского языка. - Если razdel недоступен, используется fallback-регулярное выражение. - Если tiktoken недоступен, используется грубая оценка токенов (3 символа на токен). """ import re import logging from typing import List # ---- Инициализация токенизатора tiktoken (для точного подсчёта токенов) ---- try: import tiktoken # Используем кодировку cl100k_base, которая используется в моделях GPT-4, GPT-3.5 и GigaChat TOKENIZER = tiktoken.get_encoding("cl100k_base") except ImportError: TOKENIZER = None logging.warning("tiktoken не установлен. Для точного подсчёта токенов установите: pip install tiktoken") # ---------------------------------------------------------------------------- # ---- Импорт razdel для умного разбиения по предложениям (русский язык) ---- try: from razdel import sentenize HAS_RAZDEL = True except ImportError: HAS_RAZDEL = False logging.warning("razdel не установлен. Разбиение по предложениям будет неточным. Установите: pip install razdel") # ---------------------------------------------------------------------------- logger = logging.getLogger(__name__) def count_tokens(text: str) -> int: """ Подсчитывает количество токенов в тексте, используя tiktoken (если доступен). Если tiktoken недоступен, использует грубую оценку: 1 токен ≈ 3 символа. Эта функция используется в RAG-пайплайне для обрезки истории диалога до заданного лимита токенов (max_context_tokens). Аргументы: text (str): текст для подсчёта токенов. Возвращает: int: количество токенов (или оценка). """ if not text: return 0 if TOKENIZER is not None: # Точный подсчёт через tiktoken return len(TOKENIZER.encode(text)) else: # Грубая оценка: для русского языка в среднем 1 токен ≈ 2.5–3 символа. # Используем 3 как консервативную оценку (чтобы не превысить лимит). return len(text) // 3 + 1 def split_into_chunks( text: str, chunk_size_tokens: int = 200, overlap_tokens: int = 50, approx_chunk_chars: int = 500, approx_overlap_chars: int = 100, strategy: str = "recursive_split_by_sentences" ) -> List[str]: """ Разбивает текст на чанки (фрагменты) для последующей индексации в векторной БД. Сначала пытается использовать tiktoken для точного разбиения по токенам. Если tiktoken недоступен, использует приближённое разбиение по символам. Аргументы: text (str): исходный текст. chunk_size_tokens (int): целевой размер чанка в токенах (по умолчанию 200). overlap_tokens (int): перекрытие между соседними чанками в токенах (по умолчанию 50). approx_chunk_chars (int): размер чанка в символах для fallback-режима (по умолчанию 500). approx_overlap_chars (int): перекрытие в символах для fallback-режима (по умолчанию 100). strategy (str): стратегия разбиения (пока только "recursive_split_by_sentences"). Возвращает: List[str]: список чанков (текстовых фрагментов). """ if not text or not text.strip(): return [] # Нормализуем пробелы: заменяем множественные пробелы, табуляции и переносы на один пробел. # Это упрощает дальнейшую обработку и улучшает качество разбиения. text = re.sub(r'\s+', ' ', text).strip() # Если токенизатор недоступен, используем fallback-разбиение по символам. if TOKENIZER is None: logger.debug("tiktoken недоступен, используем разбиение по символам") return _split_by_chars(text, approx_chunk_chars, approx_overlap_chars) # Основной метод: разбиение по предложениям с учётом токенов. return _split_by_tokens(text, chunk_size_tokens, overlap_tokens) def _split_by_tokens(text: str, chunk_size: int, overlap: int) -> List[str]: """ Разбивает текст на чанки, используя tiktoken для точного подсчёта токенов. Алгоритм: 1. Разбить текст на предложения (с помощью razdel для русского языка). 2. Накопить предложения, пока сумма токенов не превысит chunk_size. 3. Когда превышает – сохранить текущий чанк и создать перекрытие из последних предложений (overlap токенов). 4. Продолжить с оставшимся хвостом. 5. В конце объединить очень короткие чанки (< chunk_size/3) с предыдущим, чтобы избежать множества мелких фрагментов. Аргументы: text (str): исходный текст. chunk_size (int): максимальный размер чанка в токенах. overlap (int): количество токенов перекрытия между соседними чанками. Возвращает: List[str]: список чанков. """ # Вспомогательная функция для подсчёта токенов в строке. def _count_tokens(s: str) -> int: return len(TOKENIZER.encode(s)) # ---- 1. Разбиение на предложения ---- if HAS_RAZDEL: # razdel возвращает итератор объектов Sentence, у которых есть атрибут .text sentences = [sent.text for sent in sentenize(text)] else: # Fallback: разбиение по точкам, вопросительным и восклицательным знакам. # Это может быть неточным для аббревиатур (например, "т.е.", "и.т.д."). sentences = re.split(r'(?<=[.!?])\s+', text) # ------------------------------------- chunks = [] current_chunk = [] # Список предложений текущего чанка current_len = 0 # Суммарное количество токенов в текущем чанке for sent in sentences: sent_tokens = _count_tokens(sent) # Токены в текущем предложении # Проверяем, помещается ли предложение в текущий чанк if current_len + sent_tokens <= chunk_size: # Помещается – добавляем current_chunk.append(sent) current_len += sent_tokens else: # Не помещается – сохраняем текущий чанк if current_chunk: chunks.append(' '.join(current_chunk)) # --- Формируем перекрытие для следующего чанка --- # Берём последние предложения текущего чанка, сумма токенов которых # не превышает overlap (чтобы сохранить контекст между чанками). overlap_sentences = [] overlap_len = 0 # Идём с конца текущего чанка (от последнего предложения к первому) for prev_sent in reversed(current_chunk): prev_len = _count_tokens(prev_sent) if overlap_len + prev_len <= overlap: overlap_sentences.insert(0, prev_sent) # Вставляем в начало, чтобы сохранить порядок overlap_len += prev_len else: break # ------------------------------------------------- # Новый чанк начинается с перекрытия + текущее предложение current_chunk = overlap_sentences + [sent] current_len = overlap_len + sent_tokens # Добавляем последний чанк, если он не пуст if current_chunk: chunks.append(' '.join(current_chunk)) # ---- Склеиваем очень короткие чанки с предыдущим ---- # Это предотвращает создание большого количества мелких фрагментов, # которые могут ухудшить качество поиска. merged = [] for chunk in chunks: # Если текущий чанк слишком короткий (< 1/3 от chunk_size) и есть предыдущий чанк, # присоединяем его к предыдущему. if merged and _count_tokens(chunk) < chunk_size // 3: merged[-1] = merged[-1] + " " + chunk else: merged.append(chunk) # ------------------------------------------------- return merged def _split_by_chars(text: str, chunk_size: int, overlap: int) -> List[str]: """ Fallback-метод: разбиение по символам (при отсутствии tiktoken). Использует простое скользящее окно с фиксированной длиной chunk_size и перекрытием overlap. Аргументы: text (str): исходный текст. chunk_size (int): размер чанка в символах. overlap (int): перекрытие в символах. Возвращает: List[str]: список чанков. """ if chunk_size <= 0: return [text] chunks = [] start = 0 while start < len(text): end = min(start + chunk_size, len(text)) chunks.append(text[start:end]) start = end - overlap if start < 0: start = 0 if start >= len(text): break return chunks