whisper-translation/freeflow_polish.py

import re


# Python port of the Apache-2.0 FreeFlow polish pipeline ideas:
# https://github.com/mrinalwadhwa/freeflow

SYSTEM_PROMPT_QWEN = """Clean up this dictated text.
Fix punctuation and capitalization.
Return only the cleaned text."""

SYSTEM_PROMPT_MINIMAL = """You are a speech-to-text cleanup assistant. The user dictated text in a non-English language and a speech-to-text engine transcribed it. Your job is to clean up the transcription into polished written text.

Fix filler words, false starts, consecutive repetitions, mid-sentence corrections, punctuation, capitalization, and obvious number formatting. Keep the user's original words. Do not fabricate text. Do not translate. Preserve every <keep>...</keep> block exactly as it appears.

If the transcription is already clean, return it unchanged. Return only the cleaned text."""

SYSTEM_PROMPT_ENGLISH = """You are a speech-to-text cleanup assistant. The user dictated text and a speech-to-text engine transcribed it. Your job is to clean up the transcription into polished written text. If the transcription is already clean, return it unchanged. Do not wrap your output in quotes or add any preamble. Return only the cleaned text. If the input starts with a lowercase letter, keep it lowercase. This happens when the user is continuing a sentence.

Fix punctuation, capitalization, and spelling. Prefer periods, commas, or colons over em-dashes unless an em-dash appears inside a <keep> tag.

When the input already contains explicit punctuation like "!", "?", ":", or ";", preserve it.

When "period" appears at the end of a clause and does not make sense as a noun, treat it as a punctuation command and replace it with ".".

Remove filler words and throat-clearing preambles that add no content. Remove unintentional stuttered repetitions. When the speaker corrects themselves with phrases like "no wait", "actually", "sorry", "I mean", "let me rephrase", "never mind", "or rather", or "make that", drop the abandoned version and keep the final version.

Convert spelled-out numbers to digits where appropriate. Format times, phone numbers, ratios, and currency when the speaker clearly dictated them.

When the speaker mentions 3 or more items, format them as a vertical list. Use numbered items if the order is explicit; otherwise use bullets. Keep 2-item phrases inline.

Keep the speaker's word choices. Do not substitute synonyms, rephrase sentences, expand contractions, or add words the speaker did not say. Preserve informal words such as kinda, gonna, wanna, dunno, lemme, sorta, and gotta.

When "at" appears between a name and a domain, and "dot" appears before the TLD, format it as an email address.

The input may contain <keep>...</keep> tags around symbols inserted by preprocessing. Preserve every <keep>...</keep> block in your output exactly as-is, including the tags themselves."""


PUNCTUATION_RULES = [
    (r"\bnew paragraph\b", "[PAR]", True),
    (r"\bnew line\b", "[NL]", True),
    (r"\bnewline\b", "[NL]", True),
    (r"\bquestion mark\b", "?", False),
    (r"\bexclamation point\b", "!", False),
    (r"\bexclamation mark\b", "!", False),
    (r"\bcomma\b", ",", False),
    (r"\bcolon\b", ":", False),
    (r"\bsemicolon\b", ";", False),
    (r"\bem dash\b", "\u2014", True),
    (r"\ben dash\b", "\u2013", True),
    (r"\bhyphen\b", "-", True),
    (r"\bminus\s+(?:sign|symbol)\b", "-", True),
    (r"\bopen paren(?:t|thesis)?\b", "(", False),
    (r"\bclose paren(?:t|thesis)?\b", ")", False),
    (r"\bopen quote\b", "\u201c", False),
    (r"\b(?:close|end) quote\b", "\u201d", False),
    (r"\bunquote\b", "\u201d", False),
    (r"\b(?:apostrophe|single quote)\b", "'", False),
    (r"\bopen bracket\b", "[", False),
    (r"\bclose bracket\b", "]", False),
    (r"\b(?:open )?angle bracket\b", "<", True),
    (r"\bless[- ]than sign\b", "<", True),
    (r"\bclose angle bracket\b", ">", True),
    (r"\bgreater[- ]than sign\b", ">", True),
    (r"\bdot dot dot\b", "\u2026", True),
    (r"\bellipsis\b", "\u2026", True),
    (r"\b(?:ampersand|and sign|and symbol)\b", "&", True),
    (r"\b(?:at sign|at symbol)\b", "@", True),
    (r"\bhashtag\b", "#", True),
    (r"\b(?:back ?slash|slash en)\b", "\\", True),
    (r"\bforward slash\b", "/", True),
    (r"\b(?:asterisk|asterisk sign)\b", "*", True),
    (r"\bunderscore\b", "_", True),
    (r"\b(?:percent sign|per cent|percentage symbol)\b", "%", True),
    (r"\bdollar sign\b", "$", True),
    (r"\b(?:equals sign|equals symbol)\b", "=", True),
    (r"\b(?:plus sign|plus symbol)\b", "+", True),
    (r"\btrademark sign\b", "\u2122", True),
    (r"\btm\b", "\u2122", True),
    (r"\bcopyright sign\b", "\u00a9", True),
    (r"\bcopyright symbol\b", "\u00a9", True),
    (r"\bdegrees?\s+fahrenheit\b", "\u00b0F", True),
    (r"\bdegrees?\s+f\b", "\u00b0F", True),
    (r"\bdegrees?\s+celsius\b", "\u00b0C", True),
    (r"\bdegrees?\s+centigrade\b", "\u00b0C", True),
    (r"\b(?:degree sign|degree symbol)\b", "\u00b0", True),
]

