diff --git a/core/utils/text_utils.py b/core/utils/text_utils.py new file mode 100644 index 0000000..208ca1c --- /dev/null +++ b/core/utils/text_utils.py @@ -0,0 +1,226 @@ +# -*- 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 \ No newline at end of file