90 lines
3.2 KiB
Python
90 lines
3.2 KiB
Python
"""Server-side punctuation restoration for voice transcripts.
|
||
|
||
Uses simple rule-based approach (no external model dependency).
|
||
Can be replaced with ML-based model (e.g., silero, yandex punctuator) later.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
|
||
# Words that typically start a new sentence
|
||
_SENTENCE_STARTERS: set[str] = {
|
||
"а", "но", "и", "да", "вот", "так", "это", "тот", "кто", "что",
|
||
"как", "когда", "где", "почему", "зачем", "сколько", "чей",
|
||
"во-первых", "во-вторых", "наконец", "однако", "причем", "притом",
|
||
"кстати", "например", "между", "тем", "впрочем", "значит",
|
||
"итак", "следовательно", "кроме", "того", "помимо",
|
||
"я", "ты", "он", "она", "оно", "мы", "вы", "они",
|
||
"этот", "эта", "это", "эти", "мой", "твой", "наш", "ваш",
|
||
"сегодня", "завтра", "вчера", "сейчас", "потом", "после",
|
||
"сначала", "затем", "далее", "теперь", "тут", "там", "здесь",
|
||
}
|
||
|
||
# Interrogative words
|
||
_QUESTION_WORDS: set[str] = {
|
||
"кто", "что", "какой", "какая", "какое", "какие", "чей", "чья", "чьё", "чьи",
|
||
"сколько", "когда", "где", "куда", "откуда", "почему", "зачем", "как",
|
||
"неужели", "разве", "ли",
|
||
}
|
||
|
||
# Exclamation words
|
||
_EXCLAMATION_WORDS: set[str] = {
|
||
"ах", "ох", "ух", "ой", "эй", "ура", "браво", "караул",
|
||
"здорово", "отлично", "прекрасно", "замечательно", "классно",
|
||
"ужасно", "кошмар", "боже",
|
||
}
|
||
|
||
|
||
def restore_punctuation(text: str) -> str:
|
||
"""Add basic punctuation and capitalization to raw transcribed text.
|
||
|
||
Works on Russian text from Whisper/STT output which typically
|
||
comes without punctuation.
|
||
"""
|
||
if not text or not text.strip():
|
||
return text
|
||
|
||
text = text.strip()
|
||
|
||
# Split into clauses by common separators
|
||
clauses = re.split(r"(?<=[^.!?])[;,]\s+", text)
|
||
formatted: list[str] = []
|
||
|
||
for clause in clauses:
|
||
clause = clause.strip()
|
||
if not clause:
|
||
continue
|
||
|
||
# Ensure first word is capitalized
|
||
words = clause.split()
|
||
if words:
|
||
words[0] = words[0].capitalize()
|
||
|
||
# Detect sentence type and add punctuation
|
||
first_word = words[0].lower().rstrip(",.!?")
|
||
|
||
# Remove any trailing punctuation
|
||
last_word = words[-1].rstrip(",.!?") if words else ""
|
||
|
||
if first_word in _EXCLAMATION_WORDS:
|
||
punctuation = "!"
|
||
elif first_word in _QUESTION_WORDS:
|
||
punctuation = "?"
|
||
elif last_word in _QUESTION_WORDS:
|
||
punctuation = "?"
|
||
else:
|
||
punctuation = "."
|
||
|
||
formatted.append(" ".join(words) + punctuation)
|
||
|
||
result = " ".join(formatted)
|
||
|
||
# Clean up common issues
|
||
result = re.sub(r"\s+", " ", result)
|
||
result = re.sub(r"\s*([.!?,;:])\s*", r"\1 ", result)
|
||
result = re.sub(r"\s+\.", ".", result)
|
||
result = result.strip()
|
||
|
||
return result
|