KNOWN_TERMS = {
    "redis": "Redis",
    "postgresql": "PostgreSQL",
    "postgres": "Postgres",
    "mongodb": "MongoDB",
    "mysql": "MySQL",
    "sqlite": "SQLite",
    "kubernetes": "Kubernetes",
    "docker": "Docker",
    "terraform": "Terraform",
    "nginx": "Nginx",
    "grafana": "Grafana",
    "prometheus": "Prometheus",
    "cloudflare": "Cloudflare",
    "vercel": "Vercel",
    "rabbitmq": "RabbitMQ",
    "kafka": "Kafka",
    "javascript": "JavaScript",
    "typescript": "TypeScript",
    "python": "Python",
    "swift": "Swift",
    "rust": "Rust",
    "react": "React",
    "django": "Django",
    "fastapi": "FastAPI",
    "nextjs": "Next.js",
    "nodejs": "Node.js",
    "numpy": "NumPy",
    "pandas": "Pandas",
    "pytorch": "PyTorch",
    "github": "GitHub",
    "slack": "Slack",
    "figma": "Figma",
    "macos": "macOS",
    "ios": "iOS",
    "iphone": "iPhone",
    "xcode": "Xcode",
    "aws": "AWS",
    "gcp": "GCP",
    "api": "API",
    "sql": "SQL",
    "css": "CSS",
    "html": "HTML",
    "json": "JSON",
    "xml": "XML",
    "yaml": "YAML",
    "url": "URL",
    "cli": "CLI",
    "sdk": "SDK",
}


def ends_at_sentence_boundary(text):
    stripped = (text or "").rstrip()
    return bool(stripped) and stripped[-1] in ".?!"


def _capitalize_after_pattern(text, pattern):
    def repl(match):
        return match.group(1) + match.group(2).upper()

    return re.sub(pattern, repl, text)


def collapse_adjacent_punctuation(text):
    strength = {",": 1, ":": 2, ";": 3, ".": 4, "?": 5, "!": 6}

    def repl(match):
        marks = [ch for ch in match.group(0) if ch in strength]
        return max(marks, key=lambda ch: strength[ch]) if marks else match.group(0)

    return re.sub(r"([.,;:?!])(?:\s*[.,;:?!])+", repl, text)


def strip_noise_phrases(text):
    text = re.sub(r"(?i)\b(?:uh huh|uh-huh|mm hmm|mm-hmm)\b[,.]?\s*", "", text)
    return re.sub(r" {2,}", " ", text).strip()


def strip_filler_sounds(text):
    fillers = "um|eh|mmm|uhh|hm|umm|mm|uh|uhhh|uhm|ah|hmm|mh|ehh"
    text = re.sub(rf"(?i)\b(?:{fillers})\b[,.]?\s*", "", text)
    return re.sub(r" {2,}", " ", text).strip()


def clean_spurious_commas(text):
    text = re.sub(r",{2,}", ",", text)
    text = re.sub(r",\s*(AM|PM|a\.m\.|p\.m\.)", r" \1", text)
    text = re.sub(r",(\s*[.!?])", r"\1", text)
    text = re.sub(r",\s*$", "", text)
    text = re.sub(r",\s*,", ",", text)
    return re.sub(r" {2,}", " ", text).strip()


def substitute_dictated_punctuation(text, preceding_text=None, casual=False):
    result = text or ""
    for pattern, replacement, protect in PUNCTUATION_RULES:
        value = f"<keep>{replacement}</keep>" if protect else replacement
        result = re.sub(pattern, lambda _m, v=value: v, result, flags=re.IGNORECASE)

    result = result.replace("...", "<keep>\u2026</keep>")
    result = re.sub(r" +([.,;:?!)\]\u201d])", r"\1", result)
    result = re.sub(r"([(\[\u201c]) +", r"\1", result)
    result = re.sub(r" {2,}", " ", result)
    result = collapse_adjacent_punctuation(result)
    result = re.sub(r"[,;]\s*(<keep>\[(?:PAR|NL)\]</keep>)", r".\1", result)
    result = re.sub(r"([^.!?\s])\s*(<keep>\[(?:PAR|NL)\]</keep>)", r"\1.\2", result)
    result = re.sub(r" *\n *", "\n", result)

    if casual:
        if result[:1].isupper():
            first_word = re.match(r"[A-Za-z]+", result)
            token = first_word.group(0) if first_word else ""
            if token != "I" and not (len(token) > 1 and token.isupper()):
                result = result[0].lower() + result[1:]
    else:
        result = _capitalize_after_pattern(result, r"([.!?]\s+)(\w)")
        if not (preceding_text and not ends_at_sentence_boundary(preceding_text)):
            result = result[:1].upper() + result[1:] if result[:1].isalpha() else result

    result = clean_spurious_commas(result)
    result = strip_noise_phrases(result)
    result = strip_filler_sounds(result)
    return result.strip()


