vidconf/backend/core/summarization/chunking.py

"""Чанкинг транскрипта для map-стадии суммаризации (ТЗ §1.3).

Транскрипт — строка с одной фразой на строку в формате `[Имя MM:SS] текст`
(или `[Имя ЧЧ:MM:SS] текст` от часа, см. `workers/summarizer/transcript.py`).
Каждая строка уже атомарна и соответствует одной реконструированной фразе
(`build_phrases`) — граница фразы там уже равна смене спикера,
поэтому `chunk_transcript` закрывает чанк исключительно на границах строк и
никогда не режет фразу пополам (кроме аварийного случая монолога, см. ниже).

Правила закрытия чанка (проверяются перед добавлением очередной фразы):
- добавление фразы сделало бы охваченный чанком промежуток времени
  (от начала чанка до начала этой фразы) >= `target_chunk_minutes`;
- ЛИБО добавление фразы превысило бы `max_chunk_tokens` для чанка.
В любом из этих случаев уже накопленный чанк закрывается, а фраза уходит в
новый чанк.

Аварийный случай — монолог: единственная фраза сама по себе превышает
`max_chunk_tokens` (например, 40-минутный монолог одного участника). Такую
фразу нельзя оставить целой и нельзя просто выбросить — она делится по
границам предложений на несколько частей меньше лимита, каждая из которых
получает повторённую исходную метку `[Имя MM:SS]`, и добавляется в список
чанков как самостоятельный чанк.
"""

import re
from collections.abc import Callable

_LABEL_RE = re.compile(r"^\[(?P<speaker>.+) (?P<time>\d{1,3}:\d{2}(?::\d{2})?)\] (?P<text>.*)$")
"""Метка фразы: `[Имя MM:SS]` или `[Имя ЧЧ:MM:SS]`; имя может содержать пробелы."""

_SENTENCE_SPLIT_RE = re.compile(r"(?<=[.!?])\s+")
"""Граница предложения — пробел после `.`, `!` или `?`."""


def _parse_offset_seconds(time_str: str) -> float:
    """Разобрать метку времени `MM:SS` или `ЧЧ:MM:SS` в секунды от начала сеанса."""
    parts = [int(part) for part in time_str.split(":")]
    if len(parts) == 2:
        minutes, seconds = parts
        return float(minutes * 60 + seconds)
    hours, minutes, seconds = parts
    return float(hours * 3600 + minutes * 60 + seconds)


def _line_offset(line: str) -> float:
    """Извлечь смещение начала фразы (в секундах) из метки строки.

    Строка без распознаваемой метки считается стоящей в начале сеанса
    (0.0) — чанкер не должен падать на нестандартном входе, только терять
    точность деления по времени.
    """
    match = _LABEL_RE.match(line)
    if match is None:
        return 0.0
    return _parse_offset_seconds(match.group("time"))


def _split_sentences(text: str) -> list[str]:
    """Разбить текст фразы на предложения; пустой текст даёт список из одного элемента."""
    sentences = [s.strip() for s in _SENTENCE_SPLIT_RE.split(text) if s.strip()]
    return sentences or [text]


def _split_monologue(
    line: str,
    count_tokens: Callable[[str], int],
    max_chunk_tokens: int,
) -> list[str]:
    """Аварийно разделить фразу-монолог, превышающую лимит, по предложениям.

    Метка спикера повторяется в каждой части. Если строка не распознана как
    фраза с меткой (не должно случаться при штатном входе из
    `build_transcript`), делить не на чем — возвращаем строку одним чанком,
    даже с превышением лимита.
    """
    match = _LABEL_RE.match(line)
    if match is None:
        return [line]

    label = f"[{match.group('speaker')} {match.group('time')}]"
    sentences = _split_sentences(match.group("text"))

    pieces: list[str] = []
    current: list[str] = []
    for sentence in sentences:
        candidate = f"{label} {' '.join([*current, sentence])}"
        if current and count_tokens(candidate) > max_chunk_tokens:
            pieces.append(f"{label} {' '.join(current)}")
            current = [sentence]
        else:
            current.append(sentence)
    if current:
        pieces.append(f"{label} {' '.join(current)}")
    return pieces


def chunk_transcript(
    transcript: str,
    count_tokens: Callable[[str], int],
    *,
    max_chunk_tokens: int = 8000,
    target_chunk_minutes: int = 20,
) -> list[str]:
    """Разбить транскрипт на чанки для map-стадии суммаризации.

    Аргументы:
        transcript: строки-фразы `[Имя MM:SS] текст`, по одной на строку.
        count_tokens: подсчёт токенов для строки/чанка (в проде —
            `QwenTokenCounter`, в тестах — фейковый callable).
        max_chunk_tokens: верхняя граница токенов на чанк (ТЗ: 3–8k).
        target_chunk_minutes: целевая длительность чанка в минутах записи.

    Возвращает список строк-чанков (без изменения содержимого фраз внутри),
    пустой список для пустого транскрипта.
    """
    lines = [line for line in transcript.splitlines() if line.strip()]
    if not lines:
        return []

    target_seconds = target_chunk_minutes * 60

    chunks: list[str] = []
    current_lines: list[str] = []
    current_tokens = 0
    chunk_start_offset = 0.0

    def flush() -> None:
        nonlocal current_lines, current_tokens
        if current_lines:
            chunks.append("\n".join(current_lines))
        current_lines = []
        current_tokens = 0

    for line in lines:
        line_tokens = count_tokens(line)

        if line_tokens > max_chunk_tokens:
            # монолог одной фразой больше лимита — аварийное деление по предложениям
            flush()
            chunks.extend(_split_monologue(line, count_tokens, max_chunk_tokens))
            continue

        if current_lines:
            offset = _line_offset(line)
            prospective_span = offset - chunk_start_offset
            prospective_tokens = current_tokens + line_tokens
            if prospective_tokens > max_chunk_tokens or prospective_span >= target_seconds:
                flush()

        if not current_lines:
            chunk_start_offset = _line_offset(line)
        current_lines.append(line)
        current_tokens += line_tokens

    flush()
    return chunks