tiny-audio / asr_text.py

Promote tiny-audio-granite-qwen-500-3 (Granite NC encoder + fine-tuned Qwen3-0.6B)

26ce2ed verified 20 days ago

5.09 kB

"""Transcript text cleanup: collapse the repetition loops a decoder can fall into.
A language-model decoder that loses the audio sometimes repeats its last
character, word or phrase until it hits the token limit. These functions trim
such a trailing run back to one copy; ASRPipeline applies them to every output.
`truncate_runaway_loop` also catches a loop that starts mid-transcript and
drifts ("It's It's It it's It's"), so it never forms an exact trailing run.
"""

import re

MIN_REPEATS = 3
LOOSE_MIN_REPEATS = 6
TRAILING_CHAR_RE = re.compile(rf"(.)\1{{{MIN_REPEATS - 1},}}$")
TRAILING_WORD_RE = re.compile(rf"\b(\w+)(?:\s+\1){{{MIN_REPEATS - 1},}}\s*$", re.IGNORECASE)
NON_WORD_RE = re.compile(r"\W+")
# A 1-4 word unit repeated this many times in a row, after which at most this
# share of the words are distinct: the decoder lost the audio and never came
# back. Measured on the Open ASR leaderboard runs of 2026-10-16/17: URGENT
# -0.97 WER (13.12 -> 12.11 on an independent rerun), and no set worse in
# either run. Cutting at every 6x run regardless of what follows cost up to
# +0.05 (Monsoon, GigaSpeech): "No, no, no, no, no, no, no, sir, which..." is
# real speech that recovers.
RUNAWAY_MIN_REPEATS = 6
RUNAWAY_MAX_UNIT = 4
RUNAWAY_MAX_DISTINCT = 0.5


