feat(api-space): expose Whisper transformers decode confidence metadata.

Return avg_logprob and zlib token compression_ratio from generate with
output_scores; aggregate per chunk in _transcribe_chunked.

Apply confidence_meta rejection for all forced Italian decodes when
signals exist; add ASR_AVG_LOGPROB_REJECT_IT and gate Italian restore
with ASR_IT_RESTORE_ON_FILTER (default off) so filters are not undone.

Co-authored-by: Cursor <cursoragent@cursor.com>

main.py

@@ -14,12 +14,14 @@ import math
 import os
 import re
 import struct
+import zlib
 import threading
 import time
 from typing import Optional
 
 import numpy as np
 import torch
+import torch.nn.functional as F
 import torchaudio
 from fastapi import FastAPI, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
@@ -66,6 +68,8 @@ _NO_SPEECH_THRESHOLD_IT = float(os.environ.get("ASR_NO_SPEECH_THRESHOLD_IT", "0.
 # Lower threshold = reject repetitive / gibberish text sooner (helps tail hallucinations).
 _COMPRESSION_RATIO_THRESHOLD = float(os.environ.get("ASR_COMPRESSION_RATIO_THRESHOLD", "1.6"))
 _LOGPROB_THRESHOLD = float(os.environ.get("ASR_LOGPROB_THRESHOLD", "-1.0"))
+# Transformers forced-Italian: reject when average token logprob is below this (more negative = stricter).
+_ASR_AVG_LOGPROB_REJECT_IT = float(os.environ.get("ASR_AVG_LOGPROB_REJECT_IT", "-0.92"))
 _CPU_THREADS = max(1, int(os.environ.get("ASR_CPU_THREADS", str(os.cpu_count() or 4))))
 _CPU_INTEROP_THREADS = max(1, int(os.environ.get("ASR_CPU_INTEROP_THREADS", "1")))
 _ASR_BACKEND = (os.environ.get("ASR_BACKEND", "faster-whisper" if os.environ.get("ASR_BACKEND") is None else "").strip().lower() or "faster-whisper")
@@ -118,6 +122,10 @@ _ASR_DROP_MICRO_FRAGMENTS = (
 _ASR_BLANK_NEAR_DUP_FORCED = (
     os.environ.get("ASR_BLANK_NEAR_DUP_FORCED", "0").strip().lower() in ("1", "true", "yes")
 )
+# Off by default: restoring filtered Italian output undoes RMS/confidence rejects (e.g. weak tails).
+_ASR_IT_RESTORE_ON_FILTER = (
+    os.environ.get("ASR_IT_RESTORE_ON_FILTER", "0").strip().lower() in ("1", "true", "yes")
+)
 # Match longer repeated tails across consecutive POST /audio responses (chunk boundary stitching).
 _ASR_CROSS_CHUNK_OVERLAP_WORDS = max(3, int(os.environ.get("ASR_CROSS_CHUNK_OVERLAP_WORDS", "12")))
 # One-switch preset for HoloLens-style HTTPS chunking + forced Italian (see ASR/HOLOLENS_TUNING.md).
@@ -776,6 +784,107 @@ if DEVICE.type == "cpu":
         pass
 
 
+def _zlib_compression_ratio_from_token_ids(token_ids: list[int], vocab_size: int) -> float:
+    """Whisper-style zlib compression ratio over raw token bytes (HF generation_whisper)."""
+    if not token_ids:
+        return 1.0
+    vs = max(2, int(vocab_size))
+    length = int(math.log2(vs) / 8) + 1
+    token_bytes = b"".join([int(t).to_bytes(length, byteorder="little", signed=False) for t in token_ids])
+    compressed = zlib.compress(token_bytes)
+    if not compressed:
+        return 1.0
+    return float(len(token_bytes) / len(compressed))
+
+
+def _avg_logprob_from_generate_scores(
+    scores: tuple[torch.Tensor, ...] | None,
+    seq_tokens: torch.Tensor,
+    *,
+    temperature: float = 1.0,
+) -> float | None:
+    """Mirror HF Whisper `_retrieve_avg_logprobs` for greedy/beam batch 0."""
+    if not scores or seq_tokens.numel() == 0:
+        return None
+    stacked = torch.stack([s[0] for s in scores]).float()
+    tokens = seq_tokens.long()
+    if stacked.shape[0] > tokens.shape[0]:
+        stacked = stacked[: tokens.shape[0]]
+    else:
+        tokens = tokens[-stacked.shape[0] :]
+    if stacked.shape[0] == 0:
+        return None
+    rescale = temperature if temperature > 0 else 1.0
+    logprobs = F.log_softmax(stacked * rescale, dim=-1).to(stacked.dtype)
+    sum_lp = sum(logprobs[i][tokens[i]] for i in range(logprobs.shape[0]))
+    return float(sum_lp / logprobs.shape[0])
+
+
+def _whisper_transformers_segment_meta(
+    model: WhisperForConditionalGeneration, gen_out: object
+) -> dict[str, float | None]:
+    """Decode-side confidence from Transformers Whisper `generate` when `return_dict_in_generate=True`."""
+    meta: dict[str, float | None] = {"no_speech_prob": None, "avg_logprob": None, "compression_ratio": None}
+    if gen_out is None or isinstance(gen_out, torch.Tensor):
+        return meta
+    sequences = getattr(gen_out, "sequences", None)
+    if sequences is None and isinstance(gen_out, dict):
+        sequences = gen_out.get("sequences")
+    if sequences is None:
+        return meta
+    try:
+        if sequences.numel() == 0:
+            return meta
+    except AttributeError:
+        return meta
+    seq_row = sequences[0]
+    vocab_size = int(getattr(model.config, "vocab_size", 51866))
+    ids_list = [int(x) for x in seq_row.tolist()]
+    pad_id = getattr(model.generation_config, "pad_token_id", None)
+    eos_id = getattr(model.generation_config, "eos_token_id", None)
+    if pad_id is not None:
+        while ids_list and ids_list[-1] == pad_id:
+            ids_list.pop()
+    try:
+        meta["compression_ratio"] = float(_zlib_compression_ratio_from_token_ids(ids_list, vocab_size))
+    except Exception:
+        meta["compression_ratio"] = None
+
+    scores = getattr(gen_out, "scores", None)
+    if scores is None and isinstance(gen_out, dict):
+        scores = gen_out.get("scores")
+    avg_lp: float | None = None
+    if scores:
+        try:
+            n_gen = len(scores)
+            tail = seq_row[-n_gen:] if seq_row.shape[-1] >= n_gen else seq_row
+            avg_lp = _avg_logprob_from_generate_scores(scores, tail, temperature=1.0)
+        except Exception:
+            avg_lp = None
+    if avg_lp is None:
+        seq_scores = getattr(gen_out, "sequences_scores", None)
+        if seq_scores is None and isinstance(gen_out, dict):
+            seq_scores = gen_out.get("sequences_scores")
+        if seq_scores is not None:
+            try:
+                total_lp = float(seq_scores[0].item())
+                denom = len(scores) if scores else max(1, len(ids_list))
+                avg_lp = total_lp / float(denom)
+            except Exception:
+                avg_lp = None
+    meta["avg_logprob"] = avg_lp
+
+    nsp = getattr(gen_out, "no_speech_prob", None)
+    if nsp is None and isinstance(gen_out, dict):
+        nsp = gen_out.get("no_speech_prob")
+    if nsp is not None:
+        try:
+            meta["no_speech_prob"] = float(nsp[0]) if hasattr(nsp, "__getitem__") else float(nsp)
+        except Exception:
+            meta["no_speech_prob"] = None
+    return meta
+
+
 class WhisperASRSingle:
     def __init__(self, primary_model_name: str, fallback_model_name: str):
         self.primary_model_name = primary_model_name
@@ -855,14 +964,15 @@ class WhisperASRSingle:
         no_speech_threshold: float | None = None,
         compression_ratio_threshold: float | None = None,
         logprob_threshold: float | None = None,
-    ) -> str:
+    ) -> tuple[str, dict[str, float | None]]:
+        meta_empty: dict[str, float | None] = {"no_speech_prob": None, "avg_logprob": None, "compression_ratio": None}
         if max_new_tokens is None:
             max_new_tokens = _WHISPER_MAX_NEW
         if num_beams is None:
             num_beams = _WHISPER_BEAMS
         w = np.asarray(waveform, dtype=np.float32).reshape(-1)
         if w.size == 0 or not math.isfinite(float(np.max(np.abs(w)))):
-            return ""
+            return "", meta_empty
         sr = int(sample_rate)
         t = torch.from_numpy(w).unsqueeze(0)
         if sr != SAMPLE_RATE:
@@ -881,7 +991,7 @@ class WhisperASRSingle:
             )
         except TypeError:
             inputs = self._primary_processor(audio_16k, sampling_rate=SAMPLE_RATE, return_tensors="pt")
