diff --git a/whisperlivekit/simul_whisper/align_att_base.py b/whisperlivekit/simul_whisper/align_att_base.py index fd326c7..910be1d 100644 --- a/whisperlivekit/simul_whisper/align_att_base.py +++ b/whisperlivekit/simul_whisper/align_att_base.py @@ -120,6 +120,7 @@ class AlignAttBase(ABC): self.state.segments = [] self.state.log_segments += 1 self.state.pending_incomplete_tokens = [] + self.state.pending_retries = 0 def segments_len(self): return sum(s.shape[0] for s in self.state.segments) / 16000 @@ -223,6 +224,7 @@ class AlignAttBase(ABC): new_segment = False logits = self._apply_token_suppression(logits) + logits = self._apply_dry_penalty(logits, current_tokens) current_tokens, completed = self._update_tokens( current_tokens, logits, sum_logprobs ) @@ -326,9 +328,13 @@ class AlignAttBase(ABC): for word, word_tokens in zip(split_words, split_tokens): if replacement_char in word: - logger.warning(f"[UTF-8 Filter] Skipping: {repr(word)}") - timestamp_idx += len(word_tokens) - continue + cleaned = word.replace(replacement_char, "") + if not cleaned.strip(): + logger.debug(f"[UTF-8 Filter] Skipping: {repr(word)}") + timestamp_idx += len(word_tokens) + continue + logger.debug(f"[UTF-8 Filter] Cleaned {repr(word)} -> {repr(cleaned)}") + word = cleaned try: current_timestamp = l_absolute_timestamps[timestamp_idx] @@ -354,21 +360,84 @@ class AlignAttBase(ABC): def _handle_pending_tokens(self, split_words, split_tokens): """Handle incomplete UTF-8 tokens for next chunk.""" - self.state.pending_incomplete_tokens = [] MAX_PENDING_TOKENS = 10 + MAX_PENDING_RETRIES = 2 replacement_char = "\ufffd" + if split_words and replacement_char in split_words[-1]: - if len(split_tokens[-1]) <= MAX_PENDING_TOKENS: + self.state.pending_retries += 1 + if self.state.pending_retries > MAX_PENDING_RETRIES: + logger.warning( + f"[UTF-8 Fix] Dropping {len(split_tokens[-1])} incomplete tokens " + f"after {MAX_PENDING_RETRIES} retries (won't resolve)" + ) + self.state.pending_incomplete_tokens = [] + self.state.pending_retries = 0 + elif len(split_tokens[-1]) <= MAX_PENDING_TOKENS: self.state.pending_incomplete_tokens = split_tokens[-1] logger.debug( f"[UTF-8 Fix] Holding {len(self.state.pending_incomplete_tokens)} " - f"incomplete tokens for next chunk" + f"incomplete tokens for next chunk (retry {self.state.pending_retries})" ) else: logger.warning( f"[UTF-8 Fix] Skipping {len(split_tokens[-1])} tokens " f"(exceeds limit of {MAX_PENDING_TOKENS}, likely hallucination)" ) + self.state.pending_incomplete_tokens = [] + self.state.pending_retries = 0 + else: + self.state.pending_incomplete_tokens = [] + self.state.pending_retries = 0 + + # === Repetition penalty === + + def _apply_dry_penalty(self, logits, current_tokens): + """DRY penalty v0: penalize tokens that would extend a verbatim repetition. + See https://github.com/oobabooga/text-generation-webui/pull/5677 + + Scans the decoded sequence for positions where the current suffix already + appeared --> for each such match, the token that followed it in the past is + penalised exponentially with the match length + """ + eot = self.tokenizer.eot + seq = current_tokens[0].tolist() + if len(seq) < 5: + return logits + + last = seq[-1] + if last >= eot: + return logits + + penalties = {} + for i in range(len(seq) - 2, -1, -1): + if seq[i] != last: + continue + next_tok = seq[i + 1] + if next_tok >= eot: + continue + + length = 1 + while length < 50: + j, k = i - length, len(seq) - 1 - length + if j < 0 or k <= i: + break + if seq[j] != seq[k] or seq[j] >= eot: + break + length += 1 + + if next_tok not in penalties or length > penalties[next_tok]: + penalties[next_tok] = length + + if penalties: + max_len = max(penalties.values()) + if max_len >= 4: + logger.debug(f"[DRY] penalising {len(penalties)} tokens (longest match: {max_len})") + for tok, length in penalties.items(): + if length >= 2: + logits[:, tok] = logits[:, tok] - 1.0 * 2.0 ** (length - 2) + + return logits # === Abstract methods — subclass must implement === diff --git a/whisperlivekit/simul_whisper/decoder_state.py b/whisperlivekit/simul_whisper/decoder_state.py index 428e798..44076a3 100644 --- a/whisperlivekit/simul_whisper/decoder_state.py +++ b/whisperlivekit/simul_whisper/decoder_state.py @@ -25,7 +25,8 @@ class DecoderState: context: Any = None pending_incomplete_tokens: List[int] = field(default_factory=list) - + pending_retries: int = 0 + global_time_offset: float = 0.0 cumulative_time_offset: float = 0.0 first_timestamp: Optional[float] = None @@ -78,8 +79,9 @@ class DecoderState: self.last_attend_frame = -rewind_threshold self.cumulative_time_offset = 0.0 self.pending_incomplete_tokens = [] + self.pending_retries = 0 self.log_segments += 1 - + def full_reset(self, rewind_threshold: int = 200): """ Full reset including audio segments and tokens. diff --git a/whisperlivekit/simul_whisper/mlx/decoder_state.py b/whisperlivekit/simul_whisper/mlx/decoder_state.py index 332bc8c..5065b4c 100644 --- a/whisperlivekit/simul_whisper/mlx/decoder_state.py +++ b/whisperlivekit/simul_whisper/mlx/decoder_state.py @@ -29,7 +29,8 @@ class MLXDecoderState: context: Any = None pending_incomplete_tokens: List[int] = field(default_factory=list) - + pending_retries: int = 0 + global_time_offset: float = 0.0 cumulative_time_offset: float = 0.0 first_timestamp: Optional[float] = None @@ -59,8 +60,9 @@ class MLXDecoderState: self.last_attend_frame = -rewind_threshold self.cumulative_time_offset = 0.0 self.pending_incomplete_tokens = [] + self.pending_retries = 0 self.log_segments += 1 - + def full_reset(self, rewind_threshold: int = 200): """ Full reset including audio segments and tokens.