def truncate_repetitions(text: str) -> str:
    """Truncate repeated words/phrases/characters at end of text.
    Detects patterns like:
    - Repeated words: "the the the the" -> "the"
    - Repeated phrases: "i am sorry i am sorry i am sorry" -> "i am sorry"
    - Repeated characters: "444444" -> "4"
    Words match ignoring case. A run of `LOOSE_MIN_REPEATS` also matches
    ignoring punctuation, since the model punctuates inside a loop ("Yeah Yeah
    Yeah. Yeah. Yeah,"); shorter punctuated runs are real speech ("No, no,
    no,"), and collapsing those cost 40 errors over 47 rows of the served eval.
    Args:
        text: Input text to process
    Returns:
        Text with trailing repetitions removed
    """
    if not text:
        return text

    text = truncate_runaway_loop(text)
    text = TRAILING_CHAR_RE.sub(r"\1", text)
    while TRAILING_WORD_RE.search(text):
        text = TRAILING_WORD_RE.sub(r"\1", text)

    # 3. Truncate repeated phrases (2-20 words) at end
    # e.g., "i am sorry i am sorry i am sorry" -> "i am sorry"
    words = text.split()
    if len(words) < MIN_REPEATS * 2:
        return text

    # Cheap pre-check: trailing window must contain duplicates for any phrase repeat
    # to be possible. set(window) == window means all unique → no repetition.
    window = words[-MIN_REPEATS * 2 :]
    if len(set(window)) == len(window):
        return truncate_punctuated_loop(text)

    for phrase_len in range(2, min(21, len(words) // MIN_REPEATS + 1)):
        phrase_escaped = re.escape(" ".join(words[-phrase_len:]))
        phrase_pattern = re.compile(
            rf"(^|.*?\s)({phrase_escaped})(?:\s+{phrase_escaped}){{{MIN_REPEATS - 1},}}\s*$",
            re.IGNORECASE,
        )
        match = phrase_pattern.match(text)
        if match:
            text = (match.group(1) + match.group(2)).strip()
            break

    return truncate_punctuated_loop(text)


def truncate_punctuated_loop(text: str) -> str:
    """Collapse a trailing 1-20 word phrase run `LOOSE_MIN_REPEATS`+ times, punctuation ignored."""
    words = text.split()
    keys = [NON_WORD_RE.sub("", w).lower() for w in words]
    end = len(words)
    # Shortest phrase first, so "Yeah Yeah. Yeah," stops at "Yeah".
    for n in range(1, min(21, end // LOOSE_MIN_REPEATS + 1)):
        phrase = keys[end - n : end]
        reps = 1
        while end - (reps + 1) * n >= 0 and keys[end - (reps + 1) * n : end - reps * n] == phrase:
            reps += 1
        if reps >= LOOSE_MIN_REPEATS:
            return " ".join(words[: end - (reps - 1) * n])
    return text


def truncate_runaway_loop(text: str) -> str:
    """Cut at a loop the decoder never leaves, keeping the first copy of its unit.
    Finds the first 1-`RUNAWAY_MAX_UNIT` word unit repeated
    `RUNAWAY_MIN_REPEATS`+ times in a row (case and punctuation ignored). It
    cuts there only when the rest of the transcript, from that run on, has at
    most `RUNAWAY_MAX_DISTINCT` distinct words per word: a loop that drifts
    ("It's It's It it's It's ...") rather than speech that moves on.
    """
    words = text.split()
    keys = [NON_WORD_RE.sub("", w).lower() for w in words]
    for start in range(len(keys)):
        for size in range(1, RUNAWAY_MAX_UNIT + 1):
            unit = keys[start : start + size]
            if len(unit) < size or not all(unit):
                break
            reps = 1
            while keys[start + reps * size : start + (reps + 1) * size] == unit:
                reps += 1
            if reps < RUNAWAY_MIN_REPEATS:
                continue
            rest = keys[start:]
            if len(set(rest)) / len(rest) > RUNAWAY_MAX_DISTINCT:
                return text  # the transcript recovers after the first run
            return " ".join(words[: start + size])
    return text

	"""Transcript text cleanup: collapse the repetition loops a decoder can fall into.

	A language-model decoder that loses the audio sometimes repeats its last
	character, word or phrase until it hits the token limit. These functions trim
	such a trailing run back to one copy; ASRPipeline applies them to every output.
	`truncate_runaway_loop` also catches a loop that starts mid-transcript and
	drifts ("It's It's It it's It's"), so it never forms an exact trailing run.
	"""

	import re

	MIN_REPEATS = 3
	LOOSE_MIN_REPEATS = 6
	TRAILING_CHAR_RE = re.compile(rf"(.)\1{{{MIN_REPEATS - 1},}}$")
	TRAILING_WORD_RE = re.compile(rf"\b(\w+)(?:\s+\1){{{MIN_REPEATS - 1},}}\s*$", re.IGNORECASE)
	NON_WORD_RE = re.compile(r"\W+")
	# A 1-4 word unit repeated this many times in a row, after which at most this
	# share of the words are distinct: the decoder lost the audio and never came
	# back. Measured on the Open ASR leaderboard runs of 2026-10-16/17: URGENT
	# -0.97 WER (13.12 -> 12.11 on an independent rerun), and no set worse in
	# either run. Cutting at every 6x run regardless of what follows cost up to
	# +0.05 (Monsoon, GigaSpeech): "No, no, no, no, no, no, no, sir, which..." is
	# real speech that recovers.
	RUNAWAY_MIN_REPEATS = 6
	RUNAWAY_MAX_UNIT = 4
	RUNAWAY_MAX_DISTINCT = 0.5


	def truncate_repetitions(text: str) -> str:
	"""Truncate repeated words/phrases/characters at end of text.

	Detects patterns like:
	- Repeated words: "the the the the" -> "the"
	- Repeated phrases: "i am sorry i am sorry i am sorry" -> "i am sorry"
	- Repeated characters: "444444" -> "4"

	Words match ignoring case. A run of `LOOSE_MIN_REPEATS` also matches
	ignoring punctuation, since the model punctuates inside a loop ("Yeah Yeah
	Yeah. Yeah. Yeah,"); shorter punctuated runs are real speech ("No, no,
	no,"), and collapsing those cost 40 errors over 47 rows of the served eval.

	Args:
	text: Input text to process

	Returns:
	Text with trailing repetitions removed
	"""
	if not text:
	return text

	text = truncate_runaway_loop(text)
	text = TRAILING_CHAR_RE.sub(r"\1", text)
	while TRAILING_WORD_RE.search(text):
	text = TRAILING_WORD_RE.sub(r"\1", text)

	# 3. Truncate repeated phrases (2-20 words) at end
	# e.g., "i am sorry i am sorry i am sorry" -> "i am sorry"
	words = text.split()
	if len(words) < MIN_REPEATS * 2:
	return text

	# Cheap pre-check: trailing window must contain duplicates for any phrase repeat
	# to be possible. set(window) == window means all unique → no repetition.
	window = words[-MIN_REPEATS * 2 :]
	if len(set(window)) == len(window):
	return truncate_punctuated_loop(text)

	for phrase_len in range(2, min(21, len(words) // MIN_REPEATS + 1)):
	phrase_escaped = re.escape(" ".join(words[-phrase_len:]))
	phrase_pattern = re.compile(
	rf"(^\|.?\s)({phrase_escaped})(?:\s+{phrase_escaped}){{{MIN_REPEATS - 1},}}\s$",
	re.IGNORECASE,
	)
	match = phrase_pattern.match(text)
	if match:
	text = (match.group(1) + match.group(2)).strip()
	break

	return truncate_punctuated_loop(text)


	def truncate_punctuated_loop(text: str) -> str:
	"""Collapse a trailing 1-20 word phrase run `LOOSE_MIN_REPEATS`+ times, punctuation ignored."""
	words = text.split()
	keys = [NON_WORD_RE.sub("", w).lower() for w in words]
	end = len(words)
	# Shortest phrase first, so "Yeah Yeah. Yeah," stops at "Yeah".
	for n in range(1, min(21, end // LOOSE_MIN_REPEATS + 1)):
	phrase = keys[end - n : end]
	reps = 1
	while end - (reps + 1) * n >= 0 and keys[end - (reps + 1) * n : end - reps * n] == phrase:
	reps += 1
	if reps >= LOOSE_MIN_REPEATS:
	return " ".join(words[: end - (reps - 1) * n])
	return text


	def truncate_runaway_loop(text: str) -> str:
	"""Cut at a loop the decoder never leaves, keeping the first copy of its unit.

	Finds the first 1-`RUNAWAY_MAX_UNIT` word unit repeated
	`RUNAWAY_MIN_REPEATS`+ times in a row (case and punctuation ignored). It
	cuts there only when the rest of the transcript, from that run on, has at
	most `RUNAWAY_MAX_DISTINCT` distinct words per word: a loop that drifts
	("It's It's It it's It's ...") rather than speech that moves on.
	"""
	words = text.split()
	keys = [NON_WORD_RE.sub("", w).lower() for w in words]
	for start in range(len(keys)):
	for size in range(1, RUNAWAY_MAX_UNIT + 1):
	unit = keys[start : start + size]
	if len(unit) < size or not all(unit):
	break
	reps = 1
	while keys[start + reps * size : start + (reps + 1) * size] == unit:
	reps += 1
	if reps < RUNAWAY_MIN_REPEATS:
	continue
	rest = keys[start:]
	if len(set(rest)) / len(rest) > RUNAWAY_MAX_DISTINCT:
	return text # the transcript recovers after the first run
	return " ".join(words[: start + size])
	return text