-            
+
         input_features = inputs["input_features"].to(DEVICE)
         attention_mask = inputs.get("attention_mask")
         if attention_mask is not None:
@@ -893,7 +1003,8 @@ class WhisperASRSingle:
         requested = (forced_language or "").strip().lower()
         if requested not in {"english", "italian"}:
             requested = "auto"
-        def _decode(model, processor, language: str) -> str:
+
+        def _decode_with_meta(model: WhisperForConditionalGeneration, processor: WhisperProcessor, language: str):
             gen_common: dict = {
                 "max_new_tokens": cap,
                 "num_beams": num_beams,
@@ -911,37 +1022,58 @@ class WhisperASRSingle:
                 gen_common["compression_ratio_threshold"] = compression_ratio_threshold
             if logprob_threshold is not None and no_speech_threshold is not None:
                 gen_common["logprob_threshold"] = logprob_threshold
+
+            meta_kw = {"return_dict_in_generate": True, "output_scores": True}
             try:
-                ids_local = self._generate_ids(model, input_features, gen_common)
+                out = self._generate_ids(model, input_features, {**gen_common, **meta_kw})
             except TypeError:
-                gen_common.pop("condition_on_prev_tokens", None)
-                gen_common.pop("no_speech_threshold", None)
-                gen_common.pop("compression_ratio_threshold", None)
-                gen_common.pop("logprob_threshold", None)
-                gen_common.pop("forced_decoder_ids", None)
-                ids_local = self._generate_ids(model, input_features, gen_common)
-            return _clean_text_artifacts(processor.batch_decode(ids_local, skip_special_tokens=True)[0] or "")
+                try:
+                    out = self._generate_ids(model, input_features, dict(gen_common))
+                except TypeError:
+                    gen_common.pop("condition_on_prev_tokens", None)
+                    gen_common.pop("no_speech_threshold", None)
+                    gen_common.pop("compression_ratio_threshold", None)
+                    gen_common.pop("logprob_threshold", None)
+                    gen_common.pop("forced_decoder_ids", None)
+                    out = self._generate_ids(model, input_features, dict(gen_common))
+
+            if isinstance(out, torch.Tensor):
+                ids_local = out
+                meta = dict(meta_empty)
+            else:
+                ids_local = getattr(out, "sequences", None)
+                if ids_local is None and isinstance(out, dict):
+                    ids_local = out.get("sequences")
+                if ids_local is None:
+                    ids_local = out
+                meta = _whisper_transformers_segment_meta(model, out)
+
+            text = _clean_text_artifacts(processor.batch_decode(ids_local, skip_special_tokens=True)[0] or "")
+            return text, meta
 
         # Explicit single-language decode for speed.
         if requested == "italian":
-            return _decode(self._primary_model, self._primary_processor, "italian")
+            return _decode_with_meta(self._primary_model, self._primary_processor, "italian")
 
         if requested == "english":
             if self._ensure_fallback_loaded() and self._fallback_model is not None and self._fallback_processor is not None:
-                return _decode(self._fallback_model, self._fallback_processor, "english")
-            # Last resort if fallback model unavailable.
-            return _decode(self._primary_model, self._primary_processor, "english")
+                return _decode_with_meta(self._fallback_model, self._fallback_processor, "english")
+            return _decode_with_meta(self._primary_model, self._primary_processor, "english")
 
         # Auto mode: Italian first, English only when needed.
-        text_it = _decode(self._primary_model, self._primary_processor, "italian")
+        text_it, meta_it = _decode_with_meta(self._primary_model, self._primary_processor, "italian")
         should_fallback = _DUAL_DECODE or _should_try_english_fallback(text_it) or (not text_it and _FALLBACK_ON_EMPTY)
         if not should_fallback:
-            return text_it
+            return text_it, meta_it
 
         if not self._ensure_fallback_loaded() or self._fallback_model is None or self._fallback_processor is None:
-            return text_it
-        text_en = _decode(self._fallback_model, self._fallback_processor, "english")
-        return _pick_best_text(text_it, text_en)
+            return text_it, meta_it
+        text_en, meta_en = _decode_with_meta(self._fallback_model, self._fallback_processor, "english")
+        chosen = _pick_best_text(text_it, text_en)
+        en_st = (text_en or "").strip()
+        if en_st and chosen == en_st:
+            return chosen, meta_en
+        return chosen, meta_it
 
 
 class FasterWhisperASR:
@@ -1279,6 +1411,12 @@ def _forced_italian_speech_rms(audio_16k: np.ndarray) -> tuple[float, float]:
     return speech_rms, p90
 
 
+def _coerce_transcribe_result(out: tuple[str, dict] | str) -> tuple[str, dict[str, float | None]]:
+    if isinstance(out, tuple) and len(out) >= 2 and isinstance(out[1], dict):
+        return (out[0] or ""), out[1]
+    return (out if isinstance(out, str) else str(out or "")), {}
+
+
 def _transcribe_chunked(
     arr: np.ndarray,
     forced_language: str,
@@ -1316,31 +1454,44 @@ def _transcribe_chunked(
     prev = ""
     conf_no_speech: list[float] = []
     conf_logprob: list[float] = []
+    conf_compression: list[float] = []
     for seg in segments:
         if forced_language == "italian" and _ASR_ITALIAN_BACKEND == "transformers":
-            text = get_italian_pipe_transformers().transcribe(
-                seg,
-                SAMPLE_RATE,
-                forced_language="italian",
-                max_new_tokens=max_new_tokens,
-                num_beams=num_beams,
-                no_speech_threshold=ns_thr,
-                compression_ratio_threshold=_COMPRESSION_RATIO_THRESHOLD,
-                logprob_threshold=_LOGPROB_THRESHOLD,
+            text, seg_meta = _coerce_transcribe_result(
+                get_italian_pipe_transformers().transcribe(
+                    seg,
+                    SAMPLE_RATE,
+                    forced_language="italian",
+                    max_new_tokens=max_new_tokens,
+                    num_beams=num_beams,
+                    no_speech_threshold=ns_thr,
+                    compression_ratio_threshold=_COMPRESSION_RATIO_THRESHOLD,
+                    logprob_threshold=_LOGPROB_THRESHOLD,
+                )
             )
         else:
             pipe = get_pipe()
-            text = pipe.transcribe(
-                seg,
-                SAMPLE_RATE,
-                forced_language=forced_language,
-                max_new_tokens=max_new_tokens,
-                num_beams=num_beams,
-                no_speech_threshold=ns_thr,
-                compression_ratio_threshold=_COMPRESSION_RATIO_THRESHOLD,
-                logprob_threshold=_LOGPROB_THRESHOLD,
+            text, seg_meta = _coerce_transcribe_result(
+                pipe.transcribe(
+                    seg,
+                    SAMPLE_RATE,
+                    forced_language=forced_language,
+                    max_new_tokens=max_new_tokens,
+                    num_beams=num_beams,
+                    no_speech_threshold=ns_thr,
+                    compression_ratio_threshold=_COMPRESSION_RATIO_THRESHOLD,
+                    logprob_threshold=_LOGPROB_THRESHOLD,
+                )
             )
-            # Confidence metadata is populated when backend exposes it without extra decode.
+        nsp = seg_meta.get("no_speech_prob")
+        if nsp is not None:
+            conf_no_speech.append(float(nsp))
+        alp = seg_meta.get("avg_logprob")
+        if alp is not None:
+            conf_logprob.append(float(alp))
+        cr = seg_meta.get("compression_ratio")
+        if cr is not None:
+            conf_compression.append(float(cr))
         t = (text or "").strip()
         if not t:
             continue
@@ -1355,10 +1506,11 @@ def _transcribe_chunked(
         prev = t
 
     text_out = " ".join(out_parts).strip()
+    zlib_cr = max(conf_compression) if conf_compression else None
     conf = {
         "no_speech_prob": (sum(conf_no_speech) / len(conf_no_speech)) if conf_no_speech else None,
         "avg_logprob": (sum(conf_logprob) / len(conf_logprob)) if conf_logprob else None,
-        "compression_ratio": _text_compression_ratio(text_out),
+        "compression_ratio": float(zlib_cr) if zlib_cr is not None else _text_compression_ratio(text_out),
     }
     return text_out, conf
 
@@ -1571,6 +1723,8 @@ def startup() -> None:
         f"ASR_MAX_CHUNK_SECONDS_FORCED={_ASR_MAX_CHUNK_S_FORCED} ASR_MAX_CHUNK_SECONDS_AUTO={_ASR_MAX_CHUNK_S_AUTO} "
         f"ASR_CHUNK_OVERLAP_SECONDS={_ASR_CHUNK_OVERLAP_S} "
         f"ASR_NO_SPEECH_THRESHOLD_IT={_NO_SPEECH_THRESHOLD_IT} "
+        f"ASR_AVG_LOGPROB_REJECT_IT={_ASR_AVG_LOGPROB_REJECT_IT} "
+        f"ASR_IT_RESTORE_ON_FILTER={_ASR_IT_RESTORE_ON_FILTER} "
         f"ASR_MIN_SPEECH_RMS_FORCED_ITALIAN_ACCEPT={_ASR_MIN_SPEECH_RMS_FORCED_ITALIAN_ACCEPT} "
         f"ASR_MIN_SPEECH_RMS_FORCED_ITALIAN={_ASR_MIN_SPEECH_RMS_FORCED_ITALIAN} "
         f"ASR_MIN_SPEECH_RMS_FORCED_ITALIAN_SHORT={_ASR_MIN_SPEECH_RMS_FORCED_ITALIAN_SHORT} "
@@ -1839,22 +1993,31 @@ async def audio(request: Request) -> Response:
                 f"speech_rms={speech_rms:.5f} threshold={min_speech_thr:.5f} chunk_s={duration_s:.2f}"
             )
             raw_text = ""
-        if raw_text and short_mode_forced_it:
+        # Decode-confidence gates (Transformers Whisper now supplies avg_logprob / zlib token compression).
+        if raw_text:
             no_speech_prob = decode_conf.get("no_speech_prob")
             avg_logprob = decode_conf.get("avg_logprob")
             comp = decode_conf.get("compression_ratio")
-            weak_short = (
-                (no_speech_prob is not None and float(no_speech_prob) > 0.55)
-                or (avg_logprob is not None and float(avg_logprob) < -0.90)
-                or (comp is not None and float(comp) > (_COMPRESSION_RATIO_THRESHOLD + 0.25))
+            has_signal = no_speech_prob is not None or avg_logprob is not None or (
+                comp is not None and float(comp) > float(_COMPRESSION_RATIO_THRESHOLD)
             )
-            if weak_short:
-                print(
-                    f"[reject] lang=italian reason=short_confidence_meta "
-                    f"chunk_s={duration_s:.2f} speech_rms={speech_rms:.5f} "
-                    f"no_speech_prob={no_speech_prob} avg_logprob={avg_logprob} compression_ratio={comp}"
+            if has_signal:
+                weak_decode = (
+                    (no_speech_prob is not None and float(no_speech_prob) > 0.55)
+                    or (
+                        avg_logprob is not None
+                        and float(avg_logprob) < float(_ASR_AVG_LOGPROB_REJECT_IT)
+                    )
+                    or (comp is not None and float(comp) > (float(_COMPRESSION_RATIO_THRESHOLD) + 0.2))
                 )
-                raw_text = ""
+                if weak_decode:
+                    print(
+                        f"[reject] lang=italian reason=confidence_meta "
+                        f"chunk_s={duration_s:.2f} speech_rms={speech_rms:.5f} "
+                        f"no_speech_prob={no_speech_prob} avg_logprob={avg_logprob} "
+                        f"compression_ratio={comp}"
+                    )
+                    raw_text = ""
         # Weak-tail after strong previous speech: reject tiny tails irrespective of token content.
         if raw_text:
             now_ts = time.time()
@@ -1918,10 +2081,10 @@ async def audio(request: Request) -> Response:
         raw_text, rms, duration_s
     ):
         raw_text = ""
-    # Italian fallback only when filters blanked output but ASR text is not flagged as hallucination.
-    # (Previously we always restored here and undid Italian hallucination drops.)
+    # Optional restore: defaults off — otherwise RMS/confidence rejects get undone (e.g. weak tails).
     if (
-        forced_language == "italian"
+        _ASR_IT_RESTORE_ON_FILTER
+        and forced_language == "italian"
         and not raw_text
         and raw_text_original
         and rms >= 0.009