Tune and harden S2S generation

README.md

@@ -20,7 +20,7 @@ The model selector defaults to the 30B-A3B and can switch to 2B model.
 
 The demo includes **audio understanding, speech recognition, speech translation, text reasoning, text-to-speech, and speech-to-speech**.
 
-Reasoning has no separate token cap by default for either model. `Max new tokens` sets a total cap shared by reasoning and the final answer.
+Reasoning has no separate token cap by default for most tasks. Speech-to-speech defaults to a 1,536-token reasoning budget so the model has room for its final answer. `Max new tokens` sets a total cap shared by reasoning and the final answer.
 
 The hosted ZeroGPU runtime uses a prebuilt `mamba-ssm==2.3.2.post1` Blackwell wheel.
 

app.py

@@ -150,7 +150,7 @@ TTS_MAX_NEW_TOKENS = min(
     MAX_NEW_TOKENS,
 )
 S2S_TEXT_MAX_NEW_TOKENS = int(os.environ.get("AUDEX_S2S_TEXT_MAX_NEW_TOKENS", "2048"))
-S2S_REASONING_BUDGET = int(os.environ.get("AUDEX_S2S_REASONING_BUDGET", "0"))
+S2S_REASONING_BUDGET = int(os.environ.get("AUDEX_S2S_REASONING_BUDGET", "1536"))
 S2S_GPU_DURATION_SECONDS = int(os.environ.get("AUDEX_S2S_GPU_DURATION_SECONDS", "120"))
 S2S_SPOKEN_MAX_WORDS = 20
 S2S_TTS_MAX_NEW_TOKENS = int(os.environ.get("AUDEX_S2S_TTS_MAX_NEW_TOKENS", "2400"))
@@ -1209,8 +1209,10 @@ def _split_sentence_segments(text: str) -> list[str]:
 def _clean_transcription(text: str) -> str:
     text = text.strip()
     quoted = re.fullmatch(
-        r"(?is)(?:(?:the\s+)?(?:transcription|transcript)"
-        r"(?:\s+(?:is|reads))?\s*:?\s*)?(['\"])(.*)\1[.!]?",
+        r"(?is)(?:(?:(?:(?:source\s+)?language)\s*:[^.\n]+\.\s*)?"
+        r"(?:(?:the\s+)?(?:transcription|transcript)(?:\s+(?:is|reads))?"
+        r"|(?:the\s+)?(?:spoken\s+)?content\s+of\s+the\s+(?:input\s+)?audio\s+is)"
+        r"\s*:?\s*)?(['\"])(.*)\1[.!]?",
         text,
     )
     return quoted.group(2).strip() if quoted else text
@@ -1600,7 +1602,8 @@ with gr.Blocks(title="Nemotron-Labs-Audex") as demo:
                 gr.Markdown(
                     "Task selection restores the official defaults. ASR and AST use "
                     "greedy decoding; audio understanding and text reasoning use sampling. "
-                    "Reasoning has no separate cap unless you set one. TTS uses CFG 2.0; "
+                    "Reasoning has no separate cap unless you set one; speech-to-speech "
+                    "defaults to a 1,536-token reasoning budget. TTS uses CFG 2.0; "
                     "speech-to-speech uses CFG 1.5."
                 )
         with gr.Column():