pipeline/swedish.py

"""Swedish text normalization for QC (SPEC §7.3 step 2).

Both the script side and the ASR transcript side are normalized identically:
lowercase, punctuation stripped, digits converted to Swedish number words —
so "300 kr" and "trehundra kronor" tokenize comparably.
"""

from __future__ import annotations

import re

_ONES = [
    "noll", "ett", "två", "tre", "fyra", "fem", "sex", "sju", "åtta", "nio",
    "tio", "elva", "tolv", "tretton", "fjorton", "femton", "sexton", "sjutton",
    "arton", "nitton",
]
_TENS = ["", "", "tjugo", "trettio", "fyrtio", "femtio", "sextio", "sjuttio", "åttio", "nittio"]


def number_to_swedish(n: int) -> str:
    """Spoken Swedish for 0..999999, compounded the way it is said."""
    if n < 0 or n > 999_999:
        return str(n)  # out of scope — leave as digits rather than guess
    if n < 20:
        return _ONES[n]
    if n < 100:
        tens, ones = divmod(n, 10)
        return _TENS[tens] + (_ONES[ones] if ones else "")
    if n < 1000:
        hundreds, rest = divmod(n, 100)
        prefix = ("" if hundreds == 1 else _ONES[hundreds]) + "hundra"
        return prefix + (number_to_swedish(rest) if rest else "")
    thousands, rest = divmod(n, 1000)
    if thousands == 1:
        prefix = "tusen"
    else:
        head = number_to_swedish(thousands)
        # "...ett" + "tusen" would give three t's ("tjugoetttusen"); native Swedish
        # elides one at the boundary ("tjugoettusen").
        prefix = (head[:-1] if head.endswith("ett") else head) + "tusen"
    return prefix + (number_to_swedish(rest) if rest else "")


_DIGIT_RE = re.compile(r"\d+")
# Strip everything that is not a letter (incl. åäö/é), digit or whitespace.
_PUNCT_RE = re.compile(r"[^\w\s]|_", re.UNICODE)


def normalize(text: str) -> str:
    """Lowercase, strip punctuation, convert digit runs to Swedish words."""
    text = text.lower()
    text = _PUNCT_RE.sub(" ", text)
    text = _DIGIT_RE.sub(lambda m: " " + number_to_swedish(int(m.group())) + " ", text)
    return " ".join(text.split())


def tokens(text: str) -> list[str]:
    return normalize(text).split()