def strip_keep_tags(text, casual=False):
    result = re.sub(r"<keep>(.*?)</keep>", r"\1", text or "", flags=re.DOTALL)
    result = re.sub(r" *\[PAR\] *", "\n\n", result)
    result = re.sub(r" *\[NL\] *", "\n", result)
    result = re.sub(r" *\u00b6 *", "\n\n", result)
    result = re.sub(r" *\u21b5 *", "\n", result)
    result = re.sub(r" +([.,;:?!)\]>\u201d\u2026%\u2122\u00a9\u00b0])", r"\1", result)
    result = re.sub(r"([(\[\u201c#$<]) +", r"\1", result)
    result = re.sub(r" *([-@/\\_'\u2013\u2014]) +", r"\1", result)
    result = re.sub(r"\*[\w* ]*\*", lambda m: m.group(0).replace(" ", ""), result)
    result = re.sub(r" {2,}", " ", result)
    if not casual:
        result = _capitalize_after_pattern(result, r"(\n)(\w)")
    return result.strip()


def capitalize_known_terms(text):
    result = text
    for source, replacement in KNOWN_TERMS.items():
        result = re.sub(rf"(?<!/)\b{re.escape(source)}\b", replacement, result, flags=re.IGNORECASE)
    return result


def normalize_formatting(text, casual=False):
    result = text or ""
    result = re.sub(r" *\[PAR\] *", "\n\n", result)
    result = re.sub(r" *\[NL\] *", "\n", result)
    result = re.sub(r" *\u00b6 *", "\n\n", result)
    result = re.sub(r" *\u21b5 *", "\n", result)
    result = re.sub(r"(?<!:)//", "/", result)
    result = re.sub(r"(?<=\w)\\\\(?=\w)", r"\\", result)
    result = re.sub(r"\ba\.m\.(?= )", "AM", result)
    result = re.sub(r"\ba\.m\.(?=$|\n)", "AM.", result)
    result = re.sub(r"\bp\.m\.(?= )", "PM", result)
    result = re.sub(r"\bp\.m\.(?=$|\n)", "PM.", result)
    lines = []
    for line in result.split("\n"):
        line = line.rstrip()
        line = re.sub(r"^(\s*)-(\S)", r"\1- \2", line)
        if line.endswith(".") and re.search(r"^\s*(?:-|\d+\.)\s+", line):
            line = line[:-1]
        lines.append(line)
    result = "\n".join(lines)
    return result if casual else capitalize_known_terms(result)


def match_input_casing(text, preprocessed_input, casual=False):
    if casual or not text or not preprocessed_input:
        return text
    if not text[0].isalpha() or not preprocessed_input[0].isalpha():
        return text
    if preprocessed_input[0].islower() and text[0].isupper():
        first_word = re.match(r"[A-Za-z]+", text)
        token = first_word.group(0) if first_word else ""
        if len(token) > 1 and token.isupper():
            return text
        return text[0].lower() + text[1:]
    if preprocessed_input[0].isupper() and text[0].islower():
        return text[0].upper() + text[1:]
    return text


def guard_against_truncation(polished, preprocessed):
    if len(preprocessed or "") < 40:
        return None
    if not polished:
        return preprocessed
    return preprocessed if len(polished) / len(preprocessed) < 0.25 else None


def deterministic_polish(text, preceding_text=None, casual=False):
    substituted = substitute_dictated_punctuation(text, preceding_text=preceding_text, casual=casual)
    stripped = strip_keep_tags(substituted, casual=casual)
    return normalize_formatting(stripped, casual=casual), substituted


def build_user_prompt(text, language=None):
    parts = [f"Transcription:\n{text}"]
    if language:
        parts.append(f"Language: {language}")
    return "\n\n".join(parts)


def system_prompt_for_language(language=None, model=""):
    if "qwen" in (model or "").lower():
        return SYSTEM_PROMPT_QWEN
    if not language or language == "en":
        return SYSTEM_PROMPT_ENGLISH
    return SYSTEM_PROMPT_MINIMAL


def is_clean_enough_to_skip_llm(text):
    normalized = " ".join((text or "").split())
    if not normalized:
        return True
    if normalized[0].isalpha() and not normalized[0].isupper():
        return False
    if not ends_at_sentence_boundary(normalized):
        return False
    if re.search(r"(?i)\b(um|uh|uhh|umm|ah|hmm|you know|i mean|no wait|actually|sorry|let me rephrase)\b", normalized):
        return False
    if re.search(r"(?i)\b(\w+)(?:\s+\1\b){1,}", normalized):
        return False
    return True