truecase.py

"""Truecasing for Kazakh STT output.

The `kk-turbo` acoustic model was fine-tuned on normalized text, so it emits
lowercase with no sentence/proper-noun capitals. This restores casing:

  1. Capitalize the first word of each sentence (deterministic).
  2. Capitalize known proper names from a gazetteer (`data/kk_names.txt`, plus a
     built-in seed of common Kazakh given names).

So `сәлем қалайсың менің атым айдос` -> `Сәлем қалайсың менің атым Айдос`.

For broader proper-noun coverage, extend the gazetteer (e.g. from the KazNERD
PERSON/LOC entities) — that's the data-driven upgrade path.

    from kazakh_stt.truecase import truecase
    truecase("сәлем менің атым айдос")  # -> "Сәлем менің атым Айдос"
"""

from __future__ import annotations

import re
from pathlib import Path

# Vendored into the Space build context: kk_names.txt sits beside this
# module (not under ../data as in the upstream kazakh_stt layout).
NAMES_FILE = Path(__file__).resolve().parent / "kk_names.txt"

_WORD_RE = re.compile(r"[^\W\d_]+(?:[-’'][^\W\d_]+)*", re.UNICODE)
_SENT_END = set(".!?…")

# Built-in seed of common Kazakh given names (lowercased). Works out of the box;
# data/kk_names.txt is merged on top for user-provided additions.
_SEED_NAMES = {
    # male
    "айдос", "айбек", "айдар", "айтуар", "алихан", "арман", "аслан", "асхат",
    "азамат", "әли", "әмір", "бауыржан", "бекзат", "бекжан", "берік", "дамир",
    "данияр", "дәурен", "дәулет", "диас", "ерасыл", "ерболат", "ержан", "ерлан",
    "ернар", "есен", "жанболат", "жандос", "жасұлан", "қайрат", "қанат", "қуаныш",
    "мақсат", "медет", "мерген", "мұрат", "нұрлан", "нұржан", "нұрсұлтан", "олжас",
    "ринат", "руслан", "санжар", "серік", "сұлтан", "талғат", "темір", "тимур",
    "ұлан", "ілияс",
    # female
    "айгерім", "айгүл", "айжан", "айнұр", "айсұлу", "ақерке", "аружан", "әсел",
    "балжан", "ботагөз", "гүлназ", "гүлнұр", "дана", "динара", "дильназ", "жанар",
    "жансая", "жұлдыз", "зарина", "інжу", "камила", "лаура", "мадина", "молдір",
    "мөлдір", "нұргүл", "сәуле", "томирис", "ұлжан", "әйгерім",
}

_names_cache: set[str] | None = None


def get_names() -> set[str]:
    global _names_cache
    if _names_cache is None:
        names = set(_SEED_NAMES)
        if NAMES_FILE.exists():
            for line in NAMES_FILE.read_text(encoding="utf-8").splitlines():
                w = line.strip().lower()
                if w and not w.startswith("#"):
                    names.add(w)
        _names_cache = names
    return _names_cache


def _cap(word: str) -> str:
    return word[:1].upper() + word[1:] if word else word


def truecase(text: str, names: set[str] | None = None) -> str:
    """Restore sentence + proper-name capitalization."""
    names = names if names is not None else get_names()
    out: list[str] = []
    idx = 0
    need_cap = True  # start of text == start of a sentence
    for m in _WORD_RE.finditer(text):
        gap = text[idx:m.start()]
        out.append(gap)
        if gap and any(ch in _SENT_END for ch in gap):
            need_cap = True
        word = m.group(0)
        if need_cap or word.lower() in names:
            word = _cap(word)
        need_cap = False
        out.append(word)
        idx = m.end()
    out.append(text[idx:])
    return "".join(out)